diff --git a/.gitignore b/.gitignore index be75938ec401b1d72fa54773c85191aaac7d7f35..828bbe9bd3363853ae3f58f54a8d5f60cefad837 100644 --- a/.gitignore +++ b/.gitignore @@ -27,6 +27,7 @@ Podfile.lock /tensorflow/contrib/lite/examples/ios/simple/data/*.txt /tensorflow/contrib/lite/examples/ios/simple/data/*.tflite xcuserdata/** +/api_init_files_list.txt # Android .gradle diff --git a/CODEOWNERS b/CODEOWNERS index 007a304c3e706ce968576ec8979c08f1a3bcc552..b9f0313cc6d59d3fbdcd014e1a528126d863075a 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -45,7 +45,7 @@ # /tensorflow/contrib/session_bundle/ @nfiedel @sukritiramesh # /tensorflow/contrib/slim/ @sguada @thenbasilmanran # /tensorflow/contrib/stateless/ @girving -# /tensorflow/contrib/tensor_forest/ @gilberthendry @thomascolthurst +# /tensorflow/contrib/tensor_forest/ @gilberthendry @thomascolthurst @yupbank # /tensorflow/contrib/testing/ @dandelionmane # /tensorflow/contrib/timeseries/ @allenlavoie # /tensorflow/contrib/tpu/ @frankchn @saeta @jhseu diff --git a/README.md b/README.md index 99f4a253d99418defd7e8b21e63c03935760a4f5..e1a50c87e26d493ba3ac760f357905d89aa40dab 100644 --- a/README.md +++ b/README.md @@ -7,14 +7,14 @@ | **`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://storage.cloud.google.com/tensorflow-kokoro-build-badges/ubuntu-cc.png) | ![Build Status](https://storage.cloud.google.com/tensorflow-kokoro-build-badges/ubuntu-gpu-cc.png) | ![Build Status](https://storage.cloud.google.com/tensorflow-kokoro-build-badges/macos-py2-cc.png) | [![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](https://img.shields.io/badge/api-reference-blue.svg)](https://www.tensorflow.org/api_docs/) | ![Build Status](https://storage.googleapis.com/tensorflow-kokoro-build-badges/ubuntu-cc.png) | ![Build Status](https://storage.googleapis.com/tensorflow-kokoro-build-badges/ubuntu-gpu-cc.png) | ![Build Status](https://storage.googleapis.com/tensorflow-kokoro-build-badges/macos-py2-cc.png) | [![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 the graph edges represent the multidimensional data arrays (tensors) that flow between them. This flexible architecture enables you to deploy computation to one or more CPUs or GPUs in a desktop, server, or mobile device without rewriting -code. TensorFlow also includes TensorBoard, a data visualization toolkit. +code. TensorFlow also includes [TensorBoard](https://www.tensorflow.org/programmers_guide/summaries_and_tensorboard), a data visualization toolkit. TensorFlow was originally developed by researchers and engineers working on the Google Brain team within Google's Machine Intelligence Research diff --git a/WORKSPACE b/WORKSPACE index 11c5cdb2070e79b16540a39f13cab28608962340..4ddfb9a3832ea1ea639ace887e1d601bdd857086 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -2,11 +2,11 @@ workspace(name = "org_tensorflow") http_archive( name = "io_bazel_rules_closure", - sha256 = "6691c58a2cd30a86776dd9bb34898b041e37136f2dc7e24cadaeaf599c95c657", - strip_prefix = "rules_closure-08039ba8ca59f64248bb3b6ae016460fe9c9914f", + sha256 = "a38539c5b5c358548e75b44141b4ab637bba7c4dc02b46b1f62a96d6433f56ae", + strip_prefix = "rules_closure-dbb96841cc0a5fb2664c37822803b06dab20c7d1", urls = [ - "https://mirror.bazel.build/github.com/bazelbuild/rules_closure/archive/08039ba8ca59f64248bb3b6ae016460fe9c9914f.tar.gz", - "https://github.com/bazelbuild/rules_closure/archive/08039ba8ca59f64248bb3b6ae016460fe9c9914f.tar.gz", # 2018-01-16 + "https://mirror.bazel.build/github.com/bazelbuild/rules_closure/archive/dbb96841cc0a5fb2664c37822803b06dab20c7d1.tar.gz", + "https://github.com/bazelbuild/rules_closure/archive/dbb96841cc0a5fb2664c37822803b06dab20c7d1.tar.gz", # 2018-04-13 ], ) diff --git a/configure.py b/configure.py index 8fb8979111627b9b25be80c77c611932880e011d..b745e374a2baaffec73f9f9382e1bab322e7f0fd 100644 --- a/configure.py +++ b/configure.py @@ -226,8 +226,6 @@ def setup_python(environ_cp): # Set-up env variables used by python_configure.bzl write_action_env_to_bazelrc('PYTHON_BIN_PATH', python_bin_path) write_action_env_to_bazelrc('PYTHON_LIB_PATH', python_lib_path) - write_to_bazelrc('build --force_python=py%s' % python_major_version) - write_to_bazelrc('build --host_force_python=py%s' % python_major_version) write_to_bazelrc('build --python_path=\"%s"' % python_bin_path) environ_cp['PYTHON_BIN_PATH'] = python_bin_path diff --git a/tensorflow/BUILD b/tensorflow/BUILD index cfafffdd13012ccc29272df215bded1e6f99dd80..f2ad16fa04f5beb6616c58c28d0f0c460c3e3a17 100644 --- a/tensorflow/BUILD +++ b/tensorflow/BUILD @@ -450,11 +450,12 @@ tf_cc_shared_object( linkstatic = 1, visibility = ["//visibility:public"], deps = [ + "//tensorflow/core:core_cpu_impl", "//tensorflow/core:framework_internal_impl", + "//tensorflow/core:gpu_runtime_impl", + "//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry_impl", "//tensorflow/core:lib_internal_impl", - "//tensorflow/core:core_cpu_impl", "//tensorflow/stream_executor:stream_executor_impl", - "//tensorflow/core:gpu_runtime_impl", ] + tf_additional_binary_deps(), ) diff --git a/tensorflow/c/c_api.h b/tensorflow/c/c_api.h index fe85f8ee0ed2c58c3ba9201a9ca895c9ec48c022..c8594347451dffd465d7fa926cc53818dc9e38d4 100644 --- a/tensorflow/c/c_api.h +++ b/tensorflow/c/c_api.h @@ -72,7 +72,7 @@ limitations under the License. #ifdef SWIG #define TF_CAPI_EXPORT #else -#if defined(COMPILER_MSVC) +#if defined(_WIN32) #ifdef TF_COMPILE_LIBRARY #define TF_CAPI_EXPORT __declspec(dllexport) #else @@ -80,7 +80,7 @@ limitations under the License. #endif // TF_COMPILE_LIBRARY #else #define TF_CAPI_EXPORT __attribute__((visibility("default"))) -#endif // COMPILER_MSVC +#endif // _WIN32 #endif // SWIG #ifdef __cplusplus diff --git a/tensorflow/c/c_api_experimental.cc b/tensorflow/c/c_api_experimental.cc index 581f5743eb7057f72aa7d4e8943e0baded09ae8b..d3916bc16778a942b7eab4df93bbc19955b19e31 100644 --- a/tensorflow/c/c_api_experimental.cc +++ b/tensorflow/c/c_api_experimental.cc @@ -191,6 +191,12 @@ library { // be deleted by calling TF_DeleteFunction. static std::vector CreateImagenetDatasetFunctions( const char* file_path, std::string* dataset_name, TF_Status* status) { +#if defined(PLATFORM_WINDOWS) + status->status = tensorflow::errors::Unimplemented( + "TF_MakeFileBasedIteratorGetNextWithDatasets in the experimental C API " + "is not implemented for Windows"); + return std::vector(); +#else const char* func_def = R"PREFIX( library { function { @@ -7069,6 +7075,7 @@ library { DCHECK(found); }; return CreateFunctionsFromTextProto(func_def, &mutate_proto_func, status); +#endif } #endif @@ -7080,6 +7087,12 @@ library { static std::vector CreateMNISTDatasetFunctions( const char* file_path, int batch_size, std::string* dataset_name, TF_Status* status) { +#if defined(PLATFORM_WINDOWS) + status->status = tensorflow::errors::Unimplemented( + "TF_MakeFileBasedIteratorGetNextWithDatasets in the experimental C API " + "is not implemented for Windows"); + return nullptr; +#else const char* func_def = R"PREFIX( library { function { @@ -8209,6 +8222,7 @@ library { DCHECK(found_batch_size); }; return CreateFunctionsFromTextProto(func_def, &mutate_proto_func, status); +#endif } #endif diff --git a/tensorflow/c/c_api_experimental.h b/tensorflow/c/c_api_experimental.h index 666342974ee0a51b707955cf7468e914fead85b3..88cb173cd25f4219e32392f6722a6ea7d358a553 100644 --- a/tensorflow/c/c_api_experimental.h +++ b/tensorflow/c/c_api_experimental.h @@ -35,7 +35,7 @@ limitations under the License. #ifdef SWIG #define TF_CAPI_EXPORT #else -#if defined(COMPILER_MSVC) +#if defined(_WIN32) #ifdef TF_COMPILE_LIBRARY #define TF_CAPI_EXPORT __declspec(dllexport) #else @@ -43,7 +43,7 @@ limitations under the License. #endif // TF_COMPILE_LIBRARY #else #define TF_CAPI_EXPORT __attribute__((visibility("default"))) -#endif // COMPILER_MSVC +#endif // _WIN32 #endif // SWIG #ifdef __cplusplus diff --git a/tensorflow/c/c_api_test.cc b/tensorflow/c/c_api_test.cc index ca80db23ed3ccbbdc49c61db6cd03ff735470512..9b86425aa5fbc2be2872b3f5d2809eaa844f9d68 100644 --- a/tensorflow/c/c_api_test.cc +++ b/tensorflow/c/c_api_test.cc @@ -1700,7 +1700,7 @@ TEST_F(CApiGradientsTest, OpWithNoGradientRegistered_NoGradInputs) { TestGradientsError(false); } -// REGISTER_OP for CApiTestAttributesTest test cases. +// REGISTER_OP for CApiAttributesTest test cases. // Registers two ops, each with a single attribute called 'v'. // The attribute in one op will have a type 'type', the other // will have list(type). diff --git a/tensorflow/c/eager/BUILD b/tensorflow/c/eager/BUILD index a2d96357ac8a55be7fe03bf58e33ff1733967dd1..fae922ea3b484f668c3abaf0a1ae30235c5a2d2b 100644 --- a/tensorflow/c/eager/BUILD +++ b/tensorflow/c/eager/BUILD @@ -31,7 +31,6 @@ tf_cuda_library( "//tensorflow/core/common_runtime/eager:context", "//tensorflow/core/common_runtime/eager:eager_executor", "//tensorflow/core/common_runtime/eager:execute", - "//tensorflow/core/common_runtime/eager:execute_node", "//tensorflow/core/common_runtime/eager:kernel_and_device", "//tensorflow/core/common_runtime/eager:tensor_handle", "//tensorflow/core/common_runtime/eager:copy_to_device_node", @@ -41,6 +40,8 @@ tf_cuda_library( "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", + # TODO(b/74620627): move this here + "//tensorflow/python:cpp_shape_inference_proto_cc", ], }) + select({ "//tensorflow:with_xla_support": [ @@ -49,6 +50,7 @@ tf_cuda_library( ], "//conditions:default": [], }) + [ + "//tensorflow/core/common_runtime/eager:eager_operation", "//tensorflow/core:gpu_runtime", ], ) @@ -71,6 +73,7 @@ tf_cuda_library( "//tensorflow/core:lib_internal", "//tensorflow/core/common_runtime/eager:context", "//tensorflow/core/common_runtime/eager:eager_executor", + "//tensorflow/core/common_runtime/eager:eager_operation", "//tensorflow/core/common_runtime/eager:kernel_and_device", "//tensorflow/core/common_runtime/eager:tensor_handle", ], diff --git a/tensorflow/c/eager/c_api.cc b/tensorflow/c/eager/c_api.cc index c96a38dec3ed7bcbbd77415ec3b158390def797e..975bde7c7f30697ecdacfafdf9f1cd610b1f68d0 100644 --- a/tensorflow/c/eager/c_api.cc +++ b/tensorflow/c/eager/c_api.cc @@ -34,7 +34,6 @@ limitations under the License. #include "tensorflow/core/common_runtime/device_set.h" #include "tensorflow/core/common_runtime/eager/copy_to_device_node.h" #include "tensorflow/core/common_runtime/eager/execute.h" -#include "tensorflow/core/common_runtime/eager/execute_node.h" #include "tensorflow/core/common_runtime/function.h" #include "tensorflow/core/common_runtime/rendezvous_mgr.h" #include "tensorflow/core/framework/node_def_util.h" @@ -49,6 +48,7 @@ limitations under the License. #include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/thread_annotations.h" #include "tensorflow/core/public/version.h" +#include "tensorflow/python/framework/cpp_shape_inference.pb.h" using tensorflow::int64; using tensorflow::string; @@ -116,9 +116,7 @@ TFE_Context* TFE_NewContext(const TFE_ContextOptions* opts, TF_Status* status) { opts->async, std::move(device_mgr), r); } -void TFE_DeleteContext(TFE_Context* ctx, TF_Status* status) { - delete ctx; -} +void TFE_DeleteContext(TFE_Context* ctx, TF_Status* status) { delete ctx; } TF_DeviceList* TFE_ContextListDevices(TFE_Context* ctx, TF_Status* status) { TF_DeviceList* list = new TF_DeviceList; @@ -220,9 +218,6 @@ TF_Tensor* TFE_TensorHandleResolve(TFE_TensorHandle* h, TF_Status* status) { } return retval; } -} // extern "C" - -extern "C" { TFE_Op* TFE_NewOp(TFE_Context* ctx, const char* op_or_function_name, TF_Status* status) { @@ -242,21 +237,18 @@ TFE_Op* TFE_NewOp(TFE_Context* ctx, const char* op_or_function_name, void TFE_DeleteOp(TFE_Op* op) { delete op; } 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->context.FindDeviceByName(device_name, &d); - } - op->device = d; + status->status = op->operation.SetDevice(device_name); } const char* TFE_OpGetDevice(TFE_Op* op, TF_Status* status) { - tensorflow::Device* device = - (op->device == nullptr) ? op->ctx->context.HostCPU() : op->device; + tensorflow::Device* device = (op->operation.Device() == nullptr) + ? op->operation.EagerContext()->HostCPU() + : op->operation.Device(); return device->name().c_str(); } void TFE_OpSetXLACompilation(TFE_Op* op, unsigned char enable) { - op->use_xla = enable; + op->operation.SetUseXla(enable); #ifndef TENSORFLOW_EAGER_USE_XLA LOG(WARNING) << "This call is a no-op, as the TensorFlow library is not " "built with XLA support."; @@ -264,22 +256,20 @@ void TFE_OpSetXLACompilation(TFE_Op* op, unsigned char enable) { } void TFE_OpAddInput(TFE_Op* op, TFE_TensorHandle* h, TF_Status* status) { - h->handle->Ref(); - op->inputs.push_back(h->handle); - op->attrs.NumInputs(op->inputs.size()); + op->operation.AddInput(h->handle); } TF_AttrType TFE_OpGetAttrType(TFE_Op* op, const char* attr_name, unsigned char* is_list, TF_Status* status) { TF_AttrType ret; - if (op->is_function()) { + if (op->operation.is_function()) { status->status = tensorflow::errors::Unimplemented( "TODO(apassos): Support for attributes for TensorFlow functions is not " "ready yet."); return TF_ATTR_INT; // The compiler requires that we return something. } - status->status = - tensorflow::AttrTypeByName(*op->attr_types, attr_name, &ret, is_list); + status->status = tensorflow::AttrTypeByName(*op->operation.AttrTypes(), + attr_name, &ret, is_list); return ret; } @@ -298,23 +288,24 @@ TF_AttrType TFE_OpNameGetAttrType(TFE_Context* ctx, } void TFE_OpSetAttrString(TFE_Op* op, const char* attr_name, const char* value) { - op->attrs.Set(attr_name, value); + op->operation.MutableAttrs()->Set(attr_name, value); } void TFE_OpSetAttrInt(TFE_Op* op, const char* attr_name, int64_t value) { - op->attrs.Set(attr_name, static_cast(value)); + op->operation.MutableAttrs()->Set(attr_name, static_cast(value)); } void TFE_OpSetAttrFloat(TFE_Op* op, const char* attr_name, float value) { - op->attrs.Set(attr_name, value); + op->operation.MutableAttrs()->Set(attr_name, value); } void TFE_OpSetAttrBool(TFE_Op* op, const char* attr_name, unsigned char value) { - op->attrs.Set(attr_name, (value == 0) ? false : true); + op->operation.MutableAttrs()->Set(attr_name, (value == 0) ? false : true); } void TFE_OpSetAttrType(TFE_Op* op, const char* attr_name, TF_DataType value) { - op->attrs.Set(attr_name, static_cast(value)); + op->operation.MutableAttrs()->Set(attr_name, + static_cast(value)); } void TFE_OpSetAttrShape(TFE_Op* op, const char* attr_name, const int64_t* dims, @@ -336,23 +327,24 @@ void TFE_OpSetAttrShape(TFE_Op* op, const char* attr_name, const int64_t* dims, proto.add_dim()->set_size(dims[d]); } } - op->attrs.Set(attr_name, proto); + op->operation.MutableAttrs()->Set(attr_name, proto); } void TFE_OpSetAttrFunction(TFE_Op* op, const char* attr_name, const TFE_Op* value) { tensorflow::AttrValue attr_value; tensorflow::NameAttrList* func = attr_value.mutable_func(); - func->set_name(value->name); - value->attrs.FillAttrValueMap(func->mutable_attr()); - op->attrs.Set(attr_name, attr_value); + func->set_name(value->operation.Name()); + value->operation.Attrs().FillAttrValueMap(func->mutable_attr()); + op->operation.MutableAttrs()->Set(attr_name, attr_value); } #define TFE_OP_SET_ATTR_LIST(fn, type) \ void fn(TFE_Op* op, const char* attr_name, const type* values, \ int num_values) { \ - op->attrs.Set(attr_name, tensorflow::gtl::ArraySlice( \ - values, num_values)); \ + op->operation.MutableAttrs()->Set( \ + attr_name, \ + tensorflow::gtl::ArraySlice(values, num_values)); \ } TFE_OP_SET_ATTR_LIST(TFE_OpSetAttrStringList, char*) TFE_OP_SET_ATTR_LIST(TFE_OpSetAttrFloatList, float) @@ -360,14 +352,14 @@ TFE_OP_SET_ATTR_LIST(TFE_OpSetAttrFloatList, float) void TFE_OpSetAttrIntList(TFE_Op* op, const char* attr_name, const int64_t* values, int num_values) { - op->attrs.Set(attr_name, - tensorflow::gtl::ArraySlice( - reinterpret_cast(values), num_values)); + op->operation.MutableAttrs()->Set( + attr_name, tensorflow::gtl::ArraySlice( + reinterpret_cast(values), num_values)); } void TFE_OpSetAttrTypeList(TFE_Op* op, const char* attr_name, const TF_DataType* values, int num_values) { - op->attrs.Set( + op->operation.MutableAttrs()->Set( attr_name, tensorflow::gtl::ArraySlice( reinterpret_cast(values), num_values)); @@ -379,8 +371,8 @@ void TFE_OpSetAttrBoolList(TFE_Op* op, const char* attr_name, for (int i = 0; i < num_values; ++i) { b[i] = values[i]; } - op->attrs.Set(attr_name, - tensorflow::gtl::ArraySlice(b.get(), num_values)); + op->operation.MutableAttrs()->Set( + attr_name, tensorflow::gtl::ArraySlice(b.get(), num_values)); } void TFE_OpSetAttrShapeList(TFE_Op* op, const char* attr_name, @@ -410,9 +402,9 @@ void TFE_OpSetAttrShapeList(TFE_Op* op, const char* attr_name, } } } - op->attrs.Set(attr_name, - tensorflow::gtl::ArraySlice( - proto.get(), num_values)); + op->operation.MutableAttrs()->Set( + attr_name, tensorflow::gtl::ArraySlice( + proto.get(), num_values)); } void TFE_OpSetAttrFunctionList(TFE_Op* op, const char* attr_name, @@ -420,534 +412,25 @@ void TFE_OpSetAttrFunctionList(TFE_Op* op, const char* attr_name, 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 { - -// Initializes the step stats if needed. -void MaybeInitializeStepStats(tensorflow::StepStats* step_stats, - tensorflow::EagerContext* ctx) { - // Lazily initialize the RunMetadata with information about all devices if - // this is the first call. - while (step_stats->dev_stats_size() < ctx->devices()->size()) { - int device_idx = step_stats->dev_stats_size(); - auto* dev_stats = step_stats->add_dev_stats(); - dev_stats->set_device(ctx->devices()->at(device_idx)->name()); - } -} - -int StepStatsDeviceIndex(tensorflow::StepStats* step_stats, - tensorflow::EagerContext* ctx, - tensorflow::Device* device) { - // Find the current device's index. - if (device == nullptr) { - device = ctx->HostCPU(); - } - for (int i = 0; i < ctx->devices()->size(); ++i) { - if (ctx->devices()->at(i) == device || - ctx->devices()->at(i)->name() == device->name()) { - return i; - } - } - // TODO(apassos) do not fall back to host CPU if device is unknown. - return 0; -} - -tensorflow::Status ValidateInputTypeAndPlacement( - tensorflow::EagerContext* ctx, tensorflow::Device* op_device, TFE_Op* op, - const tensorflow::OpKernel* kernel, tensorflow::RunMetadata* run_metadata) { - tensorflow::Device* host_device = ctx->HostCPU(); - const tensorflow::MemoryTypeVector& memtypes = kernel->input_memory_types(); - if (memtypes.size() != op->inputs.size()) { - return tensorflow::errors::InvalidArgument( - "expected ", memtypes.size(), " inputs, got ", op->inputs.size()); - } - for (int i = 0; i < op->inputs.size(); ++i) { - const tensorflow::Device* expected_device = - memtypes[i] == tensorflow::HOST_MEMORY ? host_device : op_device; - tensorflow::TensorHandle* handle = op->inputs[i]; - tensorflow::Device* handle_device = nullptr; - TF_RETURN_IF_ERROR(handle->Device(&handle_device)); - const tensorflow::Device* actual_device = - handle_device == nullptr ? host_device : handle_device; - if (expected_device != actual_device) { - switch (ctx->GetDevicePlacementPolicy()) { - case tensorflow::DEVICE_PLACEMENT_SILENT_FOR_INT32: - // TODO(xpan): See if we could bubble python related error up - // to python level. - if (handle->dtype == tensorflow::DT_INT32) { - // Note: enabling silent copies of int32 tensors to match behavior - // of graph mode. - break; - } - TF_FALLTHROUGH_INTENDED; - case tensorflow::DEVICE_PLACEMENT_EXPLICIT: - return tensorflow::errors::InvalidArgument( - "Tensors on conflicting devices:" - " cannot compute ", - op->name, " as input #", i, " was expected to be on ", - expected_device->name(), " but is actually on ", - actual_device->name(), " (operation running on ", - op_device->name(), ")", - " Tensors can be copied explicitly using .gpu() or .cpu() " - "methods," - " or transparently copied by using tf.enable_eager_execution(" - "device_policy=tfe.DEVICE_PLACEMENT_SILENT). Copying tensors " - "between devices" - " may slow down your model"); - case tensorflow::DEVICE_PLACEMENT_WARN: - LOG(WARNING) << "before computing " << op->name << " input #" << i - << " was expected to be on " << expected_device->name() - << " but is actually on " << actual_device->name() - << " (operation running on " << op_device->name() - << "). This triggers a copy which can be a performance " - "bottleneck."; - break; - case tensorflow::DEVICE_PLACEMENT_SILENT: // Do nothing. - break; - } - // We are only here if the policy is warn or silent copies, so we should - // trigger a copy. - auto pre_time = tensorflow::Env::Default()->NowMicros(); - tensorflow::TensorHandle* copied_tensor = nullptr; - tensorflow::Status status = tensorflow::EagerCopyToDevice( - handle, ctx, expected_device->name().c_str(), &copied_tensor); - if (run_metadata != nullptr) { - auto* step_stats = run_metadata->mutable_step_stats(); - MaybeInitializeStepStats(step_stats, ctx); - // Record the sending on the source device for now. - int device_idx = StepStatsDeviceIndex(step_stats, ctx, handle_device); - auto* dev_stats = step_stats->mutable_dev_stats(device_idx); - auto* node_stats = dev_stats->add_node_stats(); - node_stats->set_node_name("_Send"); - node_stats->set_all_start_micros(pre_time); - node_stats->set_op_end_rel_micros( - tensorflow::Env::Default()->NowMicros() - pre_time); - } - 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()); - } - handle->Unref(); - handle = copied_tensor; - op->inputs[i] = copied_tensor; - } - 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(handle->dtype), - " tensor"); - } - } - return tensorflow::Status::OK(); -} - -tensorflow::Device* SelectDevice(const tensorflow::NodeDef& ndef, - TFE_Context* ctx, TF_Status* status) { - tensorflow::DeviceSet ds; - for (tensorflow::Device* d : *ctx->context.devices()) { - ds.AddDevice(d); - } - tensorflow::DeviceTypeVector final_devices; - status->status = tensorflow::SupportedDeviceTypesForNode( - ds.PrioritizedDeviceTypeList(), ndef, &final_devices); - if (!status->status.ok()) { - 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->context.devices()) { - if (d->device_type() == final_devices[0].type_string()) { - return d; - } - } - status->status = tensorflow::errors::Unknown( - "Could not find a device for node ", ndef.DebugString()); - return nullptr; -} - - -#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; - { - status->status = ctx->context.FindFunctionOpData(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(); - - status->status = ctx->context.AddFunctionDef(fdef); - if (!status->status.ok()) return nullptr; - const auto ret = ctx->context.FindFunctionDef(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; - { - fdef = op->ctx->context.FindFunctionDef(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 (tensorflow::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())); + funcs[i].set_name(value[i]->operation.Name()); + value[i]->operation.Attrs().FillAttrValueMap(funcs[i].mutable_attr()); } - 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; + op->operation.MutableAttrs()->Set( + attr_name, tensorflow::gtl::ArraySlice( + funcs.get(), num_values)); } -#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->context.GetStatus(); + tensorflow::gtl::InlinedVector handle_retvals( + *num_retvals); + status->status = + tensorflow::EagerExecute(&op->operation, &handle_retvals, num_retvals); 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; - VLOG(2) << "for op " << op->name << " input " << i << " " - << tensorflow::DataTypeString(op->inputs[i]->dtype) << " " - << (input_op_device == nullptr ? "cpu" : input_op_device->name()) - << " " << (op->device == nullptr ? "cpu" : op->device->name()); - if (op->inputs[i]->dtype == tensorflow::DT_RESOURCE && - (input_op_device != op->device || input_op_device == nullptr)) { - tensorflow::Device* d = - input_op_device == nullptr ? ctx->context.HostCPU() : 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; - - tensorflow::Fprint128 cache_key = - op->attrs.CacheKey(device == nullptr ? "unspecified" : device->name()); - tensorflow::KernelAndDevice* kernel = ctx->context.GetCachedKernel(cache_key); - if (kernel == nullptr) { - const tensorflow::NodeDef& ndef = op->attrs.BuildNodeDef(); - if (device == nullptr) { - device = SelectDevice(ndef, ctx, status); - if (!status->status.ok()) { - return; - } - } - CHECK(device != nullptr); - if (ctx->context.LogDevicePlacement()) { - LOG(INFO) << "Executing op " << ndef.op() << " in device " - << device->name(); - } - kernel = new tensorflow::KernelAndDevice(ctx->context.GetRendezvous()); - // 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->context.FunctionsMu()); - status->status = tensorflow::KernelAndDevice::Init( - ndef, ctx->context.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->context.FuncLibDef()->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; - } - ctx->context.AddKernelToCache(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->context, device, op, kernel->kernel(), - ctx->context.ShouldStoreMetadata() ? ctx->context.RunMetadataProto() - : nullptr); - if (!status->status.ok()) return; - std::unique_ptr maybe_stats; - if (ctx->context.ShouldStoreMetadata()) { - 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->context.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. - tensorflow::gtl::InlinedVector handle_retvals( - *num_retvals); - tensorflow::uint64 id = op->ctx->context.NextId(); - for (int i = 0; i < *num_retvals; ++i) { - tensorflow::TensorHandle* h = - new tensorflow::TensorHandle(id, output_dtypes[i], &op->ctx->context); - retvals[i] = new TFE_TensorHandle(h); - handle_retvals[i] = h; - } - tensorflow::EagerNode* node = new tensorflow::ExecuteNode( - id, &op->ctx->context, op->device, op->inputs, kernel, - maybe_stats.release(), output_dtypes, handle_retvals); - ctx->context.ExecutorAdd(node); - } else { - // Execute checks if retvals[i] is nullptr or not to figure if it needs to - // allocate it. - std::vector handle_retvals(*num_retvals, - nullptr); - status->status = tensorflow::EagerExecute( - &op->ctx->context, op->device, op->inputs, kernel, maybe_stats.get(), - handle_retvals.data(), *num_retvals); - for (int i = 0; i < *num_retvals; ++i) { - retvals[i] = new TFE_TensorHandle(handle_retvals[i]); - } + for (int i = 0; i < *num_retvals; ++i) { + retvals[i] = new TFE_TensorHandle(handle_retvals[i]); } } @@ -1020,6 +503,62 @@ void TFE_ContextExportRunMetadata(TFE_Context* ctx, TF_Buffer* buf, ctx->context.RunMetadataProto()->Clear(); } +void TFE_GetResourceHandleShapeAndType(TF_Graph* graph, TF_Output output, + TF_Buffer* output_proto, + TF_Status* status) { + tensorflow::Node* node = &output.oper->node; + tensorflow::CppShapeInferenceResult::HandleData handle_data; + handle_data.set_is_set(true); + { + tensorflow::mutex_lock l(graph->mu); + tensorflow::shape_inference::InferenceContext* ic = + graph->refiner.GetContext(node); + CHECK(ic != nullptr); + CHECK_LT(output.index, ic->num_outputs()); + const auto* shapes_and_types = + ic->output_handle_shapes_and_types(output.index); + if (shapes_and_types == nullptr) { + output_proto->data = nullptr; + output_proto->length = 0; + output_proto->data_deallocator = nullptr; + return; + } + + for (const auto& p : *shapes_and_types) { + auto* out_shape_and_type = handle_data.add_shape_and_type(); + ic->ShapeHandleToProto(p.shape, out_shape_and_type->mutable_shape()); + out_shape_and_type->set_dtype(p.dtype); + } + } + status->status = MessageToBuffer(handle_data, output_proto); +} + +void TFE_SetResourceHandleShapeAndType(TF_Graph* graph, TF_Output output, + const void* proto, size_t proto_len, + TF_Status* status) { + tensorflow::CppShapeInferenceResult::HandleData handle_data; + if (!handle_data.ParseFromArray(proto, proto_len)) { + status->status = tensorflow::errors::InvalidArgument( + "Couldn't deserialize HandleData proto"); + return; + } + DCHECK(handle_data.is_set()); + + tensorflow::mutex_lock l(graph->mu); + tensorflow::shape_inference::InferenceContext* ic = + graph->refiner.GetContext(&output.oper->node); + + std::vector shapes_and_types; + for (const auto& shape_and_type_proto : handle_data.shape_and_type()) { + tensorflow::shape_inference::ShapeHandle shape; + status->status = + ic->MakeShapeFromShapeProto(shape_and_type_proto.shape(), &shape); + if (status->status.ok()) return; + shapes_and_types.emplace_back(shape, shape_and_type_proto.dtype()); + } + ic->set_output_handle_shapes_and_types(output.index, shapes_and_types); +} + namespace { TFE_Op* GetFunc(TFE_Context* ctx, const tensorflow::NameAttrList& func, TF_Status* status) { @@ -1090,10 +629,3 @@ void SetOpAttrValueScalar(TFE_Context* ctx, TFE_Op* op, } } } // namespace tensorflow - - -TFE_Op::~TFE_Op() { - for (tensorflow::TensorHandle* h : inputs) { - h->Unref(); - } -} diff --git a/tensorflow/c/eager/c_api.h b/tensorflow/c/eager/c_api.h index 3926c22ce1f9e194b1452c796c83944d10cfdc64..ba77f3cd07ff686782ad7b5c5751811aa3336555 100644 --- a/tensorflow/c/eager/c_api.h +++ b/tensorflow/c/eager/c_api.h @@ -30,7 +30,7 @@ limitations under the License. #ifdef SWIG #define TF_CAPI_EXPORT #else -#if defined(COMPILER_MSVC) +#if defined(_WIN32) #ifdef TF_COMPILE_LIBRARY #define TF_CAPI_EXPORT __declspec(dllexport) #else @@ -38,7 +38,7 @@ limitations under the License. #endif // TF_COMPILE_LIBRARY #else #define TF_CAPI_EXPORT __attribute__((visibility("default"))) -#endif // COMPILER_MSVC +#endif // _WIN32 #endif // SWIG #ifdef __cplusplus @@ -329,6 +329,20 @@ TF_CAPI_EXPORT extern void TFE_ContextExportRunMetadata(TFE_Context* ctx, TF_Buffer* buf, TF_Status* status); +// Returns the serialized CppShapeInferenceResult::HandleData proto for +// `output` if its a resource tensor, or otherwise returns an empty buffer. +TF_CAPI_EXPORT extern void TFE_GetResourceHandleShapeAndType( + TF_Graph* graph, TF_Output output, TF_Buffer* output_proto, + TF_Status* status); + +// Sets `output` based on `proto`, which should be a serialized +// CppShapeInferenceResult::HandleData proto. +TF_CAPI_EXPORT extern void TFE_SetResourceHandleShapeAndType(TF_Graph* graph, + TF_Output output, + const void* proto, + size_t proto_len, + TF_Status* status); + #ifdef __cplusplus } /* end extern "C" */ #endif diff --git a/tensorflow/c/eager/c_api_internal.h b/tensorflow/c/eager/c_api_internal.h index 05dc64f521735f944559392f470a37590e93f17c..49e1aab1cef9577256d9b081858cf094c788caf8 100644 --- a/tensorflow/c/eager/c_api_internal.h +++ b/tensorflow/c/eager/c_api_internal.h @@ -32,6 +32,7 @@ limitations under the License. #include "tensorflow/core/common_runtime/device_factory.h" #include "tensorflow/core/common_runtime/eager/context.h" #include "tensorflow/core/common_runtime/eager/eager_executor.h" +#include "tensorflow/core/common_runtime/eager/eager_operation.h" #include "tensorflow/core/common_runtime/eager/kernel_and_device.h" #include "tensorflow/core/common_runtime/eager/tensor_handle.h" #include "tensorflow/core/common_runtime/function.h" @@ -45,7 +46,6 @@ limitations under the License. #include "tensorflow/core/platform/thread_annotations.h" #include "tensorflow/core/public/version.h" - struct TFE_ContextOptions { TF_SessionOptions session_options; // true if async execution is enabled. @@ -85,19 +85,9 @@ 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; } + : operation(&ctx->context, op, t) {} - TFE_Context* ctx; // Must outlive the TFE_Op. - const tensorflow::string name; - tensorflow::AttrBuilder attrs; - const tensorflow::AttrTypeMap* attr_types; - tensorflow::gtl::InlinedVector inputs; - tensorflow::Device* device; - bool use_xla = false; + tensorflow::EagerOperation operation; }; namespace tensorflow { diff --git a/tensorflow/c/eager/runtime.cc b/tensorflow/c/eager/runtime.cc index abe2793ce894ad07c252575c5d55d98342916eac..e6c51ab17a867a0697f15d7683d8ca52c062035d 100644 --- a/tensorflow/c/eager/runtime.cc +++ b/tensorflow/c/eager/runtime.cc @@ -184,8 +184,7 @@ void CombineUnordered(const tensorflow::Fprint128& a, inline tensorflow::Fprint128 CacheKeyHelper(StringPiece s, const tensorflow::Fprint128& b) { - // TODO(agarwal): avoid ToString(). - tensorflow::Fprint128 a = tensorflow::Fingerprint128(s.ToString()); + tensorflow::Fprint128 a = tensorflow::Fingerprint128(s); return FingerprintCat128(a, b); } @@ -213,10 +212,8 @@ tensorflow::Fprint128 AttrBuilder::CacheKey(const string& device) const { if (node_def_finalized_) return f; } for (const auto& p : string_attrs_) { - // TODO(agarwal): avoid ToString(). - CombineUnordered(CacheKeyHelper(p.first, tensorflow::Fingerprint128( - p.second.ToString())), - &f); + CombineUnordered( + CacheKeyHelper(p.first, tensorflow::Fingerprint128(p.second)), &f); } for (const auto& p : int_attrs_) { CombineUnordered(CacheKeyHelper(p.first, static_cast(p.second)), diff --git a/tensorflow/cc/gradients/array_grad.cc b/tensorflow/cc/gradients/array_grad.cc index 6545e4ee3eb406436937a43ddac66d017af8e108..ff348fadb24e29a83bd6c8853aa67931f6df4182 100644 --- a/tensorflow/cc/gradients/array_grad.cc +++ b/tensorflow/cc/gradients/array_grad.cc @@ -385,6 +385,42 @@ Status MirrorPadGradGrad(const Scope& scope, const Operation& op, } REGISTER_GRADIENT_OP("MirrorPadGrad", MirrorPadGradGrad); +Status StridedSliceGradHelper(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + Input x = Shape(scope, op.input(0)); + Input begin = op.input(1); + Input end = op.input(2); + Input strides = op.input(3); + int64 begin_mask; + int64 end_mask; + int64 ellipsis_mask; + int64 new_axis_mask; + int64 shrink_axis_mask; + TF_RETURN_IF_ERROR( + GetNodeAttr(op.node()->attrs(), "begin_mask", &begin_mask)); + TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "end_mask", &end_mask)); + TF_RETURN_IF_ERROR( + GetNodeAttr(op.node()->attrs(), "ellipsis_mask", &ellipsis_mask)); + TF_RETURN_IF_ERROR( + GetNodeAttr(op.node()->attrs(), "new_axis_mask", &new_axis_mask)); + TF_RETURN_IF_ERROR( + GetNodeAttr(op.node()->attrs(), "shrink_axis_mask", &shrink_axis_mask)); + grad_outputs->push_back( + StridedSliceGrad(scope, x, begin, end, strides, grad_inputs[0], + StridedSliceGrad::BeginMask(begin_mask) + .EndMask(end_mask) + .EllipsisMask(ellipsis_mask) + .NewAxisMask(new_axis_mask) + .ShrinkAxisMask(shrink_axis_mask))); + // No gradients returned for begin, end and strides + grad_outputs->push_back(NoGradient()); + grad_outputs->push_back(NoGradient()); + grad_outputs->push_back(NoGradient()); + return scope.status(); +} +REGISTER_GRADIENT_OP("StridedSlice", StridedSliceGradHelper); + } // anonymous namespace } // namespace ops } // namespace tensorflow diff --git a/tensorflow/cc/gradients/array_grad_test.cc b/tensorflow/cc/gradients/array_grad_test.cc index 4a215fcc9299cf8b8da04cbf151640631ed0d449..de3bd0fc9e2493f8ff76163f5be6bd4327c58c5a 100644 --- a/tensorflow/cc/gradients/array_grad_test.cc +++ b/tensorflow/cc/gradients/array_grad_test.cc @@ -354,5 +354,29 @@ TEST_F(ArrayGradTest, MirrorPadGradGrad_Symmetric) { RunTest(x, x_shape, y, y_shape); } +TEST_F(ArrayGradTest, StridedSliceGrad) { + TensorShape x_shape({6, 4, 4}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)); + + // y = x[2:6:2, 1:3, 1:3] + auto y = StridedSlice(scope_, x, {2, 1, 1}, {6, 3, 3}, {2, 1, 1}); + // y.shape = [2, 2, 2]; + RunTest(x, x_shape, y, {2, 2, 2}); + + // y = x[2:6:2, 1:3, 1:3] + // begin_mask = 1<<1 (ignore begin_index = 1) + // end_mask = 1<<2 (ignore end_index = 2) + y = StridedSlice(scope_, x, {2, 1, 1}, {6, 3, 3}, {2, 1, 1}, + StridedSlice::BeginMask(1 << 1).EndMask(1 << 2)); + // y.shape = [2, 3, 3]; + RunTest(x, x_shape, y, {2, 3, 3}); + + // y = [tf.newaxis, 2:6:2, 1:3, 1:3] + y = StridedSlice(scope_, x, {0, 2, 1, 1}, {0, 6, 3, 3}, {1, 2, 1, 1}, + StridedSlice::NewAxisMask(1 << 0)); + // y.shape = [1, 2, 2, 2]; + RunTest(x, x_shape, y, {1, 2, 2, 2}); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/aot/BUILD b/tensorflow/compiler/aot/BUILD index fa03b1f3c2dfc334d4a3871e6a1bf5503fa8d5f8..19e6bf68e77725bb3cae4e1d338c52dff472cb18 100644 --- a/tensorflow/compiler/aot/BUILD +++ b/tensorflow/compiler/aot/BUILD @@ -60,6 +60,7 @@ cc_library( "//tensorflow/compiler/tf2xla:tf2xla_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla/kernels:xla_cpu_only_ops", + "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:statusor", diff --git a/tensorflow/compiler/aot/embedded_protocol_buffers.cc b/tensorflow/compiler/aot/embedded_protocol_buffers.cc index 6489929a576d6469c4ff1358ca5ee9d27fb578bb..0048eec93bbe10271d9aa535203f19473a38b342 100644 --- a/tensorflow/compiler/aot/embedded_protocol_buffers.cc +++ b/tensorflow/compiler/aot/embedded_protocol_buffers.cc @@ -19,7 +19,6 @@ limitations under the License. #include #include "llvm/ADT/Triple.h" -#include "llvm/ExecutionEngine/ObjectMemoryBuffer.h" #include "llvm/IR/GlobalVariable.h" #include "llvm/IR/LLVMContext.h" #include "llvm/IR/LegacyPassManager.h" diff --git a/tensorflow/compiler/aot/runtime.cc b/tensorflow/compiler/aot/runtime.cc index 5772776666129ed55a479c8917e69df3f3ce2fc0..5e74079fc158379b8977ada6412141e39142c3d3 100644 --- a/tensorflow/compiler/aot/runtime.cc +++ b/tensorflow/compiler/aot/runtime.cc @@ -31,7 +31,7 @@ namespace { inline void* aligned_malloc(size_t size, int minimum_alignment) { #if defined(__ANDROID__) || defined(OS_ANDROID) || defined(OS_CYGWIN) return memalign(minimum_alignment, size); -#elif defined(COMPILER_MSVC) +#elif defined(_WIN32) return _aligned_malloc(size, minimum_alignment); #else // !__ANDROID__ && !OS_ANDROID && !OS_CYGWIN void* ptr = nullptr; @@ -48,7 +48,7 @@ inline void* aligned_malloc(size_t size, int minimum_alignment) { } inline void aligned_free(void* aligned_memory) { -#if defined(COMPILER_MSVC) +#if defined(_WIN32) _aligned_free(aligned_memory); #else free(aligned_memory); diff --git a/tensorflow/compiler/aot/tests/BUILD b/tensorflow/compiler/aot/tests/BUILD index b053dad1b57c258b7cb0d6831923e6a0f30f5e7e..bb73cb19c57a654058af5bbb4535c76b0aca8e8c 100644 --- a/tensorflow/compiler/aot/tests/BUILD +++ b/tensorflow/compiler/aot/tests/BUILD @@ -14,6 +14,7 @@ test_suite( ":test_graph_tfadd_test", ":test_graph_tfadd_with_ckpt_saver_test", ":test_graph_tfadd_with_ckpt_test", + ":test_graph_tfassert_eq_test", ":test_graph_tffunction_test", ":test_graph_tfgather_test", ":test_graph_tfmatmul_test", @@ -33,6 +34,7 @@ py_binary( "//tensorflow/python", # TODO(b/34059704): remove when fixed "//tensorflow/python:array_ops", "//tensorflow/python:client", + "//tensorflow/python:control_flow_ops", "//tensorflow/python:framework_for_generated_wrappers", "//tensorflow/python:math_ops", "//tensorflow/python:platform", @@ -52,6 +54,7 @@ genrule( "test_graph_tfadd_with_ckpt_saver.ckpt", "test_graph_tfadd_with_ckpt_saver.pb", "test_graph_tfadd_with_ckpt_saver.saver", + "test_graph_tfassert_eq.pb", "test_graph_tffunction.pb", "test_graph_tfgather.pb", "test_graph_tfmatmul.pb", @@ -104,6 +107,17 @@ tf_library( ], ) +tf_library( + name = "test_graph_tfassert_eq", + testonly = 1, + config = "test_graph_tfassert_eq.config.pbtxt", + cpp_class = "AssertComp", + graph = "test_graph_tfassert_eq.pb", + tags = [ + "manual", + ], +) + tf_library( name = "test_graph_tffunction", testonly = 1, @@ -170,6 +184,7 @@ tf_cc_test( ":test_graph_tfadd", ":test_graph_tfadd_with_ckpt", ":test_graph_tfadd_with_ckpt_saver", + ":test_graph_tfassert_eq", ":test_graph_tffunction", ":test_graph_tfgather", ":test_graph_tfmatmul", diff --git a/tensorflow/compiler/aot/tests/make_test_graphs.py b/tensorflow/compiler/aot/tests/make_test_graphs.py index 89c7cd4507cbd476104a039d6083d8f89de11278..67767f55dae9b15aafbd8b129328bde2c59a9ef3 100644 --- a/tensorflow/compiler/aot/tests/make_test_graphs.py +++ b/tensorflow/compiler/aot/tests/make_test_graphs.py @@ -29,6 +29,7 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import function from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops +from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import variables from tensorflow.python.platform import app @@ -125,6 +126,14 @@ def tfsplits(_): array_ops.identity(y, name='result') +def tfassert_eq(_): + x = array_ops.placeholder(dtypes.int32, name='x_hold') + y = array_ops.placeholder(dtypes.int32, name='y_hold') + control_flow_ops.Assert( + math_ops.equal(x, y), ['Expected x == y.'], name='assert_eq') + math_ops.add(x, math_ops.negative(y), name='x_y_diff') + + def write_graph(build_graph, out_dir): """Build a graph using build_graph and write it out.""" g = ops.Graph() @@ -144,6 +153,7 @@ def main(_): write_graph(tfmatmulandadd, FLAGS.out_dir) write_graph(tffunction, FLAGS.out_dir) write_graph(tfsplits, FLAGS.out_dir) + write_graph(tfassert_eq, FLAGS.out_dir) if __name__ == '__main__': diff --git a/tensorflow/compiler/aot/tests/test_graph_tfassert_eq.config.pbtxt b/tensorflow/compiler/aot/tests/test_graph_tfassert_eq.config.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..8732d1709e809bb47d3769c483483c2c4f350e1c --- /dev/null +++ b/tensorflow/compiler/aot/tests/test_graph_tfassert_eq.config.pbtxt @@ -0,0 +1,16 @@ +# Text form of tensorflow.tf2xla.Config proto. +feed { + id { node_name: "x_hold" } + shape { + dim { size: 1 } + } +} +feed { + id { node_name: "y_hold" } + shape { + dim { size: 1 } + } +} +fetch { + id { node_name: "x_y_diff" } +} diff --git a/tensorflow/compiler/aot/tests/tfcompile_test.cc b/tensorflow/compiler/aot/tests/tfcompile_test.cc index 413efd9cea3b6f71574615ad9ca92471ff925781..67dbd643bfc7bf2c214e7eb5ae8bd2cc7d6e164b 100644 --- a/tensorflow/compiler/aot/tests/tfcompile_test.cc +++ b/tensorflow/compiler/aot/tests/tfcompile_test.cc @@ -20,6 +20,7 @@ limitations under the License. #include "tensorflow/compiler/aot/tests/test_graph_tfadd.h" #include "tensorflow/compiler/aot/tests/test_graph_tfadd_with_ckpt.h" #include "tensorflow/compiler/aot/tests/test_graph_tfadd_with_ckpt_saver.h" +#include "tensorflow/compiler/aot/tests/test_graph_tfassert_eq.h" #include "tensorflow/compiler/aot/tests/test_graph_tffunction.h" #include "tensorflow/compiler/aot/tests/test_graph_tfgather.h" #include "tensorflow/compiler/aot/tests/test_graph_tfmatmul.h" @@ -413,6 +414,23 @@ TEST(TFCompileTest, Splits) { EXPECT_NEAR(expected[3], fn.result0(1, 1), 1e4); } +TEST(TFCompileTest, AssertEqAndReturnDiff) { + // Assert is converted into a no-op in XLA, so there is no failure even if the + // two args are different. + AssertComp assert; + EXPECT_EQ(assert.arg0_data(), assert.args()[0]); + EXPECT_EQ(assert.arg1_data(), assert.args()[1]); + + assert.arg0() = 2; + assert.arg1() = 1; + const int32 expected_result = assert.arg0() - assert.arg1(); + EXPECT_TRUE(assert.Run()); + EXPECT_EQ(assert.error_msg(), ""); + EXPECT_EQ(assert.result0(), expected_result); + EXPECT_EQ(assert.result0_data()[0], expected_result); + EXPECT_EQ(assert.result0_data(), assert.results()[0]); +} + TEST(TFCompileTest, LookupNameIndex) { // add doesn't have any names defined in its config. AddComp add; diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD index a492fc6b9bb83dd64f1528bd4d8f13d640c5fbed..50fa95c4f322e85c22f7be2d63f2bcd194ee419e 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -183,6 +183,13 @@ cc_library( ], ) +cc_library( + name = "shape_inference_helpers", + srcs = ["shape_inference_helpers.cc"], + hdrs = ["shape_inference_helpers.h"], + deps = ["//tensorflow/core:graph"], +) + # Internal targets below this point. cc_library( @@ -293,6 +300,7 @@ cc_library( deps = [ ":common", ":graph_to_functiondef", + ":shape_inference_helpers", ":union_find", "//tensorflow/compiler/jit/graphcycles", "//tensorflow/compiler/jit/kernels:parallel_check_op", @@ -310,6 +318,7 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", + "//tensorflow/core/kernels:bounds_check", ], ) @@ -318,6 +327,25 @@ cc_library( hdrs = ["union_find.h"], ) +cc_library( + name = "producer_consumer_queue", + hdrs = ["producer_consumer_queue.h"], + deps = ["//tensorflow/core:lib"], +) + +tf_cc_test( + name = "producer_consumer_queue_test", + size = "small", + srcs = ["producer_consumer_queue_test.cc"], + deps = [ + ":producer_consumer_queue", + "//tensorflow/core:lib", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + "//tensorflow/core:testlib", + ], +) + tf_cc_test( name = "graph_to_functiondef_test", size = "small", diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc index b04b333141a616e7c4db2751c14ec6eb0b7725b5..7507e193b56d7200876fb938001d06666784b905 100644 --- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc +++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc @@ -23,8 +23,10 @@ limitations under the License. #include #include "tensorflow/compiler/jit/graph_to_functiondef.h" +#include "tensorflow/compiler/jit/graphcycles/graphcycles.h" #include "tensorflow/compiler/jit/legacy_flags/encapsulate_subgraphs_pass_flags.h" #include "tensorflow/compiler/jit/mark_for_compilation_pass.h" +#include "tensorflow/compiler/jit/shape_inference_helpers.h" #include "tensorflow/compiler/tf2xla/const_analysis.h" #include "tensorflow/compiler/tf2xla/dump_graph.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -36,6 +38,7 @@ limitations under the License. #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/control_flow.h" #include "tensorflow/core/graph/graph.h" #include "tensorflow/core/graph/graph_def_builder.h" #include "tensorflow/core/graph/tensor_id.h" @@ -158,6 +161,11 @@ class Encapsulator { std::move(outside_compilation_attribute)), graph_in_(graph_in) {} + // Find dependencies between subgraphs and outside_compilation clusters that + // only manifest via edges between outside_compilation clusters in the outer + // (non-compiled) graph. + Status FindClusterDependencies(); + // Find subgraphs marked with 'group_attribute', and build a new // subgraph, one for each value of 'group_attribute'. Status SplitIntoSubgraphs(); @@ -228,6 +236,19 @@ class Encapsulator { // 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. + // + // Once edges between compiled and outside_compilation clusters have been + // replaced by send/recv ops, some dependencies may no longer be apparent. + // A clustering pass finds all the dependencies between HC nodes that are only + // present as a result of edges between nodes in outside_compilaton clusters. + // Suppose there is a path from outside_compilation cluster C in subgraph S + // to outside_compilation cluster D in subgraph T. If S != T then a control + // edge is added from the call node for S to the call node for T, which + // ensures that C will execute before D because S executes before T. If S==T + // then a control dependency is added between the HC nodes for C and D in S, + // and the HC node for C is added to an 'ancestors' attr in the HC node for D + // so that during compilation of the HC node for D, an XLA control dependency + // can be added to ensure C's SendToHost executes before D's RecvFromHost. class Subgraph { public: // Creates a graph to build the subgraph in, if it doesn't already exist, @@ -322,6 +343,18 @@ class Encapsulator { void RecordOutsideCompilationOutputOrControl( const string& outside_compilation_id, const Edge* edge); + // Records the fact that there is a path from a node in outside_compilation + // cluster ancestor to node in cluster successor that does not go through + // the subgraph. + void RecordOutsideCompilationDependency(const string& successor, + const string& ancestor); + + // Returns the mapping from outside_compilation cluster C to the set of + // outside_compilation clusters that have a path to C entirely outside + // compiled subgraphs. + const std::unordered_map> + OutsideCompilationAncestorMap() const; + // Adds the HostCompute nodes for each outside_compilation subgraph. Status AddHostComputes( const string& subgraph_name, @@ -404,6 +437,13 @@ class Encapsulator { Status AddHostComputeKeyPlaceholder(OutsideCompilationSubgraph* oc_subgraph, Graph* graph_out); + // Get the set of outside_compilation clusters and the dependency edges + // between them. + void GetActiveClusterDependencyGraph( + std::unordered_set* clusters, + std::unordered_set* has_successor, + std::unordered_map>* ancestors_map); + // 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& group_attribute, @@ -466,6 +506,14 @@ class Encapsulator { // The outside_compilation clusters in this subgraph. std::unordered_map outside_compilation_subgraphs_; + // For each outside_compilation cluster C, the outside_compilation clusters + // that have a path to C outside the compiled graph. + std::unordered_map> + outside_compilation_ancestors_; + // For each outside_compilation cluster C, the outside_compilation clusters + // that have a path from C outside the compiled graph. + std::unordered_map> + outside_compilation_successors_; // NoOp node in the output graph that is sequenced after the call node and // used to prevent host-side outside_compilation sends and recvs from being @@ -554,6 +602,10 @@ class Encapsulator { std::unordered_set, NodeSlot::PairHasher>* edges_added); + // Adds control dependencies between subgraph call nodes that have + // dependencies via outside_compilation edges. + Status AddCallNodeDependencies(Graph* graph_out); + // Adds all edges to the output graph. Status AddEdgesToOutputGraph( const std::unordered_map& node_images, @@ -576,7 +628,8 @@ class Encapsulator { // 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 Graph& graph_in, const BackEdgeHelper& back_edge_helper, + const ShapeRefiner& shape_refiner, const std::unordered_set& recv_at_host_nodes, Node* send_node, FunctionLibraryDefinition* library, std::vector* static_shape_out, @@ -599,7 +652,7 @@ class Encapsulator { // to nodes in pruned_graph. Status MakeGraphForOutsideCompilationSends( const Graph& graph, std::unique_ptr* pruned_graph, - ShapeRefiner* shape_refiner, + BackEdgeHelper* back_edge_helper, ShapeRefiner* shape_refiner, std::unordered_map* node_images, FunctionLibraryDefinition* library); @@ -617,10 +670,65 @@ class Encapsulator { const Graph* graph_in_; std::unordered_map subgraphs_; + // For each subgraph S the subgraphs S' such that there is a path in some + // outside_compilation cluster C in S to some outside_compilation cluster C' + // in S', that goes only through the uncompiled graph. + std::unordered_map> subgraph_ancestors_; TF_DISALLOW_COPY_AND_ASSIGN(Encapsulator); }; +namespace { + +// Return in 'sorted' a topological sort of clusters according to the +// dependencies encoded in ancestors. clusters is the list of all clusters +// including clusters that are not present in the ancestors map. has_successors +// is the set of clusters that are ancestors of some other cluster. +void TopologicalClusterSort( + const std::unordered_set& clusters, + const std::unordered_set& has_successors, + const std::unordered_map>& ancestors, + std::vector* sorted) { + // The nodes are placed in 'sorted' in topological order. + sorted->clear(); + // We don't use the standard DFS because we are not operating on Node* + // objects. + struct Work { + string cluster; + bool leave; + }; + std::set visited; + std::vector stack; + // Seed the processing list with clusters that have no successors. + for (const auto& cluster : clusters) { + if (has_successors.find(cluster) == has_successors.end()) { + stack.push_back({cluster, false}); + } + } + while (!stack.empty()) { + const Work item = stack.back(); + stack.pop_back(); + if (item.leave) { + sorted->push_back(item.cluster); + continue; + } + + if (visited.find(item.cluster) != visited.end()) continue; + visited.insert(item.cluster); + + stack.push_back({item.cluster, true}); + const auto& iter = ancestors.find(item.cluster); + if (iter != ancestors.end()) { + for (const auto& ancestor : iter->second) { + stack.push_back({ancestor, false}); + } + } + } + CHECK(sorted->size() == clusters.size()); +} + +} // namespace + Node* Encapsulator::Subgraph::GetCallNodeForInputs() const { return call_node_inputs_; } @@ -783,12 +891,71 @@ void Encapsulator::Subgraph::RecordOutsideCompilationOutputOrControl( } } +void Encapsulator::Subgraph::RecordOutsideCompilationDependency( + const string& successor, const string& ancestor) { + outside_compilation_ancestors_[successor].insert(ancestor); + outside_compilation_successors_[ancestor].insert(successor); +} + +const std::unordered_map> +Encapsulator::Subgraph::OutsideCompilationAncestorMap() const { + return outside_compilation_ancestors_; +} + +void Encapsulator::Subgraph::GetActiveClusterDependencyGraph( + std::unordered_set* clusters, + std::unordered_set* has_successor, + std::unordered_map>* ancestors_map) { + // During initial clustering the ancestor and successor datastructures may + // have been built including oc_cluster names that never turned into subgraphs + // because they had no edges into or out of the compiled cluster. Remove them + // before proceeding to simplify the logic. Get the set of clusters that was + // actually added, then remove references to the others. + for (const auto& oc_subgraph : outside_compilation_subgraphs_) { + clusters->insert(oc_subgraph.first); + } + for (const auto& cluster : outside_compilation_successors_) { + if (clusters->find(cluster.first) != clusters->end()) { + for (const auto& successor : cluster.second) { + if (clusters->find(successor) != clusters->end()) { + has_successor->insert(cluster.first); + break; + } + } + } + } + for (const auto& cluster : outside_compilation_ancestors_) { + if (clusters->find(cluster.first) != clusters->end()) { + std::unordered_set& ancestors = (*ancestors_map)[cluster.first]; + for (const auto& ancestor : cluster.second) { + if (clusters->find(ancestor) != clusters->end()) { + ancestors.insert(ancestor); + } + } + } + } +} + 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; + // Get the set of outside_compilation clusters and the dependency edges + // between them. + std::unordered_set clusters; + std::unordered_set has_successor; + std::unordered_map> ancestors_map; + GetActiveClusterDependencyGraph(&clusters, &has_successor, &ancestors_map); + // Topologically sort the outside_compilation clusters according to their + // dependency relation. + std::vector sorted_clusters; + TopologicalClusterSort(clusters, has_successor, ancestors_map, + &sorted_clusters); + + // The host compute nodes added for each outside_compilation_cluster; + std::unordered_map host_compute_node; + for (const string& oc_subgraph_name : sorted_clusters) { + OutsideCompilationSubgraph& oc_subgraph = + outside_compilation_subgraphs_[oc_subgraph_name]; if (!oc_subgraph.inputs.empty() || !oc_subgraph.control_inputs.empty() || !oc_subgraph.outputs_by_src.empty() || !oc_subgraph.control_outputs.empty()) { @@ -808,13 +975,22 @@ Status Encapsulator::Subgraph::AddHostComputes( 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; } + std::vector host_compute_ancestors; + const auto iter = ancestors_map.find(oc_subgraph_name); + if (iter != ancestors_map.end()) { + for (const string& ancestor_cluster : iter->second) { + host_compute_ancestors.push_back( + outside_compilation_subgraphs_[ancestor_cluster] + .host_compute_name); + } + } + NodeDef host_compute_def; NodeDefBuilder builder(strings::StrCat("outside_compilation_", oc_subgraph_name, "_host_compute"), @@ -822,6 +998,7 @@ Status Encapsulator::Subgraph::AddHostComputes( builder.Input(inputs); builder.Attr("Tinputs", input_dtypes); builder.Attr("Toutputs", output_dtypes); + builder.Attr("ancestors", host_compute_ancestors); builder.Attr("key", strings::StrCat("host_compute_channel_", subgraph_name, "_", oc_subgraph_name)); @@ -831,6 +1008,7 @@ Status Encapsulator::Subgraph::AddHostComputes( Node* host_compute = graph_->AddNode(host_compute_def, &s); if (!s.ok()) return s; + host_compute_node[host_compute->name()] = host_compute; oc_subgraph.host_compute_name = host_compute->name(); // Connect the _HostCompute node to its producers in the subgraph. @@ -849,6 +1027,12 @@ Status Encapsulator::Subgraph::AddHostComputes( graph_->AddControlEdge(src_image, host_compute); } + // Connect the _HostCompute node to its ancestor host compute nodes. + for (const auto& ancestor_name : host_compute_ancestors) { + Node* ancestor = host_compute_node[ancestor_name]; + graph_->AddControlEdge(ancestor, host_compute); + } + // 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; @@ -1651,6 +1835,17 @@ Status Encapsulator::CopyEdgeToOutputGraph( return Status::OK(); } +Status Encapsulator::AddCallNodeDependencies(Graph* graph_out) { + for (const auto& ancestors : subgraph_ancestors_) { + const string& subgraph = ancestors.first; + for (const string& ancestor : ancestors.second) { + graph_out->AddControlEdge(subgraphs_[ancestor].GetCallNodeForOutputs(), + subgraphs_[subgraph].GetCallNodeForInputs()); + } + } + return Status::OK(); +} + Status Encapsulator::AddEdgesToOutputGraph( const std::unordered_map& node_images, bool parallel_checking, Graph* graph_out) { @@ -1700,6 +1895,7 @@ Status Encapsulator::AddEdgesToOutputGraph( Subgraph& subgraph = subgraph_entry.second; subgraph.ConnectSequencerToCallNode(graph_out); } + TF_RETURN_IF_ERROR(AddCallNodeDependencies(graph_out)); return Status::OK(); } @@ -1712,9 +1908,13 @@ namespace { // 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) { +// are added, and is not necessary.) If the node being replaced was within a +// control flow frame, adds appropriate Enter nodes so that the use of the Const +// is well-formed. +Node* AddDummyShapedNode(const Node* src_node, int src_port, + const std::vector& control_flow_info, + const TensorShapeProto& shape, Graph* graph_out) { + DataType data_type = src_node->output_type(src_port); TensorProto dummy_proto; dummy_proto.set_dtype(data_type); *dummy_proto.mutable_tensor_shape() = shape; @@ -1725,7 +1925,23 @@ Node* AddDummyShapedNode(DataType data_type, const TensorShapeProto& shape, 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); + Node* node = options.FinalizeBuilder(&node_builder); + // Add any Enter nodes required to bring the constant to the correct control + // flow frame. + while (!control_flow_info[src_node->id()].frame_name.empty()) { + NodeBuilder enter_builder(options.GetNameForOp("Enter"), "Enter", + options.op_registry()); + enter_builder.Attr("frame_name", + control_flow_info[src_node->id()].frame_name); + enter_builder.Attr("is_constant", true); + enter_builder.Input(node, 0); + Node* enter_node = options.FinalizeBuilder(&enter_builder); + // Adopt the new Enter node as the value in the current frame. + node = enter_node; + // Recurse to the parent frame to see if more Enter nodes need to be added. + src_node = control_flow_info[src_node->id()].parent_frame; + } + return node; } // Adds a copy of node_in to graph_out and adds the mapping to @@ -1767,17 +1983,30 @@ Status CopyShapeInferenceNodeToGraph( } } } + // Work around the fact that Enter nodes refuse to propagate shape information + // unless they are marked loop invariant. Since we are never going to execute + // this graph, marking them all loop invariant is fine. + if (node_out->type_string() == "Enter") { + node_out->ClearAttr("is_constant"); + node_out->AddAttr("is_constant", true); + } return Status::OK(); } } // namespace Status Encapsulator::DoStaticShapeInferenceForOutsideCompilationSend( - const Graph& graph_in, const ShapeRefiner& shape_refiner, + const Graph& graph_in, const BackEdgeHelper& back_edge_helper, + 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* graph_out) { + // Get the control flow structure of the input graph so we can build + // well-formed output graphs. + std::vector control_flow_info; + TF_RETURN_IF_ERROR(BuildControlFlowInfo(&graph_in, &control_flow_info)); + // 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 @@ -1802,7 +2031,6 @@ Status Encapsulator::DoStaticShapeInferenceForOutsideCompilationSend( // 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? @@ -1840,8 +2068,9 @@ Status Encapsulator::DoStaticShapeInferenceForOutsideCompilationSend( 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()); + dummy_node_images[src_node] = + AddDummyShapedNode(src_node, src_port, control_flow_info, + proto, graph_out->get()); } // The final input to the send node is the dynamic key, which we // don't include in the static shapes. @@ -1889,6 +2118,214 @@ Status Encapsulator::DoStaticShapeInferenceForOutsideCompilationSend( } } + for (const auto edge : back_edge_helper.RemovedEdges()) { + if (copied_node_images.find(edge.dst) != copied_node_images.end()) { + // The destination of this back edge was added to the inference graph, so + // fix it up. + Node* dst = copied_node_images[edge.dst]; + if (dst->type_string() != "Merge") { + return errors::InvalidArgument( + "outside_compilation cluster contains a back-edge to node ", + dst->name(), " of type ", dst->type_string(), + ". The analysis pass only supports back-edges to Merge nodes."); + } + const Edge* existing_input_edge; + if (edge.dst_input != 1 || dst->num_inputs() != 2 || + !dst->input_edge(0, &existing_input_edge).ok()) { + // TODO(misard) if we see graphs built with a different structure, relax + // this constraint. Leaving it here for now to avoid writing unnecessary + // complex code since we believe graphs generated by front ends all have + // the back edge as the second input to the merge node. + return errors::Internal( + "Internal assumption failed while rewriting an outside_compilation " + "cluster that contains a while loop. Logic assumes back-edge is to " + "port 1 of a 2-input " + "Merge node."); + } + // Connect the existing edge to both inputs of the Merge node so that the + // graph will be well-formed. + (*graph_out) + ->AddEdge(existing_input_edge->src(), + existing_input_edge->src_output(), dst, edge.dst_input); + } + } + + return Status::OK(); +} + +namespace { + +// Helper struct for building cluster dependencies and also debugging cycles in +// the dependencies. While computing dependencies we construct a mapping from +// Node* to PathDetails. +struct PathDetails { + struct SubgraphAndCluster { + string subgraph; + string outside_compilation_cluster; + bool operator==(const SubgraphAndCluster& other) const { + return subgraph == other.subgraph && + outside_compilation_cluster == other.outside_compilation_cluster; + } + }; + + struct SubgraphAndClusterHash { + inline std::size_t operator()(const SubgraphAndCluster& v) const { + return hash()( + strings::StrCat(v.subgraph, v.outside_compilation_cluster)); + } + }; + + typedef std::unordered_set + SubgraphAndClusterSet; + + // Returns the set of (subgraph, oc_cluster) pairs that should be recorded as + // ancestors for any successor of this node. If the node is in the outer + // graph, it returns the transitive union of the ancestors of the node's + // inputs. If the node is in an outside_compilation cluster, it returns just + // that cluster. If the node is compiled, it returns the empty set. + SubgraphAndClusterSet AncestorsForSuccessor() { + if (subgraph.empty()) { + return ancestor_clusters; + } else if (outside_compilation_cluster.empty()) { + return SubgraphAndClusterSet(); + } else { + SubgraphAndCluster entry; + entry.subgraph = subgraph; + entry.outside_compilation_cluster = outside_compilation_cluster; + return SubgraphAndClusterSet({entry}); + } + } + + // The transitive union of the ancestor's of this node's inputs. This is only + // saved for debugging in order to print out enough information to debug a + // discovered cycle. + SubgraphAndClusterSet ancestor_clusters; + // The subgraph attr on this node. + string subgraph; + // The outside_compilation attr on this node. + string outside_compilation_cluster; +}; + +// Adds an edge from ancestor to successor to the cycle detector, and returns an +// error if that edge causes the formation of a cycle. In the error case, logs +// the contents of the node_ancestors_map to facilitate debugging. +Status CheckClusterDependencyForCycles( + const string& ancestor, const string& successor, + const std::unordered_map>& ancestors, + const std::unordered_map& node_ancestors_map, + GraphCycles* cycle_detector, std::map* cycle_detector_map) { + if (cycle_detector_map->find(ancestor) == cycle_detector_map->end()) { + (*cycle_detector_map)[ancestor] = cycle_detector->NewNode(); + } + if (cycle_detector_map->find(successor) == cycle_detector_map->end()) { + (*cycle_detector_map)[successor] = cycle_detector->NewNode(); + } + + if (!cycle_detector->InsertEdge((*cycle_detector_map)[ancestor], + (*cycle_detector_map)[successor])) { + LOG(ERROR) << "Cycle in outside_compilation clusters"; + for (const auto& cluster : ancestors) { + LOG(ERROR) << "Cluster " << cluster.first << " depends on:"; + for (const auto& ancestor : cluster.second) { + LOG(ERROR) << " " << ancestor; + } + } + for (const auto& node_ancestors : node_ancestors_map) { + LOG(ERROR) << "Node " << node_ancestors.first->name() << " (" + << node_ancestors.second.subgraph << ";" + << node_ancestors.second.outside_compilation_cluster + << ") has ancestor clusters:"; + for (const auto& ancestor : node_ancestors.second.ancestor_clusters) { + LOG(ERROR) << " " << ancestor.subgraph << ";" + << ancestor.outside_compilation_cluster; + } + } + return errors::InvalidArgument( + "Can't compile outside_compilation clusters because there is a " + "dependency cycle: see error log for details."); + } + return Status::OK(); +} + +} // namespace + +Status Encapsulator::FindClusterDependencies() { + // Map from nodes to ancestor details. A node is entered into the map if it is + // in a compilation subgraph, and outside_compilation cluster, or appears on a + // path in the outer graph leading from an outside_compilation subgraph. + std::unordered_map node_ancestors_map; + // We check that clusters are acyclic using this cycle detector. + GraphCycles cycle_detector; + // Map from cluster name to cycle detector node id. + std::map cycle_detector_map; + // Process the nodes in topologically-sorted order. + std::vector nodes; + GetReversePostOrder(*graph_in_, &nodes); + for (Node* node : nodes) { + string subgraph_name; + string oc_cluster; + TF_RETURN_IF_ERROR(GetFunctionNameAttr(node, &subgraph_name, &oc_cluster)); + // First create an entry in the ancestors map if the node is in a compiled + // subgraph or outside_compilation cluster, or if any incoming edge is from + // a node with an ancestor map entry; and find the union of all the + // ancestors. + if (!subgraph_name.empty()) { + node_ancestors_map[node].subgraph = subgraph_name; + node_ancestors_map[node].outside_compilation_cluster = oc_cluster; + } + for (Node* src : node->in_nodes()) { + const auto iter = node_ancestors_map.find(src); + if (iter != node_ancestors_map.end()) { + const auto& ancestors_to_follow = iter->second.AncestorsForSuccessor(); + for (const auto& ancestor : ancestors_to_follow) { + if (ancestor.subgraph != subgraph_name || + ancestor.outside_compilation_cluster != oc_cluster) { + node_ancestors_map[node].ancestor_clusters.insert(ancestor); + } + } + } + } + if (!subgraph_name.empty()) { + // The node is in a compiled subgraph or an outside_compilation cluster. + if (oc_cluster.empty()) { + // The node is not in an outside_compilation cluster. Record the + // subgraph's ancestor dependencies. + for (const auto& cluster : node_ancestors_map[node].ancestor_clusters) { + if (cluster.subgraph != subgraph_name) { + subgraph_ancestors_[subgraph_name].insert(cluster.subgraph); + TF_RETURN_IF_ERROR(CheckClusterDependencyForCycles( + cluster.subgraph, subgraph_name, subgraph_ancestors_, + node_ancestors_map, &cycle_detector, &cycle_detector_map)); + } + } + } else { + Subgraph& subgraph = subgraphs_[subgraph_name]; + // The node is in an outside_compilation cluster. Record the cluster + // and/or subgraph ancestor dependencies. + for (const auto& cluster : node_ancestors_map[node].ancestor_clusters) { + if (cluster.subgraph == subgraph_name) { + // The ancestor is in the same subgraph. + if (cluster.outside_compilation_cluster != oc_cluster) { + // But not in the same oc_cluster, so record the dependency. + subgraph.RecordOutsideCompilationDependency( + oc_cluster, cluster.outside_compilation_cluster); + TF_RETURN_IF_ERROR(CheckClusterDependencyForCycles( + cluster.outside_compilation_cluster, oc_cluster, + subgraph.OutsideCompilationAncestorMap(), node_ancestors_map, + &cycle_detector, &cycle_detector_map)); + } + } else { + // The ancestor is in a different subgraph, so record the + // dependency. + subgraph_ancestors_[subgraph_name].insert(cluster.subgraph); + TF_RETURN_IF_ERROR(CheckClusterDependencyForCycles( + cluster.subgraph, subgraph_name, subgraph_ancestors_, + node_ancestors_map, &cycle_detector, &cycle_detector_map)); + } + } + } + } + } return Status::OK(); } @@ -1956,7 +2393,7 @@ Status Encapsulator::MakePrunedGraphCopyAndInline( Status Encapsulator::MakeGraphForOutsideCompilationSends( const Graph& graph, std::unique_ptr* pruned_graph, - ShapeRefiner* shape_refiner, + BackEdgeHelper* back_edge_helper, 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 @@ -1978,10 +2415,15 @@ Status Encapsulator::MakeGraphForOutsideCompilationSends( // nodes, inlining any functions as needed. TF_RETURN_IF_ERROR(MakePrunedGraphCopyAndInline( graph, send_from_host_nodes, pruned_graph, node_images, library)); + FixupSourceAndSinkEdges(pruned_graph->get()); + + // Remove back edges from any cycles in the pruned graph to simplify shape + // inference traversal. They will be fixed up in the per-subgraph shape + // inference graphs stored in the function library. + TF_RETURN_IF_ERROR(back_edge_helper->Remove(pruned_graph->get())); // 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) { @@ -1999,11 +2441,13 @@ Status Encapsulator::MakeGraphForOutsideCompilationSends( Status Encapsulator::GetShapeInfoForOutsideCompilationSends( Graph* graph_out, FunctionLibraryDefinition* library) { + BackEdgeHelper back_edge_helper; 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)); + *graph_out, &pruned_graph, &back_edge_helper, &shape_refiner, + &node_images, library)); if (VLOG_IS_ON(1)) { dump_graph::DumpGraphToFile("pruned_graph_for_shape_inference", @@ -2033,7 +2477,7 @@ Status Encapsulator::GetShapeInfoForOutsideCompilationSends( std::unique_ptr graph; if (send_node != nullptr) { TF_RETURN_IF_ERROR(DoStaticShapeInferenceForOutsideCompilationSend( - *pruned_graph, shape_refiner, recv_at_host_names, + *pruned_graph, back_edge_helper, shape_refiner, recv_at_host_names, node_images[send_node], library, &static_shape, &graph)); if (graph == nullptr) { VLOG(2) << "Send node " << send_node->name() << " shapes"; @@ -2091,6 +2535,7 @@ Status EncapsulateSubgraphsInFunctions( Encapsulator encapsulator(std::move(group_attribute), std::move(outside_compilation_attribute), &graph_in); + TF_RETURN_IF_ERROR(encapsulator.FindClusterDependencies()); TF_RETURN_IF_ERROR(encapsulator.SplitIntoSubgraphs()); TF_RETURN_IF_ERROR(encapsulator.BuildFunctionDefs( diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc index 8599a7038af9663e5af6f3231429cb7f6ea5f69b..3502d1bb4596bc8e8a57523d0d77dd3f12b9bec6 100644 --- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc +++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc @@ -74,7 +74,7 @@ bool EqualProtoMap(const ::tensorflow::protobuf::Map& a, 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 '", + key_to_string(elt_a.first), "' has value '", value_to_string(elt_a.second), "' got: '", value_to_string(iter->second), "'"); } @@ -121,8 +121,22 @@ bool EqualFunctionNodeDef(const NodeDef& a, const NodeDef& b, } return false; } + std::unordered_set control_input_a; + std::unordered_set control_input_b; for (int i = 0; i < a.input_size(); ++i) { - if (a.input(i) != b.input(i)) { + if (str_util::StartsWith(a.input(i), "^")) { + if (!str_util::StartsWith(b.input(i), "^")) { + if (diff) { + *diff = strings::StrCat( + diff_preamble, " mismatch for node ", a.name(), " input ", i, + ", expected control input ", a.input(i), " got ", b.input(i), + " expected:\n", a.DebugString(), "\ngot:\n", b.DebugString()); + } + return false; + } + control_input_a.insert(a.input(i)); + control_input_b.insert(b.input(i)); + } else 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), @@ -132,11 +146,29 @@ bool EqualFunctionNodeDef(const NodeDef& a, const NodeDef& b, return false; } } + if (control_input_a != control_input_b) { + if (diff) { + *diff = strings::StrCat(diff_preamble, " mismatch for node ", a.name(), + " control inputs differ 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) { - return av.DebugString() == bv.DebugString(); + if (key == "ancestors") { + // The ancestors are added from a set so the order is unpredictable; + // just compare set equality not list equality. + std::unordered_set a_set(av.list().s().begin(), + av.list().s().end()); + std::unordered_set b_set(bv.list().s().begin(), + bv.list().s().end()); + return a_set == b_set; + } else { + return av.DebugString() == bv.DebugString(); + } }, strings::StrCat(diff_preamble, " attr mismatch for node ", a.name()), diff); @@ -261,6 +293,7 @@ REGISTER_OP("XlaHostCompute") .Output("outputs: Toutputs") .Attr("Tinputs: list(type) >= 0") .Attr("Toutputs: list(type) >= 0") + .Attr("ancestors: list(string) >= 0") .Attr("key: string") .Attr("shape_inference_graph: string = ''") .Attr("shapes: list(shape) >= 0") @@ -899,6 +932,7 @@ TEST(EncapsulateSubgraphsTest, OneFunctionOneOutside) { {"C:o:0", "c:o:0"}, {{"Tinputs", gtl::ArraySlice({DT_FLOAT, DT_FLOAT})}, {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"ancestors", gtl::ArraySlice({})}, {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", "_outside_compilation_shape_inference_F1_O1"}, @@ -1044,17 +1078,20 @@ TEST(EncapsulateSubgraphsTest, OneFunctionTwoOutside) { {"D:o:0", "F:o:0"}, {{"Tinputs", gtl::ArraySlice({DT_FLOAT, DT_FLOAT})}, {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"ancestors", + gtl::ArraySlice({"outside_compilation_O1_host_compute"})}, {"key", "host_compute_channel_F1_O2"}, {"shape_inference_graph", "_outside_compilation_shape_inference_F1_O2"}, {"shapes", gtl::ArraySlice({})}, {"_outside_compilation_subgraph", "O2"}}, - {"F"}}, + {"F", "outside_compilation_O1_host_compute"}}, {{"outside_compilation_O1_host_compute"}, "XlaHostCompute", {"C:o:0", "D:o:0"}, {{"Tinputs", gtl::ArraySlice({DT_FLOAT, DT_FLOAT})}, {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"ancestors", gtl::ArraySlice({})}, {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", "_outside_compilation_shape_inference_F1_O1"}, @@ -1193,6 +1230,7 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { {"C:o:0", "D:o:0"}, {{"Tinputs", gtl::ArraySlice({DT_FLOAT, DT_FLOAT})}, {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"ancestors", gtl::ArraySlice({})}, {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", "_outside_compilation_shape_inference_F1_O1"}, @@ -1215,6 +1253,7 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { {"G:o:0"}, {{"Tinputs", gtl::ArraySlice({DT_FLOAT})}, {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"ancestors", gtl::ArraySlice({})}, {"key", "host_compute_channel_F2_O1"}, {"shape_inference_graph", ""}, {"shapes", @@ -1279,6 +1318,179 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { TF_EXPECT_FUNCTIONDEFLIBRARY_EQ(library_expected, library); } +// Test with two functions to transform, each with one outside_compilation +// cluster, with the dependency between them purely from an outside_compilation +// edge. +TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutsideDependencyFromOutside) { + FunctionDefLibrary library; + GraphDef graphdef; + + { + GraphDefBuilder b1(GraphDefBuilder::kFailImmediately); + 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")); + Node* e = Binary(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")); + Node* g = + Binary(a, b, b1.opts().WithName("G").WithAttr("_encapsulate", "F2")); + Node* h = Unary(g, b1.opts() + .WithName("H") + .WithAttr("_encapsulate", "F2") + .WithAttr("_outside", "O1") + .WithControlInput(e)); + Node* i = Unary(h, b1.opts().WithName("I").WithAttr("_encapsulate", "F2")); + Binary(f, i, b1.opts().WithName("J")); + TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); + } + + TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + + FunctionDefLibrary library_expected; + GraphDef graphdef_expected; + + { + GraphDefBuilder shape(GraphDefBuilder::kFailImmediately); + Node* key_constant = + KeyPlaceholderShape(shape.opts().WithName("KnownShape/_0")); + Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", + {DT_FLOAT, DT_FLOAT}, shape.opts()); + Node* e = Binary(ops::NodeOut(recv, 0), ops::NodeOut(recv, 1), + shape.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, shape.opts()); + TF_EXPECT_OK( + AddGraphDefToFunctionLibrary(shape, "F1_O1", &library_expected)); + } + + { + GraphDefBuilder shape(GraphDefBuilder::kFailImmediately); + Node* key_constant = + KeyPlaceholderShape(shape.opts().WithName("KnownShape/_0")); + Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F2", "O1", + {DT_FLOAT}, shape.opts()); + Node* h = Unary(recv, shape.opts() + .WithName("H") + .WithAttr("_encapsulate", "F2") + .WithAttr("_outside", "O1")); + SendFromHost(ops::NodeOut(key_constant, 0), "F2", "O1", {h}, shape.opts()); + TF_EXPECT_OK( + AddGraphDefToFunctionLibrary(shape, "F2_O1", &library_expected)); + } + + *library_expected.add_function() = FunctionDefHelper::Create( + "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval:float"}, {}, + { + {{"C"}, "UnaryTest", {"a_0_arg"}}, + {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, + {{"F"}, + "BinaryTest", + {"C:o:0", "outside_compilation_O1_host_compute:outputs:0"}, + {}, + {"outside_compilation_O1_host_compute"}}, + {{"outside_compilation_O1_host_compute"}, + "XlaHostCompute", + {"C:o:0", "D:o:0"}, + {{"Tinputs", gtl::ArraySlice({DT_FLOAT, DT_FLOAT})}, + {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"ancestors", gtl::ArraySlice({})}, + {"key", "host_compute_channel_F1_O1"}, + {"shape_inference_graph", + "_outside_compilation_shape_inference_F1_O1"}, + {"shapes", gtl::ArraySlice({})}, + {"_outside_compilation_subgraph", "O1"}}, + {"D"}}, + }, + {{"f_0_retval", "F:o:0"}}); + + *library_expected.add_function() = FunctionDefHelper::Create( + "F2", {"a_0_arg:float", "b_0_arg:float"}, {"i_0_retval:float"}, {}, + { + {{"G"}, "BinaryTest", {"a_0_arg", "b_0_arg"}}, + {{"I"}, + "UnaryTest", + {"outside_compilation_O1_host_compute:outputs:0"}}, + {{"outside_compilation_O1_host_compute"}, + "XlaHostCompute", + {"G:o:0"}, + {{"Tinputs", gtl::ArraySlice({DT_FLOAT})}, + {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"ancestors", gtl::ArraySlice({})}, + {"key", "host_compute_channel_F2_O1"}, + {"shape_inference_graph", + "_outside_compilation_shape_inference_F2_O1"}, + {"shapes", gtl::ArraySlice({})}, + {"_outside_compilation_subgraph", "O1"}}}, + }, + {{"i_0_retval", "I: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 = InputShaped(b2.opts().WithName("B")); + + Node* key_constant1 = + KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); + Node* recv1 = RecvAtHost(ops::NodeOut(key_constant1, 0), "F1", "O1", + {DT_FLOAT, DT_FLOAT}, b2.opts()); + Node* e = Binary(ops::NodeOut(recv1, 0), ops::NodeOut(recv1, 1), + b2.opts() + .WithName("E") + .WithControlInputs({recv1, b}) + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + Node* send1 = SendFromHost(ops::NodeOut(key_constant1, 0), "F1", "O1", {e}, + b2.opts().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().WithControlInput(s1).FinalizeBuilder(&node_builder1); + + Node* key_constant2 = + KeyPlaceholder("F2", b2.opts().WithName("F2_key_placeholder")); + Node* recv2 = RecvAtHost(ops::NodeOut(key_constant2, 0), "F2", "O1", + {DT_FLOAT}, b2.opts()); + Node* h = Unary(recv2, b2.opts() + .WithName("H") + .WithAttr("_encapsulate", "F2") + .WithAttr("_outside", "O1") + .WithControlInput(e)); + Node* send2 = SendFromHost(ops::NodeOut(key_constant2, 0), "F2", "O1", {h}, + b2.opts()); + + Node* s2 = Sequencer( + b2.opts().WithName("F2_sequencer").WithControlInputs({recv2, send2}), + "F2"); + NodeBuilder node_builder2("F2", "F2", lib_def.get()); + node_builder2.Input(a).Input(b); + Node* call2 = b2.opts() + .WithControlInputs({s2, call1}) + .FinalizeBuilder(&node_builder2); + Binary(call1, call2, b2.opts().WithName("J")); + TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); + } + + TF_EXPECT_GRAPH_EQ(graphdef_expected, graphdef); + TF_EXPECT_FUNCTIONDEFLIBRARY_EQ(library_expected, library); +} + // Test with one outside_compilation cluster that has no inputs from the // compiled subgraph. TEST(EncapsulateSubgraphsTest, OutsideCompilationNoInputs) { @@ -1323,6 +1535,7 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoInputs) { {}, {{"Tinputs", gtl::ArraySlice({})}, {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"ancestors", gtl::ArraySlice({})}, {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", ""}, {"shapes", @@ -1406,6 +1619,7 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlInput) { {}, {{"Tinputs", gtl::ArraySlice({})}, {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"ancestors", gtl::ArraySlice({})}, {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", ""}, {"shapes", @@ -1487,6 +1701,7 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoOutputs) { {"D:o:0"}, {{"Tinputs", gtl::ArraySlice({DT_FLOAT})}, {"Toutputs", gtl::ArraySlice({})}, + {"ancestors", gtl::ArraySlice({})}, {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", ""}, {"shapes", gtl::ArraySlice({})}, @@ -1567,6 +1782,7 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlOutput) { {"D:o:0"}, {{"Tinputs", gtl::ArraySlice({DT_FLOAT})}, {"Toutputs", gtl::ArraySlice({})}, + {"ancestors", gtl::ArraySlice({})}, {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", ""}, {"shapes", gtl::ArraySlice({})}, @@ -1607,6 +1823,371 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlOutput) { TF_EXPECT_FUNCTIONDEFLIBRARY_EQ(library_expected, library); } +// Test with two outside_compilation clusters that interact outside the compiled +// subgraph, where the ancestor has no HostCompute Op. +TEST(EncapsulateSubgraphsTest, + OutsideCompilationClusterDependencyNoSrcCluster) { + FunctionDefLibrary library; + GraphDef graphdef; + + { + GraphDefBuilder b1(GraphDefBuilder::kFailImmediately); + Node* a = Input(b1.opts().WithName("A")); + Node* b = Input(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")); + Node* e = Unary(a, b1.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + Node* f = Unary(d, b1.opts().WithName("F").WithAttr("_encapsulate", "F1")); + Node* g = Unary(f, b1.opts() + .WithName("G") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O2") + .WithControlInput(e)); + Node* h = Unary(g, b1.opts().WithName("H").WithAttr("_encapsulate", "F1")); + Binary(e, h, b1.opts().WithName("I")); + TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); + } + + TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + + FunctionDefLibrary library_expected; + GraphDef graphdef_expected; + + { + GraphDefBuilder shape2(GraphDefBuilder::kFailImmediately); + Node* key_constant = + KeyPlaceholderShape(shape2.opts().WithName("KnownShape/_0")); + Node* recv2 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2", + {DT_FLOAT}, shape2.opts()); + Node* g = Unary(ops::NodeOut(recv2, 0), shape2.opts() + .WithName("G") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O2")); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {g}, shape2.opts()); + TF_EXPECT_OK( + AddGraphDefToFunctionLibrary(shape2, "F1_O2", &library_expected)); + } + + *library_expected.add_function() = FunctionDefHelper::Create( + "F1", {"a_0_arg:float", "b_0_arg:float"}, {"h_0_retval:float"}, {}, + { + {{"C"}, "UnaryTest", {"a_0_arg"}}, + {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, + {{"F"}, "UnaryTest", {"D:o:0"}}, + {{"H"}, + "UnaryTest", + {"outside_compilation_O2_host_compute:outputs:0"}}, + {{"outside_compilation_O2_host_compute"}, + "XlaHostCompute", + {"F:o:0"}, + {{"Tinputs", gtl::ArraySlice({DT_FLOAT})}, + {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"ancestors", gtl::ArraySlice({})}, + {"key", "host_compute_channel_F1_O2"}, + {"shape_inference_graph", + "_outside_compilation_shape_inference_F1_O2"}, + {"shapes", gtl::ArraySlice({})}, + {"_outside_compilation_subgraph", "O2"}}}, + }, + {{"h_0_retval", "H:o:0"}}); + + { + 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* e = Unary(a, b2.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + Node* key_constant = + KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); + Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2", + {DT_FLOAT}, b2.opts()); + Node* g = Unary(recv, b2.opts() + .WithName("G") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O2") + .WithControlInput(e)); + Node* send = + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {g}, b2.opts()); + Node* s1 = Sequencer( + b2.opts().WithName("F1_sequencer").WithControlInputs({recv, send}), + "F1"); + NodeBuilder node_builder1("F1", "F1", lib_def.get()); + node_builder1.Input(a).Input(b).ControlInput(s1); + Node* call1 = b2.opts().FinalizeBuilder(&node_builder1); + + Binary(e, call1, b2.opts().WithName("I")); + TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); + } + + TF_EXPECT_GRAPH_EQ(graphdef_expected, graphdef); + TF_EXPECT_FUNCTIONDEFLIBRARY_EQ(library_expected, library); +} + +// Test with two outside_compilation clusters that interact outside the compiled +// subgraph, where the successor has no HostCompute Op. +TEST(EncapsulateSubgraphsTest, + OutsideCompilationClusterDependencyNoDstCluster) { + FunctionDefLibrary library; + GraphDef graphdef; + + { + GraphDefBuilder b1(GraphDefBuilder::kFailImmediately); + Node* a = Input(b1.opts().WithName("A")); + Node* b = Input(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")); + Node* e = Unary(d, b1.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + Node* f = Unary(e, b1.opts().WithName("F").WithAttr("_encapsulate", "F1")); + /*Node* g =*/Unary(a, b1.opts() + .WithName("G") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O2") + .WithControlInput(e)); + Node* h = Unary(f, b1.opts().WithName("H").WithAttr("_encapsulate", "F1")); + Binary(e, h, b1.opts().WithName("I")); + TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); + } + + TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + + FunctionDefLibrary library_expected; + GraphDef graphdef_expected; + + { + GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately); + Node* key_constant = + KeyPlaceholderShape(shape1.opts().WithName("KnownShape/_0")); + Node* recv2 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", + {DT_FLOAT}, shape1.opts()); + Node* e = Unary(ops::NodeOut(recv2, 0), shape1.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, shape1.opts()); + TF_EXPECT_OK( + AddGraphDefToFunctionLibrary(shape1, "F1_O1", &library_expected)); + } + + *library_expected.add_function() = FunctionDefHelper::Create( + "F1", {"a_0_arg:float", "b_0_arg:float"}, {"h_0_retval:float"}, {}, + { + {{"C"}, "UnaryTest", {"a_0_arg"}}, + {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, + {{"F"}, + "UnaryTest", + {"outside_compilation_O1_host_compute:outputs:0"}}, + {{"H"}, "UnaryTest", {"F:o:0"}}, + {{"outside_compilation_O1_host_compute"}, + "XlaHostCompute", + {"D:o:0"}, + {{"Tinputs", gtl::ArraySlice({DT_FLOAT})}, + {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"ancestors", gtl::ArraySlice({})}, + {"key", "host_compute_channel_F1_O1"}, + {"shape_inference_graph", + "_outside_compilation_shape_inference_F1_O1"}, + {"shapes", gtl::ArraySlice({})}, + {"_outside_compilation_subgraph", "O1"}}}, + }, + {{"h_0_retval", "H:o:0"}}); + + { + 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* key_constant = + KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); + Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", + {DT_FLOAT}, b2.opts()); + Node* e = Unary(recv, b2.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + Node* send = + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, b2.opts()); + /*Node* g =*/Unary(a, b2.opts() + .WithName("G") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O2") + .WithControlInput(e)); + Node* s1 = Sequencer( + b2.opts().WithName("F1_sequencer").WithControlInputs({recv, send}), + "F1"); + NodeBuilder node_builder1("F1", "F1", lib_def.get()); + node_builder1.Input(a).Input(b).ControlInput(s1); + Node* call1 = b2.opts().FinalizeBuilder(&node_builder1); + + Binary(e, call1, b2.opts().WithName("I")); + TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); + } + + TF_EXPECT_GRAPH_EQ(graphdef_expected, graphdef); + TF_EXPECT_FUNCTIONDEFLIBRARY_EQ(library_expected, library); +} + +// Test with two outside_compilation clusters that interact outside the compiled +// subgraph. +TEST(EncapsulateSubgraphsTest, OutsideCompilationClusterDependency) { + FunctionDefLibrary library; + GraphDef graphdef; + + { + GraphDefBuilder b1(GraphDefBuilder::kFailImmediately); + Node* a = Input(b1.opts().WithName("A")); + Node* b = Input(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")); + Node* e = Unary(d, b1.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + Node* f = Unary(e, b1.opts().WithName("F").WithAttr("_encapsulate", "F1")); + Node* g = Unary(d, b1.opts() + .WithName("G") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O2") + .WithControlInput(e)); + Node* h = Unary(f, b1.opts().WithName("H").WithAttr("_encapsulate", "F1")); + /*Node* i =*/Binary(d, e, + b1.opts() + .WithName("I") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O3") + .WithControlInput(g)); + Binary(e, h, b1.opts().WithName("J")); + TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); + } + + TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + + FunctionDefLibrary library_expected; + GraphDef graphdef_expected; + + { + GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately); + Node* key_constant = + KeyPlaceholderShape(shape1.opts().WithName("KnownShape/_0")); + Node* recv2 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", + {DT_FLOAT}, shape1.opts()); + Node* e = Unary(ops::NodeOut(recv2, 0), shape1.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, shape1.opts()); + TF_EXPECT_OK( + AddGraphDefToFunctionLibrary(shape1, "F1_O1", &library_expected)); + } + + *library_expected.add_function() = FunctionDefHelper::Create( + "F1", {"a_0_arg:float", "b_0_arg:float"}, {"h_0_retval:float"}, {}, + {{{"C"}, "UnaryTest", {"a_0_arg"}}, + {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, + {{"F"}, "UnaryTest", {"outside_compilation_O1_host_compute:outputs:0"}}, + {{"H"}, "UnaryTest", {"F:o:0"}}, + {{"outside_compilation_O1_host_compute"}, + "XlaHostCompute", + {"D:o:0"}, + {{"Tinputs", gtl::ArraySlice({DT_FLOAT})}, + {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"ancestors", gtl::ArraySlice({})}, + {"key", "host_compute_channel_F1_O1"}, + {"shape_inference_graph", + "_outside_compilation_shape_inference_F1_O1"}, + {"shapes", gtl::ArraySlice({})}, + {"_outside_compilation_subgraph", "O1"}}}, + {{"outside_compilation_O2_host_compute"}, + "XlaHostCompute", + {"D:o:0"}, + {{"Tinputs", gtl::ArraySlice({DT_FLOAT})}, + {"Toutputs", gtl::ArraySlice({})}, + {"ancestors", + gtl::ArraySlice({"outside_compilation_O1_host_compute"})}, + {"key", "host_compute_channel_F1_O2"}, + {"shape_inference_graph", ""}, + {"shapes", gtl::ArraySlice({})}, + {"_outside_compilation_subgraph", "O2"}}, + {"outside_compilation_O1_host_compute"}}, + {{"outside_compilation_O3_host_compute"}, + "XlaHostCompute", + {"D:o:0"}, + {{"Tinputs", gtl::ArraySlice({DT_FLOAT})}, + {"Toutputs", gtl::ArraySlice({})}, + {"ancestors", + gtl::ArraySlice({"outside_compilation_O1_host_compute", + "outside_compilation_O2_host_compute"})}, + {"key", "host_compute_channel_F1_O3"}, + {"shape_inference_graph", ""}, + {"shapes", gtl::ArraySlice({})}, + {"_outside_compilation_subgraph", "O3"}}, + {"outside_compilation_O1_host_compute", + "outside_compilation_O2_host_compute"}}}, + {{"h_0_retval", "H:o:0"}}); + + { + 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* key_constant = + KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); + Node* recv1 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", + {DT_FLOAT}, b2.opts()); + Node* e = Unary(recv1, b2.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + Node* send = + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, b2.opts()); + Node* recv2 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2", + {DT_FLOAT}, b2.opts()); + Node* g = Unary(recv2, b2.opts() + .WithName("G") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O2") + .WithControlInput(e)); + Node* recv3 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O3", + {DT_FLOAT}, b2.opts()); + /*Node* i =*/Binary(recv3, e, + b2.opts() + .WithName("I") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O3") + .WithControlInput(g)); + Node* s1 = Sequencer(b2.opts() + .WithName("F1_sequencer") + .WithControlInputs({recv1, send, recv2, recv3}), + "F1"); + NodeBuilder node_builder1("F1", "F1", lib_def.get()); + node_builder1.Input(a).Input(b).ControlInput(s1); + Node* call1 = b2.opts().FinalizeBuilder(&node_builder1); + + Binary(e, call1, b2.opts().WithName("J")); + TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); + } + + TF_EXPECT_GRAPH_EQ(graphdef_expected, graphdef); + TF_EXPECT_FUNCTIONDEFLIBRARY_EQ(library_expected, library); +} + // Test with one outside_compilation cluster that has no outputs from the // compiled subgraph. TEST(EncapsulateSubgraphsTest, OutsideCompilationNoInputsOrOutputs) { @@ -1731,6 +2312,7 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationShapeInference) { {"c:o:0"}, {{"Tinputs", gtl::ArraySlice({DT_FLOAT})}, {"Toutputs", gtl::ArraySlice({DT_FLOAT})}, + {"ancestors", gtl::ArraySlice({})}, {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", "_outside_compilation_shape_inference_F1_O1"}, diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc index 386240ff8d1a562ff4894c40ac79022b8c61fb15..8e2ee0f1d71bc17b4c12c792c38002af4f9eb5eb 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc @@ -51,6 +51,15 @@ bool HasXLAKernel(const Node& node, const DeviceType& jit_device_type) { // is really a kind of function call and will be handled by // IsCompilableCall(). if (node.type_string() == "SymbolicGradient") return false; + if (node.type_string() == "Const") { + // Skip Const op with type DT_STRING, since XLA doesn't support it, but the + // registered Const KernelDef says that it does, to support no-op Assert for + // tfcompile. + const AttrValue* attr = node.attrs().Find("dtype"); + if (attr != nullptr && attr->type() == DT_STRING) { + return false; + } + } return FindKernelDef(jit_device_type, node.def(), nullptr, nullptr).ok(); } diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc index 80edaf28b83348e3a8071a6e9696bc3ebad5d70f..703d8825d74ced8d4d69c31ccd730adc89a8bffe 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc @@ -609,5 +609,29 @@ TEST(XlaCompilationTest, DontCountIdentityOpsWithLocalJit) { EXPECT_TRUE(clusters.empty()); } +TEST(XlaCompilationTest, ConstOp) { + // valid data type + { + std::unique_ptr graph(new Graph(OpRegistry::Global())); + Scope root = Scope::NewRootScope().ExitOnError(); + auto c = ops::Const(root.WithOpName("const"), 0.5f); + c.node()->AddAttr(kXlaCompileAttr, true); + TF_ASSERT_OK(root.ToGraph(graph.get())); + TF_ASSERT_OK(MarkForCompilation(&graph)); + EXPECT_EQ(1, GetClusters(*graph).size()); + } + + // invalid data type + { + std::unique_ptr graph(new Graph(OpRegistry::Global())); + Scope root = Scope::NewRootScope().ExitOnError(); + auto c = ops::Const(root.WithOpName("const"), string("string")); + c.node()->AddAttr(kXlaCompileAttr, true); + TF_ASSERT_OK(root.ToGraph(graph.get())); + TF_ASSERT_OK(MarkForCompilation(&graph)); + EXPECT_TRUE(GetClusters(*graph).empty()); + } +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/jit/producer_consumer_queue.h b/tensorflow/compiler/jit/producer_consumer_queue.h new file mode 100644 index 0000000000000000000000000000000000000000..7c8c04152d2f3a0fd46711df24756b7e68b967ea --- /dev/null +++ b/tensorflow/compiler/jit/producer_consumer_queue.h @@ -0,0 +1,132 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_JIT_PRODUCER_CONSUMER_QUEUE_H_ +#define TENSORFLOW_COMPILER_JIT_PRODUCER_CONSUMER_QUEUE_H_ + +#include +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/mutex.h" + +namespace tensorflow { + +// A thread-safe, first-in-first-out queue. +template +class ProducerConsumerQueue { + public: + ProducerConsumerQueue() + : capacity_(std::numeric_limits::max()) {} + ~ProducerConsumerQueue() = default; + + // Wait until the queue is non-full, then append a copy of v. + void Put(const T &v); + + // Wait until the queue is non-empty, then remove and return the head value. + T Get(); + + // If the queue is non-empty, remove the head value, placing it in *pv, and + // return true; otherwise return false. + bool TryGet(T *pv); + + // Set the capacity of the queue; the queue is full whenever count() >= + // capacity(). The initial value is the maximum size_t. Requires size > 0. + void set_capacity(std::size_t size); + + // Return the capacity of the queue. + std::size_t capacity() const; + + // Return the number of elements in the queue. + std::size_t count() const; + + // Implementation details follow. Clients should ignore. + private: + mutable tensorflow::mutex mu_; // protects all fields below + tensorflow::condition_variable non_empty_ GUARDED_BY(mu_); + tensorflow::condition_variable non_full_ GUARDED_BY(mu_); + std::size_t capacity_ GUARDED_BY(mu_); + std::deque queue_ GUARDED_BY(mu_); + + TF_DISALLOW_COPY_AND_ASSIGN(ProducerConsumerQueue); +}; + +// ------------------------------------------------------ +// Implementation details follow. Clients should ignore. + +// Wait until the queue is non-full, then append a copy of v. +template +void ProducerConsumerQueue::Put(const T &v) { + mutex_lock lock(mu_); + while (queue_.size() >= capacity_) { + non_full_.wait(lock); + } + queue_.push_back(v); + non_empty_.notify_one(); +} + +// Wait until the queue is non-empty, then remove and return the head value. +template +T ProducerConsumerQueue::Get() { + mutex_lock lock(mu_); + while (queue_.empty()) { + non_empty_.wait(lock); + } + non_full_.notify_one(); + T result_value = queue_.front(); + queue_.pop_front(); + return result_value; +} + +// If the queue is non-empty, remove the head value, placing it in *pv, and +// return true; otherwise return false. +template +bool ProducerConsumerQueue::TryGet(T *pv) { + mutex_lock lock(mu_); + bool got_element = !queue_.empty(); + if (got_element) { + non_full_.notify_one(); + *pv = queue_.front(); + queue_.pop_front(); + } + return got_element; +} + +// Set the capacity of the queue; the queue is full whenever count() >= +// capacity(). The initial value is the maximum size_t. Requires size > 0. +template +void ProducerConsumerQueue::set_capacity(std::size_t size) { + mutex_lock lock(mu_); + CHECK_NE(size, 0); + capacity_ = size; + non_full_.notify_all(); +} + +// Return the capacity of the queue. +template +std::size_t ProducerConsumerQueue::capacity() const { + mutex_lock lock(mu_); + std::size_t max_elements = capacity_; + return max_elements; +} + +// Return the number of elements in the queue. +template +std::size_t ProducerConsumerQueue::count() const { + mutex_lock lock(mu_); + std::size_t num_elements = queue_.size(); + return num_elements; +} +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_JIT_PRODUCER_CONSUMER_QUEUE_H_ diff --git a/tensorflow/compiler/jit/producer_consumer_queue_test.cc b/tensorflow/compiler/jit/producer_consumer_queue_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..f61260c6e52756ee039829afdc7452f5f760c221 --- /dev/null +++ b/tensorflow/compiler/jit/producer_consumer_queue_test.cc @@ -0,0 +1,139 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/jit/producer_consumer_queue.h" + +#include "tensorflow/core/lib/core/threadpool.h" +#include "tensorflow/core/platform/env.h" +#include "tensorflow/core/platform/mutex.h" +#include "tensorflow/core/platform/test.h" + +namespace tensorflow { +namespace { + +typedef ProducerConsumerQueue IntQueue; + +// Insert integers between low inclusive and high exclusive into q. +void PushRange(IntQueue *q, int low, int high) { + while (low != high) { + q->Put(low); + VLOG(2) << "Pushing " << low; + ++low; + } +} + +// Push the numbers between 0 and 999 inclusive from several threads in the +// pool. +void PushRanges(IntQueue *queue, thread::ThreadPool *pool) { + VLOG(1) << "Adding 20-36"; + pool->Schedule([queue] { PushRange(queue, 20, 36); }); + VLOG(1) << "Adding 7-20"; + pool->Schedule([queue] { PushRange(queue, 7, 20); }); + VLOG(1) << "Adding 36-501"; + pool->Schedule([queue] { PushRange(queue, 36, 501); }); + VLOG(1) << "Adding 501-1000"; + pool->Schedule([queue] { PushRange(queue, 501, 1000); }); + VLOG(1) << "Adding 0-5"; + pool->Schedule([queue] { PushRange(queue, 0, 5); }); + VLOG(1) << "Adding 5-7"; + pool->Schedule([queue] { PushRange(queue, 5, 7); }); +} + +// Pop elements from queue using Get(). Make sure that exactly elements +// were present and their values are all integers between 0 and high-1 +// inclusive. +void GetRange(IntQueue *queue, int high) { + VLOG(1) << "Testing Wait"; + std::vector results; + for (int i = 0; i != high; ++i) { + int r = queue->Get(); + VLOG(2) << "Waited and got " << r; + results.push_back(r); + } + CHECK_EQ(queue->count(), 0); + std::sort(results.begin(), results.end()); + for (int i = 0; i != high; ++i) { + CHECK(results[i] == i); + } +} + +// Pop elements from queue using TryGet(). Make sure that exactly +// elements were present and their values are all integers between 0 and high-1 +// inclusive. +void TryGetRange(IntQueue *queue, int high) { + std::vector results; + // Give up if we don't get all the elements back from the queue + // in 10 seconds. + int timeout = 10; + int r; + for (int i = 0; i != high; ++i) { + while (!queue->TryGet(&r)) { + if (!timeout--) { + LOG(FATAL) << "Can't find all elements in the queue"; + } + VLOG(1) << "Sleeping for a second..."; + sleep(1); + } + VLOG(2) << "Popped " << r; + results.push_back(r); + } + CHECK_EQ(queue->count(), 0); + CHECK(!queue->TryGet(&r)); + std::sort(results.begin(), results.end()); + for (int i = 0; i != high; ++i) { + CHECK_EQ(i, results[i]); + } +} + +const int kNumThreads = 15; + +TEST(ProducerConsumerQueue, GetRange) { + IntQueue queue; + { + thread::ThreadPool pool(Env::Default(), "test", kNumThreads); + PushRanges(&queue, &pool); + } + GetRange(&queue, 1000); +} + +TEST(ProducerConsumerQueue, TryGetRange) { + IntQueue queue; + { + thread::ThreadPool pool(Env::Default(), "test", kNumThreads); + PushRanges(&queue, &pool); + } + TryGetRange(&queue, 1000); +} + +TEST(ProducerConsumerQueue, ParallelGetRange) { + IntQueue queue; + { + thread::ThreadPool pool(Env::Default(), "test", kNumThreads); + pool.Schedule([&queue] { GetRange(&queue, 1000); }); + PushRanges(&queue, &pool); + } +} + +TEST(ProducerConsumerQueue, ParallelTryGetRange) { + IntQueue queue; + { + thread::ThreadPool pool(Env::Default(), "test", kNumThreads); + pool.Schedule([&queue] { TryGetRange(&queue, 1000); }); + PushRanges(&queue, &pool); + } +} + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/jit/shape_inference_helpers.cc b/tensorflow/compiler/jit/shape_inference_helpers.cc new file mode 100644 index 0000000000000000000000000000000000000000..d9cfa16526bc5d809942a35e86075b4ec6e88a59 --- /dev/null +++ b/tensorflow/compiler/jit/shape_inference_helpers.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. +==============================================================================*/ + +// Contains helpers for use in shape inference. + +#include "tensorflow/compiler/jit/shape_inference_helpers.h" + +#include + +#include "tensorflow/core/graph/graph.h" + +namespace tensorflow { + +Status BackEdgeHelper::Remove(Graph* graph) { + if (graph_ != nullptr) { + return errors::Internal("BackEdgeHelper duplicate call to Remove."); + } + graph_ = graph; + for (Node* n : graph_->nodes()) { + if (n->IsMerge()) { + for (const Edge* e : n->in_edges()) { + if (e->src()->IsNextIteration()) { + back_edges_.push_back( + BackEdge{e, e->src(), e->src_output(), e->dst(), e->dst_input()}); + } + } + } + } + for (const BackEdge& be : back_edges_) { + graph_->RemoveEdge(be.edge); + } + return Status::OK(); +} + +const std::vector& BackEdgeHelper::RemovedEdges() + const { + return back_edges_; +} + +Status BackEdgeHelper::Replace() { + if (graph_ == nullptr) { + return errors::Internal("BackEdgeHelper Replace called before Remove."); + } + if (replaced_) { + return errors::Internal("BackEdgeHelper Replace called more than once."); + } + replaced_ = true; + for (const BackEdge& be : back_edges_) { + graph_->AddEdge(be.src, be.src_output, be.dst, be.dst_input); + } + return Status::OK(); +} + +} // namespace tensorflow diff --git a/tensorflow/compiler/jit/shape_inference_helpers.h b/tensorflow/compiler/jit/shape_inference_helpers.h new file mode 100644 index 0000000000000000000000000000000000000000..2f053c9a45dd47ca1b056634d2248d6181e77d68 --- /dev/null +++ b/tensorflow/compiler/jit/shape_inference_helpers.h @@ -0,0 +1,65 @@ +/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_JIT_SHAPE_INFERENCE_HELPERS_H_ +#define TENSORFLOW_COMPILER_JIT_SHAPE_INFERENCE_HELPERS_H_ + +#include + +#include "tensorflow/core/graph/graph.h" + +namespace tensorflow { + +// Helper class to temporarily remove, then replace, the back edges in a +// graph. Simple algorithms for shape inference don't work with cycles, and this +// class can be used to remove cycles before running inference and replace them +// after. Correct usage requires exactly one call to Remove(), followed by any +// number of calls to RemovedEdges() and at most one call to Replace(). The call +// to Replace() is optional if the graph will be discarded without being +// executed, e.g., if it is being used purely for a shape inference pass. +class BackEdgeHelper { + public: + struct BackEdge { + const Edge* edge; + Node* src; + int src_output; + Node* dst; + int dst_input; + }; + + BackEdgeHelper() = default; + // Disallows copy and assign. + BackEdgeHelper(const BackEdgeHelper& other) = delete; + BackEdgeHelper& operator=(const BackEdgeHelper& other) = delete; + + // Temporarily removes all the back edges in graph. + Status Remove(Graph* graph); + + // Gets the list of removed edges. + const std::vector& RemovedEdges() const; + + // Replaces the back edges removed by a prior call to Remove. + Status Replace(); + + private: + Graph* graph_ = nullptr; // not owned + std::vector back_edges_; + // Set once Replace has been called. + bool replaced_ = false; +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_JIT_SHAPE_INFERENCE_HELPERS_H_ diff --git a/tensorflow/compiler/jit/xla_compile_on_demand_op.cc b/tensorflow/compiler/jit/xla_compile_on_demand_op.cc index 682d6ea8ccc4a54912ccad4666cf0a7a03a7a698..6c2782e28e99df0801ab8d938fa3f961e4e1574c 100644 --- a/tensorflow/compiler/jit/xla_compile_on_demand_op.cc +++ b/tensorflow/compiler/jit/xla_compile_on_demand_op.cc @@ -67,6 +67,7 @@ Status XlaCompileOnDemandOp::Run(OpKernelContext* ctx, run_options.set_stream(stream); run_options.set_allocator(client->backend().memory_allocator()); run_options.set_intra_op_thread_pool(&ctx->eigen_cpu_device()); + run_options.set_rng_seed(ctx->step_id()); auto run_result = executable->Run(launch_context.arguments(), run_options); TF_RETURN_IF_ERROR(run_result.status()); diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc index 50b0061d692f2a8c5ea475c0b00c4cb42a1a84e6..3520501c1a39d243cc84583c47cbb0663bb201a8 100644 --- a/tensorflow/compiler/jit/xla_launch_util.cc +++ b/tensorflow/compiler/jit/xla_launch_util.cc @@ -32,10 +32,13 @@ limitations under the License. #include "tensorflow/core/framework/types.h" #include "tensorflow/core/util/stream_executor_util.h" +namespace { namespace gpu = perftools::gputools; +using xla::ScopedShapedBuffer; +using xla::ShapedBuffer; +} // anonymous namespace namespace tensorflow { - std::map SnapshotResourceVariables(OpKernelContext* ctx, int num_variables) { std::map snapshot; @@ -80,17 +83,17 @@ namespace { // Return the 'index''th subtree of the given ShapedBuffer as a // ScopedShapedBuffer. The returned ScopedShapedBuffer takes ownership of the // subtree, and sets the input's buffer pointers to nullptr for the subtree. -std::unique_ptr ExtractSubShapedBuffer( - xla::ShapedBuffer* shaped_buffer, int index, +ScopedShapedBuffer ExtractSubShapedBuffer( + ShapedBuffer* shaped_buffer, int index, xla::DeviceMemoryAllocator* allocator) { 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()); + 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(); @@ -102,8 +105,7 @@ std::unique_ptr ExtractSubShapedBuffer( index_to_buffer.second = gpu::DeviceMemoryBase(nullptr, 0); } } - return xla::ScopedShapedBuffer::MakeScoped(&sub_shaped_buffer, allocator) - .ValueOrDie(); + return ScopedShapedBuffer(std::move(sub_shaped_buffer), allocator); } } // namespace @@ -118,10 +120,10 @@ XlaComputationLaunchContext::XlaComputationLaunchContext( void XlaComputationLaunchContext::PopulateInputs( OpKernelContext* ctx, const XlaCompiler::CompilationResult* kernel, const std::map& variables) { - // Build xla::ShapedBuffers that point directly to the Tensor buffers. + // Build 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()); + arg_ptrs_ = std::vector(arg_buffers_.size()); // Pass remaining parameters. const Tensor* t; @@ -140,8 +142,7 @@ void XlaComputationLaunchContext::PopulateInputs( if (xla::ShapeUtil::IsTuple(on_device_shape)) { const XlaTensor* xla_tensor = XlaTensor::FromTensor(t); CHECK(xla_tensor && xla_tensor->has_shaped_buffer()); - arg_ptrs_[i] = - const_cast(&xla_tensor->shaped_buffer()); + arg_ptrs_[i] = const_cast(&xla_tensor->shaped_buffer()); } else { CHECK(xla::ShapeUtil::Equal(shape, on_device_shape)) << "On-device shape " @@ -149,7 +150,7 @@ void XlaComputationLaunchContext::PopulateInputs( << " not the same as on-host shape " << xla::ShapeUtil::HumanStringWithLayout(shape); gpu::DeviceMemoryBase dmem = XlaTensor::DeviceMemoryFromTensor(*t); - arg_buffers_[i] = xla::MakeUnique( + 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=*/{}); @@ -160,15 +161,15 @@ void XlaComputationLaunchContext::PopulateInputs( void XlaComputationLaunchContext::PopulateOutputs( OpKernelContext* ctx, const XlaCompiler::CompilationResult* kernel, - std::unique_ptr output) { + ScopedShapedBuffer 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(); + VLOG(2) << "Result tuple shape: " << output.on_host_shape().DebugString(); VLOG(2) << "Result tuple shape (on device): " - << output->on_device_shape().DebugString(); + << output.on_device_shape().DebugString(); } CHECK_EQ(ctx->num_outputs(), kernel->outputs.size()); @@ -226,18 +227,18 @@ void XlaComputationLaunchContext::PopulateOutputs( const TensorShape& shape = kernel->outputs[i].shape; VLOG(2) << "Retval " << i << " shape " << shape.DebugString(); - gpu::DeviceMemoryBase buffer = output->buffer({output_num}); + gpu::DeviceMemoryBase buffer = output.buffer({output_num}); if (allocate_xla_tensors_) { Tensor* output_tensor; OP_REQUIRES_OK(ctx, ctx->allocate_output(i, shape, &output_tensor)); XlaTensor* xla_tensor = XlaTensor::FromTensor(output_tensor); CHECK(xla_tensor); - xla_tensor->set_shaped_buffer( - ExtractSubShapedBuffer(output.get(), output_num, xla_allocator_)); + xla_tensor->set_shaped_buffer(ScopedShapedBuffer( + ExtractSubShapedBuffer(&output, output_num, xla_allocator_))); } else { Tensor output_tensor = XlaTensorBuffer::MakeTensor( ctx->expected_output_dtype(i), shape, buffer, allocator); - output->set_buffer(gpu::DeviceMemoryBase(nullptr, 0), {output_num}); + output.set_buffer(gpu::DeviceMemoryBase(nullptr, 0), {output_num}); ctx->set_output(i, output_tensor); } ++output_num; @@ -257,7 +258,7 @@ void XlaComputationLaunchContext::PopulateOutputs( 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}); + gpu::DeviceMemoryBase buffer = output.buffer({output_num}); Var* variable = nullptr; // TODO(b/35625933): tensorflow::Var should contain a PersistentTensor, @@ -282,12 +283,12 @@ void XlaComputationLaunchContext::PopulateOutputs( XlaTensor* xla_tensor = XlaTensor::FromTensor(&output_tensor); CHECK(xla_tensor); xla_tensor->set_shaped_buffer( - ExtractSubShapedBuffer(output.get(), output_num, xla_allocator_)); + ExtractSubShapedBuffer(&output, output_num, xla_allocator_)); *variable->tensor() = output_tensor; } else { Tensor output_tensor = XlaTensorBuffer::MakeTensor( write.type, write.shape, buffer, allocator); - output->set_buffer(gpu::DeviceMemoryBase(nullptr, 0), {output_num}); + output.set_buffer(gpu::DeviceMemoryBase(nullptr, 0), {output_num}); *variable->tensor() = output_tensor; } ++output_num; diff --git a/tensorflow/compiler/jit/xla_launch_util.h b/tensorflow/compiler/jit/xla_launch_util.h index 14f70fe35891040ff3460567adb223be0f1c910f..26dcaa8a51d90f6706106ab8a4d762ddfec6a88d 100644 --- a/tensorflow/compiler/jit/xla_launch_util.h +++ b/tensorflow/compiler/jit/xla_launch_util.h @@ -87,7 +87,7 @@ class XlaComputationLaunchContext { // Given the XLA output in `output`, populate all outputs of `ctx`. void PopulateOutputs(OpKernelContext* ctx, const XlaCompiler::CompilationResult* kernel, - std::unique_ptr output); + xla::ScopedShapedBuffer output); // Return the argument list. Only valid after PopulateInputs() has been // called. diff --git a/tensorflow/compiler/jit/xla_tensor.cc b/tensorflow/compiler/jit/xla_tensor.cc index 956328e6757f4c903e3995a54635682d19052794..84b2835c4066005c194f73becb841fbcc48399b3 100644 --- a/tensorflow/compiler/jit/xla_tensor.cc +++ b/tensorflow/compiler/jit/xla_tensor.cc @@ -65,10 +65,8 @@ Status XlaTensor::AllocateShapedBuffer(DataType dtype, const TensorShape& shape, device_ordinal, size, /*retry_on_failure=*/false)); } - TF_ASSIGN_OR_RETURN(auto scoped_buffer, - xla::ScopedShapedBuffer::MakeScoped( - &buffer, client->backend().memory_allocator())); - set_shaped_buffer(std::move(scoped_buffer)); + set_shaped_buffer(xla::ScopedShapedBuffer( + std::move(buffer), client->backend().memory_allocator())); return Status::OK(); } diff --git a/tensorflow/compiler/jit/xla_tensor.h b/tensorflow/compiler/jit/xla_tensor.h index 5ff2fb08f03548260215c6aeded2c124f8d28f43..2334fd272beb9e8185de10048421369179a0e475 100644 --- a/tensorflow/compiler/jit/xla_tensor.h +++ b/tensorflow/compiler/jit/xla_tensor.h @@ -64,9 +64,9 @@ class XlaTensor { return *shaped_buffer_; } // Mutates the TensorInfo to set the ShapedBuffer. - void set_shaped_buffer( - std::unique_ptr shaped_buffer) { - shaped_buffer_ = std::move(shaped_buffer); + void set_shaped_buffer(xla::ScopedShapedBuffer shaped_buffer) { + shaped_buffer_ = + xla::MakeUnique(std::move(shaped_buffer)); } // Some tensors on the device may have known values on the host. We use these diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index a7a8d2d1ffe4ab4a8017bb377a2da7145b63dab3..46b86c53aa6870c9290e25da27f82a06302f6865 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -203,6 +203,7 @@ tf_xla_py_test( tags = [ # Allocates very large amounts of memory and does not work under TSAN. "notsan", + "optonly", # Times out frequently in fastbuild. ], deps = [ ":xla_test", @@ -270,6 +271,18 @@ tf_xla_py_test( ], ) +tf_xla_py_test( + name = "dynamic_slice_ops_test", + size = "small", + srcs = ["dynamic_slice_ops_test.py"], + deps = [ + "//tensorflow/compiler/tests:xla_test", + "//tensorflow/compiler/tf2xla/python:xla", + "//tensorflow/python:array_ops", + "//tensorflow/python:dtypes", + ], +) + tf_xla_py_test( name = "dynamic_stitch_test", size = "small", @@ -327,7 +340,7 @@ tf_xla_py_test( tf_xla_py_test( name = "ftrl_test", - size = "small", + size = "medium", srcs = ["ftrl_test.py"], deps = [ ":xla_test", @@ -496,6 +509,22 @@ tf_xla_py_test( ], ) +tf_xla_py_test( + name = "reduce_window_test", + size = "small", + srcs = ["reduce_window_test.py"], + disabled_backends = ["cpu_ondemand"], + deps = [ + ":xla_test", + "//tensorflow/compiler/tf2xla/python:xla", + "//tensorflow/python:array_ops", + "//tensorflow/python:errors", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:math_ops", + "//tensorflow/python:platform_test", + ], +) + tf_xla_py_test( name = "reverse_ops_test", size = "medium", @@ -688,6 +717,21 @@ tf_xla_py_test( ], ) +tf_xla_py_test( + name = "while_test", + size = "small", + srcs = ["while_test.py"], + disabled_backends = ["cpu_ondemand"], + deps = [ + ":xla_test", + "//tensorflow/compiler/tf2xla/python:xla", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:platform_test", + "//tensorflow/python:training", + ], +) + tf_xla_py_test( name = "gather_test", size = "medium", diff --git a/tensorflow/compiler/tests/binary_ops_test.py b/tensorflow/compiler/tests/binary_ops_test.py index d1d7379c0a32eff4ff96e791dacbe800bbd70b7d..1e4dd32916c3a40282735fb8f75670b0e9ef0dc9 100644 --- a/tensorflow/compiler/tests/binary_ops_test.py +++ b/tensorflow/compiler/tests/binary_ops_test.py @@ -360,11 +360,13 @@ class BinaryOpsTest(XLATestCase): np.array([2, -1], dtype=dtype), expected=np.array([[[[3, 1], [5, 3]]]], 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)) + if np.int64 in self.numeric_types: + 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)) def testComplexOps(self): for dtype in self.complex_types: diff --git a/tensorflow/compiler/tests/build_defs.bzl b/tensorflow/compiler/tests/build_defs.bzl index 45b6a6eb86f8ec3d20e06efa103bec2944b9b095..7b114d4f85d3a5cadc6af25b55c5a21f90d2a768 100644 --- a/tensorflow/compiler/tests/build_defs.bzl +++ b/tensorflow/compiler/tests/build_defs.bzl @@ -56,7 +56,7 @@ def tf_xla_py_test(name, srcs=[], deps=[], tags=[], data=[], main=None, elif backend == "gpu": backend_args += [ "--test_device=XLA_GPU", - "--types=DT_FLOAT,DT_DOUBLE,DT_INT32,DT_INT64,DT_BOOL,DT_COMPLEX64,DT_BFLOAT16" + "--types=DT_HALF,DT_FLOAT,DT_DOUBLE,DT_INT32,DT_INT64,DT_BOOL,DT_COMPLEX64,DT_BFLOAT16" ] backend_tags += ["requires-gpu-sm35"] elif backend in plugins: diff --git a/tensorflow/compiler/tests/dynamic_slice_ops_test.py b/tensorflow/compiler/tests/dynamic_slice_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..6a46d2ec3e7aee3a4ecfbf1ab9f622d8eb659e3c --- /dev/null +++ b/tensorflow/compiler/tests/dynamic_slice_ops_test.py @@ -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. +# ============================================================================== +"""Tests for XLA dynamic slicing ops.""" + +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.compiler.tf2xla.python import xla +from tensorflow.python.framework import dtypes +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import test + + +class DynamicUpdateSliceOpsTest(XLATestCase): + + def _assertOpOutputMatchesExpected(self, op, args, expected): + with self.test_session() as session: + with self.test_scope(): + placeholders = [ + array_ops.placeholder(dtypes.as_dtype(arg.dtype), arg.shape) + for arg in args + ] + feeds = {placeholders[i]: args[i] for i in range(0, len(args))} + output = op(*placeholders) + result = session.run(output, feeds) + self.assertAllClose(result, expected, rtol=1e-3) + + def testUpdateSlice(self): + for dtype in self.numeric_types: + self._assertOpOutputMatchesExpected( + xla.dynamic_update_slice, [ + np.array([], dtype=dtype), + np.array([], dtype=dtype), + np.array([0], dtype=np.int32) + ], + expected=np.array([], dtype=dtype)) + + self._assertOpOutputMatchesExpected( + xla.dynamic_update_slice, [ + np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], dtype=dtype), + np.array([-1, -2, -3], dtype=dtype), + np.array([6], dtype=np.int32) + ], + expected=np.array([1, 2, 3, 4, 5, 6, -1, -2, -3, 10], dtype=dtype)) + + self._assertOpOutputMatchesExpected( + xla.dynamic_update_slice, [ + np.array( + [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], dtype=dtype), + np.array([[42, 43], [44, 45]], dtype=dtype), + np.array([1, 2], dtype=np.int32) + ], + expected=np.array( + [[1, 2, 3, 4], [5, 6, 42, 43], [9, 10, 44, 45]], dtype=dtype)) + + self._assertOpOutputMatchesExpected( + xla.dynamic_update_slice, [ + np.array( + [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], dtype=dtype), + np.array([[], []], dtype=dtype), + np.array([1, 2], dtype=np.int32) + ], + expected=np.array( + [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], dtype=dtype)) + + self._assertOpOutputMatchesExpected( + xla.dynamic_update_slice, [ + np.array( + [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], dtype=dtype), + np.ones([3, 4], dtype=dtype), + np.array([0, 0], dtype=np.int32) + ], + expected=np.ones([3, 4], dtype=dtype)) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/compiler/tests/ftrl_test.py b/tensorflow/compiler/tests/ftrl_test.py index f9db4cf2017c0b4b6dc0cfeeda6dca7bb9d14f19..8e6407dffdac3adbcda8cbca2109ef9196defa8c 100644 --- a/tensorflow/compiler/tests/ftrl_test.py +++ b/tensorflow/compiler/tests/ftrl_test.py @@ -134,9 +134,15 @@ class FtrlOptimizerTest(XLATestCase): # Validate updated params self.assertAllCloseAccordingToType( - np.array([-2.60260963, -4.29698515]), var0.eval(), float_rtol=1e-5) + np.array([-2.60260963, -4.29698515]), + var0.eval(), + float_rtol=1e-5, + half_rtol=1e-2) self.assertAllCloseAccordingToType( - np.array([-0.28432083, -0.56694895]), var1.eval(), float_rtol=1e-5) + np.array([-0.28432083, -0.56694895]), + var1.eval(), + float_rtol=1e-5, + half_rtol=1e-2) def testFtrlwithoutRegularization2(self): for dtype in self.float_types: @@ -272,8 +278,8 @@ class FtrlOptimizerTest(XLATestCase): with self.test_session(), self.test_scope(): val2, val3 = self.equivAdagradTest_AdagradPart(steps, dtype) - self.assertAllCloseAccordingToType(val0, val2, rtol=1e-4) - self.assertAllCloseAccordingToType(val1, val3, rtol=1e-4) + self.assertAllCloseAccordingToType(val0, val2, rtol=1e-4, half_rtol=1e-2) + self.assertAllCloseAccordingToType(val1, val3, rtol=1e-4, half_rtol=1e-2) def testEquivGradientDescentwithoutRegularization(self): steps = 5 diff --git a/tensorflow/compiler/tests/function_test.py b/tensorflow/compiler/tests/function_test.py index 11d8a99ffe1a136a54b16e20f1792062203f7969..fbc3c994d163a504351fcccd1ba71a0997e6516f 100644 --- a/tensorflow/compiler/tests/function_test.py +++ b/tensorflow/compiler/tests/function_test.py @@ -105,6 +105,28 @@ class FunctionTest(XLATestCase): result = sess.run(call_f) self.assertAllClose(result, expected, rtol=1e-3) + def testCompileTimeConstantsInDefun(self): + """Tests that XLA handles compile-time constants in defuns.""" + with self.test_session() as sess: + + @function.Defun(dtypes.float32, dtypes.int32, dtypes.int32) + def Foo(a, c, d): + # c and d must be known at compile time + x = array_ops.slice(a, c, d) + return x + + a = array_ops.placeholder(dtypes.float32) + c = array_ops.placeholder(dtypes.int32, shape=[4]) + d = array_ops.placeholder(dtypes.int32, shape=[4]) + with self.test_scope(): + call_f = Foo(a, c, d) + result = sess.run(call_f, feed_dict={ + a: np.ones([1, 4, 4, 1]), + c: [0, 0, 0, 0], + d: [1, 2, 2, 1]}) + + self.assertAllEqual(np.ones([1, 2, 2, 1]), result) + # TODO(b/36139787): Re-enable this test when noinline works again. def DISABLED_testFunctionsNoInline(self): diff --git a/tensorflow/compiler/tests/image_ops_test.py b/tensorflow/compiler/tests/image_ops_test.py index 3bc41b7cfd72bec7572097f8c53eef314a4369f6..42e637734c578fcc70473060cb156e172a0a1995 100644 --- a/tensorflow/compiler/tests/image_ops_test.py +++ b/tensorflow/compiler/tests/image_ops_test.py @@ -34,6 +34,13 @@ from tensorflow.python.ops import image_ops from tensorflow.python.platform import test +def GenerateNumpyRandomRGB(shape): + # Only generate floating points that are fractions like n / 256, since they + # are RGB pixels. Some low-precision floating point types in this test can't + # handle arbitrary precision floating points well. + return np.random.randint(0, 256, shape) / 256. + + class RGBToHSVTest(XLATestCase): def testBatch(self): @@ -43,7 +50,7 @@ class RGBToHSVTest(XLATestCase): shape = (batch_size, 2, 7, 3) for nptype in self.float_types: - inp = np.random.rand(*shape).astype(nptype) + inp = GenerateNumpyRandomRGB(shape).astype(nptype) # Convert to HSV and back, as a batch and individually with self.test_session() as sess: @@ -65,7 +72,8 @@ class RGBToHSVTest(XLATestCase): # Verify that processing batch elements together is the same as separate self.assertAllClose(batch1, join1) self.assertAllClose(batch2, join2) - self.assertAllCloseAccordingToType(batch2, inp, bfloat16_atol=0.03) + self.assertAllCloseAccordingToType( + batch2, inp, bfloat16_atol=0.03, half_rtol=0.02) def testRGBToHSVRoundTrip(self): data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1] @@ -82,7 +90,7 @@ class RGBToHSVTest(XLATestCase): def testRGBToHSVNumpy(self): """Tests the RGB to HSV conversion matches a reference implementation.""" for nptype in self.float_types: - rgb_flat = np.random.random(64 * 3).reshape((64, 3)).astype(nptype) + rgb_flat = GenerateNumpyRandomRGB((64, 3)).astype(nptype) rgb_np = rgb_flat.reshape(4, 4, 4, 3) hsv_np = np.array([ colorsys.rgb_to_hsv( diff --git a/tensorflow/compiler/tests/reduce_window_test.py b/tensorflow/compiler/tests/reduce_window_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e78a63465b80644d8810d9fa7433653bc4639fed --- /dev/null +++ b/tensorflow/compiler/tests/reduce_window_test.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. +# ============================================================================== +"""Tests for xla.reduce_window.""" + +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.compiler.tf2xla.python import xla +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import function +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import googletest + + +class ReduceWindowTest(XLATestCase): + """Test cases for xla.reduce_window.""" + + def _reduce_window(self, operand, init, reducer, **kwargs): + with self.test_session(): + placeholder = array_ops.placeholder(operand.dtype) + with self.test_scope(): + output = xla.reduce_window(placeholder, init, reducer, **kwargs) + return output.eval(feed_dict={placeholder: operand}) + + def testReduceWindow(self): + + # TODO(b/77644762): float16 and float64 ReduceWindow are unimplemented. + for dtype in set(self.numeric_types).intersection( + set([dtypes.bfloat16.as_numpy_dtype, np.float32])): + + @function.Defun(dtype, dtype) + def sum_reducer(x, y): + return x + y + + @function.Defun(dtype, dtype) + def mul_reducer(x, y): + return x * y + + self.assertAllClose( + np.array([3, 5, 7, 9, 11, 13], dtype=dtype), + self._reduce_window( + np.array([1, 2, 3, 4, 5, 6, 7], dtype=dtype), + 0.0, + sum_reducer, + window_dimensions=[2])) + + self.assertAllClose( + np.array([3, 7, 11], dtype=dtype), + self._reduce_window( + np.array([1, 2, 3, 4, 5, 6, 7], dtype=dtype), + 0.0, + sum_reducer, + window_dimensions=[2], + window_strides=[2])) + + self.assertAllClose( + np.array([1, 4, 7], dtype=dtype), + self._reduce_window( + np.array([1, 2, 3, 4, 5, 6, 7], dtype=dtype), + 0.0, + sum_reducer, + window_dimensions=[1], + window_strides=[3])) + + self.assertAllClose( + np.array([[24, 36, 24], [96, 0, 0]], dtype=dtype), + self._reduce_window( + np.array([[1, 2, 3, 4], [4, 3, 2, 1], [2, 4, 0, 1]], dtype=dtype), + 1.0, + mul_reducer, + window_dimensions=[2, 2], + window_strides=[1, 1])) + + self.assertAllClose( + np.array([[0, 0, 0], [5, 10, 5], [2, 4, 1], [0, 0, 0]], dtype=dtype), + self._reduce_window( + np.array([[1, 2, 3, 4], [4, 3, 2, 1], [2, 4, 0, 1]], dtype=dtype), + 0.0, + sum_reducer, + window_dimensions=[2, 2], + window_strides=[2, 2], + padding=[[2, 3], [1, 2]])) + + +if __name__ == '__main__': + googletest.main() diff --git a/tensorflow/compiler/tests/spacetobatch_op_test.py b/tensorflow/compiler/tests/spacetobatch_op_test.py index ef47187477545d019ed6e8f12ebe4a771248e607..f37c34156f96761632247be4bc1b62fca54f666e 100644 --- a/tensorflow/compiler/tests/spacetobatch_op_test.py +++ b/tensorflow/compiler/tests/spacetobatch_op_test.py @@ -163,17 +163,26 @@ class SpaceToBatchNDTest(XLATestCase): # error. if dtype == dtypes.bfloat16.as_numpy_dtype: continue - # TODO(b/77694432): Half test failed on CPU, last ran on 04-06-2018. - if dtype == np.float16 and self.device == "XLA_CPU": - continue + if dtype == np.float16: + actual_inputs = np.array(inputs).astype(dtype) + actual_paddings = np.array(paddings).astype(dtype) + expected_outputs = np.array(outputs).astype(dtype) + else: + actual_inputs = inputs + actual_paddings = paddings + expected_outputs = outputs placeholder = array_ops.placeholder(dtype) # outputs = space_to_batch(inputs) - x_tf = array_ops.space_to_batch_nd(placeholder, block_shape, paddings) - self.assertAllEqual(sess.run(x_tf, {placeholder: inputs}), outputs) + x_tf = array_ops.space_to_batch_nd(placeholder, block_shape, + actual_paddings) + self.assertAllEqual( + sess.run(x_tf, {placeholder: actual_inputs}), expected_outputs) # inputs = batch_to_space(outputs) placeholder = array_ops.placeholder(dtype) - x_tf = array_ops.batch_to_space_nd(placeholder, block_shape, paddings) - self.assertAllEqual(sess.run(x_tf, {placeholder: outputs}), inputs) + x_tf = array_ops.batch_to_space_nd(placeholder, block_shape, + actual_paddings) + self.assertAllEqual( + sess.run(x_tf, {placeholder: expected_outputs}), actual_inputs) def _testDirect(self, input_shape, block_shape, paddings): inputs = np.arange(np.prod(input_shape), dtype=np.float32) diff --git a/tensorflow/compiler/tests/ternary_ops_test.py b/tensorflow/compiler/tests/ternary_ops_test.py index ba5f829936fd82ca0cc53eda34aefbca6d80482b..75a2cf07c5adc262cdbd9c10d317a5d00f370e3d 100644 --- a/tensorflow/compiler/tests/ternary_ops_test.py +++ b/tensorflow/compiler/tests/ternary_ops_test.py @@ -23,6 +23,7 @@ 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_math_ops from tensorflow.python.ops import math_ops from tensorflow.python.platform import googletest @@ -119,6 +120,23 @@ class TernaryOpsTest(XLATestCase): np.array([2, 1], dtype=np.int32), expected=np.array([[2], [5]], dtype=dtype)) + def testClipByValue(self): + # TODO(b/78258593): enable integer types here too. + for dtype in self.float_types: + test_cases = [ + (np.array([2, 4, 5], dtype=dtype), dtype(7)), # + (dtype(1), np.array([2, 4, 5], dtype=dtype)), # + (np.array([-2, 7, 7], dtype=dtype), np.array([-2, 9, 8], dtype=dtype)) + ] + x = np.array([-2, 10, 6], dtype=dtype) + for lower, upper in test_cases: + self._testTernary( + gen_math_ops._clip_by_value, + x, + lower, + upper, + expected=np.minimum(np.maximum(x, lower), upper)) + if __name__ == "__main__": googletest.main() diff --git a/tensorflow/compiler/tests/while_test.py b/tensorflow/compiler/tests/while_test.py new file mode 100644 index 0000000000000000000000000000000000000000..f79eb27435cc954cebde4357c1d946a320f4ed75 --- /dev/null +++ b/tensorflow/compiler/tests/while_test.py @@ -0,0 +1,130 @@ +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for while loops in XLA.""" + +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.compiler.tf2xla.python import xla +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import function +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import test + + +class WhileTest(XLATestCase): + + def testSingletonLoopHandrolled(self): + # Define a function for the loop body + @function.Defun(dtypes.int32) + def loop_body(step): + step_out = step + constant_op.constant(1, dtype=dtypes.int32) + return step_out + + # Define a function for the loop condition + @function.Defun(dtypes.int32) + def loop_cond(step): + return step < 10 + + with self.test_session() as sess: + init_index = array_ops.placeholder(dtypes.int32, []) + with self.test_scope(): + loop_outputs = xla.while_loop([init_index], loop_cond, loop_body) + + result = sess.run(loop_outputs, {init_index: 0}) + self.assertAllClose(result, [10], rtol=1e-3) + + def testCountingLoopHandrolled(self): + # Define a function for the loop body + @function.Defun(dtypes.int32, dtypes.float32) + def loop_body(step, rsum): + step_out = step + constant_op.constant(1, dtype=dtypes.int32) + sum_out = rsum + constant_op.constant(1.5, dtype=dtypes.float32) + return step_out, sum_out + + # Define a function for the loop condition + @function.Defun(dtypes.int32, dtypes.float32) + def loop_cond(step, rsum): + del rsum + return step < 10 + + with self.test_session() as sess: + init_index = array_ops.placeholder(dtypes.int32, []) + init_sum = array_ops.placeholder(dtypes.float32, []) + with self.test_scope(): + loop_outputs = xla.while_loop([init_index, init_sum], loop_cond, + loop_body) + + result = sess.run(loop_outputs, {init_index: 0, init_sum: 0.0}) + self.assertAllClose(result, [10, 15.0], rtol=1e-3) + no_iters_result = sess.run(loop_outputs, {init_index: 10, init_sum: 0.0}) + self.assertAllClose(no_iters_result, [10, 0.0], rtol=1e-3) + + def testCountingLoopHandrolledC64(self): + # Define a function for the loop body + @function.Defun(dtypes.int32, dtypes.complex64) + def loop_body(step, rsum): + step_out = step + constant_op.constant(1, dtype=dtypes.int32) + sum_out = rsum + constant_op.constant(1.5 + 2j, dtype=dtypes.complex64) + return step_out, sum_out + + # Define a function for the loop condition + @function.Defun(dtypes.int32, dtypes.complex64) + def loop_cond(step, rsum): + del rsum + return step < 10 + + with self.test_session() as sess: + init_index = array_ops.placeholder(dtypes.int32, []) + init_sum = array_ops.placeholder(dtypes.complex64, []) + with self.test_scope(): + loop_outputs = xla.while_loop([init_index, init_sum], loop_cond, + loop_body) + + result = sess.run(loop_outputs, {init_index: 0, init_sum: 0.0}) + self.assertAllClose(result[1], np.complex64(15 + 20j), rtol=1e-3) + no_iters_result = sess.run(loop_outputs, {init_index: 10, init_sum: 0.0}) + self.assertAllClose(no_iters_result[1], np.complex64(0), rtol=1e-3) + + def testLoopWithConstantOutput(self): + # Define a function for the loop body + @function.Defun(dtypes.int32, dtypes.int32) + def loop_body(step, x): + del x + step_out = step + constant_op.constant(1, dtype=dtypes.int32) + return (step_out, 7) + + # Define a function for the loop condition + @function.Defun(dtypes.int32, dtypes.int32) + def loop_cond(step, x): + del x + return step < 10 + + with self.test_session() as sess: + init_index = array_ops.placeholder(dtypes.int32, []) + with self.test_scope(): + loop_outputs = xla.while_loop([init_index, 42], loop_cond, loop_body) + + result = sess.run(loop_outputs, {init_index: 0}) + self.assertAllClose(result, [10, 7], rtol=1e-3) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD index e7daf4e01c45c3705216fce7dd3db5baa0c261fc..ba5c3a14849cefcb680b03425232724ff32375a8 100644 --- a/tensorflow/compiler/tf2xla/BUILD +++ b/tensorflow/compiler/tf2xla/BUILD @@ -415,7 +415,7 @@ cc_library( "//tensorflow/compiler/jit:graph_to_functiondef", "//tensorflow/compiler/jit:union_find", "//tensorflow/compiler/tf2xla:dump_graph", - "//tensorflow/compiler/tf2xla/ops:functional_ops", + "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:util", "//tensorflow/core:core_cpu", @@ -437,7 +437,7 @@ tf_cc_test( "//tensorflow/cc:function_ops", "//tensorflow/cc:ops", "//tensorflow/cc:resource_variable_ops", - "//tensorflow/compiler/tf2xla/cc:functional_ops", + "//tensorflow/compiler/tf2xla/cc:xla_ops", "//tensorflow/compiler/xla:status_macros", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", diff --git a/tensorflow/compiler/tf2xla/cc/BUILD b/tensorflow/compiler/tf2xla/cc/BUILD index c30bb9cacd48fb93ac359a6a25699ba6a74183c5..4f8bb8ad743afe69a6544c2ae0dc7309891b2df3 100644 --- a/tensorflow/compiler/tf2xla/cc/BUILD +++ b/tensorflow/compiler/tf2xla/cc/BUILD @@ -7,44 +7,20 @@ licenses(["notice"]) # Apache 2.0 load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_cc") tf_gen_op_wrapper_cc( - name = "functional_ops_gen", - include_internal_ops = 1, - out_ops_file = "ops/functional_ops", - deps = ["//tensorflow/compiler/tf2xla/ops:functional_ops"], + name = "xla_ops_gen", + out_ops_file = "ops/xla_ops", + deps = ["//tensorflow/compiler/tf2xla/ops:xla_ops"], ) cc_library( - name = "functional_ops", - srcs = ["ops/functional_ops.cc"], - hdrs = ["ops/functional_ops.h"], + name = "xla_ops", + srcs = ["ops/xla_ops.cc"], + hdrs = ["ops/xla_ops.h"], deps = [ "//tensorflow/cc:const_op", "//tensorflow/cc:ops", "//tensorflow/cc:scope", - "//tensorflow/compiler/tf2xla/ops:functional_ops", - "//tensorflow/core:core_cpu", - "//tensorflow/core:framework", - "//tensorflow/core:lib", - "//tensorflow/core:protos_all_cc", - ], -) - -tf_gen_op_wrapper_cc( - name = "sendrecv_ops_gen", - include_internal_ops = 1, - out_ops_file = "ops/sendrecv_ops", - deps = ["//tensorflow/compiler/tf2xla/ops:sendrecv_ops"], -) - -cc_library( - name = "sendrecv_ops", - srcs = ["ops/sendrecv_ops.cc"], - hdrs = ["ops/sendrecv_ops.h"], - deps = [ - "//tensorflow/cc:const_op", - "//tensorflow/cc:ops", - "//tensorflow/cc:scope", - "//tensorflow/compiler/tf2xla/ops:sendrecv_ops", + "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:lib", diff --git a/tensorflow/compiler/tf2xla/functionalize_control_flow.cc b/tensorflow/compiler/tf2xla/functionalize_control_flow.cc index 16b9142cbf7d2afe99c22acbc32fb17c09b00081..23629d85aed316764b8ddd775c51c3890e5c16e3 100644 --- a/tensorflow/compiler/tf2xla/functionalize_control_flow.cc +++ b/tensorflow/compiler/tf2xla/functionalize_control_flow.cc @@ -870,6 +870,9 @@ FunctionalizeCond::DeterminePredicateSwitchOrder() { // 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()) { + // Only consider the data inputs to the Switch node. + if (e->IsControlEdge()) continue; + Node* src = e->src(); UnionFind* src_cluster = find_output_cluster(src); int src_cluster_depth = switch_depth[src_cluster->Get().representative]; diff --git a/tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc b/tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc index bc7276c3afd5060d6faeceb4d479416299ecc5da..e494f42e8ed254ac0c7c7a23a13728d3f015e9d3 100644 --- a/tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc +++ b/tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc @@ -20,7 +20,7 @@ limitations under the License. #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/cc/ops/functional_ops.h" +#include "tensorflow/compiler/tf2xla/cc/ops/xla_ops.h" #include "tensorflow/compiler/tf2xla/test_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/core/common_runtime/function.h" diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD index f1bc7d6af49a09f84ef251eaa1c3d684792d0c1e..00fd08b1a0750739445a124adc7ccf436a4a9b71 100644 --- a/tensorflow/compiler/tf2xla/kernels/BUILD +++ b/tensorflow/compiler/tf2xla/kernels/BUILD @@ -21,6 +21,7 @@ tf_kernel_library( "cast_op.cc", "categorical_op.cc", "cholesky_op.cc", + "clip_by_value_op.cc", "concat_op.cc", "const_op.cc", "conv_ops.cc", @@ -29,6 +30,7 @@ tf_kernel_library( "cwise_ops.h", "depthtospace_op.cc", "diag_op.cc", + "dynamic_slice_ops.cc", "dynamic_stitch_op.cc", "elu_op.cc", "extract_image_patches_op.cc", @@ -56,6 +58,7 @@ tf_kernel_library( "pooling_ops.cc", "quantize_and_dequantize_op.cc", "random_ops.cc", + "reduce_window_op.cc", "reduction_ops.cc", "reduction_ops.h", "reduction_ops_common.cc", @@ -103,7 +106,7 @@ tf_kernel_library( "//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/tf2xla/ops:xla_ops", "//tensorflow/compiler/xla:array4d", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", @@ -146,7 +149,7 @@ tf_kernel_library( deps = [ "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:xla_compiler", - "//tensorflow/compiler/tf2xla/ops:functional_ops", + "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/core:framework", @@ -162,7 +165,7 @@ tf_kernel_library( deps = [ "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:xla_compiler", - "//tensorflow/compiler/tf2xla/ops:functional_ops", + "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/core:framework", @@ -171,6 +174,23 @@ tf_kernel_library( ], ) +# Kernels that have a dummy (no-op) implementation. +tf_kernel_library( + name = "xla_dummy_ops", + srcs = [ + "assert_op.cc", + "check_numerics_op.cc", + ], + deps = [ + "//tensorflow/compiler/tf2xla:xla_compiler", + "//tensorflow/core:array_ops_op_lib", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:logging_ops_op_lib", + ], + alwayslink = 1, +) + # Kernels that only work on CPU, because they use XLA custom calls. # Only link this when using the CPU backend for XLA. tf_kernel_library( diff --git a/tensorflow/compiler/tf2xla/kernels/assert_op.cc b/tensorflow/compiler/tf2xla/kernels/assert_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..af4ab5e8ef6e268226edc90515706405ac36858c --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/assert_op.cc @@ -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. +==============================================================================*/ + +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/mutex.h" + +namespace tensorflow { + +namespace { + +// This TensorFlow op supports the Assert primitve. +class AssertOp : public XlaOpKernel { + public: + explicit AssertOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + ~AssertOp() override {} + + void Compile(XlaOpKernelContext* ctx) override { + static mutex mu(tensorflow::LINKER_INITIALIZED); + static int log_counter = 0; + + mutex_lock l(mu); + if (log_counter < 20) { + ++log_counter; + LOG(WARNING) << "Ignoring Assert operator " << name(); + } + } + + private: + TF_DISALLOW_COPY_AND_ASSIGN(AssertOp); +}; + +REGISTER_XLA_OP(Name("Assert"), AssertOp); + +} // anonymous namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/check_numerics_op.cc b/tensorflow/compiler/tf2xla/kernels/check_numerics_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..6061e822d8d9c6c807a63aad4e9e9526a49e456c --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/check_numerics_op.cc @@ -0,0 +1,50 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/mutex.h" + +namespace tensorflow { +namespace { + +class CheckNumericsOp : public XlaOpKernel { + public: + explicit CheckNumericsOp(OpKernelConstruction* context) + : XlaOpKernel(context) {} + + void Compile(XlaOpKernelContext* ctx) override { + // TODO(b/32223192): add a real implementation of CheckNumerics + { + static mutex mu(tensorflow::LINKER_INITIALIZED); + static int log_counter = 0; + mutex_lock l(mu); + if (log_counter < 20) { + ++log_counter; + LOG(WARNING) << "Ignoring CheckNumerics operator " << name(); + } + } + ctx->SetOutput(0, ctx->Input(0)); + } + + private: + TF_DISALLOW_COPY_AND_ASSIGN(CheckNumericsOp); +}; + +REGISTER_XLA_OP(Name("CheckNumerics"), CheckNumericsOp); + +} // anonymous namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/clip_by_value_op.cc b/tensorflow/compiler/tf2xla/kernels/clip_by_value_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..fdf75be7b1156540d762e3bc04a51f2478f00f46 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/clip_by_value_op.cc @@ -0,0 +1,61 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/core/framework/tensor_shape.h" + +namespace tensorflow { +namespace { + +class ClipByValueOp : public XlaOpKernel { + public: + explicit ClipByValueOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + + void Compile(XlaOpKernelContext* ctx) override { + const TensorShape shape = ctx->InputShape(0); + const TensorShape min_shape = ctx->InputShape(1); + const TensorShape max_shape = ctx->InputShape(2); + + xla::ComputationBuilder* builder = ctx->builder(); + auto input = ctx->Input(0); + auto min = ctx->Input(1); + auto max = ctx->Input(2); + + auto shape_error = [&]() -> tensorflow::Status { + return errors::InvalidArgument( + "clip_value_min and clip_value_max must be either of " + "the same shape as input, or a scalar. ", + "Input shape: ", shape.DebugString(), + " clip_value_min shape: ", min_shape.DebugString(), + " clip_value_max shape: ", max_shape.DebugString()); + }; + + if (shape != min_shape) { + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(min_shape), shape_error()); + min = builder->Broadcast(min, shape.dim_sizes()); + } + if (shape != max_shape) { + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(max_shape), shape_error()); + max = builder->Broadcast(max, shape.dim_sizes()); + } + ctx->SetOutput(0, builder->Clamp(min, input, max)); + } +}; + +REGISTER_XLA_OP(Name("ClipByValue"), ClipByValueOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/dynamic_slice_ops.cc b/tensorflow/compiler/tf2xla/kernels/dynamic_slice_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..800ef5ab98d70ad822c6efffb33db28b46ae50fe --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/dynamic_slice_ops.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 "tensorflow/compiler/tf2xla/shape_util.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/core/framework/op_kernel.h" + +#include "tensorflow/compiler/tf2xla/type_util.h" +#include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "tensorflow/core/framework/kernel_def_builder.h" + +namespace tensorflow { +namespace { + +class DynamicUpdateSliceOp : public XlaOpKernel { + public: + explicit DynamicUpdateSliceOp(OpKernelConstruction* context) + : XlaOpKernel(context) {} + + void Compile(XlaOpKernelContext* ctx) override { + VLOG(3) << "DynamicUpdateSliceOp::Compile"; + + DataType index_type = input_type(2); + OP_REQUIRES(ctx, index_type == DT_INT32 || index_type == DT_INT64, + errors::InvalidArgument("index must be int32 or int64")); + + const TensorShape input_shape = ctx->InputShape(0); + const TensorShape update_shape = ctx->InputShape(1); + const TensorShape index_shape = ctx->InputShape(2); + + OP_REQUIRES( + ctx, + TensorShapeUtils::IsVector(index_shape) && + index_shape.num_elements() == input_shape.dims(), + errors::InvalidArgument("index must be a vector with length equal to " + "the number of input dimensions")); + OP_REQUIRES( + ctx, input_shape.dims() == update_shape.dims(), + errors::InvalidArgument("input and update must have the same rank," + " input shape is ", + input_shape.DebugString(), "; update shape is ", + update_shape.DebugString())); + + xla::ComputationDataHandle result = ctx->builder()->DynamicUpdateSlice( + ctx->Input(0), ctx->Input(1), ctx->Input(2)); + ctx->SetOutput(0, result); + } +}; + +REGISTER_XLA_OP(Name("XlaDynamicUpdateSlice"), DynamicUpdateSliceOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op.cc b/tensorflow/compiler/tf2xla/kernels/gather_op.cc index 7945c05af40df21a798a2cff51fe7f8e935793f6..0b79cb0916ee8a7d0e26c5dc12557639336f8ab1 100644 --- a/tensorflow/compiler/tf2xla/kernels/gather_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/gather_op.cc @@ -29,52 +29,54 @@ namespace tensorflow { 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, + const TensorShape& indices_shape, int64 axis, + bool indices_are_nd, DataType dtype, DataType index_type, xla::ComputationBuilder* builder, xla::ComputationDataHandle* gather_output) { + // There is no deep reason why we need this precondition, but this is the only + // combination that is used and tested today. + CHECK(!indices_are_nd || axis == 0); + + // num_index_dims is the number of components in each index in the indices + // tensor. + // + // num_indices is the total number of (n dimensional or scalar) indices in the + // indices tensor. + // // 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; + int64 num_index_dims; + int64 num_indices = 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); + for (int64 i = 0, e = indices_shape.dims() - 1; i < e; i++) { + num_indices *= indices_shape.dim_size(i); + } + } else { + num_index_dims = 1; + for (int64 i = 0, e = indices_shape.dims(); i < e; i++) { + num_indices *= indices_shape.dim_size(i); + } } - // 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 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 + num_index_dims); - // Each slice of the input tensor has shape: - // [, 1, ..., 1, ] - TensorShape slice_shape(input_shape); - for (int64 i = 0; i < num_index_dims; ++i) { - slice_shape.set_dim(axis + i, 1); - } + // Degenerate case: empty indices. + if (num_indices == 0) { + 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 + num_index_dims); - TensorShape loop_out_shape; - loop_out_shape.AppendShape(input_shape_pre_axis); - loop_out_shape.AddDim(num_indices); - loop_out_shape.AppendShape(input_shape_post_axis); - TensorShape loop_out_slice_shape; - loop_out_slice_shape.AppendShape(input_shape_pre_axis); - loop_out_slice_shape.AddDim(1); - loop_out_slice_shape.AppendShape(input_shape_post_axis); + TensorShape indices_shape_no_index_vectors{indices_shape}; + if (indices_are_nd) { + indices_shape_no_index_vectors.RemoveLastDims(1); + } - TensorShape out_shape; - out_shape.AppendShape(input_shape_pre_axis); - out_shape.AppendShape(indices_shape); - out_shape.AppendShape(input_shape_post_axis); + TensorShape out_shape; + out_shape.AppendShape(input_shape_pre_axis); + out_shape.AppendShape(indices_shape_no_index_vectors); + out_shape.AppendShape(input_shape_post_axis); - // Degenerate case: empty indices. - if (num_indices == 0) { *gather_output = builder->Broadcast(XlaHelpers::Zero(builder, dtype), out_shape.dim_sizes()); return Status::OK(); @@ -88,76 +90,61 @@ Status XlaGather(const xla::ComputationDataHandle& input, } } - // 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. - - // Construct the initial values of the loop-carried Tensors. - auto flat_indices = builder->Reshape(indices, flat_indices_shape); - auto init_out = builder->Broadcast(XlaHelpers::Zero(builder, dtype), - loop_out_shape.dim_sizes()); - 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) - 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}); + // Example of a 1-D gather with axis=1, pulling two [3,1] tensors out of a + // tensor of shape [3,3]. + // + // operand = s32[3,3] parameter(0) + // indices = s32[2] parameter(1) + // 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} + // + // + // Example of an N-D gather pulling out slices of shape [1,1,2] out of a + // tensor of shape [3,3,2]. + // + // operand = s32[3,3,2] parameter(0) + // indices = s32[2,2] parameter(1) + // 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} + + xla::GatherDimensionNumbers dim_numbers; + std::vector window_bounds; + window_bounds.reserve(input_shape.dims()); + for (int64 i = 0; i < input_shape.dims(); i++) { + int64 window_bound; + if (axis <= i && i < (axis + num_index_dims)) { + dim_numbers.add_elided_window_dims(i); + window_bound = 1; } else { - index = bodyb->DynamicSlice(indices, indices_offset, {1}); + window_bound = input_shape.dim_size(i); + } + + window_bounds.push_back(window_bound); + + if (i < axis) { + dim_numbers.add_output_window_dims(i); + } else if (i >= (axis + num_index_dims)) { + int64 indices_rank = + indices_are_nd ? (indices_shape.dims() - 1) : indices_shape.dims(); + dim_numbers.add_output_window_dims(i + indices_rank - num_index_dims); } + } + + dim_numbers.set_index_vector_dim(indices_are_nd ? (indices_shape.dims() - 1) + : indices_shape.dims()); + for (int64 i = axis; i < axis + num_index_dims; i++) { + dim_numbers.add_gather_dims_to_operand_dims(i); + } - // 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()), - 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(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); - - return std::vector{input, indices, - updated_output}; - }; - - // Construct the While loop, extract and reshape the output. - 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()); + *gather_output = builder->Gather(input, indices, dim_numbers, window_bounds); return Status::OK(); } diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h b/tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h index bd8b92c22d71fe89ab8951ec79f411feef6505e3..f9376f0eabdc0f0c565eb4b9f86425de96b5aa22 100644 --- a/tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h +++ b/tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h @@ -36,8 +36,8 @@ namespace tensorflow { 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, + const TensorShape& indices_shape, int64 axis, + bool indices_are_nd, DataType dtype, DataType index_type, xla::ComputationBuilder* builder, xla::ComputationDataHandle* gather_output); diff --git a/tensorflow/compiler/tf2xla/kernels/reduce_window_op.cc b/tensorflow/compiler/tf2xla/kernels/reduce_window_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..cb144bea9e429b7c8bcc3d07f688ed6a254c3be0 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/reduce_window_op.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/compiler/tf2xla/kernels/while_op.h" + +#include "tensorflow/compiler/tf2xla/shape_util.h" +#include "tensorflow/compiler/tf2xla/xla_compiler.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/core/framework/function.h" +#include "tensorflow/core/framework/op_kernel.h" + +namespace tensorflow { +namespace { + +class ReduceWindowOp : public XlaOpKernel { + public: + explicit ReduceWindowOp(OpKernelConstruction* context) + : XlaOpKernel(context) { + OP_REQUIRES_OK(context, context->GetAttr("computation", &computation_)); + OP_REQUIRES_OK(context, + context->GetAttr("window_dimensions", &window_dimensions_)); + OP_REQUIRES_OK(context, + context->GetAttr("window_strides", &window_strides_)); + OP_REQUIRES_OK(context, context->GetAttr("padding_low", &padding_low_)); + OP_REQUIRES_OK(context, context->GetAttr("padding_high", &padding_high_)); + } + + void Compile(XlaOpKernelContext* context) override { + const TensorShape input_shape = context->InputShape(0); + const DataType dtype = context->input_type(0); + + const int rank = input_shape.dims(); + OP_REQUIRES(context, rank == window_dimensions_.size(), + errors::InvalidArgument( + "The size of window_dimensions must be equal to the input " + "rank (", + window_dimensions_.size(), " vs. ", rank, ")")); + OP_REQUIRES(context, rank == window_strides_.size(), + errors::InvalidArgument( + "The size of window_strides must be equal to the input " + "rank (", + window_strides_.size(), " vs. ", rank, ")")); + OP_REQUIRES(context, rank == padding_low_.size(), + errors::InvalidArgument( + "The size of padding_low must be equal to the input " + "rank (", + padding_low_.size(), " vs. ", rank, ")")); + OP_REQUIRES(context, rank == padding_high_.size(), + errors::InvalidArgument( + "The size of padding_high must be equal to the input " + "rank (", + padding_high_.size(), " vs. ", rank, ")")); + + xla::ComputationBuilder* builder = context->builder(); + + // Build the reducer function. + XlaCompiler::Argument reducer_arg; + reducer_arg.kind = XlaCompiler::Argument::kParameter; + reducer_arg.type = dtype; + reducer_arg.shape = TensorShape(); + + XlaCompiler::CompileOptions compile_options; + compile_options.use_tuple_arg = false; + compile_options.resolve_compile_time_constants = false; + compile_options.is_entry_computation = false; + XlaCompiler::CompilationResult reducer; + OP_REQUIRES_OK(context, context->compiler()->CompileFunction( + compile_options, *computation_, + {reducer_arg, reducer_arg}, &reducer)); + + xla::Shape scalar_shape; + OP_REQUIRES_OK(context, + TensorShapeToXLAShape(dtype, TensorShape(), &scalar_shape)); + OP_REQUIRES(context, + xla::ShapeUtil::Compatible( + reducer.xla_output_shape, + xla::ShapeUtil::MakeTupleShape({scalar_shape})), + errors::InvalidArgument( + "Invalid output shape of ReduceWindow reducer. Expected ", + xla::ShapeUtil::HumanString(scalar_shape), " got ", + xla::ShapeUtil::HumanString(reducer.xla_output_shape))); + + // Wraps the reducer in a computation that unpacks the output tuple. + xla::Computation wrapper; + { + std::unique_ptr cb = + builder->CreateSubBuilder("wrapper"); + auto x = cb->Parameter(0, scalar_shape, "x"); + auto y = cb->Parameter(1, scalar_shape, "y"); + auto outputs = cb->Call(*reducer.computation, {x, y}); + cb->GetTupleElement(outputs, 0); + xla::StatusOr result = cb->Build(); + OP_REQUIRES_OK(context, result.status()); + wrapper = std::move(result.ValueOrDie()); + } + + std::vector> padding(rank); + for (int i = 0; i < rank; ++i) { + padding[i] = {padding_low_[i], padding_high_[i]}; + } + + xla::ComputationDataHandle output = builder->ReduceWindowWithGeneralPadding( + context->Input(0), context->Input(1), wrapper, window_dimensions_, + window_strides_, padding); + context->SetOutput(0, output); + } + + private: + const NameAttrList* computation_; + std::vector window_dimensions_; + std::vector window_strides_; + std::vector padding_low_; + std::vector padding_high_; + + TF_DISALLOW_COPY_AND_ASSIGN(ReduceWindowOp); +}; + +REGISTER_XLA_OP(Name("XlaReduceWindow"), ReduceWindowOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/sendrecv_ops.cc b/tensorflow/compiler/tf2xla/kernels/sendrecv_ops.cc index 5172781c0d05b6682fe92086654e3b86961949ee..d079b89861817a5639ac72b5ee49d76cb4506ae8 100644 --- a/tensorflow/compiler/tf2xla/kernels/sendrecv_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/sendrecv_ops.cc @@ -48,7 +48,7 @@ void SendOp::Compile(XlaOpKernelContext* ctx) { ctx->builder()->Send(ctx->Input(0), channel); } -REGISTER_XLA_OP(Name("_XLASend"), SendOp); +REGISTER_XLA_OP(Name("XlaSend"), SendOp); class RecvOp : public XlaOpKernel { public: @@ -68,7 +68,7 @@ RecvOp::RecvOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { TensorShape tensor_shape; DataType dtype; OP_REQUIRES_OK(ctx, ctx->GetAttr("shape", &tensor_shape)); - OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("dtype", &dtype)); OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(dtype, tensor_shape, &shape_)); } @@ -79,7 +79,7 @@ void RecvOp::Compile(XlaOpKernelContext* ctx) { ctx->SetOutput(0, ctx->builder()->Recv(shape_, channel)); } -REGISTER_XLA_OP(Name("_XLARecv"), RecvOp); +REGISTER_XLA_OP(Name("XlaRecv"), RecvOp); } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/lib/BUILD b/tensorflow/compiler/tf2xla/lib/BUILD index 344773c8c5f8e1a552d585d0317c62c56d9f9d46..fde1977c1b1834156b87b4fb3516f7bf8df435d7 100644 --- a/tensorflow/compiler/tf2xla/lib/BUILD +++ b/tensorflow/compiler/tf2xla/lib/BUILD @@ -39,6 +39,7 @@ cc_library( ":batch_dot", ":triangular_solve", ":util", + ":while_loop", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", @@ -126,6 +127,30 @@ cc_library( ], ) +xla_test( + name = "util_test", + srcs = ["util_test.cc"], + deps = [ + ":batch_dot", + ":util", + "//tensorflow/compiler/xla:array2d", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/compiler/xla/client:global_data", + "//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:xla_internal_test_main", + "//tensorflow/core:lib", + "//tensorflow/core:test", + ], +) + cc_library( name = "while_loop", srcs = ["while_loop.cc"], diff --git a/tensorflow/compiler/tf2xla/lib/cholesky.cc b/tensorflow/compiler/tf2xla/lib/cholesky.cc index e795701181dd80a2ff544743d513bffd52fd2399..203365e2ab07e0da1abfac5452a8ec41a4ddf406 100644 --- a/tensorflow/compiler/tf2xla/lib/cholesky.cc +++ b/tensorflow/compiler/tf2xla/lib/cholesky.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. @@ -21,6 +21,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/lib/batch_dot.h" #include "tensorflow/compiler/tf2xla/lib/triangular_solve.h" #include "tensorflow/compiler/tf2xla/lib/util.h" +#include "tensorflow/compiler/tf2xla/lib/while_loop.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -31,68 +32,122 @@ namespace tensorflow { namespace { +// The Cholesky–Banachiewicz algorithm. See +// https://en.wikipedia.org/wiki/Cholesky_decomposition#The_Cholesky–Banachiewicz_and_Cholesky–Crout_algorithms +// for a description. +// // def cholesky_unblocked(a): // assert len(a.shape) == 2 and a.shape[-2] == a.shape[-1] // n = a.shape[-2] // l = np.zeros_like(a) // for j in xrange(n): -// r = l[..., j, :j] -// l[..., j, j] = np.sqrt(a[..., j, j] - np.dot(r, r)) -// l[..., j+1:, j] = (a[..., j+1:, j] - np.dot(l[..., j+1:, :j], -// np.transpose(r))) / l[..., j, j] +// row = l[..., j, :j] +// row_t = np.swapaxes(row, -1, -2) +// l[..., j, j] = np.sqrt(a[..., j, j] - np.dot(row, row_t)) +// l[..., j+1:, j] = (a[..., j+1:, j] - np.dot(l[..., j+1:, :j], row_t)) / +// l[..., j, j] // return l xla::StatusOr CholeskyUnblocked( xla::ComputationBuilder* builder, const xla::ComputationDataHandle& a) { - TF_ASSIGN_OR_RETURN(std::unique_ptr shape, builder->GetShape(a)); - xla::ComputationDataHandle l = Zeros(builder, *shape); - const int64 n = xla::ShapeUtil::GetDimension(*shape, -2); - for (int j = 0; j < n; ++j) { - // Picture of block structure: - // ... \ - // \ - // -- r -- d - // |\ - // B c \ - // | \ - // | ... - // - // ^ - // column j - TF_ASSIGN_OR_RETURN(auto d, - SliceInMinorDims(builder, a, {j, j}, {j + 1, j + 1})); - TF_ASSIGN_OR_RETURN(auto c, - SliceInMinorDims(builder, a, {j + 1, j}, {n, j + 1})); - xla::ComputationDataHandle new_d_squared = d; - xla::ComputationDataHandle br; - if (j > 0) { - TF_ASSIGN_OR_RETURN(auto r, - SliceInMinorDims(builder, l, {j, 0}, {j + 1, j})); - TF_ASSIGN_OR_RETURN(auto b, - SliceInMinorDims(builder, l, {j + 1, 0}, {n, j})); - TF_ASSIGN_OR_RETURN(auto r_squared, - BatchDot(builder, r, r, /*transpose_x=*/false, - /*transpose_y=*/true, /*conjugate_x=*/false, - /*conjugate_y=*/false)); - new_d_squared = builder->Sub(new_d_squared, r_squared); + TF_ASSIGN_OR_RETURN(std::unique_ptr a_shape, + builder->GetShape(a)); + const int n_dims = xla::ShapeUtil::Rank(*a_shape); + const int64 n = xla::ShapeUtil::GetDimension(*a_shape, -1); + gtl::ArraySlice major_dims(xla::AsInt64Slice(a_shape->dimensions()), + /*pos=*/0, + /*len=*/n_dims - 2); - TF_ASSIGN_OR_RETURN(br, BatchDot(builder, b, r, /*transpose_x=*/false, - /*transpose_y=*/true, - /*conjugate_x=*/false, - /*conjugate_y=*/false)); - } - auto new_d_inv = builder->Pow( - new_d_squared, FloatLiteral(builder, shape->element_type(), -0.5)); - auto new_d = builder->Mul(new_d_inv, new_d_squared); - TF_ASSIGN_OR_RETURN(l, UpdateSliceInMinorDims(builder, l, new_d, {j, j})); + xla::ComputationDataHandle l = Zeros(builder, *a_shape); - if (j > 0) { - c = builder->Sub(c, br); + // Construct the for loop body to iterate over rows. + auto body_fn = [&](xla::ComputationDataHandle i, + gtl::ArraySlice loop_vars, + xla::ComputationBuilder* body_builder) + -> xla::StatusOr> { + xla::Shape col_shape; + xla::Shape row_shape; + for (int64 d : major_dims) { + row_shape.add_dimensions(d); + col_shape.add_dimensions(d); } - auto new_c = builder->Mul(c, new_d_inv); - TF_ASSIGN_OR_RETURN(l, - UpdateSliceInMinorDims(builder, l, new_c, {j + 1, j})); - } - return l; + row_shape.add_dimensions(1); + row_shape.add_dimensions(n); + row_shape.set_element_type(a_shape->element_type()); + auto mask_zeros_row = Zeros(body_builder, row_shape); + + col_shape.add_dimensions(n); + col_shape.add_dimensions(1); + col_shape.set_element_type(a_shape->element_type()); + auto mask_zeros_col = Zeros(body_builder, col_shape); + + std::vector mask_vector(n); + std::iota(mask_vector.begin(), mask_vector.end(), 0); + auto mask_range = body_builder->ConstantR1(mask_vector); + auto mask_range_row = body_builder->Broadcast( + body_builder->Reshape(mask_range, {0}, {1, n}), major_dims); + auto mask_range_col = body_builder->Broadcast( + body_builder->Reshape(mask_range, {0}, {n, 1}), major_dims); + auto body_a = loop_vars[0]; + auto body_l = loop_vars[1]; + + // row = l[..., i, :i] + // select the whole i-th row, then mask out all columns past i-1 + auto zero = body_builder->ConstantR0(0); + TF_ASSIGN_OR_RETURN(auto l_i, DynamicSliceInMinorDims(body_builder, body_l, + {i, zero}, {1, n})); + auto row = body_builder->Select(body_builder->Ge(mask_range_row, i), + mask_zeros_row, l_i); + // a[..., i, i] + TF_ASSIGN_OR_RETURN(auto a_ii, DynamicSliceInMinorDims(body_builder, body_a, + {i, i}, {1, 1})); + // np.dot(row, np.swapaxes(row, -1, -2)) + xla::ComputationDataHandle diag_dot; + TF_ASSIGN_OR_RETURN(diag_dot, BatchDot(body_builder, row, row, + /*transpose_x=*/false, + /*transpose_y=*/true)); + // l[..., i, i] = np.sqrt(a[..., i, i] - np.dot(row, + // np.swapaxes(row, -1, -2))) + auto l_ii = body_builder->Pow( + body_builder->Sub(a_ii, diag_dot), + FloatLiteral(body_builder, a_shape->element_type(), 0.5)); + + // a[..., i+1:, i] + auto ip1 = body_builder->Add(i, body_builder->ConstantR0(1)); + // select the whole i-th column, then mask out all rows above i+1 + TF_ASSIGN_OR_RETURN( + auto a_0i, DynamicSliceInMinorDims(body_builder, body_a, {i}, {1})); + auto a_ip1i = body_builder->Select(body_builder->Le(mask_range_col, i), + mask_zeros_col, a_0i); + + // l[..., i+1:, i] = (a[..., i+1:, i] - np.dot(l[..., i+1:, :i], r.T)) / + // l[..., i, i] + // The columns in [i, n] are zeroed out in `row`, so we just have to + // zero out rows above i+1 after the BatchDot. np.dot(l[..., :, :i], + // r.T) + TF_ASSIGN_OR_RETURN(auto dot, BatchDot(body_builder, body_l, row, + /*transpose_x=*/false, + /*transpose_y=*/true)); + // np.dot(l[..., i+1:, :i], r.T) + auto dot_ip1 = body_builder->Select(body_builder->Le(mask_range_col, i), + mask_zeros_col, dot); + + auto col_update = + body_builder->Div(body_builder->Sub(a_ip1i, dot_ip1), l_ii); + TF_ASSIGN_OR_RETURN(body_l, DynamicUpdateSliceInMinorDims( + body_builder, body_l, col_update, {i})); + // Assign the diagonal after the rest of the column because otherwise the + // column assign will wrap around and overwrite the diagonal assign. + TF_ASSIGN_OR_RETURN(body_l, DynamicUpdateSliceInMinorDims( + body_builder, body_l, l_ii, {i, i})); + + return std::vector{body_a, body_l}; + }; + + TF_ASSIGN_OR_RETURN( + auto cholesky_while, + XlaForEachIndex(n, xla::S32, body_fn, {a, l}, "unblocked", builder)); + + return cholesky_while[1]; } } // namespace diff --git a/tensorflow/compiler/tf2xla/lib/cholesky.h b/tensorflow/compiler/tf2xla/lib/cholesky.h index e083a383be4be0d1b556b63214fe5f70323b4149..17da8d8b22d107701ce768ac945c1404df6d47e8 100644 --- a/tensorflow/compiler/tf2xla/lib/cholesky.h +++ b/tensorflow/compiler/tf2xla/lib/cholesky.h @@ -1,4 +1,4 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -29,7 +29,7 @@ namespace tensorflow { // the block size to use. // TODO(phawkins): check for negative values on the diagonal and return an // error, instead of silently yielding NaNs. -// TODO(mattjj): handle the complex Hermitian case +// TODO(znado): handle the complex Hermitian case xla::StatusOr Cholesky( xla::ComputationBuilder* builder, xla::ComputationDataHandle a, int64 block_size = 256); diff --git a/tensorflow/compiler/tf2xla/lib/util.cc b/tensorflow/compiler/tf2xla/lib/util.cc index f579669bbd852b514e021ce71d635f8ce5e4fe4d..31d823ca336039f691f2c16e37028c0de98b1ee5 100644 --- a/tensorflow/compiler/tf2xla/lib/util.cc +++ b/tensorflow/compiler/tf2xla/lib/util.cc @@ -1,4 +1,4 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -140,13 +140,47 @@ xla::StatusOr SliceInMinorDims( return builder->Slice(x, padded_start, padded_end, strides); } +std::vector PrependMajorDims(xla::ComputationBuilder* builder, + const gtl::ArraySlice& major_dims, + const gtl::ArraySlice& indices) { + std::vector output(indices.size() + major_dims.size()); + std::copy(major_dims.begin(), major_dims.end(), output.begin()); + std::copy(indices.begin(), indices.end(), output.begin() + major_dims.size()); + return output; +} + +xla::StatusOr DynamicSliceInMinorDims( + xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x, + const std::vector& starts, + const gtl::ArraySlice& sizes) { + TF_ASSIGN_OR_RETURN(std::unique_ptr shape, builder->GetShape(x)); + const int64 n_dims = xla::ShapeUtil::Rank(*shape); + int64 n_minor_dims = starts.size(); + TF_RET_CHECK(n_minor_dims == sizes.size()); + TF_RET_CHECK(n_minor_dims <= n_dims); + gtl::ArraySlice major_dims(xla::AsInt64Slice(shape->dimensions()), + /*pos=*/0, + /*len=*/n_dims - sizes.size()); + TF_ASSIGN_OR_RETURN(auto padded_starts, + PrependZerosInMajorDims(builder, x, starts)); + auto padded_sizes = PrependMajorDims(builder, major_dims, sizes); + return builder->DynamicSlice(x, padded_starts, padded_sizes); +} + xla::StatusOr UpdateSlice( xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x, const xla::ComputationDataHandle& update, gtl::ArraySlice start) { // TODO(phawkins): make int64 work on all backends, remove the int32 cast. std::vector start_as_int32(start.begin(), start.end()); - return builder->DynamicUpdateSlice( - x, update, builder->ConstantR1(start_as_int32)); + auto start_constant = builder->ConstantR1(start_as_int32); + TF_ASSIGN_OR_RETURN(std::unique_ptr shape, builder->GetShape(x)); + const int64 n_dims = xla::ShapeUtil::Rank(*shape); + TF_ASSIGN_OR_RETURN(std::unique_ptr start_constant_shape, + builder->GetShape(start_constant)); + const int64 start_length = + xla::ShapeUtil::GetDimension(*start_constant_shape, -1); + TF_RET_CHECK(start_length == n_dims); + return builder->DynamicUpdateSlice(x, update, start_constant); } xla::StatusOr UpdateSliceInMinorDims( @@ -162,6 +196,29 @@ xla::StatusOr UpdateSliceInMinorDims( return UpdateSlice(builder, x, update, padded_start); } +xla::StatusOr DynamicUpdateSliceInMinorDims( + xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x, + const xla::ComputationDataHandle& update, + const std::vector& starts) { + TF_ASSIGN_OR_RETURN(auto padded_starts, + PrependZerosInMajorDims(builder, x, starts)); + return builder->DynamicUpdateSlice(x, update, padded_starts); +} + +xla::StatusOr PrependZerosInMajorDims( + xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x, + const std::vector& starts) { + TF_ASSIGN_OR_RETURN(std::unique_ptr shape, builder->GetShape(x)); + const int64 n_dims = xla::ShapeUtil::Rank(*shape); + auto zero = builder->Reshape(builder->ConstantR0(0), {1}); + std::vector padded_starts(n_dims, zero); + for (int i = 0; i < starts.size(); ++i) { + padded_starts[n_dims - starts.size() + i] = + builder->Reshape(starts[i], {1}); + } + return builder->ConcatInDim(padded_starts, 0); +} + xla::StatusOr TransposeInMinorDims( xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x) { TF_ASSIGN_OR_RETURN(std::unique_ptr shape, builder->GetShape(x)); diff --git a/tensorflow/compiler/tf2xla/lib/util.h b/tensorflow/compiler/tf2xla/lib/util.h index 51f8baaf00bd8fd25baa1a87be8cb0089dfb22b5..b684123f1363cff9e6ac4314cc3a8ae7630cbdf3 100644 --- a/tensorflow/compiler/tf2xla/lib/util.h +++ b/tensorflow/compiler/tf2xla/lib/util.h @@ -1,4 +1,4 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -32,16 +32,39 @@ xla::ComputationDataHandle Zeros(xla::ComputationBuilder* builder, xla::ComputationDataHandle FloatLiteral(xla::ComputationBuilder* builder, xla::PrimitiveType type, double value); +// Makes a 1D tensor [0, ..., x, y] from two tensors x and y with zeros +// prepended until the array is length n_dims. +xla::ComputationDataHandle PrependZerosInMajorDims( + xla::ComputationBuilder* builder, + gtl::ArraySlice starts); + // 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); +// Builds a vector of zeros of length rank(x) with the last two values being +// those in `starts`. +xla::StatusOr PrependZerosInMajorDims( + xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x, + const std::vector& starts); + // Performs a slice in the minor dimensions of a Tensor. xla::StatusOr SliceInMinorDims( xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x, gtl::ArraySlice start, gtl::ArraySlice end); +// Builds a 1-d vector out of a concatenation of `major_dims` and `starts`. +std::vector PrependMajorDims(xla::ComputationBuilder* builder, + const gtl::ArraySlice& major_dims, + const gtl::ArraySlice& indices); + +// Performs a dynamic slice in the minor dimensions of a Tensor. +xla::StatusOr DynamicSliceInMinorDims( + xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x, + const std::vector& starts, + const gtl::ArraySlice& sizes); + // Updates a slice of 'x', i.e., // x[start[0], ..., start[n]] = update xla::StatusOr UpdateSlice( @@ -54,6 +77,11 @@ xla::StatusOr UpdateSliceInMinorDims( xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x, const xla::ComputationDataHandle& update, gtl::ArraySlice start); +xla::StatusOr DynamicUpdateSliceInMinorDims( + xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x, + const xla::ComputationDataHandle& update, + const std::vector& starts); + // Transposes a stack of matrices `x` by swapping the last two dimensions. xla::StatusOr TransposeInMinorDims( xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x); diff --git a/tensorflow/compiler/tf2xla/lib/util_test.cc b/tensorflow/compiler/tf2xla/lib/util_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..b6bd33af2e42a4ab93a22528fd49ef53c46bb479 --- /dev/null +++ b/tensorflow/compiler/tf2xla/lib/util_test.cc @@ -0,0 +1,145 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT 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/util.h" + +#include +#include +#include + +#include "tensorflow/compiler/tf2xla/lib/batch_dot.h" +#include "tensorflow/compiler/xla/array2d.h" +#include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/tests/client_library_test_base.h" +#include "tensorflow/compiler/xla/tests/literal_test_util.h" +#include "tensorflow/compiler/xla/tests/test_macros.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/core/lib/core/status_test_util.h" + +namespace tensorflow { +namespace { + +using UtilTest = xla::ClientLibraryTestBase; +using UtilLeftLookingTest = xla::ClientLibraryTestBase; + +xla::Array2D BValsRight() { + return {{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}; +} + +xla::Array2D BValsLeft() { + return {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {10, 11, 12}}; +} + +xla::Array2D AValsFull() { + return {{2, 0, 1, 2}, {3, 6, 0, 1}, {4, 7, 9, 0}, {5, 8, 10, 11}}; +} + +xla::Array3D BatchedAValsFull() { + return {{ + {2, 0, 1, 2}, + {3, 6, 0, 1}, + {4, 7, 9, 0}, + {5, 8, 10, 11}, + }, + { + {16, 24, 8, 12}, + {24, 61, 82, 48}, + {8, 82, 456, 106}, + {12, 48, 106, 62}, + }}; +} + +XLA_TEST_F(UtilTest, Simple2dLookup) { + xla::ComputationBuilder builder(client_, TestName()); + + xla::ComputationDataHandle a, x, y; + auto a_data = CreateR2Parameter(BValsRight(), 0, "a", &builder, &a); + auto x_data = CreateR0Parameter(2, 1, "x", &builder, &x); + auto y_data = CreateR0Parameter(1, 2, "y", &builder, &y); + auto result = DynamicSliceInMinorDims(&builder, a, {x, y}, {1, 1}); + TF_ASSERT_OK(result.status()); + + ComputeAndCompareR2(&builder, {{10}}, + {a_data.get(), x_data.get(), y_data.get()}, + xla::ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(UtilTest, Simple3dLookup) { + xla::ComputationBuilder builder(client_, TestName()); + + xla::ComputationDataHandle a, index; + auto a_data = + CreateR3Parameter(BatchedAValsFull(), 0, "a", &builder, &a); + auto index_data = CreateR0Parameter(1, 1, "index", &builder, &index); + + TF_ASSERT_OK_AND_ASSIGN( + auto l_index, + DynamicSliceInMinorDims(&builder, a, + {index, builder.ConstantR0(0)}, {1, 4})); + + ComputeAndCompareR3(&builder, {{{3, 6, 0, 1}}, {{24, 61, 82, 48}}}, + {a_data.get(), index_data.get()}); +} + +XLA_TEST_F(UtilTest, SimpleSliceUpdate) { + xla::ComputationBuilder builder(client_, TestName()); + + xla::ComputationDataHandle a, b, x, y; + auto a_data = CreateR2Parameter(AValsFull(), 0, "a", &builder, &a); + auto b_data = CreateR2Parameter({{9, 1, -10}}, 1, "b", &builder, &b); + auto x_data = CreateR0Parameter(2, 2, "x", &builder, &x); + auto y_data = CreateR0Parameter(1, 3, "y", &builder, &y); + + auto result = DynamicUpdateSliceInMinorDims(&builder, a, b, {x, y}); + TF_ASSERT_OK(result.status()); + + xla::Array2D expected( + {{{2, 0, 1, 2}, {3, 6, 0, 1}, {4, 9, 1, -10}, {5, 8, 10, 11}}}); + + ComputeAndCompareR2( + &builder, expected, + {a_data.get(), b_data.get(), x_data.get(), y_data.get()}); +} + +XLA_TEST_F(UtilTest, RowBatchDot) { + xla::ComputationBuilder builder(client_, TestName()); + + int n = 4; + + xla::ComputationDataHandle a, row, index; + auto a_data = + CreateR3Parameter(BatchedAValsFull(), 0, "a", &builder, &a); + auto row_data = CreateR3Parameter({{{9, 1, 0, 0}}, {{2, 4, 0, 0}}}, 1, + "row", &builder, &row); + // Select {{3, 6, 0, 1}, {24, 61, 82, 48}} out of BatchedAValsFull(). + auto index_data = CreateR0Parameter(1, 2, "index", &builder, &index); + + TF_ASSERT_OK_AND_ASSIGN( + auto l_index, + DynamicSliceInMinorDims(&builder, a, + {index, builder.ConstantR0(0)}, {1, n})); + TF_ASSERT_OK_AND_ASSIGN( + auto dot, BatchDot(&builder, l_index, row, + /*transpose_x=*/false, /*transpose_y=*/true)); + + ComputeAndCompareR3(&builder, {{{33}}, {{292}}}, + {a_data.get(), row_data.get(), index_data.get()}); +} + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/ops/BUILD b/tensorflow/compiler/tf2xla/ops/BUILD index aeb743a6634673f2e8c4dee9ae1e5017944aae2c..bb9168fa358154f3db9dab87bacc9bf28dd16406 100644 --- a/tensorflow/compiler/tf2xla/ops/BUILD +++ b/tensorflow/compiler/tf2xla/ops/BUILD @@ -7,17 +7,13 @@ licenses(["notice"]) # Apache 2.0 load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py") cc_library( - name = "functional_ops", - srcs = ["functional_ops.cc"], - deps = [ - "//tensorflow/core:framework", + name = "xla_ops", + srcs = [ + "dynamic_slice_ops.cc", + "functional_ops.cc", + "reduce_window_op.cc", + "sendrecv_ops.cc", ], - alwayslink = 1, -) - -cc_library( - name = "sendrecv_ops", - srcs = ["sendrecv_ops.cc"], deps = [ "//tensorflow/core:framework", ], @@ -25,17 +21,9 @@ cc_library( ) tf_gen_op_wrapper_py( - name = "gen_functional_ops", - out = "gen_functional_ops.py", - deps = [ - ":functional_ops", - ], -) - -tf_gen_op_wrapper_py( - name = "gen_sendrecv_ops", - out = "gen_sendrecv_ops.py", + name = "gen_xla_ops", + out = "gen_xla_ops.py", deps = [ - ":sendrecv_ops", + ":xla_ops", ], ) diff --git a/tensorflow/compiler/tf2xla/ops/dynamic_slice_ops.cc b/tensorflow/compiler/tf2xla/ops/dynamic_slice_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..d6c0edbb889b1751ac9d9d47d0c9534b543196ff --- /dev/null +++ b/tensorflow/compiler/tf2xla/ops/dynamic_slice_ops.cc @@ -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. +==============================================================================*/ + +#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("XlaDynamicUpdateSlice") + .Input("input: T") + .Input("update: T") + .Input("indices: Tindices") + .Output("output: T") + .Attr("T: type") + .Attr("Tindices: {int32, int64}") + .SetShapeFn(shape_inference::UnchangedShape) + .Doc(R"doc( +Wraps the XLA DynamicUpdateSlice operator, documented at + https://www.tensorflow.org/performance/xla/operation_semantics#dynamicupdateslice +. + +XlaDynamicUpdateSlice generates a result which is the value of the `input` +operand, with a slice update overwritten at `indices`. The shape of `update` +determines the shape of the sub-array of the result which is updated. The shape +of indices must be rank == 1, with dimension size equal to the rank of `input`. + +Handling of out-of-bounds slice indices is implementation-defined. + +input: A `Tensor` of type T. +indices: A vector of indices into `input`. Must have length equal to the rank of + `input`. +update: A `Tensor` of type T. Same rank as `input`. +output: A `Tensor` of type T. +)doc"); + +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/ops/reduce_window_op.cc b/tensorflow/compiler/tf2xla/ops/reduce_window_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..d9af982adc090ea78c711fd4656ba429c53b18c9 --- /dev/null +++ b/tensorflow/compiler/tf2xla/ops/reduce_window_op.cc @@ -0,0 +1,45 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/core/framework/common_shape_fns.h" +#include "tensorflow/core/framework/op.h" + +namespace tensorflow { + +REGISTER_OP("XlaReduceWindow") + .Input("input: T") + .Input("init_value: T") + .Attr("T: numbertype") + .Attr("computation: func") + .Attr("window_dimensions: list(int)") + .Attr("window_strides: list(int)") + .Attr("padding_low: list(int)") + .Attr("padding_high: list(int)") + .Output("output: T") + .SetShapeFn(shape_inference::UnknownShape) + .Doc(R"doc( +Wraps the XLA ReduceWindow operator, documented at + https://www.tensorflow.org/performance/xla/operation_semantics#reducewindow . + +input: the input tensor +init_value: a scalar representing the initial value for the reduction +computation: a reducer function to apply +window_dimensions: the shape of the window +window_strides: the inter-window strides +padding_low: the padding to apply at the start of each input dimensions +padding_high: the padding to apply at the end of each input dimension. +)doc"); + +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/ops/sendrecv_ops.cc b/tensorflow/compiler/tf2xla/ops/sendrecv_ops.cc index 4b41c16a8b3fdc0c3412c76d29d3ec2b7bdfd0aa..7ec7b50e905a6cbdecea4543dcb87322b5a7e844 100644 --- a/tensorflow/compiler/tf2xla/ops/sendrecv_ops.cc +++ b/tensorflow/compiler/tf2xla/ops/sendrecv_ops.cc @@ -18,22 +18,24 @@ limitations under the License. namespace tensorflow { -REGISTER_OP("_XLASend") +REGISTER_OP("XlaSend") .Input("tensor: T") .Attr("T: type") .Attr("tensor_name: string") .SetIsStateful() .SetShapeFn(shape_inference::UnknownShape) .Doc(R"doc( -Sends the named tensor to another XLA computation. +Sends the named tensor to another XLA computation. Wraps the XLA Send operator +documented at + https://www.tensorflow.org/performance/xla/operation_semantics#send . tensor: The tensor to send. -tensor_name: The name of the tensor to send. +tensor_name: A string key that identifies the channel. )doc"); -REGISTER_OP("_XLARecv") - .Output("tensor: T") - .Attr("T: type") +REGISTER_OP("XlaRecv") + .Output("tensor: dtype") + .Attr("dtype: type") .Attr("tensor_name: string") .Attr("shape: shape") .SetIsStateful() @@ -46,11 +48,14 @@ REGISTER_OP("_XLARecv") return Status::OK(); }) .Doc(R"doc( -Receives the named tensor from another XLA computation. +Receives the named tensor from another XLA computation. Wraps the XLA Recv +operator documented at + https://www.tensorflow.org/performance/xla/operation_semantics#recv . tensor: The tensor to receive. -tensor_name: The name of the tensor to receive. -shape: The shape of the input tensor. +dtype: The type of the tensor. +tensor_name: A string key that identifies the channel. +shape: The shape of the tensor. )doc"); } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/python/BUILD b/tensorflow/compiler/tf2xla/python/BUILD index f0a2ef0651ff6115bd201a3b1c34b3c061a22a3d..42b6292f79ffddd155c05758a1420a2a583eb0c6 100644 --- a/tensorflow/compiler/tf2xla/python/BUILD +++ b/tensorflow/compiler/tf2xla/python/BUILD @@ -22,3 +22,11 @@ tf_py_clif_cc( "//tensorflow/compiler/tf2xla:xla_compiler", ], ) + +py_library( + name = "xla", + srcs = ["xla.py"], + deps = [ + "//tensorflow/compiler/tf2xla/ops:gen_xla_ops", + ], +) diff --git a/tensorflow/compiler/tf2xla/python/xla.py b/tensorflow/compiler/tf2xla/python/xla.py new file mode 100644 index 0000000000000000000000000000000000000000..e5ce65bec950fdfd38c3ca5bc62ac745ef8ca4a7 --- /dev/null +++ b/tensorflow/compiler/tf2xla/python/xla.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. +# ============================================================================== +"""Experimental library that exposes XLA operations directly in TensorFlow. + +It is sometimes useful to be able to build HLO programs directly from +TensorFlow. This file provides Tensorflow operators that map as closely as +possible to HLO operators. + +There is no promise of backward or forward compatibility for operators defined +in this module. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.compiler.tf2xla.ops import gen_xla_ops + +# TODO(phawkins): provide wrappers for all XLA operators. + +dynamic_update_slice = gen_xla_ops.xla_dynamic_update_slice + + +def reduce_window(operand, + init, + reducer, + window_dimensions, + window_strides=None, + padding=None, + name=None): + """Wraps the XLA ReduceWindow operator. + + ReduceWindow is documented at + https://www.tensorflow.org/performance/xla/operation_semantics#reducewindow . + + Args: + operand: the input tensor + init: a scalar tensor representing the initial value for the reduction + reducer: a reduction function that combines a pair of scalars. + window_dimensions: shape of the window, as a list of integers + window_strides: inter-window strides, as a list of integers. Optional; + if omitted, defaults to strides of 1. + padding: padding to apply to 'operand'. List of (low, high) pairs of + integers that specify the padding to apply before and after each + dimension. Optional; if omitted, defaults to no padding. + name: the operator name, or None. + Returns: + A tensor that represents the output of the reduce_window operator. + """ + window_strides = window_strides or [1] * len(window_dimensions) + padding = padding or [(0, 0)] * len(window_dimensions) + padding_low = [x for (x, _) in padding] + padding_high = [y for (_, y) in padding] + return gen_xla_ops.xla_reduce_window( + operand, + init, + reducer, + window_dimensions, + window_strides, + padding_low, + padding_high, + name=name) + + +recv = gen_xla_ops.xla_recv +send = gen_xla_ops.xla_send + +while_loop = gen_xla_ops.xla_while diff --git a/tensorflow/compiler/tf2xla/tf2xla_util.cc b/tensorflow/compiler/tf2xla/tf2xla_util.cc index 2fc77cc4bc536a9e0f48f3933615985650c2cd84..7ec85aa3cdec622cae509f45c5ba7740222025f9 100644 --- a/tensorflow/compiler/tf2xla/tf2xla_util.cc +++ b/tensorflow/compiler/tf2xla/tf2xla_util.cc @@ -288,4 +288,13 @@ Status SetNodeShardingFromNeighbors(Node* n, bool out_edges) { return Status::OK(); } +void AddDtypeToKernalDefConstraint(StringPiece name, DataType dtype, + KernelDef* kdef) { + for (KernelDef::AttrConstraint& constraint : *kdef->mutable_constraint()) { + if (constraint.name() == name) { + constraint.mutable_allowed_values()->mutable_list()->add_type(dtype); + } + } +} + } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/tf2xla_util.h b/tensorflow/compiler/tf2xla/tf2xla_util.h index e5fba8ede7745febbb42c572a7b52247213afc95..745beb39c1d917cd0d1cd219536ee26a96253ec9 100644 --- a/tensorflow/compiler/tf2xla/tf2xla_util.h +++ b/tensorflow/compiler/tf2xla/tf2xla_util.h @@ -20,6 +20,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/tf2xla.pb.h" #include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/framework/kernel_def.pb.h" #include "tensorflow/core/framework/op.h" #include "tensorflow/core/graph/graph.h" #include "tensorflow/core/lib/core/status.h" @@ -51,6 +52,10 @@ string TensorIdToString(const tf2xla::TensorId& id); // edges are considered. Status SetNodeShardingFromNeighbors(Node* n, bool out_edges); +// Add an allowed data type to the AttrConstraint with the given name. +void AddDtypeToKernalDefConstraint(StringPiece name, DataType dtype, + KernelDef* kdef); + } // namespace tensorflow #endif // TENSORFLOW_COMPILER_TF2XLA_TF2XLA_UTIL_H_ diff --git a/tensorflow/compiler/tf2xla/xla_compiler.cc b/tensorflow/compiler/tf2xla/xla_compiler.cc index 86263d847ae02d50e70dafb0129b2664c522f2a3..c0e996768491a6315c21021ce874b8a11557de6e 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler.cc @@ -813,4 +813,29 @@ Status XlaCompiler::SetHostToDeviceMetadata( return Status::OK(); } +Status XlaCompiler::GetHostComputeControlDependency( + const string& host_compute_name, xla::ComputationDataHandle* handle) { + const auto iter = host_compute_control_output_.find(host_compute_name); + if (iter == host_compute_control_output_.end()) { + return errors::InvalidArgument( + "No registered control handle for host compute Op '", host_compute_name, + "'"); + } else { + *handle = iter->second; + } + return Status::OK(); +} + +Status XlaCompiler::SetHostComputeControlDependency( + const string& host_compute_name, const xla::ComputationDataHandle& handle) { + if (host_compute_control_output_.find(host_compute_name) != + host_compute_control_output_.end()) { + return errors::InvalidArgument( + "Duplicate control handles registered for for host compute Op ", + host_compute_name); + } + host_compute_control_output_[host_compute_name] = handle; + return Status::OK(); +} + } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_compiler.h b/tensorflow/compiler/tf2xla/xla_compiler.h index a6747bbe72e161b2ece55697825cce0e71145a5c..8f564f35ec81765e8998513dfd4805d221200c6c 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.h +++ b/tensorflow/compiler/tf2xla/xla_compiler.h @@ -325,6 +325,23 @@ class XlaCompiler { gtl::ArraySlice types, gtl::ArraySlice shapes); + // In order to avoid deadlocks from dependencies in host computations, it can + // be necessary to enforce a partial order on the execution of HostCompute + // Ops. In particular it may be necessary to constrain the SendToHost for one + // HostCompute to run before blocking on the RecvAtHost for another + // HostCompute. The compiler maintains a mapping from 'host_compute_name' to + // handle, where the handle is an 'output' of the HostCompute Op corresponding + // to 'host_compute_name'. Another HostCompute Op that needs to be sequenced + // later can add the handle as an 'input' to enforce the constraints. + // 'host_compute_name' can be any string the client wishes to use to identify + // a given HostCompute Op as long as the names are unique within the + // compilation. + Status GetHostComputeControlDependency(const string& host_compute_name, + xla::ComputationDataHandle* handle); + Status SetHostComputeControlDependency( + const string& host_compute_name, + const xla::ComputationDataHandle& handle); + const Options& options() const { return options_; } xla::Client* client() const { return options_.client; } FunctionLibraryRuntime* flib_runtime() const { return flib_runtime_; } @@ -391,6 +408,9 @@ class XlaCompiler { std::unordered_map host_compute_sends_; std::unordered_map host_compute_recvs_; + std::unordered_map + host_compute_control_output_; + TF_DISALLOW_COPY_AND_ASSIGN(XlaCompiler); }; diff --git a/tensorflow/compiler/tf2xla/xla_cpu_backend.cc b/tensorflow/compiler/tf2xla/xla_cpu_backend.cc index 8286480e0ea07429adbe31ec4f16d043e321df0a..ead229aaccc292d4944db0c1eaf98c82583533cd 100644 --- a/tensorflow/compiler/tf2xla/xla_cpu_backend.cc +++ b/tensorflow/compiler/tf2xla/xla_cpu_backend.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "tensorflow/compiler/tf2xla/tf2xla_util.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/core/framework/kernel_def.pb.h" @@ -30,6 +31,12 @@ bool CpuOpFilter(KernelDef* kdef) { DT_FLOAT); return true; } + if (kdef->op() == "Const") { + AddDtypeToKernalDefConstraint("dtype", DT_STRING, kdef); + } + if (kdef->op() == "Assert") { + AddDtypeToKernalDefConstraint("T", DT_STRING, kdef); + } return true; } diff --git a/tensorflow/compiler/tf2xla/xla_gpu_backend.cc b/tensorflow/compiler/tf2xla/xla_gpu_backend.cc index 8ca757e72355d890c13b8b448d35c327d3986696..62168b648331844bfe2db1a4d5dcad895c8726f3 100644 --- a/tensorflow/compiler/tf2xla/xla_gpu_backend.cc +++ b/tensorflow/compiler/tf2xla/xla_gpu_backend.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "tensorflow/compiler/tf2xla/tf2xla_util.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/core/framework/kernel_def.pb.h" @@ -25,6 +26,12 @@ bool GpuOpFilter(KernelDef* kdef) { kdef->op() == "RandomUniformInt" || kdef->op() == "TruncatedNormal") { return false; } + if (kdef->op() == "Const") { + AddDtypeToKernalDefConstraint("dtype", DT_STRING, kdef); + } + if (kdef->op() == "Assert") { + AddDtypeToKernalDefConstraint("T", DT_STRING, kdef); + } return true; } diff --git a/tensorflow/compiler/xla/BUILD b/tensorflow/compiler/xla/BUILD index 751777222fcc7ec073958349aa2677d5b4e6757d..88f37433a552e0dff1587075f1d31fd5d050302b 100644 --- a/tensorflow/compiler/xla/BUILD +++ b/tensorflow/compiler/xla/BUILD @@ -443,6 +443,9 @@ cc_library( srcs = ["executable_run_options.cc"], hdrs = ["executable_run_options.h"], visibility = ["//visibility:public"], + deps = [ + ":types", + ], ) cc_library( diff --git a/tensorflow/compiler/xla/client/BUILD b/tensorflow/compiler/xla/client/BUILD index a299c2afd45aa6b785964b8a8e1400ddf54083a4..286d06d12ffca7410067f2d33398497576986807 100644 --- a/tensorflow/compiler/xla/client/BUILD +++ b/tensorflow/compiler/xla/client/BUILD @@ -130,6 +130,7 @@ cc_library( "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/client/xla_client:xla_computation", "//tensorflow/compiler/xla/service:compile_only_service", "//tensorflow/compiler/xla/service:compiler", "//tensorflow/core:stream_executor_no_cuda", diff --git a/tensorflow/compiler/xla/client/client.cc b/tensorflow/compiler/xla/client/client.cc index 3f45167fcb77cd3085c9645fba0b2901329c4bb2..f0f94298a05f7c4bdc41cbfb8572454fbedd371d 100644 --- a/tensorflow/compiler/xla/client/client.cc +++ b/tensorflow/compiler/xla/client/client.cc @@ -193,6 +193,34 @@ StatusOr> Client::ExecuteAndTransfer( return Transfer(*data, shape_with_output_layout); } +StatusOr> Client::ComputeConstant( + const XlaComputation& computation, const Layout* output_layout) const { + ComputeConstantGraphRequest request; + *request.mutable_computation() = computation.proto(); + if (output_layout != nullptr) { + *request.mutable_output_layout() = *output_layout; + } + + ComputeConstantResponse response; + + VLOG(2) << "making compute-constant-graph request"; + Status s = stub_->ComputeConstantGraph(&request, &response); + VLOG(2) << "done with request"; + + if (!s.ok()) { + return s; + } + + VLOG(3) << "ComputeConstant: {" << response.DebugString() << "}"; + + if (!response.has_literal()) { + return InternalError( + "no computed literal in the provided response in ComputeConstantGraph " + "request"); + } + return Literal::CreateFromProto(response.literal()); +} + StatusOr Client::LoadSnapshot(const SessionModule& module) { LoadComputationSnapshotRequest request; *request.mutable_module() = module; diff --git a/tensorflow/compiler/xla/client/client.h b/tensorflow/compiler/xla/client/client.h index 05d707dab1533f44ce827157e888720e218d4c9c..14c685d94ea31c382d84223ca4e2eba544420d78 100644 --- a/tensorflow/compiler/xla/client/client.h +++ b/tensorflow/compiler/xla/client/client.h @@ -194,6 +194,27 @@ class Client { const ExecutionOptions* execution_options = nullptr, ExecutionProfile* execution_profile = nullptr); + // Computes the value of the given computation using a non-optimized + // interpreter on the host. + // + // The computation must not depend on any parameters, or on stateful operators + // such as `RngNormal` or `Infeed`. + // + // This functionality can be useful when translating a computation into XLA + // where something that looked dynamic is required by XLA to be specified as a + // constant. E.g. the source computation (outside of XLA) may include a + // dynamic computation of the shape of something and ComputeConstant lets you + // determine what the value of that computation is in the case where the value + // can be determined at compile time. + // + // If output_layout is non-null, then the output of the computation will be + // stored using that layout. + // + // TODO(b/74197823): This is a part of a NOT YET ready refactor. + StatusOr> ComputeConstant( + const XlaComputation& computation, + const Layout* output_layout = nullptr) const; + // Unregister the memory for the given GlobalData on the device. Status Unregister(const GlobalData& data); diff --git a/tensorflow/compiler/xla/client/client_library.cc b/tensorflow/compiler/xla/client/client_library.cc index b1663bc815719c3da75b37593ac665b1f3493db8..803a9e40094391ba47ed27713f4538caf875c4f6 100644 --- a/tensorflow/compiler/xla/client/client_library.cc +++ b/tensorflow/compiler/xla/client/client_library.cc @@ -23,22 +23,19 @@ limitations under the License. namespace xla { -LocalClientOptions::LocalClientOptions(perftools::gputools::Platform* platform, +LocalClientOptions::LocalClientOptions(se::Platform* platform, int number_of_replicas, int intra_op_parallelism_threads) : platform_(platform), number_of_replicas_(number_of_replicas), intra_op_parallelism_threads_(intra_op_parallelism_threads) {} -LocalClientOptions& LocalClientOptions::set_platform( - perftools::gputools::Platform* platform) { +LocalClientOptions& LocalClientOptions::set_platform(se::Platform* platform) { platform_ = platform; return *this; } -perftools::gputools::Platform* LocalClientOptions::platform() const { - return platform_; -} +se::Platform* LocalClientOptions::platform() const { return platform_; } LocalClientOptions& LocalClientOptions::set_number_of_replicas( int number_of_replicas) { @@ -69,7 +66,7 @@ ClientLibrary::ClientLibrary() = default; ClientLibrary::~ClientLibrary() = default; /* static */ StatusOr ClientLibrary::GetOrCreateLocalClient( - perftools::gputools::Platform* platform) { + se::Platform* platform) { LocalClientOptions default_options; default_options.set_platform(platform); return GetOrCreateLocalClient(default_options); @@ -77,7 +74,7 @@ ClientLibrary::~ClientLibrary() = default; /* static */ StatusOr ClientLibrary::GetOrCreateLocalClient( const LocalClientOptions& options) { - perftools::gputools::Platform* platform = options.platform(); + se::Platform* platform = options.platform(); int replica_count = options.number_of_replicas(); ClientLibrary& client_library = Singleton(); tensorflow::mutex_lock lock(client_library.service_mutex_); @@ -115,7 +112,7 @@ ClientLibrary::~ClientLibrary() = default; } /* static */ LocalService* ClientLibrary::GetXlaService( - perftools::gputools::Platform* platform) { + se::Platform* platform) { ClientLibrary& client_library = Singleton(); tensorflow::mutex_lock lock(client_library.service_mutex_); auto it = client_library.local_instances_.find(platform->id()); @@ -124,8 +121,7 @@ ClientLibrary::~ClientLibrary() = default; } /* static */ StatusOr -ClientLibrary::GetOrCreateCompileOnlyClient( - perftools::gputools::Platform* platform) { +ClientLibrary::GetOrCreateCompileOnlyClient(se::Platform* platform) { ClientLibrary& client_library = Singleton(); tensorflow::mutex_lock lock(client_library.service_mutex_); diff --git a/tensorflow/compiler/xla/client/client_library.h b/tensorflow/compiler/xla/client/client_library.h index a6f30d82e43587135697e76e8bc7d122edc0f602..3ad558fa532931937fab898f7b855f0a3370eaec 100644 --- a/tensorflow/compiler/xla/client/client_library.h +++ b/tensorflow/compiler/xla/client/client_library.h @@ -43,13 +43,13 @@ namespace xla { // Options to configure the local client when it is created. class LocalClientOptions { public: - LocalClientOptions(perftools::gputools::Platform* platform = nullptr, + LocalClientOptions(se::Platform* platform = nullptr, int number_of_replicas = 1, int intra_op_parallelism_threads = -1); // Set the platform backing the service, or nullptr for the default platform. - LocalClientOptions& set_platform(perftools::gputools::Platform* platform); - perftools::gputools::Platform* platform() const; + LocalClientOptions& set_platform(se::Platform* platform); + se::Platform* platform() const; // Set the number of replicas to use when compiling replicated // programs. @@ -61,7 +61,7 @@ class LocalClientOptions { int intra_op_parallelism_threads() const; private: - perftools::gputools::Platform* platform_; + se::Platform* platform_; int number_of_replicas_; int intra_op_parallelism_threads_; }; @@ -74,7 +74,7 @@ class ClientLibrary { // platform : The platform the underlying XLA service should target. If // null then default platform is used. static StatusOr GetOrCreateLocalClient( - perftools::gputools::Platform* platform = nullptr); + se::Platform* platform = nullptr); static StatusOr GetOrCreateLocalClient( const LocalClientOptions& options); @@ -84,14 +84,14 @@ class ClientLibrary { // Returns the service from the service thread. Only used in unit tests to // access user computations from client. - static LocalService* GetXlaService(perftools::gputools::Platform* platform); + static LocalService* GetXlaService(se::Platform* platform); // Singleton constructor-or-accessor for compile-only clients. Arguments: // // platform : The platform the underlying XLA service should target. If // null then default platform is used. static StatusOr GetOrCreateCompileOnlyClient( - perftools::gputools::Platform* platform = nullptr); + se::Platform* platform = nullptr); // Clears the local instance and compile only instance caches. The client // pointers returned by the previous GetOrCreateLocalClient() or @@ -120,12 +120,10 @@ class ClientLibrary { }; tensorflow::mutex service_mutex_; // Guards the singleton creation state. - std::unordered_map> + std::unordered_map> local_instances_ GUARDED_BY(service_mutex_); - std::unordered_map> + std::unordered_map> compile_only_instances_ GUARDED_BY(service_mutex_); TF_DISALLOW_COPY_AND_ASSIGN(ClientLibrary); diff --git a/tensorflow/compiler/xla/client/compile_only_client.cc b/tensorflow/compiler/xla/client/compile_only_client.cc index 59662c95ac15e7c23790c5b5ff5d75a694613aeb..96e38bca01087991943aff40ed1cb3e21f9e6cba 100644 --- a/tensorflow/compiler/xla/client/compile_only_client.cc +++ b/tensorflow/compiler/xla/client/compile_only_client.cc @@ -39,6 +39,24 @@ CompileOnlyClient::CompileAheadOfTime( return compiler_service_->CompileAheadOfTime(service_instances, options); } +StatusOr>> +CompileOnlyClient::CompileAheadOfTime( + const tensorflow::gtl::ArraySlice computations, + const AotCompilationOptions& options) { + std::vector service_instances; + service_instances.reserve(computations.size()); + for (const AotXlaComputationInstance& instance : computations) { + service_instances.emplace_back(); + CompileOnlyService::AotXlaComputationInstance& service_instance = + service_instances.back(); + TF_RET_CHECK(instance.computation != nullptr); + service_instance.computation = instance.computation->proto(); + service_instance.argument_layouts = instance.argument_layouts; + service_instance.result_layout = instance.result_layout; + } + return compiler_service_->CompileAheadOfTime(service_instances, options); +} + int64 CompileOnlyClient::PointerSizeForTriple(tensorflow::StringPiece triple) { llvm::Triple llvm_triple( llvm::Triple::normalize(llvm::StringRef(triple.data(), triple.size()))); diff --git a/tensorflow/compiler/xla/client/compile_only_client.h b/tensorflow/compiler/xla/client/compile_only_client.h index 5900048711384e0240a3cd502260eb388eb40f51..c8725b8517484acdaf093bc3b34adb00f69155b1 100644 --- a/tensorflow/compiler/xla/client/compile_only_client.h +++ b/tensorflow/compiler/xla/client/compile_only_client.h @@ -18,6 +18,7 @@ limitations under the License. #include "tensorflow/compiler/xla/client/client.h" #include "tensorflow/compiler/xla/client/computation.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h" #include "tensorflow/compiler/xla/service/compile_only_service.h" #include "tensorflow/compiler/xla/service/compiler.h" #include "tensorflow/compiler/xla/statusor.h" @@ -54,6 +55,27 @@ class CompileOnlyClient : public Client { const tensorflow::gtl::ArraySlice computations, const AotCompilationOptions& options); + // A description of an xla computation to compile using CompileAheadOfTime. + // + // TODO(b/74197823): This is a part of a NOT YET ready refactor. + struct AotXlaComputationInstance { + const XlaComputation* computation; + // Inform the compiler of the expected layout for arguments. + std::vector argument_layouts; + // Specifies the expected result layout. + const Shape* result_layout; + }; + + // Compiles a list of xla computations for ahead-of-time execution. This is + // intended for use in static compilation. The |options| parameter describes + // the target for which the compiler should emit code. + // + // TODO(b/74197823): This is a part of a NOT YET ready refactor. + StatusOr>> + CompileAheadOfTime( + const tensorflow::gtl::ArraySlice computations, + const AotCompilationOptions& options); + // Returns the size of a pointer in bytes for a given triple. static int64 PointerSizeForTriple(tensorflow::StringPiece triple); diff --git a/tensorflow/compiler/xla/client/computation_builder.cc b/tensorflow/compiler/xla/client/computation_builder.cc index 4d3b0ee0d6e9ba82cfa09af0fbff0ae1efa0ac64..83c7cb174402133706fbde6a734a29afd8edfe80 100644 --- a/tensorflow/compiler/xla/client/computation_builder.cc +++ b/tensorflow/compiler/xla/client/computation_builder.cc @@ -1046,6 +1046,11 @@ ComputationDataHandle ComputationBuilder::Neg( return UnaryOp(UNOP_NEGATE, operand); } +ComputationDataHandle ComputationBuilder::Clz( + const ComputationDataHandle& operand) { + return UnaryOp(UNOP_CLZ, operand); +} + ComputationDataHandle ComputationBuilder::Clamp( const ComputationDataHandle& min, const ComputationDataHandle& operand, const ComputationDataHandle& max) { diff --git a/tensorflow/compiler/xla/client/computation_builder.h b/tensorflow/compiler/xla/client/computation_builder.h index 019c6f3afb5d57bfe453988ded19120a4483cf36..9431c2c459a564e3cf509d9dae16e71fc27ee2c0 100644 --- a/tensorflow/compiler/xla/client/computation_builder.h +++ b/tensorflow/compiler/xla/client/computation_builder.h @@ -657,6 +657,9 @@ class ComputationBuilder { // Enqueues a negate instruction onto the computation. ComputationDataHandle Neg(const ComputationDataHandle& operand); + // Enqueues a count-leading-zeros instruction onto the computation. + ComputationDataHandle Clz(const ComputationDataHandle& operand); + // Enqueues a transpose instruction onto the computation. ComputationDataHandle Transpose( const ComputationDataHandle& operand, diff --git a/tensorflow/compiler/xla/client/lib/BUILD b/tensorflow/compiler/xla/client/lib/BUILD index f4673a8204f27e93441c73f6dcc9130d96cfcebc..59c4a53c05a45490a7c8e732840a4e70767c46c2 100644 --- a/tensorflow/compiler/xla/client/lib/BUILD +++ b/tensorflow/compiler/xla/client/lib/BUILD @@ -46,6 +46,7 @@ cc_library( "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", + "//tensorflow/compiler/xla/client/xla_client:xla_computation", "//tensorflow/compiler/xla/tests:test_utils", "//tensorflow/core:lib", ], diff --git a/tensorflow/compiler/xla/client/lib/testing.cc b/tensorflow/compiler/xla/client/lib/testing.cc index b63a1465ea755b906853860d47768ecbeaa0dcdd..311dc4bdd72cfd7999e83a26e11614d6ca005bce 100644 --- a/tensorflow/compiler/xla/client/lib/testing.cc +++ b/tensorflow/compiler/xla/client/lib/testing.cc @@ -111,4 +111,20 @@ std::vector> MakeFakeArgumentsOrDie( return fake_arguments; } +std::vector> MakeFakeArgumentsOrDie( + const XlaComputation& computation, Client* client) { + CHECK(computation.proto().has_program_shape()) + << "Computation should have progran shape."; + auto program_shape = computation.proto().program_shape(); + + // For every (unbound) parameter that the computation wants, we manufacture + // some arbitrary data so that we can invoke the computation. + std::vector> fake_arguments; + for (const Shape& parameter : program_shape.parameters()) { + fake_arguments.push_back(MakeFakeDataOrDie(parameter, client)); + } + + return fake_arguments; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/testing.h b/tensorflow/compiler/xla/client/lib/testing.h index 7e640d1307edcc3e2c021f4391c456f578a015ee..1dc2622972d5fd3da6991d70b800cc3fd5a638f4 100644 --- a/tensorflow/compiler/xla/client/lib/testing.h +++ b/tensorflow/compiler/xla/client/lib/testing.h @@ -22,6 +22,7 @@ limitations under the License. #include "tensorflow/compiler/xla/client/client.h" #include "tensorflow/compiler/xla/client/computation.h" #include "tensorflow/compiler/xla/client/global_data.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h" #include "tensorflow/compiler/xla/xla_data.pb.h" namespace xla { @@ -38,6 +39,12 @@ std::unique_ptr MakeFakeDataOrDie(const Shape& shape, std::vector> MakeFakeArgumentsOrDie( const Computation& computation, Client* client); +// Returns vector of GlobalData handles of fake data (created using +// MakeFakeDataOrDie) that are correctly shaped arguments for the given +// xla computation. +std::vector> MakeFakeArgumentsOrDie( + const XlaComputation& computation, Client* client); + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_TESTING_H_ diff --git a/tensorflow/compiler/xla/client/local_client.cc b/tensorflow/compiler/xla/client/local_client.cc index 30594243dcf51d2b5312b9dcb2bea7d0cd78524d..d0e945b70fdb3c7abdcf94cb03d3b962f8556576 100644 --- a/tensorflow/compiler/xla/client/local_client.cc +++ b/tensorflow/compiler/xla/client/local_client.cc @@ -24,8 +24,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/source_map_util.h" #include "tensorflow/compiler/xla/status_macros.h" -namespace se = ::perftools::gputools; - using xla::source_map_util::InvalidParameterArgument; namespace xla { @@ -136,7 +134,7 @@ tensorflow::Status LocalExecutable::ValidateExecutionOptions( return Status::OK(); } -StatusOr> LocalExecutable::Run( +StatusOr LocalExecutable::Run( const tensorflow::gtl::ArraySlice arguments, ExecutableRunOptions run_options) { TF_RETURN_IF_ERROR( @@ -169,27 +167,26 @@ StatusOr> LocalExecutable::Run( return ExecuteAndDump(&service_options, arguments); } TF_ASSIGN_OR_RETURN( - std::unique_ptr result, + ShapedBuffer result, executable_->ExecuteOnStreamWrapper( &service_options, run_options.execution_profile(), arguments)); - return MakeUnique(std::move(*result), - run_options.allocator()); + return ScopedShapedBuffer(std::move(result), run_options.allocator()); } -StatusOr> LocalExecutable::ExecuteAndDump( +StatusOr LocalExecutable::ExecuteAndDump( const ServiceExecutableRunOptions* run_options, const tensorflow::gtl::ArraySlice arguments) { executable_->session_module()->set_execution_platform( backend_->platform()->Name()); TF_RETURN_IF_ERROR(RecordArguments(arguments, executable_->session_module())); TF_ASSIGN_OR_RETURN( - std::unique_ptr result, + ShapedBuffer result, executable_->ExecuteOnStream(run_options, arguments, /*hlo_execution_profile=*/nullptr)); - TF_RETURN_IF_ERROR(RecordResult(result.get(), executable_->session_module())); + TF_RETURN_IF_ERROR(RecordResult(&result, executable_->session_module())); TF_RETURN_IF_ERROR(executable_->DumpSessionModule()); - return ScopedShapedBuffer::MakeScoped(result.get(), run_options->allocator()); + return ScopedShapedBuffer(std::move(result), run_options->allocator()); } tensorflow::Status LocalExecutable::RecordArguments( @@ -283,9 +280,9 @@ StatusOr> LocalClient::Compile( updated_options)); } -StatusOr> -LocalClient::LiteralToShapedBuffer(const Literal& literal, int device_ordinal, - DeviceMemoryAllocator* allocator) { +StatusOr LocalClient::LiteralToShapedBuffer( + const Literal& literal, int device_ordinal, + DeviceMemoryAllocator* allocator) { if (allocator == nullptr) { allocator = backend().memory_allocator(); } @@ -295,7 +292,7 @@ LocalClient::LiteralToShapedBuffer(const Literal& literal, int device_ordinal, TF_ASSIGN_OR_RETURN(se::StreamExecutor * executor, backend().stream_executor(device_ordinal)); TF_RETURN_IF_ERROR(backend().transfer_manager()->TransferLiteralToDevice( - executor, literal, *scoped_buffer)); + executor, literal, scoped_buffer)); return std::move(scoped_buffer); } diff --git a/tensorflow/compiler/xla/client/local_client.h b/tensorflow/compiler/xla/client/local_client.h index 98ee7c62c94be7c618cedd3dc12ecbfc812ee180..f306c520ede0014be52d1b952849c8894b092baf 100644 --- a/tensorflow/compiler/xla/client/local_client.h +++ b/tensorflow/compiler/xla/client/local_client.h @@ -38,7 +38,7 @@ class LocalExecutable { public: // Run the compiled computation with the given arguments and options and // return the result. - StatusOr> Run( + StatusOr Run( const tensorflow::gtl::ArraySlice arguments, ExecutableRunOptions run_options); @@ -73,7 +73,7 @@ class LocalExecutable { // Records the computation in a SessionModule proto with the arguments used to // invoke it, and the result. Enabled by flag: --tla_dump_executions_to. - StatusOr> ExecuteAndDump( + StatusOr ExecuteAndDump( const ServiceExecutableRunOptions* run_options, const tensorflow::gtl::ArraySlice arguments); @@ -136,7 +136,7 @@ class LocalClient : public Client { // ScopedShapedBuffer. If non-null the given memory allocator is used for // device memory allocation. If null, the default memory allocator for the // device is used. - StatusOr> LiteralToShapedBuffer( + StatusOr LiteralToShapedBuffer( const Literal& literal, int device_ordinal, DeviceMemoryAllocator* allocator = nullptr); @@ -167,7 +167,7 @@ class LocalClient : public Client { StatusOr ReplicaNumberToDeviceOrdinal(int replica_number); // Returns the platform that the underlying service targets. - perftools::gputools::Platform* platform() const; + se::Platform* platform() const; // Returns the number of devices on the system of the service platform // type. Not all devices may be supported by the service (see diff --git a/tensorflow/compiler/xla/client/xla_client/BUILD b/tensorflow/compiler/xla/client/xla_client/BUILD index b1dba168565cca86cba0403604736fecd00d6f29..31fa1241ee474a31575c45cf7652063dfc818fac 100644 --- a/tensorflow/compiler/xla/client/xla_client/BUILD +++ b/tensorflow/compiler/xla/client/xla_client/BUILD @@ -44,6 +44,7 @@ cc_library( hdrs = ["xla_builder.h"], deps = [ ":xla_computation", + "//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/client/xla_client/xla_builder.cc b/tensorflow/compiler/xla/client/xla_client/xla_builder.cc index ed9f994d39c36a5e7b8e58e532ebcf75b1943442..1899983e442116d3ebf8a3e79b0515653cd624cb 100644 --- a/tensorflow/compiler/xla/client/xla_client/xla_builder.cc +++ b/tensorflow/compiler/xla/client/xla_client/xla_builder.cc @@ -17,12 +17,15 @@ limitations under the License. #include #include +#include #include #include +#include "tensorflow/compiler/xla/execution_options_util.h" #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/gtl/flatset.h" #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/mutex.h" @@ -82,7 +85,7 @@ StatusOr XlaOp::GetShape() const { } XlaBuilder::XlaBuilder(const string& computation_name) - : name_(computation_name), unique_id_(GetUniqueId()) {} + : name_(computation_name) {} XlaBuilder::~XlaBuilder() {} @@ -111,10 +114,11 @@ XlaOp XlaBuilder::NoteErrorOrReturn( return op.ConsumeValueOrDie(); } -StatusOr XlaBuilder::GetProgramShape(int64* root_id) { +StatusOr XlaBuilder::GetProgramShape(int64* root_id) const { TF_RETURN_IF_ERROR(first_error_); TF_RET_CHECK(root_id != nullptr); + ProgramShape program_shape; // Not all instructions can be roots. Walk backwards from the last added @@ -155,9 +159,56 @@ StatusOr XlaBuilder::GetProgramShape(int64* root_id) { return program_shape; } -StatusOr XlaBuilder::GetProgramShape() { - int64 root_id; - return GetProgramShape(&root_id); +StatusOr XlaBuilder::GetProgramShape() const { + int64 root; + return GetProgramShape(&root); +} + +void XlaBuilder::IsConstantVisitor(const int64 op_handle, + std::set* visited, + bool* is_constant) const { + if (visited->count(op_handle) != 0 || !*is_constant) { + return; + } + + CHECK(op_handle < instructions_.size() && op_handle >= 0); + + const HloInstructionProto& instr = instructions_[op_handle]; + const HloOpcode opcode = StringToHloOpcode(instr.opcode()).ValueOrDie(); + switch (opcode) { + default: + for (const int64 operand_id : instr.operand_ids()) { + IsConstantVisitor(operand_id, visited, is_constant); + } + // TODO(b/32495713): We aren't checking the called computations. + break; + + // Non functional ops. + case HloOpcode::kRng: + case HloOpcode::kCrossReplicaSum: + // TODO(b/33009255): Implmement constant folding for cross replica sum. + case HloOpcode::kInfeed: + case HloOpcode::kOutfeed: + case HloOpcode::kHostCompute: + case HloOpcode::kCall: + // TODO(b/32495713): We aren't checking the to_apply computation itself, + // so we conservatively say that computations containing the Call op + // cannot be constant. We cannot set is_functional=false in other similar + // cases since we're already relying on IsConstant to return true. + case HloOpcode::kCustomCall: + case HloOpcode::kWhile: + // TODO(b/32495713): We aren't checking the condition and body + // computations themselves. + case HloOpcode::kSend: + case HloOpcode::kRecv: + case HloOpcode::kParameter: + *is_constant = false; + break; + } + if (!*is_constant) { + VLOG(1) << "Non-constant: " << instr.name(); + } + visited->insert(op_handle); } XlaComputation XlaBuilder::BuildAndNoteError() { @@ -180,21 +231,24 @@ StatusOr XlaBuilder::Build() { } HloComputationProto entry; + entry.set_id(GetUniqueId()); // Give the computation a global unique id. + entry.set_name(StrCat(name_, entry.id())); // Ensure that the name is unique. { int64 root_id; - ProgramShape program_shape; - TF_ASSIGN_OR_RETURN(program_shape, GetProgramShape(&root_id)); - entry.mutable_program_shape()->Swap(&program_shape); + TF_ASSIGN_OR_RETURN(*entry.mutable_program_shape(), + GetProgramShape(&root_id)); entry.set_root_id(root_id); } for (auto& instruction : instructions_) { + // Ensures that the instruction names are unique among the whole graph. + const string& new_name = + StrCat(instruction.name(), ".", entry.id(), ".", instruction.id()); + instruction.set_name(new_name); entry.add_instructions()->Swap(&instruction); } - entry.set_id(unique_id_); - entry.set_name(StrCat(name_, entry.id())); // Ensure that the name is unique. XlaComputation computation(entry.id()); HloModuleProto* module = computation.mutable_proto(); module->set_name(entry.name()); @@ -417,11 +471,10 @@ XlaOp XlaBuilder::Parameter(int64 parameter_number, const Shape& shape, const string& name) { return NoteErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; - if (parameter_numbers_.find(parameter_number) != parameter_numbers_.end()) { + if (!parameter_numbers_.insert(parameter_number).second) { 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; @@ -485,7 +538,17 @@ XlaOp XlaBuilder::Slice(const XlaOp& operand, XlaOp XlaBuilder::SliceInDim(const XlaOp& operand, int64 start_index, int64 limit_index, int64 stride, int64 dimno) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(const Shape& shape, GetShape(operand)); + std::vector starts(ShapeUtil::Rank(shape), 0); + std::vector limits(shape.dimensions().begin(), + shape.dimensions().end()); + std::vector strides(ShapeUtil::Rank(shape), 1); + starts[dimno] = start_index; + limits[dimno] = limit_index; + strides[dimno] = stride; + return Slice(operand, starts, limits, strides); + }); } XlaOp XlaBuilder::DynamicSlice(const XlaOp& operand, const XlaOp& start_indices, @@ -737,24 +800,101 @@ XlaOp XlaBuilder::DotGeneral(const XlaOp& lhs, const XlaOp& rhs, }); } +Status XlaBuilder::VerifyConvolution( + const Shape& lhs_shape, const Shape& rhs_shape, + const ConvolutionDimensionNumbers& dimension_numbers) const { + if (ShapeUtil::Rank(lhs_shape) != ShapeUtil::Rank(rhs_shape)) { + return InvalidArgument( + "Convolution arguments must have same number of " + "dimensions. Got: %s and %s", + ShapeUtil::HumanString(lhs_shape).c_str(), + ShapeUtil::HumanString(rhs_shape).c_str()); + } + int num_dims = ShapeUtil::Rank(lhs_shape); + if (num_dims < 2) { + return InvalidArgument( + "Convolution expects argument arrays with >= 3 dimensions. " + "Got: %s and %s", + ShapeUtil::HumanString(lhs_shape).c_str(), + ShapeUtil::HumanString(rhs_shape).c_str()); + } + int num_spatial_dims = num_dims - 2; + + const auto check_spatial_dimensions = + [&](const char* const field_name, + const tensorflow::protobuf::RepeatedField& + numbers) { + if (numbers.size() != num_spatial_dims) { + return InvalidArgument("Expected %d elements for %s, but got %d.", + num_spatial_dims, field_name, numbers.size()); + } + for (int i = 0; i < numbers.size(); ++i) { + if (numbers.Get(i) < 0 || numbers.Get(i) >= num_dims) { + return InvalidArgument("Convolution %s[%d] is out of bounds: %lld", + field_name, i, numbers.Get(i)); + } + } + return Status::OK(); + }; + TF_RETURN_IF_ERROR( + check_spatial_dimensions("input_spatial_dimensions", + dimension_numbers.input_spatial_dimensions())); + TF_RETURN_IF_ERROR( + check_spatial_dimensions("kernel_spatial_dimensions", + dimension_numbers.kernel_spatial_dimensions())); + return check_spatial_dimensions( + "output_spatial_dimensions", + dimension_numbers.output_spatial_dimensions()); +} + XlaOp XlaBuilder::Conv(const XlaOp& lhs, const XlaOp& rhs, tensorflow::gtl::ArraySlice window_strides, Padding padding) { - return UnimplementedOp(); + return ConvWithGeneralDimensions( + lhs, rhs, window_strides, padding, + CreateDefaultConvDimensionNumbers(window_strides.size())); } XlaOp XlaBuilder::ConvWithGeneralPadding( const XlaOp& lhs, const XlaOp& rhs, tensorflow::gtl::ArraySlice window_strides, tensorflow::gtl::ArraySlice> padding) { - return UnimplementedOp(); + return ConvGeneral(lhs, rhs, window_strides, padding, + CreateDefaultConvDimensionNumbers(window_strides.size())); } XlaOp XlaBuilder::ConvWithGeneralDimensions( const XlaOp& lhs, const XlaOp& rhs, tensorflow::gtl::ArraySlice window_strides, Padding padding, const ConvolutionDimensionNumbers& dimension_numbers) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(const Shape& lhs_shape, GetShape(lhs)); + TF_ASSIGN_OR_RETURN(const Shape& rhs_shape, GetShape(rhs)); + + TF_RETURN_IF_ERROR( + VerifyConvolution(lhs_shape, rhs_shape, dimension_numbers)); + + std::vector base_area_dimensions( + dimension_numbers.input_spatial_dimensions_size()); + for (std::vector::size_type i = 0; i < base_area_dimensions.size(); + ++i) { + base_area_dimensions[i] = + lhs_shape.dimensions(dimension_numbers.input_spatial_dimensions(i)); + } + + std::vector window_dimensions( + dimension_numbers.kernel_spatial_dimensions_size()); + for (std::vector::size_type i = 0; i < window_dimensions.size(); + ++i) { + window_dimensions[i] = + rhs_shape.dimensions(dimension_numbers.kernel_spatial_dimensions(i)); + } + + return ConvGeneral(lhs, rhs, window_strides, + MakePadding(base_area_dimensions, window_dimensions, + window_strides, padding), + dimension_numbers); + }); } XlaOp XlaBuilder::ConvGeneral( @@ -762,7 +902,8 @@ XlaOp XlaBuilder::ConvGeneral( tensorflow::gtl::ArraySlice window_strides, tensorflow::gtl::ArraySlice> padding, const ConvolutionDimensionNumbers& dimension_numbers) { - return UnimplementedOp(); + return ConvGeneralDilated(lhs, rhs, window_strides, padding, {}, {}, + dimension_numbers); } XlaOp XlaBuilder::ConvGeneralDilated( @@ -772,12 +913,107 @@ XlaOp XlaBuilder::ConvGeneralDilated( tensorflow::gtl::ArraySlice lhs_dilation, tensorflow::gtl::ArraySlice rhs_dilation, const ConvolutionDimensionNumbers& dimension_numbers) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + TF_ASSIGN_OR_RETURN(const Shape& lhs_shape, GetShape(lhs)); + TF_ASSIGN_OR_RETURN(const Shape& rhs_shape, GetShape(rhs)); + TF_RETURN_IF_ERROR( + VerifyConvolution(lhs_shape, rhs_shape, dimension_numbers)); + + std::vector window_dimensions( + dimension_numbers.kernel_spatial_dimensions_size()); + for (std::vector::size_type i = 0; i < window_dimensions.size(); + ++i) { + window_dimensions[i] = + rhs_shape.dimensions(dimension_numbers.kernel_spatial_dimensions(i)); + } + TF_ASSIGN_OR_RETURN(*instr.mutable_window(), + MakeWindow(window_dimensions, window_strides, padding, + lhs_dilation, rhs_dilation)); + + TF_ASSIGN_OR_RETURN( + *instr.mutable_shape(), + ShapeInference::InferConvolveShape(lhs_shape, rhs_shape, instr.window(), + dimension_numbers)); + + *instr.mutable_convolution_dimension_numbers() = dimension_numbers; + + return AddInstruction(std::move(instr), HloOpcode::kConvolution, + {lhs, rhs}); + }); +} + +StatusOr XlaBuilder::MakeWindow( + tensorflow::gtl::ArraySlice window_dimensions, + tensorflow::gtl::ArraySlice window_strides, + tensorflow::gtl::ArraySlice> padding, + tensorflow::gtl::ArraySlice lhs_dilation, + tensorflow::gtl::ArraySlice rhs_dilation) const { + const auto verify_size = [&](const size_t x, const char* x_name) { + if (x == 0 || x == window_dimensions.size()) { + return Status::OK(); + } else { + return InvalidArgument( + "%s", tensorflow::strings::StrCat( + "Window has different number of window dimensions than of ", + x_name, + "\nNumber of window dimensions: ", window_dimensions.size(), + "\nNumber of ", x_name, ": ", x, "\n") + .c_str()); + } + }; + TF_RETURN_IF_ERROR(verify_size(window_strides.size(), "window strides")); + TF_RETURN_IF_ERROR(verify_size(padding.size(), "padding entries")); + TF_RETURN_IF_ERROR(verify_size(lhs_dilation.size(), "lhs dilation factors")); + TF_RETURN_IF_ERROR(verify_size(rhs_dilation.size(), "rhs dilation factors")); + + Window window; + for (size_t i = 0; i < window_dimensions.size(); i++) { + auto dim = window.add_dimensions(); + dim->set_size(window_dimensions[i]); + if (!window_strides.empty()) { + dim->set_stride(window_strides[i]); + } else { + dim->set_stride(1); + } + if (!padding.empty()) { + dim->set_padding_low(padding[i].first); + dim->set_padding_high(padding[i].second); + } else { + dim->set_padding_low(0); + dim->set_padding_high(0); + } + if (!lhs_dilation.empty()) { + dim->set_base_dilation(lhs_dilation[i]); + } else { + dim->set_base_dilation(1); + } + if (!rhs_dilation.empty()) { + dim->set_window_dilation(rhs_dilation[i]); + } else { + dim->set_window_dilation(1); + } + dim->set_window_reversal(false); + } + return window; } XlaOp XlaBuilder::Fft(const XlaOp& operand, const FftType fft_type, const tensorflow::gtl::ArraySlice fft_length) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + TF_ASSIGN_OR_RETURN(const Shape& operand_shape, operand.GetShape()); + TF_ASSIGN_OR_RETURN( + *instr.mutable_shape(), + ShapeInference::InferFftShape(operand_shape, fft_type, fft_length)); + + instr.set_fft_type(fft_type); + for (int64 i : fft_length) { + instr.add_fft_length(i); + } + + return AddInstruction(std::move(instr), HloOpcode::kFft, {operand}); + }); } XlaOp XlaBuilder::Infeed(const Shape& shape, const string& config) { @@ -821,13 +1057,30 @@ void XlaBuilder::Outfeed(const XlaOp& operand, const Shape& shape_with_layout, XlaOp XlaBuilder::CustomCall(const string& call_target_name, tensorflow::gtl::ArraySlice operands, const Shape& shape) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + if (tensorflow::str_util::StartsWith(call_target_name, "$")) { + return InvalidArgument( + "Invalid custom_call_target \"%s\": Call targets that start with '$' " + "are reserved for internal use.", + call_target_name.c_str()); + } + *instr.mutable_shape() = shape; + instr.set_custom_call_target(call_target_name); + return AddInstruction(std::move(instr), HloOpcode::kCustomCall, operands); + }); } XlaOp XlaBuilder::HostCompute(tensorflow::gtl::ArraySlice operands, const string& channel_name, int64 cost_estimate_ns, const Shape& shape) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + *instr.mutable_shape() = shape; + instr.set_channel_name(channel_name); + instr.set_cost_estimate_ns(cost_estimate_ns); + return AddInstruction(std::move(instr), HloOpcode::kHostCompute, operands); + }); } XlaOp XlaBuilder::Complex( @@ -836,7 +1089,9 @@ XlaOp XlaBuilder::Complex( return BinaryOp(HloOpcode::kComplex, real, imag, broadcast_dimensions); } -XlaOp XlaBuilder::Conj(const XlaOp& operand) { return UnimplementedOp(); } +XlaOp XlaBuilder::Conj(const XlaOp& operand) { + return Complex(Real(operand), Neg(Imag(operand))); +} XlaOp XlaBuilder::Sub(const XlaOp& lhs, const XlaOp& rhs, tensorflow::gtl::ArraySlice broadcast_dimensions) { @@ -938,6 +1193,10 @@ XlaOp XlaBuilder::Sign(const XlaOp& operand) { return UnaryOp(HloOpcode::kSign, operand); } +XlaOp XlaBuilder::Clz(const XlaOp& operand) { + return UnaryOp(HloOpcode::kClz, operand); +} + XlaOp XlaBuilder::Cos(const XlaOp& operand) { return UnaryOp(HloOpcode::kCos, operand); } @@ -979,7 +1238,17 @@ XlaOp XlaBuilder::Transpose(const XlaOp& operand, XlaOp XlaBuilder::Rev(const XlaOp& operand, tensorflow::gtl::ArraySlice dimensions) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); + TF_ASSIGN_OR_RETURN( + *instr.mutable_shape(), + ShapeInference::InferReverseShape(operand_shape, dimensions)); + for (int64 dim : dimensions) { + instr.add_dimensions(dim); + } + return AddInstruction(std::move(instr), HloOpcode::kReverse, {operand}); + }); } XlaOp XlaBuilder::Sort(const XlaOp& operand) { @@ -1010,7 +1279,15 @@ XlaOp XlaBuilder::ConvertElementType(const XlaOp& operand, XlaOp XlaBuilder::BitcastConvertType(const XlaOp& operand, PrimitiveType new_element_type) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); + TF_ASSIGN_OR_RETURN( + *instr.mutable_shape(), + ShapeInference::InferConvertShape(operand_shape, new_element_type)); + return AddInstruction(std::move(instr), HloOpcode::kBitcastConvert, + {operand}); + }); } XlaOp XlaBuilder::SquareF32(const XlaOp& operand) { @@ -1036,7 +1313,28 @@ XlaOp XlaBuilder::Map(tensorflow::gtl::ArraySlice operands, const XlaComputation& computation, tensorflow::gtl::ArraySlice dimensions, tensorflow::gtl::ArraySlice static_operands) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + if (!static_operands.empty()) { + return Unimplemented("static_operands is not supported in Map"); + } + + HloInstructionProto instr; + + std::vector operand_shape_ptrs; + TF_ASSIGN_OR_RETURN(const auto& operand_shapes, GetOperandShapes(operands)); + c_transform(operand_shapes, std::back_inserter(operand_shape_ptrs), + [](const Shape& shape) { return &shape; }); + TF_ASSIGN_OR_RETURN(const ProgramShape& called_program_shape, + computation.GetProgramShape()); + TF_ASSIGN_OR_RETURN( + *instr.mutable_shape(), + ShapeInference::InferMapShape(operand_shape_ptrs, called_program_shape, + dimensions)); + + AddCalledComputation(computation, &instr); + + return AddInstruction(std::move(instr), HloOpcode::kMap, operands); + }); } XlaOp XlaBuilder::RngOp(RandomDistribution distribution, @@ -1102,14 +1400,57 @@ XlaOp XlaBuilder::While(const XlaComputation& condition, XlaOp XlaBuilder::Gather(const XlaOp& input, const XlaOp& gather_indices, const GatherDimensionNumbers& dimension_numbers, tensorflow::gtl::ArraySlice window_bounds) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + + TF_ASSIGN_OR_RETURN(const Shape& input_shape, GetShape(input)); + TF_ASSIGN_OR_RETURN(const Shape& gather_indices_shape, + GetShape(gather_indices)); + TF_ASSIGN_OR_RETURN( + *instr.mutable_shape(), + ShapeInference::InferGatherShape(input_shape, gather_indices_shape, + dimension_numbers, window_bounds)); + + *instr.mutable_gather_dimension_numbers() = dimension_numbers; + for (int64 bound : window_bounds) { + instr.add_gather_window_bounds(bound); + } + + return AddInstruction(std::move(instr), HloOpcode::kGather, + {input, gather_indices}); + }); } XlaOp XlaBuilder::Conditional(const XlaOp& predicate, const XlaOp& true_operand, const XlaComputation& true_computation, const XlaOp& false_operand, const XlaComputation& false_computation) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + + TF_ASSIGN_OR_RETURN(const Shape& predicate_shape, GetShape(predicate)); + TF_ASSIGN_OR_RETURN(const Shape& true_operand_shape, + GetShape(true_operand)); + TF_ASSIGN_OR_RETURN(const ProgramShape& true_computation_shape, + true_computation.GetProgramShape()); + TF_ASSIGN_OR_RETURN(const Shape& false_operand_shape, + GetShape(false_operand)); + TF_ASSIGN_OR_RETURN(const ProgramShape& false_computation_shape, + false_computation.GetProgramShape()); + TF_ASSIGN_OR_RETURN( + *instr.mutable_shape(), + ShapeInference::InferConditionalShape( + predicate_shape, true_operand_shape, false_operand_shape, + true_computation_shape, false_computation_shape)); + + // The index of true_computation must be 0 and that of false computation + // must be 1. + AddCalledComputation(true_computation, &instr); + AddCalledComputation(false_computation, &instr); + + return AddInstruction(std::move(instr), HloOpcode::kConditional, + {predicate, true_operand, false_operand}); + }); } XlaOp XlaBuilder::Reduce( @@ -1141,7 +1482,12 @@ XlaOp XlaBuilder::Reduce( XlaOp XlaBuilder::ReduceAll(const XlaOp& operand, const XlaOp& init_value, const XlaComputation& computation) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); + std::vector all_dimnos(ShapeUtil::Rank(operand_shape)); + std::iota(all_dimnos.begin(), all_dimnos.end(), 0); + return Reduce(operand, init_value, computation, all_dimnos); + }); } XlaOp XlaBuilder::ReduceWindow( @@ -1149,7 +1495,21 @@ XlaOp XlaBuilder::ReduceWindow( const XlaComputation& computation, tensorflow::gtl::ArraySlice window_dimensions, tensorflow::gtl::ArraySlice window_strides, Padding padding) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + + TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); + TF_RETURN_IF_ERROR( + ValidatePaddingValues(AsInt64Slice(operand_shape.dimensions()), + window_dimensions, window_strides)); + + std::vector> padding_values = + MakePadding(AsInt64Slice(operand_shape.dimensions()), window_dimensions, + window_strides, padding); + return ReduceWindowWithGeneralPadding(operand, init_value, computation, + window_dimensions, window_strides, + padding_values); + }); } XlaOp XlaBuilder::ReduceWindowWithGeneralPadding( @@ -1158,31 +1518,111 @@ XlaOp XlaBuilder::ReduceWindowWithGeneralPadding( tensorflow::gtl::ArraySlice window_dimensions, tensorflow::gtl::ArraySlice window_strides, tensorflow::gtl::ArraySlice> padding) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + + TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); + TF_ASSIGN_OR_RETURN(const Shape& init_shape, GetShape(init_value)); + TF_ASSIGN_OR_RETURN(const ProgramShape& to_apply_shape, + computation.GetProgramShape()); + TF_ASSIGN_OR_RETURN(*instr.mutable_window(), + MakeWindow(window_dimensions, window_strides, padding, + /*lhs_dilation=*/{}, /*rhs_dilation=*/{})); + TF_ASSIGN_OR_RETURN( + *instr.mutable_shape(), + ShapeInference::InferReduceWindowShape(operand_shape, init_shape, + instr.window(), to_apply_shape)); + + AddCalledComputation(computation, &instr); + return AddInstruction(std::move(instr), HloOpcode::kReduceWindow, + {operand, init_value}); + }); } XlaOp XlaBuilder::BatchNormTraining(const XlaOp& operand, const XlaOp& scale, const XlaOp& offset, float epsilon, int64 feature_index) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + + TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); + TF_ASSIGN_OR_RETURN(const Shape& scale_shape, GetShape(scale)); + TF_ASSIGN_OR_RETURN(const Shape& offset_shape, GetShape(offset)); + TF_ASSIGN_OR_RETURN( + *instr.mutable_shape(), + ShapeInference::InferBatchNormTrainingShape( + operand_shape, scale_shape, offset_shape, feature_index)); + + instr.set_epsilon(epsilon); + instr.set_feature_index(feature_index); + + return AddInstruction(std::move(instr), HloOpcode::kBatchNormTraining, + {operand, scale, offset}); + }); } XlaOp XlaBuilder::BatchNormInference(const XlaOp& operand, const XlaOp& scale, const XlaOp& offset, const XlaOp& mean, const XlaOp& variance, float epsilon, int64 feature_index) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + + TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); + TF_ASSIGN_OR_RETURN(const Shape& scale_shape, GetShape(scale)); + TF_ASSIGN_OR_RETURN(const Shape& offset_shape, GetShape(offset)); + TF_ASSIGN_OR_RETURN(const Shape& mean_shape, GetShape(mean)); + TF_ASSIGN_OR_RETURN(const Shape& variance_shape, GetShape(variance)); + TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), + ShapeInference::InferBatchNormInferenceShape( + operand_shape, scale_shape, offset_shape, + mean_shape, variance_shape, feature_index)); + + instr.set_epsilon(epsilon); + instr.set_feature_index(feature_index); + + return AddInstruction(std::move(instr), HloOpcode::kBatchNormInference, + {operand, scale, offset, mean, variance}); + }); } XlaOp XlaBuilder::BatchNormGrad(const XlaOp& operand, const XlaOp& scale, const XlaOp& batch_mean, const XlaOp& batch_var, const XlaOp& grad_output, float epsilon, int64 feature_index) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + + TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); + TF_ASSIGN_OR_RETURN(const Shape& scale_shape, GetShape(scale)); + TF_ASSIGN_OR_RETURN(const Shape& batch_mean_shape, GetShape(batch_mean)); + TF_ASSIGN_OR_RETURN(const Shape& batch_var_shape, GetShape(batch_var)); + TF_ASSIGN_OR_RETURN(const Shape& grad_output_shape, GetShape(grad_output)); + TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), + ShapeInference::InferBatchNormGradShape( + operand_shape, scale_shape, batch_mean_shape, + batch_var_shape, grad_output_shape, feature_index)); + + instr.set_epsilon(epsilon); + instr.set_feature_index(feature_index); + + return AddInstruction(std::move(instr), HloOpcode::kBatchNormGrad, + {operand, scale, batch_mean, batch_var, grad_output}); + }); } XlaOp XlaBuilder::CrossReplicaSum(const XlaOp& operand) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + + TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); + TF_ASSIGN_OR_RETURN( + *instr.mutable_shape(), + ShapeInference::InferCrossReplicaSumShape({&operand_shape})); + + return AddInstruction(std::move(instr), HloOpcode::kCrossReplicaSum, + {operand}); + }); } XlaOp XlaBuilder::SelectAndScatter( @@ -1191,7 +1631,14 @@ XlaOp XlaBuilder::SelectAndScatter( tensorflow::gtl::ArraySlice window_strides, Padding padding, const XlaOp& source, const XlaOp& init_value, const XlaComputation& scatter) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); + return SelectAndScatterWithGeneralPadding( + operand, select, window_dimensions, window_strides, + MakePadding(AsInt64Slice(operand_shape.dimensions()), window_dimensions, + window_strides, padding), + source, init_value, scatter); + }); } XlaOp XlaBuilder::SelectAndScatterWithGeneralPadding( @@ -1201,12 +1648,45 @@ XlaOp XlaBuilder::SelectAndScatterWithGeneralPadding( tensorflow::gtl::ArraySlice> padding, const XlaOp& source, const XlaOp& init_value, const XlaComputation& scatter) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + + TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); + TF_ASSIGN_OR_RETURN(const Shape& source_shape, GetShape(source)); + TF_ASSIGN_OR_RETURN(const Shape& init_shape, GetShape(init_value)); + TF_ASSIGN_OR_RETURN(const ProgramShape& select_shape, + select.GetProgramShape()); + TF_ASSIGN_OR_RETURN(const ProgramShape& scatter_shape, + scatter.GetProgramShape()); + TF_ASSIGN_OR_RETURN(*instr.mutable_window(), + MakeWindow(window_dimensions, window_strides, padding, + /*lhs_dilation=*/{}, /*rhs_dilation=*/{})); + TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), + ShapeInference::InferSelectAndScatterShape( + operand_shape, select_shape, instr.window(), + source_shape, init_shape, scatter_shape)); + + AddCalledComputation(select, &instr); + AddCalledComputation(scatter, &instr); + + return AddInstruction(std::move(instr), HloOpcode::kSelectAndScatter, + {operand, source, init_value}); + }); } XlaOp XlaBuilder::ReducePrecision(const XlaOp& operand, const int exponent_bits, const int mantissa_bits) { - return UnimplementedOp(); + return NoteErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); + TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), + ShapeInference::InferReducePrecisionShape( + operand_shape, exponent_bits, mantissa_bits)); + instr.set_exponent_bits(exponent_bits); + instr.set_mantissa_bits(mantissa_bits); + return AddInstruction(std::move(instr), HloOpcode::kReducePrecision, + {operand}); + }); } void XlaBuilder::Send(const XlaOp& operand, const ChannelHandle& handle) { @@ -1249,15 +1729,98 @@ XlaOp XlaBuilder::Recv(const Shape& shape, const ChannelHandle& handle) { }); } -StatusOr XlaBuilder::IsConstant(const XlaOp& operand, - int64 num_parameters) { - return Unimplemented("IsConstant is not implemented."); +StatusOr XlaBuilder::IsConstant(const XlaOp& operand) const { + TF_RETURN_IF_ERROR(first_error_); + + // Verify that the handle is valid. + TF_RETURN_IF_ERROR(LookUpInstruction(operand).status()); + + bool is_constant = true; + std::set visited; + IsConstantVisitor(operand.handle(), &visited, &is_constant); + return is_constant; } -StatusOr> XlaBuilder::ComputeConstant( - const XlaOp& operand, const Layout* output_layout, - tensorflow::gtl::ArraySlice parameters) { - return Unimplemented("ComputeConstant is not implemented"); +StatusOr XlaBuilder::BuildConstantSubGraph( + const XlaOp& root_op) const { + TF_ASSIGN_OR_RETURN(bool is_constant, IsConstant(root_op)); + if (!is_constant) { + auto op_status = LookUpInstruction(root_op); + string op_string = + op_status.ok() ? op_status.ValueOrDie()->name() : ""; + return InvalidArgument( + "Operand to BuildConstantSubGraph depends on a parameter.\n\n" + " op requested for constant subgraph: %s\n\n" + "This is an internal error that typically happens when the XLA user " + "(e.g. TensorFlow) is attempting to determine a value that must be a " + "compile-time constant (e.g. an array dimension) but it is not capable " + "of being evaluated at XLA compile time.\n\n" + "Please file a usability bug with the framework being used (e.g. " + "TensorFlow).", + op_string.c_str()); + } + + TF_ASSIGN_OR_RETURN(const HloInstructionProto* root, + LookUpInstruction(root_op)); + TF_ASSIGN_OR_RETURN(HloOpcode opcode, StringToHloOpcode(root->opcode())); + if (!CanBeRoot(opcode)) { + return InvalidArgument("the operand with opcode %s cannot be root", + root->opcode().c_str()); + } + + HloComputationProto entry; + entry.set_id(GetUniqueId()); // Give the computation a global unique id. + entry.set_name(StrCat(name_, entry.id(), "_compute_constant")); + entry.set_root_id(root->id()); + ProgramShape* program_shape = entry.mutable_program_shape(); + *program_shape->mutable_result() = root->shape(); + + // We use std::set to keep the instruction ids in ascending order (which is + // also a valid denpendency order). The related ops will be added to the + // subgraph in the same order. + std::set related_ops; + tensorflow::gtl::FlatSet related_calls; // Related computations. + std::queue worklist; + worklist.push(root->id()); + related_ops.insert(root->id()); + while (!worklist.empty()) { + int64 node = worklist.front(); + worklist.pop(); + for (int64 id : instructions_[node].operand_ids()) { + if (related_ops.insert(id).second) { + worklist.push(id); + } + } + for (int64 called_id : instructions_[node].called_computation_ids()) { + related_calls.insert(called_id); + } + } + + // Add related ops to the computation. + for (int64 id : related_ops) { + auto* instr = entry.add_instructions(); + *instr = instructions_[id]; + // Ensures that the instruction names are unique among the graph. + const string& new_name = + StrCat(instr->name(), ".", entry.id(), ".", instr->id()); + instr->set_name(new_name); + } + + XlaComputation computation(entry.id()); + HloModuleProto* module = computation.mutable_proto(); + module->set_name(entry.name()); + module->set_id(entry.id()); + module->set_entry_computation_name(entry.name()); + module->set_entry_computation_id(entry.id()); + *module->mutable_program_shape() = *program_shape; + for (auto& e : embedded_) { + if (related_calls.find(e.second.id()) != related_calls.end()) { + *module->add_computations() = e.second; + } + } + *module->add_computations() = std::move(entry); + + return std::move(computation); } std::unique_ptr XlaBuilder::CreateSubBuilder( @@ -1268,10 +1831,6 @@ std::unique_ptr XlaBuilder::CreateSubBuilder( return sub_builder; } -Status XlaBuilder::SetReturnValue(const XlaOp& operand) { - return Unimplemented("SetReturnValue is not implemented."); -} - /* static */ ConvolutionDimensionNumbers XlaBuilder::CreateDefaultConvDimensionNumbers(int num_spatial_dims) { ConvolutionDimensionNumbers dimension_numbers; @@ -1351,10 +1910,7 @@ StatusOr XlaBuilder::AddInstruction( instr.set_id(handle); instr.set_opcode(HloOpcodeString(opcode)); if (instr.name().empty()) { - instr.set_name(StrCat(instr.opcode(), ".", unique_id_, ".", handle)); - } else { - // Append the handle to make sure the name is unique. - instr.set_name(StrCat(instr.name(), ".", unique_id_, ".", handle)); + instr.set_name(StrCat(instr.opcode())); } for (const auto& operand : operands) { if (operand.builder_ == nullptr) { diff --git a/tensorflow/compiler/xla/client/xla_client/xla_builder.h b/tensorflow/compiler/xla/client/xla_client/xla_builder.h index 0673b86646eeecae45b1076baf0002ed94242846..4955f1515d66af00ddf72e4c7621292a590e662c 100644 --- a/tensorflow/compiler/xla/client/xla_client/xla_builder.h +++ b/tensorflow/compiler/xla/client/xla_client/xla_builder.h @@ -53,14 +53,31 @@ class XlaBuilder; class XlaOp { public: XlaOp() : handle_(0), builder_(nullptr) {} + ~XlaOp() {} StatusOr GetShape() const; + const XlaBuilder* builder() const { return builder_; } + + bool operator==(const XlaOp& rhs) const { + return handle_ == rhs.handle_ && builder_ == rhs.builder_; + } + + bool operator!=(const XlaOp& rhs) const { + return handle_ != rhs.handle_ || builder_ != rhs.builder_; + } + + friend std::ostream& operator<<(std::ostream& out, const XlaOp& op) { + out << op.handle(); + return out; + } + private: XlaOp(int64 handle, XlaBuilder* builder) : handle_(handle), builder_(builder) {} int64 handle() const { return handle_; } + friend class XlaBuilder; int64 handle_; @@ -570,6 +587,9 @@ class XlaBuilder { // Enqueues a sign instruction onto the computation. XlaOp Sign(const XlaOp& operand); + // Enqueues a count leading zeros instruction onto the computation. + XlaOp Clz(const XlaOp& operand); + // Enqueues a cosine instruction onto the computation. XlaOp Cos(const XlaOp& operand); @@ -687,11 +707,12 @@ class XlaBuilder { XlaOp Recv(const Shape& shape, const ChannelHandle& handle); // Returns true if 'operand' is a compile-time constant. A compile-time - // constant does not depend on parameters with index greater than or equal to - // `num_parameters`, or on stateful operators such as `RngNormal` or `Infeed`. - // Unlike `ComputeConstant`, `IsConstant` tests whether a computation is a - // compile-time constant without evaluating the computation. - StatusOr IsConstant(const XlaOp& operand, int64 num_parameters = 0); + // constant does not depend on any parameters, or on stateful operators such + // as `RngNormal` or `Infeed`. + // + // This tests whether a computation is a compile-time constant without + // evaluating the computation. + StatusOr IsConstant(const XlaOp& operand) const; // Normalizes operand across spatial and batch dimensions for each feature. // @@ -731,47 +752,14 @@ class XlaBuilder { const XlaOp& grad_output, float epsilon, int64 feature_index); - // Computes the value of a constant indicated by a XlaOp using a non-optimized - // interpreter on the host. - // - // The operand must represent a constant value, which in this case - // means that it must not statically depend on any parameter of the - // computation that is being built other then the ones specified on the - // parameter list. The parameters in the list will be indexed by their - // parameter id property so the number of parameters specified should be at - // least as many as the largest used parameter index. - // - // `IsConstant` can be used to test whether a computation is a compile-time - // constant without evaluation it. `ComputeConstant` only succeeds for - // computations where `IsConstant` returns true. - // - // This functionality can be useful when translating a computation - // into XLA where something that looked dynamic is required by - // XLA to be specified as a constant. E.g. the source - // computation (outside of XLA) may include a dynamic - // computation of the shape of something and ComputeConstant lets - // you determine what the value of that computation is in the case - // where the value can be determined at compile time. - // - // If output_layout is non-null, then the output of the computation - // will be stored using that layout. - StatusOr> ComputeConstant( - const XlaOp& operand, const Layout* output_layout = nullptr, - tensorflow::gtl::ArraySlice parameters = {}); - // Returns a new XlaBuilder whose resultant Computation is used only by this // XlaBuilder. The sub-XlaBuilder has the same die_immediately_on_error // behavior as the parent. std::unique_ptr CreateSubBuilder(const string& computation_name); - // Modifies the computation being built so that executions of it will return - // the value associated with operand, rather than the last expression enqueued - // on the XlaBuilder. Any subsequent operations added to the XlaBuilder will - // not have any effect unless SetReturnValue is called again. - Status SetReturnValue(const XlaOp& operand); - // Builds the computation with the requested operations, or returns a non-ok - // status. + // status. Note that all ops that have been enqueued will be moved to the + // computation being returned. StatusOr Build(); // Builds the computation with the requested operations, or notes an error in @@ -784,6 +772,12 @@ class XlaBuilder { // instead. XlaComputation BuildAndNoteError(); + // Returns a subgraph that roots on the given root. If the root is not a + // compile-time constant (see `IsConstant`), returns an error. + // + // This will copy the needed ops/computations to the subgraph. + StatusOr BuildConstantSubGraph(const XlaOp& root_op) const; + // Returns the first error that was encountered while building the // computation. When an error is encountered, by default we return a vacuous // XlaOp and inform the user of the error that occurred while @@ -796,7 +790,7 @@ class XlaBuilder { StatusOr GetShape(const XlaOp& op) const; // Returns the (inferred) result for the current computation's shape. - StatusOr GetProgramShape(); + StatusOr GetProgramShape() const; private: StatusOr AddInstruction( @@ -851,10 +845,31 @@ class XlaBuilder { // Returns the (inferred) result for the program shape for the current // computation and fills the root_id in the pointer. - StatusOr GetProgramShape(int64* root_id); + StatusOr GetProgramShape(int64* root_id) const; + + // A visitor which checks whether an operation is a compile-time constant, + // meaning that it doesn't depend on any parameters, or on any stateful + // operation such as `RngNormal` or `Infeed`. The visitor walks the + // computation starting at a given operation and sets is_constant to false iff + // a parameter or stateful operation is encountered. + void IsConstantVisitor(const int64 op_handle, std::set* visited, + bool* is_constant) const; + + // Checks bounds for convolution parameters. + Status VerifyConvolution( + const Shape& lhs_shape, const Shape& rhs_shape, + const ConvolutionDimensionNumbers& dimension_numbers) const; + + // Helper function for creating a Window proto from user-supplied data. + // Returns error if the user-supplied data was invalid. + StatusOr MakeWindow( + tensorflow::gtl::ArraySlice window_dimensions, + tensorflow::gtl::ArraySlice window_strides, + tensorflow::gtl::ArraySlice> padding, + tensorflow::gtl::ArraySlice lhs_dilation, + tensorflow::gtl::ArraySlice rhs_dilation) const; - string name_; // Name to use for the built computation. - int64 unique_id_; // The unique id for the built computation. + 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. @@ -964,6 +979,37 @@ XlaOp XlaBuilder::ConstantR4FromArray4D(const Array4D& values) { return ConstantFromArray(values); } +// RAII-style object: sets the current sharding assignment in builder on +// construction, and sets back to the previous assignment on destruction. +// +// TODO(b/74197823): This is a part of a NOT YET ready refactor. +class XlaScopedShardingAssignment { + public: + XlaScopedShardingAssignment(xla::XlaBuilder* builder, + tensorflow::gtl::optional sharding) + : builder_(builder), prev_sharding_(builder->sharding()) { + SetSharding(sharding); + } + + XlaScopedShardingAssignment(const XlaScopedShardingAssignment&) = delete; + XlaScopedShardingAssignment& operator=(const XlaScopedShardingAssignment&) = + delete; + + ~XlaScopedShardingAssignment() { SetSharding(prev_sharding_); } + + private: + void SetSharding(const tensorflow::gtl::optional& sharding) { + if (sharding.has_value()) { + builder_->SetSharding(sharding.value()); + } else { + builder_->ClearSharding(); + } + } + + xla::XlaBuilder* const builder_; + tensorflow::gtl::optional prev_sharding_; +}; + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_CLIENT_XLA_CLIENT_XLA_BUILDER_H_ diff --git a/tensorflow/compiler/xla/client/xla_client/xla_computation.h b/tensorflow/compiler/xla/client/xla_client/xla_computation.h index 2a3c6952667a434b68ca0c5e4e9874397da173d3..7ad212aa24cd32d104cc4db7aa164c22c9f5be8f 100644 --- a/tensorflow/compiler/xla/client/xla_client/xla_computation.h +++ b/tensorflow/compiler/xla/client/xla_client/xla_computation.h @@ -30,6 +30,10 @@ namespace xla { class XlaComputation { public: XlaComputation() : unique_id_(-1) {} + XlaComputation(const HloModuleProto& proto) + : unique_id_(proto.id()), proto_(proto) {} + + ~XlaComputation() {} XlaComputation(const XlaComputation&) = delete; XlaComputation& operator=(const XlaComputation&) = delete; @@ -44,6 +48,9 @@ class XlaComputation { const HloModuleProto& proto() const { return proto_; } + // Returns true if this object is a null Computation. + bool IsNull() const { return unique_id_ == -1; } + private: XlaComputation(const int64 unique_id) : unique_id_(unique_id) {} HloModuleProto* mutable_proto() { return &proto_; } diff --git a/tensorflow/compiler/xla/device_util.h b/tensorflow/compiler/xla/device_util.h index 23a622b1ad0e2f3b220645f62767271f28df24e9..1a51fdee680721a4a03fa5de79a81746d92af76b 100644 --- a/tensorflow/compiler/xla/device_util.h +++ b/tensorflow/compiler/xla/device_util.h @@ -29,7 +29,7 @@ namespace xla { // Returns a string that represents the device in terms of platform and ordinal; // e.g. the first CUDA device will be "cuda:0" -string DeviceIdentifier(perftools::gputools::StreamExecutor* stream_exec) { +string DeviceIdentifier(se::StreamExecutor* stream_exec) { return tensorflow::strings::StrCat(stream_exec->platform()->Name(), ":", stream_exec->device_ordinal()); } diff --git a/tensorflow/compiler/xla/executable_run_options.cc b/tensorflow/compiler/xla/executable_run_options.cc index 1700c977189a9e4aedf6a6a75923c13678dae667..99b8f0558e6e39649c0e2e6ef345332e2aa55736 100644 --- a/tensorflow/compiler/xla/executable_run_options.cc +++ b/tensorflow/compiler/xla/executable_run_options.cc @@ -36,12 +36,12 @@ DeviceMemoryAllocator* ExecutableRunOptions::allocator() const { } ExecutableRunOptions& ExecutableRunOptions::set_stream( - perftools::gputools::Stream* stream) { + stream_executor::Stream* stream) { stream_ = stream; return *this; } -perftools::gputools::Stream* ExecutableRunOptions::stream() const { +stream_executor::Stream* ExecutableRunOptions::stream() const { return stream_; } diff --git a/tensorflow/compiler/xla/executable_run_options.h b/tensorflow/compiler/xla/executable_run_options.h index 2c1d9ffff10ed26410898ad258aa6b5b2cd37518..a306ae16ba4aee87e565ec3d119b84d90bc69c6c 100644 --- a/tensorflow/compiler/xla/executable_run_options.h +++ b/tensorflow/compiler/xla/executable_run_options.h @@ -16,26 +16,27 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_EXECUTABLE_RUN_OPTIONS_H_ #define TENSORFLOW_COMPILER_XLA_EXECUTABLE_RUN_OPTIONS_H_ -// Intentionally forward declared so that ExecutableRunOptions can be linked +// Pulls in the ::stream_executor -> ::xla::se namespace alias. +#include "tensorflow/compiler/xla/types.h" + +// These classes are forward declared so that ExecutableRunOptions can be linked // into an XLA-compiled binary without having to link all of the pointed-to // objects (e.g., for an ahead-of-time compiled CPU binary, the gpu tools don't // need to be linked). -namespace perftools { -namespace gputools { +namespace stream_executor { class Stream; class Platform; -} -} +} // namespace stream_executor namespace tensorflow { namespace thread { class ThreadPool; -} -} +} // namespace thread +} // namespace tensorflow namespace Eigen { struct ThreadPoolDevice; -} +} // namespace Eigen namespace xla { @@ -61,8 +62,8 @@ class ExecutableRunOptions { // If set, this is the stream to run the computation on. The platform of the // stream must match the platform the executable was built for. A value of // nullptr indicates the option has not been set. - ExecutableRunOptions& set_stream(perftools::gputools::Stream* stream); - perftools::gputools::Stream* stream() const; + ExecutableRunOptions& set_stream(stream_executor::Stream* stream); + stream_executor::Stream* stream() const; // Sets the thread pool on which to run parallel CPU backend // computations. Does not take ownership. @@ -91,7 +92,7 @@ class ExecutableRunOptions { DeviceMemoryAllocator* allocator_ = nullptr; int device_ordinal_ = -1; DeviceAssignment* device_assignment_ = nullptr; - perftools::gputools::Stream* stream_ = nullptr; + stream_executor::Stream* stream_ = nullptr; tensorflow::thread::ThreadPool* inter_op_thread_pool_ = nullptr; const Eigen::ThreadPoolDevice* intra_op_thread_pool_ = nullptr; ExecutionProfile* execution_profile_ = nullptr; diff --git a/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc b/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc index f037663e3f296861d0f9ae69175fd9f7b20b55e8..70ae95bf47398589e3c20f72c1f2084a738f253a 100644 --- a/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc +++ b/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc @@ -43,7 +43,7 @@ void SetDebugOptionsDefaults(DebugOptions* flags) { #ifdef INTEL_MKL flags->set_xla_cpu_use_mkl_dnn(true); #endif // INTEL_MKL - + flags->set_xla_gpu_max_kernel_unroll_factor(1); // Set cudnn batchnorm off by default; it does not provide a performance win // on average. flags->set_xla_gpu_use_cudnn_batchnorm(false); @@ -223,6 +223,11 @@ void AllocateFlags() { bool_setter_for(&DebugOptions::set_xla_gpu_disable_multi_streaming), flag_values->xla_gpu_disable_multi_streaming(), "If true, multi-streaming in the GPU backend is disabled."), + tensorflow::Flag( + "xla_gpu_max_kernel_unroll_factor", + int32_setter_for(&DebugOptions::set_xla_gpu_max_kernel_unroll_factor), + flag_values->xla_gpu_max_kernel_unroll_factor(), + "Specify the maximum kernel unroll factor for the GPU backend."), tensorflow::Flag( "xla_dump_optimized_hlo_proto_to", flag_values->mutable_xla_dump_optimized_hlo_proto_to(), diff --git a/tensorflow/compiler/xla/literal_util.cc b/tensorflow/compiler/xla/literal_util.cc index c2950c1faa7ac109d4cd90ababa82cb4877ecb57..bb6dd4f9098aefc1c2bbb1b1c41b3cee856b67de 100644 --- a/tensorflow/compiler/xla/literal_util.cc +++ b/tensorflow/compiler/xla/literal_util.cc @@ -44,8 +44,16 @@ namespace { constexpr bool kLittleEndian = __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__; -// Converts between little and big endian, assuming elements in the array are 16 -// bits long. +// Converts between little and big endian. +// +// Precondition: size % 2 == 0 (elements in the array are 16 bits long) +void ConvertEndianShort(string* bytes) { + CHECK_EQ(bytes->size() / 2, 0); + for (int64 i = 0; i < bytes->size(); i += 2) { + std::swap((*bytes)[i], (*bytes)[i + 1]); + } +} + void ConvertEndianShort(char* bytes, int64 size) { CHECK_EQ(size / 2, 0); for (int64 i = 0; i < size; i += 2) { @@ -97,11 +105,18 @@ Literal::Literal(const Shape& shape, bool allocate_arrays) const Shape& subshape = piece.subshape(); if (ShapeUtil::IsArray(subshape)) { if (allocate_arrays) { - piece.set_buffer(new char[piece.size_bytes()]); if (LayoutUtil::IsSparseArray(subshape)) { + // For sparse arrays, the buffer must be of the size of the maximum + // number of sparse elements possible. + const int64 max_sparse_elements = + LayoutUtil::MaxSparseElements(subshape.layout()); + piece.set_buffer( + new char[max_sparse_elements * ShapeUtil::ByteSizeOfPrimitiveType( + subshape.element_type())]); piece.set_sparse_indices(new SparseIndexArray( - LayoutUtil::MaxSparseElements(subshape.layout()), - ShapeUtil::Rank(subshape))); + max_sparse_elements, ShapeUtil::Rank(subshape))); + } else { + piece.set_buffer(new char[piece.size_bytes()]); } } else { piece.set_buffer(nullptr); @@ -1923,16 +1938,14 @@ void Literal::Piece::WriteToProto(LiteralProto* proto) const { *proto->mutable_f16s() = string( reinterpret_cast(data().data()), size_bytes()); if (!kLittleEndian) { - ConvertEndianShort(const_cast(proto->mutable_f16s()->data()), - proto->f16s().size()); + ConvertEndianShort(proto->mutable_f16s()); } break; case BF16: *proto->mutable_bf16s() = string( reinterpret_cast(data().data()), size_bytes()); if (!kLittleEndian) { - ConvertEndianShort(const_cast(proto->mutable_bf16s()->data()), - proto->bf16s().size()); + ConvertEndianShort(proto->mutable_bf16s()); } break; case F32: diff --git a/tensorflow/compiler/xla/literal_util.h b/tensorflow/compiler/xla/literal_util.h index 66ff39ecbb4a8e524981c339eca2d1d3a672637c..8aa19222dc4b9175ec72128dfdad448f65c23e91 100644 --- a/tensorflow/compiler/xla/literal_util.h +++ b/tensorflow/compiler/xla/literal_util.h @@ -741,7 +741,13 @@ class Literal { int64 size_bytes() const { return ShapeUtil::ByteSizeOf(subshape()); } // Returns the number of elements in this piece's array. - int64 element_count() const { return ShapeUtil::ElementsIn(subshape()); } + int64 element_count() const { + // If this is a sparse array, use the number of elements represented by + // the indices in the associated SparseIndexArray. + return LayoutUtil::IsSparseArray(subshape()) + ? sparse_indices()->index_count() + : ShapeUtil::ElementsIn(subshape()); + } // Copy the data from 'src' into this piece's buffer. Shapes of this piece // and src must be compatible. @@ -853,8 +859,7 @@ tensorflow::gtl::ArraySlice Literal::Piece::data() const { << " type, but literal element type is " << PrimitiveType_Name(subshape().element_type()); return tensorflow::gtl::ArraySlice( - reinterpret_cast(buffer()), - ShapeUtil::ElementsIn(subshape())); + reinterpret_cast(buffer()), element_count()); } template @@ -867,7 +872,7 @@ tensorflow::gtl::MutableArraySlice Literal::Piece::data() { << " type, but literal element type is " << PrimitiveType_Name(subshape().element_type()); return tensorflow::gtl::MutableArraySlice( - reinterpret_cast(buffer()), ShapeUtil::ElementsIn(subshape())); + reinterpret_cast(buffer()), element_count()); } template @@ -1282,12 +1287,13 @@ void Literal::PopulateSparse(SparseIndexArray indices, CHECK_LE(num_elements, max_elements); CHECK_EQ(num_elements, indices.index_count()); auto root_data = root_piece().data(); - root_data.remove_suffix(max_elements - values.size()); + // Piece::data() returns an ArraySlice of size equal to the number of indices + // in the SparseIndexArray. So there is no need to adjust the size of the data + // here. It is enough to just copy the incoming values into the data buffer. std::copy(values.begin(), values.end(), root_data.begin()); *this->root_piece().sparse_indices() = std::move(indices); if (sort) { auto root_data = this->root_piece().data(); - root_data.remove_suffix(root_data.size() - num_elements); this->root_piece().sparse_indices()->SortWithValues(root_data); } DCHECK(this->root_piece().sparse_indices()->Validate(shape())); diff --git a/tensorflow/compiler/xla/literal_util_test.cc b/tensorflow/compiler/xla/literal_util_test.cc index be4f2bc5ce466315253a05b81898ad0dace8c7c8..61046784e05623cd3117c24ecc6d6c474739bbd5 100644 --- a/tensorflow/compiler/xla/literal_util_test.cc +++ b/tensorflow/compiler/xla/literal_util_test.cc @@ -218,9 +218,7 @@ TEST_F(LiteralUtilTest, CreateSparse) { EXPECT_EQ(literal->sparse_indices()->data(), ArraySlice(expected_indices.data(), expected_indices.num_elements())); - EXPECT_EQ( - ArraySlice(literal->data().data(), expected_values.size()), - ArraySlice(expected_values)); + EXPECT_EQ(literal->data(), ArraySlice(expected_values)); } TEST_F(LiteralUtilTest, LiteralR4F32ProjectedStringifies) { diff --git a/tensorflow/compiler/xla/python/BUILD b/tensorflow/compiler/xla/python/BUILD index 0517a5502e686def4ffea59f929aef225186a8aa..ecb87bd8893276fbb9ecffaa0f8a3233d2e0043f 100644 --- a/tensorflow/compiler/xla/python/BUILD +++ b/tensorflow/compiler/xla/python/BUILD @@ -20,6 +20,7 @@ py_test( srcs = ["xla_client_test.py"], main = "xla_client_test.py", srcs_version = "PY2AND3", + tags = ["no_oss"], deps = [ ":xla_client", "//tensorflow/python:platform_test", diff --git a/tensorflow/compiler/xla/python/local_computation_builder.cc b/tensorflow/compiler/xla/python/local_computation_builder.cc index 2bacc6a9142971f6d14b3929fb1a69e2a40052e2..24e17abbe06197d2a090421508132e611da3cfa0 100644 --- a/tensorflow/compiler/xla/python/local_computation_builder.cc +++ b/tensorflow/compiler/xla/python/local_computation_builder.cc @@ -89,17 +89,16 @@ StatusOr> TransferFromOutfeedLocalReplica( return client->TransferFromOutfeedLocal(shape, device_ordinal); } -LocalShapedBuffer::LocalShapedBuffer( - std::unique_ptr shaped_buffer) +LocalShapedBuffer::LocalShapedBuffer(ScopedShapedBuffer shaped_buffer) : shaped_buffer_(std::move(shaped_buffer)) {} -const std::unique_ptr& LocalShapedBuffer::shaped_buffer() - const { - return shaped_buffer_; +const ScopedShapedBuffer* LocalShapedBuffer::shaped_buffer() const { + return &shaped_buffer_; } -static StatusOr> ToBuffer( - LocalClient* client, int device_ordinal, const Literal& arg) { +static StatusOr ToBuffer(LocalClient* client, + int device_ordinal, + const Literal& arg) { return client->LiteralToShapedBuffer(arg, device_ordinal, client->backend().memory_allocator()); } @@ -109,14 +108,15 @@ LocalShapedBuffer* LocalShapedBuffer::FromLiteral( const Literal& argument, const tensorflow::gtl::optional& shape_with_layout) { LocalClient* client = GetOrCreateLocalClient(); - std::unique_ptr buf; - if (shape_with_layout) { - std::unique_ptr relaid = - argument.Relayout(shape_with_layout.value()); - buf = ToBuffer(client, /*device_ordinal=*/0, *relaid).ConsumeValueOrDie(); - } else { - buf = ToBuffer(client, /*device_ordinal=*/0, argument).ConsumeValueOrDie(); - } + ScopedShapedBuffer buf = [&] { + if (shape_with_layout) { + std::unique_ptr relaid = + argument.Relayout(shape_with_layout.value()); + return ToBuffer(client, /*device_ordinal=*/0, *relaid) + .ConsumeValueOrDie(); + } + return ToBuffer(client, /*device_ordinal=*/0, argument).ConsumeValueOrDie(); + }(); return new LocalShapedBuffer(std::move(buf)); } @@ -158,14 +158,14 @@ StatusOr> CompiledLocalComputation::Execute( << device_ordinal; // Transfer arguments in - std::vector> scoped_buffers; + std::vector scoped_buffers; scoped_buffers.reserve(arguments.size()); for (int i = 0; i < arguments.size(); ++i) { const Literal& argument = arguments[i]; const tensorflow::gtl::optional& shape_with_layout = shapes_with_layout[i]; - StatusOr> pushed; + StatusOr pushed; if (shape_with_layout) { std::unique_ptr relaid = argument.Relayout(shape_with_layout.value()); @@ -185,7 +185,7 @@ StatusOr> CompiledLocalComputation::Execute( std::vector argument_buffers; argument_buffers.reserve(scoped_buffers.size()); for (auto& buffer : scoped_buffers) { - argument_buffers.push_back(buffer.get()); + argument_buffers.push_back(&buffer); } DeviceAssignment device_assignment = @@ -202,7 +202,7 @@ StatusOr> CompiledLocalComputation::Execute( options.set_intra_op_thread_pool( client->backend().eigen_intra_op_thread_pool_device()); options.set_device_assignment(&device_assignment); - StatusOr> result_buffer_status = + StatusOr result_buffer_status = executable_->Run(argument_buffers, options); if (!result_buffer_status.ok()) { results[replica] = result_buffer_status.status(); @@ -210,8 +210,8 @@ StatusOr> CompiledLocalComputation::Execute( } // Transfer result out - results[replica] = - client->ShapedBufferToLiteral(*result_buffer_status.ValueOrDie()); + results[replica] = client->ShapedBufferToLiteral( + std::move(result_buffer_status).ValueOrDie()); }); } } @@ -236,7 +236,7 @@ LocalShapedBuffer* CompiledLocalComputation::ExecuteWithShapedBuffers( std::vector argument_buffers; argument_buffers.reserve(argument_handles.size()); for (auto& handle : argument_handles) { - argument_buffers.push_back(handle->shaped_buffer().get()); + argument_buffers.push_back(handle->shaped_buffer()); } // Execute @@ -245,7 +245,7 @@ LocalShapedBuffer* CompiledLocalComputation::ExecuteWithShapedBuffers( options.set_inter_op_thread_pool(client->backend().inter_op_thread_pool()); options.set_intra_op_thread_pool( client->backend().eigen_intra_op_thread_pool_device()); - std::unique_ptr result_buffer = + ScopedShapedBuffer result_buffer = executable_->Run(argument_buffers, options).ConsumeValueOrDie(); return new LocalShapedBuffer(std::move(result_buffer)); diff --git a/tensorflow/compiler/xla/python/local_computation_builder.h b/tensorflow/compiler/xla/python/local_computation_builder.h index 31046e60f11af9cc89ddec4c5fd16babfc8eb231..e1048909ab29c2147a37ed72844391400d99e90d 100644 --- a/tensorflow/compiler/xla/python/local_computation_builder.h +++ b/tensorflow/compiler/xla/python/local_computation_builder.h @@ -62,12 +62,12 @@ class LocalShapedBuffer { 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; + LocalShapedBuffer(ScopedShapedBuffer shaped_buffer); + const ScopedShapedBuffer* shaped_buffer() const; std::unique_ptr ToLiteral() const; private: - std::unique_ptr shaped_buffer_; + ScopedShapedBuffer shaped_buffer_; }; // Wraps a LocalExecutable produced by compiling a diff --git a/tensorflow/compiler/xla/python/numpy_bridge.cc b/tensorflow/compiler/xla/python/numpy_bridge.cc index eec48479c929ab0823fef342fc284bfdc4b1f339..dc6f5fe5fcc067c99ced01812f9f2388a00766d0 100644 --- a/tensorflow/compiler/xla/python/numpy_bridge.cc +++ b/tensorflow/compiler/xla/python/numpy_bridge.cc @@ -181,16 +181,6 @@ StatusOr XlaShapeFromPyShape(PyObject* o) { PyObjectCppRepr(o).c_str()); }; - auto get_attr = [o, &error](const string& field) -> StatusOr { - PyObject* result = - PyObject_GetAttrString(o, const_cast(field.c_str())); - if (result == nullptr) { - return error(tensorflow::strings::StrCat( - "Failed to get attribute of Shape object:", field)); - } - return result; - }; - auto call_method = [o, &error](const string& method) -> StatusOr { PyObject* result = PyObject_CallMethod(o, const_cast(method.c_str()), nullptr); @@ -202,12 +192,16 @@ StatusOr XlaShapeFromPyShape(PyObject* o) { }; PyObject* np_type; - TF_ASSIGN_OR_RETURN(np_type, get_attr("np_dtype")); + TF_ASSIGN_OR_RETURN(np_type, call_method("numpy_dtype")); if (np_type->ob_type != &PyArrayDescr_Type) { - return error("Shape attribute np_dtype is not an integer numpy dtype"); + return error( + "Return value of shape method numpy_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"); + return error( + "Return value of shape method numpy_dtype " + "is not a valid integer numpy dtype"); } const PrimitiveType element_type = NumpyTypeToPrimitiveType(NumpyTypenum(np_type)); diff --git a/tensorflow/compiler/xla/python/xla_client.py b/tensorflow/compiler/xla/python/xla_client.py index 9c81f6439d0d9f0a0f0d1d3402e9c1ada46e8691..f6809b6b871d7e246dd43811c7e8c08378d53989 100644 --- a/tensorflow/compiler/xla/python/xla_client.py +++ b/tensorflow/compiler/xla/python/xla_client.py @@ -166,14 +166,14 @@ class LocalBuffer(object): self._delete = c_api.DeleteLocalShapedBuffer @staticmethod - def from_py(npval, layout_fn=None): - npval = require_numpy_array_layout(npval) + def from_pyval(pyval, layout_fn=None): + pyval = require_numpy_array_layout(pyval) if layout_fn: - shape = Shape.from_numpy(npval) + shape = Shape.from_pyval(pyval) shape = shape.map_leaves(layout_fn) else: shape = None - return LocalBuffer(c_api.LocalShapedBuffer.FromLiteral(npval, shape)) + return LocalBuffer(c_api.LocalShapedBuffer.FromLiteral(pyval, shape)) def to_py(self): return self.c_local_shaped_buffer.ToLiteral() @@ -191,53 +191,104 @@ class LocalBuffer(object): class Shape(object): - """XLA shape. + """Represents an XLA shape. - Represents an XLA shape by a corresponding Python/Numpy type and a - list of dimensions, which are themselves Shapes in case this one - represents an XLA tuple. + A shape is either an array shape, having rank-many integer + dimensions and an element type (represented by a Numpy dtype), or it + is a tuple shape, having a shape for every tuple component: + + type shape = + TupleShape of shape list + | ArrayShape of { dimensions: int list; element_type: dtype } + + Callers are expected to instantiate this class only via the static + constructors: tuple_shape, array_shape, and from_pyval. """ - def __init__(self, np_dtype, dimensions, minor_to_major=None): + @staticmethod + def tuple_shape(tuple_shapes): + """Construct a tuple shape.""" + if (not isinstance(tuple_shapes, (tuple, list)) or + not all(isinstance(t, Shape) for t in tuple_shapes)): + raise TypeError('tuple_shapes must be a tuple of Shapes') + return Shape(tuple_shapes, tuple) + + @staticmethod + def array_shape(element_type, dimensions, minor_to_major=None): + """Construct an array shape.""" + if (not isinstance(dimensions, tuple) or + not all(isinstance(i, int) for i in dimensions)): + dimensions = tuple(int(i) for i in dimensions) + return Shape(dimensions, np.dtype(element_type), + minor_to_major=minor_to_major) + + @staticmethod + def from_pyval(pyval): + def convert(pyval): + if isinstance(pyval, tuple): + return Shape.tuple_shape(tuple(convert(elt) for elt in pyval)) + else: + pyval = require_numpy_array_layout(pyval) + return Shape.array_shape(pyval.dtype, np.shape(pyval)) + return convert(pyval) + + def __init__(self, dimensions, dtype, minor_to_major=None): assert isinstance(dimensions, tuple) - self.np_dtype = np_dtype self._dimensions = dimensions + self._dtype = dtype + self._is_tuple = dtype == tuple 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 + return (self._dtype == other._dtype and self._dimensions == other._dimensions and self._minor_to_major == other._minor_to_major) def __repr__(self): - return ('xla_client.Shape(np_dtype={!r}, dimensions={!r}, ' - 'minor_to_major={!r})').format(self.np_dtype, self._dimensions, - self._minor_to_major) - - def element_type(self): - return DTYPE_TO_XLA_ELEMENT_TYPE[str(self.np_dtype)] + return ('xla_client.Shape(_dtype={!r}, _dimensions={!r}, ' + '_is_tuple={!r}), _minor_to_major={!r}').format( + self._dtype, self._dimensions, self._is_tuple, + self._minor_to_major) def is_tuple(self): - return self.element_type() == xla_data_pb2.TUPLE + return self._is_tuple - def dimensions(self): - if self.is_tuple(): - raise ValueError('Tuple shape has no dimensions') - return self._dimensions - - def minor_to_major(self): - return self._minor_to_major + def is_array(self): + return not self._is_tuple def tuple_shapes(self): if not self.is_tuple(): - raise ValueError('Shape is not a tuple shape') + raise ValueError('not a tuple shape') + return self._dimensions + + def numpy_dtype(self): + """Like element_type(), but returns dtype('O') in case of a tuple shape.""" + if self.is_tuple(): + return np.dtype(np.object) + else: + return self.element_type() + + def xla_element_type(self): + return DTYPE_TO_XLA_ELEMENT_TYPE[str(self.numpy_dtype())] + + def element_type(self): + if not self.is_array(): + raise ValueError('not an array shape') + return self._dtype + + def dimensions(self): + if not self.is_array(): + raise ValueError('not an array shape') return self._dimensions def rank(self): return len(self.dimensions()) + def minor_to_major(self): + return self._minor_to_major + def map_leaves(self, f): """Map f over each leaf-level array subshape. @@ -250,7 +301,7 @@ class Shape(object): """ if self.is_tuple(): children = tuple(child.map_leaves(f) for child in self.tuple_shapes()) - return Shape(np.dtype('O'), children) + return Shape.tuple_shape(children) else: mapped = f(self) return self if mapped is None else mapped @@ -264,30 +315,24 @@ class Shape(object): assert sorted(mtm) == range(len(mtm)), self def update_minor_to_major(self, minor_to_major): + if not self.is_array(): + raise ValueError('not an array shape') 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 = Shape.array_shape( + self.element_type(), self.dimensions(), minor_to_major) updated._check_minor_to_major() # pylint: disable=protected-access return updated - @staticmethod - def from_numpy(npval): - - def convert(npval): - if isinstance(npval, tuple): - return Shape(np.dtype('O'), tuple(convert(elt) for elt in npval)) - else: - return Shape(npval.dtype, np.shape(npval)) - - return convert(require_numpy_array_layout(npval)) - 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 = tuple(_wrap_shape(subshape_info) for subshape_info in dims) - return Shape(dtype, dims) + shapes = tuple(_wrap_shape(subshape_info) for subshape_info in dims) + return Shape.tuple_shape(shapes) + else: + return Shape.array_shape(dtype, dims) def _wrap_data_handle(handle): @@ -420,7 +465,7 @@ class LocalComputation(object): compile_options=None, layout_fn=None): return self.Compile( - argument_shapes=[Shape.from_numpy(arg) for arg in arguments], + argument_shapes=[Shape.from_pyval(arg) for arg in arguments], compile_options=compile_options, layout_fn=layout_fn) @@ -428,7 +473,7 @@ class LocalComputation(object): """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] + argument_shapes = [Shape.from_pyval(arg) for arg in arguments] if layout_fn: argument_shapes = [ shape.map_leaves(layout_fn) for shape in argument_shapes @@ -607,7 +652,7 @@ class ComputationBuilder(object): A ComputationDataHandle message. """ return self.ParameterWithShape( - Shape.from_numpy(value), name=name, parameter_num=parameter_num) + Shape.from_pyval(value), name=name, parameter_num=parameter_num) def Broadcast(self, operand, sizes): """Enqueues a broadcast operation onto the computation. @@ -968,7 +1013,7 @@ class ComputationBuilder(object): Returns: a ComputationDataHandle to the generated array of F32 values. """ - shape = Shape(self.GetShape(mu).np_dtype, dims) + shape = Shape.array_shape(self.GetShape(mu).element_type(), dims) return _wrap_data_handle( self._client.RngNormal( _unwrap_data_handle(mu), _unwrap_data_handle(sigma), shape)) @@ -988,7 +1033,7 @@ class ComputationBuilder(object): Returns: a ComputationDataHandle to the generated array of values with the same numeric type (F32, S32, or U32) as the arguments a and b. """ - shape = Shape(self.GetShape(a).np_dtype, dims) + shape = Shape.array_shape(self.GetShape(a).element_type(), dims) return _wrap_data_handle( self._client.RngUniform( _unwrap_data_handle(a), _unwrap_data_handle(b), shape)) diff --git a/tensorflow/compiler/xla/python/xla_client_test.py b/tensorflow/compiler/xla/python/xla_client_test.py index d97264ea640787ab865f3cd64867addedd73cc1d..c073c02040e4d260cf760ea2b25f70d60ddd41a1 100644 --- a/tensorflow/compiler/xla/python/xla_client_test.py +++ b/tensorflow/compiler/xla/python/xla_client_test.py @@ -319,7 +319,7 @@ class LocalBufferTest(LocalComputationTest): def _Execute(self, c, arguments): compiled_c = c.Build().CompileWithExampleArguments(arguments) - arg_buffers = [xla_client.LocalBuffer.from_py(arg) for arg in arguments] + arg_buffers = [xla_client.LocalBuffer.from_pyval(arg) for arg in arguments] result_buffer = compiled_c.ExecuteWithLocalBuffers(arg_buffers) return result_buffer.to_py() @@ -350,7 +350,7 @@ class LocalBufferTest(LocalComputationTest): c.Add(c.ParameterFromNumpy(NumpyArrayF32(0.)), c.ConstantF32Scalar(3.14)) arg = NumpyArrayF32(1.11) compiled_c = c.Build().CompileWithExampleArguments([arg]) - arg_buffer = xla_client.LocalBuffer.from_py(arg) + arg_buffer = xla_client.LocalBuffer.from_pyval(arg) arg_buffer.delete() with self.assertRaises(ValueError): compiled_c.ExecuteWithLocalBuffers([arg_buffer]) @@ -1160,7 +1160,6 @@ class EmbeddedComputationsTest(LocalComputationTest): self._ExecuteAndCompareClose( c, expected=np.sum(input_array, axis=tuple(dims))) - _ReduceAndTest(0) _ReduceAndTest(0) _ReduceAndTest(0, 1) _ReduceAndTest(0, 2) @@ -1288,7 +1287,7 @@ class EmbeddedComputationsTest(LocalComputationTest): def testInfeedS32Values(self): to_infeed = NumpyArrayS32([1, 2, 3, 4]) c = self._NewComputation() - c.Infeed(xla_client.Shape.from_numpy(to_infeed[0])) + c.Infeed(xla_client.Shape.from_pyval(to_infeed[0])) compiled_c = c.Build().CompileWithExampleArguments() for item in to_infeed: xla_client.transfer_to_infeed(item) @@ -1300,7 +1299,7 @@ class EmbeddedComputationsTest(LocalComputationTest): def testInfeedThenOutfeedS32(self): to_round_trip = NumpyArrayS32([1, 2, 3, 4]) c = self._NewComputation() - x = c.Infeed(xla_client.Shape.from_numpy(to_round_trip[0])) + x = c.Infeed(xla_client.Shape.from_pyval(to_round_trip[0])) c.Outfeed(x) compiled_c = c.Build().CompileWithExampleArguments() @@ -1310,7 +1309,7 @@ class EmbeddedComputationsTest(LocalComputationTest): execution.start() xla_client.transfer_to_infeed(want) got = xla_client.transfer_from_outfeed( - xla_client.Shape.from_numpy(to_round_trip[0])) + xla_client.Shape.from_pyval(to_round_trip[0])) execution.join() self.assertEqual(want, got) diff --git a/tensorflow/compiler/xla/rpc/BUILD b/tensorflow/compiler/xla/rpc/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..977f8637873a4b6555798f533010a28ff36e8679 --- /dev/null +++ b/tensorflow/compiler/xla/rpc/BUILD @@ -0,0 +1,79 @@ +licenses(["notice"]) # Apache 2.0 + +package(default_visibility = ["//tensorflow:internal"]) + +load("//tensorflow:tensorflow.bzl", "tf_cc_test") +load("//tensorflow:tensorflow.bzl", "tf_cc_binary") +load( + "//tensorflow/compiler/xla:xla.bzl", + "xla_proto_library", + "xla_py_grpc_library", +) + +xla_proto_library( + name = "xla_service_proto", + srcs = ["xla_service.proto"], + use_grpc_plugin = True, + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla:xla_proto", + ], +) + +cc_library( + name = "grpc_stub", + srcs = ["grpc_stub.cc"], + hdrs = ["grpc_stub.h"], + deps = [ + ":xla_service_proto", + "//tensorflow/compiler/xla:service_interface", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/core:lib", + "//tensorflow/core/distributed_runtime/rpc:grpc_util", + ], +) + +tf_cc_binary( + name = "grpc_service_main_cpu", + srcs = ["grpc_service_main.cc"], + deps = [ + ":grpc_service", + "//tensorflow/compiler/xla/service:cpu_plugin", + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + "@grpc//:grpc++_unsecure", + ], +) + +tf_cc_test( + name = "grpc_client_test", + srcs = ["grpc_client_test.cc"], + data = [ + "//tensorflow/compiler/xla/rpc:grpc_service_main_cpu", + ], + deps = [ + ":grpc_stub", + "//tensorflow/compiler/xla/client", + "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/compiler/xla/tests:literal_test_util", + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + "@grpc//:grpc++_unsecure", + ], +) + +cc_library( + name = "grpc_service", + srcs = ["grpc_service.cc"], + hdrs = ["grpc_service.h"], + deps = [ + ":xla_service_proto", + "//tensorflow/compiler/xla/service", + "//tensorflow/compiler/xla/service:platform_util", + "//tensorflow/core/distributed_runtime/rpc:grpc_util", + "@grpc//:grpc++_unsecure", + ], +) diff --git a/tensorflow/compiler/xla/rpc/grpc_client_test.cc b/tensorflow/compiler/xla/rpc/grpc_client_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..b559ee4b5a345dbb2cc481b571562a0a630b3294 --- /dev/null +++ b/tensorflow/compiler/xla/rpc/grpc_client_test.cc @@ -0,0 +1,109 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// Simple C++ test to exercise the GRPC capabilities of XLA. +// +// Launches an RPC service in a subprocess and connects to it over a socket +// using an RPCStub. +#include +#include + +#include "grpc++/create_channel.h" +#include "grpc++/security/credentials.h" + +#include "tensorflow/compiler/xla/client/client.h" +#include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/rpc/grpc_stub.h" +#include "tensorflow/compiler/xla/tests/literal_test_util.h" +#include "tensorflow/core/lib/io/path.h" +#include "tensorflow/core/lib/strings/stringprintf.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/net.h" +#include "tensorflow/core/platform/subprocess.h" +#include "tensorflow/core/platform/test.h" + +namespace xla { +namespace { + +class GRPCClientTestBase : public ::testing::Test { + protected: + GRPCClientTestBase() { + string test_srcdir = tensorflow::testing::TensorFlowSrcRoot(); + string service_main_path = tensorflow::io::JoinPath( + test_srcdir, "compiler/xla/rpc/grpc_service_main_cpu"); + int port = tensorflow::internal::PickUnusedPortOrDie(); + subprocess_.SetProgram( + service_main_path, + {service_main_path, tensorflow::strings::Printf("--port=%d", port)}); + subprocess_.SetChannelAction(tensorflow::CHAN_STDOUT, + tensorflow::ACTION_DUPPARENT); + subprocess_.SetChannelAction(tensorflow::CHAN_STDERR, + tensorflow::ACTION_DUPPARENT); + CHECK(subprocess_.Start()); + LOG(INFO) << "Launched subprocess"; + + auto channel = + ::grpc::CreateChannel(tensorflow::strings::Printf("localhost:%d", port), + ::grpc::InsecureChannelCredentials()); + channel->WaitForConnected(gpr_time_add( + gpr_now(GPR_CLOCK_REALTIME), gpr_time_from_seconds(10, GPR_TIMESPAN))); + LOG(INFO) << "Channel to server is connected on port " << port; + + xla_service_ = grpc::XlaService::NewStub(channel); + stub_.reset(new GRPCStub(xla_service_.get())); + client_.reset(new Client(stub_.get())); + } + + ~GRPCClientTestBase() override { + LOG(INFO) << "Killing subprocess"; + subprocess_.Kill(SIGKILL); + } + + tensorflow::SubProcess subprocess_; + std::unique_ptr xla_service_; + std::unique_ptr stub_; + std::unique_ptr client_; +}; + +TEST_F(GRPCClientTestBase, ItsAlive) { + ASSERT_NE(xla_service_, nullptr); + ASSERT_NE(stub_, nullptr); + ASSERT_NE(client_, nullptr); +} + +TEST_F(GRPCClientTestBase, AxpyTenValues) { + ComputationBuilder builder(client_.get(), "axpy_10"); + auto alpha = builder.ConstantR0(3.1415926535); + auto x = builder.ConstantR1( + {-1.0, 1.0, 2.0, -2.0, -3.0, 3.0, 4.0, -4.0, -5.0, 5.0}); + auto y = builder.ConstantR1( + {5.0, -5.0, -4.0, 4.0, 3.0, -3.0, -2.0, 2.0, 1.0, -1.0}); + auto ax = builder.Mul(alpha, x); + auto axpy = builder.Add(ax, y); + + std::vector expected = { + 1.85840735, -1.85840735, 2.28318531, -2.28318531, -6.42477796, + 6.42477796, 10.56637061, -10.56637061, -14.70796327, 14.70796327}; + std::unique_ptr expected_literal = + Literal::CreateR1(expected); + TF_ASSERT_OK_AND_ASSIGN(auto computation, builder.Build()); + TF_ASSERT_OK_AND_ASSIGN(auto result_literal, client_->ExecuteAndTransfer( + computation, {}, nullptr)); + LiteralTestUtil::ExpectNear(*expected_literal, *result_literal, + ErrorSpec(0.0001)); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/rpc/grpc_service.cc b/tensorflow/compiler/xla/rpc/grpc_service.cc new file mode 100644 index 0000000000000000000000000000000000000000..0b100bd108e239964483ed5ba279dff61bce0023 --- /dev/null +++ b/tensorflow/compiler/xla/rpc/grpc_service.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/xla/rpc/grpc_service.h" +#include "tensorflow/compiler/xla/service/platform_util.h" +#include "tensorflow/core/distributed_runtime/rpc/grpc_util.h" + +namespace xla { + +/* static */ StatusOr> GRPCService::NewService( + se::Platform* platform) { + std::unique_ptr grpc_service(new GRPCService()); + TF_ASSIGN_OR_RETURN(grpc_service->service_, + ::xla::Service::NewService(platform)); + return std::move(grpc_service); +} + +::grpc::Status DelegateRPC(std::function op) { + tensorflow::Status s = op(); + return tensorflow::ToGrpcStatus(s); +} + +::grpc::Status GRPCService::Computation(::grpc::ServerContext* context, + const ComputationRequest* arg, + ComputationResponse* result) { + return DelegateRPC( + [this, arg, result]() { return service_->Computation(arg, result); }); +} + +::grpc::Status GRPCService::CreateOp(::grpc::ServerContext* context, + const OpRequest* arg, OpResponse* result) { + return DelegateRPC( + [this, arg, result]() { return service_->Op(arg, result); }); +} + +::grpc::Status GRPCService::Unregister(::grpc::ServerContext* context, + const UnregisterRequest* arg, + UnregisterResponse* result) { + return DelegateRPC( + [this, arg, result]() { return service_->Unregister(arg, result); }); +} + +::grpc::Status GRPCService::DeconstructTuple(::grpc::ServerContext* context, + const DeconstructTupleRequest* arg, + DeconstructTupleResponse* result) { + return DelegateRPC([this, arg, result]() { + return service_->DeconstructTuple(arg, result); + }); +} + +::grpc::Status GRPCService::SetReturnValue(::grpc::ServerContext* context, + const SetReturnValueRequest* arg, + SetReturnValueResponse* results) { + return DelegateRPC([this, arg, results]() { + return service_->SetReturnValue(arg, results); + }); +} + +::grpc::Status GRPCService::Execute(::grpc::ServerContext* context, + const ExecuteRequest* arg, + ExecuteResponse* result) { + return DelegateRPC( + [this, arg, result]() { return service_->Execute(arg, result); }); +} + +::grpc::Status GRPCService::ExecuteAsync(::grpc::ServerContext* context, + const ExecuteAsyncRequest* arg, + ExecuteAsyncResponse* result) { + return DelegateRPC( + [this, arg, result]() { return service_->ExecuteAsync(arg, result); }); +} + +::grpc::Status GRPCService::WaitForExecution(::grpc::ServerContext* context, + const WaitForExecutionRequest* arg, + WaitForExecutionResponse* result) { + return DelegateRPC([this, arg, result]() { + return service_->WaitForExecution(arg, result); + }); +} + +::grpc::Status GRPCService::TransferToClient(::grpc::ServerContext* context, + const TransferToClientRequest* arg, + TransferToClientResponse* result) { + return DelegateRPC([this, arg, result]() { + return service_->TransferToClient(arg, result); + }); +} + +::grpc::Status GRPCService::TransferToServer(::grpc::ServerContext* context, + const TransferToServerRequest* arg, + TransferToServerResponse* result) { + return DelegateRPC([this, arg, result]() { + return service_->TransferToServer(arg, result); + }); +} + +::grpc::Status GRPCService::TransferToInfeed(::grpc::ServerContext* context, + const TransferToInfeedRequest* arg, + TransferToInfeedResponse* result) { + return DelegateRPC([this, arg, result]() { + return service_->TransferToInfeed(arg, result); + }); +} + +::grpc::Status GRPCService::TransferFromOutfeed( + ::grpc::ServerContext* context, const TransferFromOutfeedRequest* arg, + TransferFromOutfeedResponse* result) { + return DelegateRPC([this, arg, result]() { + return service_->TransferFromOutfeed(arg, result); + }); +} + +::grpc::Status GRPCService::ResetDevice(::grpc::ServerContext* context, + const ResetDeviceRequest* arg, + ResetDeviceResponse* result) { + return DelegateRPC( + [this, arg, result]() { return service_->ResetDevice(arg, result); }); +} + +::grpc::Status GRPCService::IsConstant(::grpc::ServerContext* context, + const IsConstantRequest* arg, + IsConstantResponse* result) { + return DelegateRPC( + [this, arg, result]() { return service_->IsConstant(arg, result); }); +} + +::grpc::Status GRPCService::ComputeConstant(::grpc::ServerContext* context, + const ComputeConstantRequest* arg, + ComputeConstantResponse* result) { + return DelegateRPC( + [this, arg, result]() { return service_->ComputeConstant(arg, result); }); +} + +::grpc::Status GRPCService::GetShape(::grpc::ServerContext* context, + const GetShapeRequest* arg, + GetShapeResponse* result) { + return DelegateRPC( + [this, arg, result]() { return service_->GetShape(arg, result); }); +} + +::grpc::Status GRPCService::GetComputationShape( + ::grpc::ServerContext* context, const GetComputationShapeRequest* arg, + GetComputationShapeResponse* result) { + return DelegateRPC([this, arg, result]() { + return service_->GetComputationShape(arg, result); + }); +} + +::grpc::Status GRPCService::GetLocalShape(::grpc::ServerContext* context, + const GetLocalShapeRequest* arg, + GetLocalShapeResponse* result) { + return DelegateRPC( + [this, arg, result]() { return service_->GetLocalShape(arg, result); }); +} + +::grpc::Status GRPCService::GetComputationStats( + ::grpc::ServerContext* context, const ComputationStatsRequest* arg, + ComputationStatsResponse* result) { + return DelegateRPC([this, arg, result]() { + return service_->GetComputationStats(arg, result); + }); +} + +::grpc::Status GRPCService::SnapshotComputation( + ::grpc::ServerContext* context, const SnapshotComputationRequest* arg, + SnapshotComputationResponse* result) { + return DelegateRPC([this, arg, result]() { + return service_->SnapshotComputation(arg, result); + }); +} + +::grpc::Status GRPCService::LoadComputationSnapshot( + ::grpc::ServerContext* context, const LoadComputationSnapshotRequest* arg, + LoadComputationSnapshotResponse* result) { + return DelegateRPC([this, arg, result]() { + return service_->LoadComputationSnapshot(arg, result); + }); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/rpc/grpc_service.h b/tensorflow/compiler/xla/rpc/grpc_service.h new file mode 100644 index 0000000000000000000000000000000000000000..fad74375bd59f7254d97c4adbc6b3d2f5fbf6b29 --- /dev/null +++ b/tensorflow/compiler/xla/rpc/grpc_service.h @@ -0,0 +1,126 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT 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_RPC_GRPC_SERVICE_H_ +#define TENSORFLOW_COMPILER_XLA_RPC_GRPC_SERVICE_H_ + +#include "grpc++/server_context.h" +#include "tensorflow/compiler/xla/rpc/xla_service.grpc.pb.h" +#include "tensorflow/compiler/xla/service/service.h" + +namespace xla { + +// Service implementation which wraps a XLA Service with a GRPC interface. +class GRPCService : public grpc::XlaService::Service { + public: + // Factory for creating a RPCService. The parameter platform is the platform + // that the service should target. If platform is null then the default + // platform is used. + static StatusOr> NewService( + se::Platform* platform = nullptr); + + ::grpc::Status Computation(::grpc::ServerContext* context, + const ComputationRequest* arg, + ComputationResponse* result) override; + + ::grpc::Status CreateOp(::grpc::ServerContext* context, const OpRequest* arg, + OpResponse* result) override; + + ::grpc::Status Unregister(::grpc::ServerContext* context, + const UnregisterRequest* arg, + UnregisterResponse* result) override; + + ::grpc::Status DeconstructTuple(::grpc::ServerContext* context, + const DeconstructTupleRequest* arg, + DeconstructTupleResponse* result) override; + + ::grpc::Status SetReturnValue(::grpc::ServerContext* context, + const SetReturnValueRequest* arg, + SetReturnValueResponse* results) override; + + ::grpc::Status Execute(::grpc::ServerContext* context, + const ExecuteRequest* arg, + ExecuteResponse* result) override; + + ::grpc::Status ExecuteAsync(::grpc::ServerContext* context, + const ExecuteAsyncRequest* arg, + ExecuteAsyncResponse* result) override; + + ::grpc::Status WaitForExecution(::grpc::ServerContext* context, + const WaitForExecutionRequest* arg, + WaitForExecutionResponse* result) override; + + ::grpc::Status TransferToClient(::grpc::ServerContext* context, + const TransferToClientRequest* arg, + TransferToClientResponse* result) override; + + ::grpc::Status TransferToServer(::grpc::ServerContext* context, + const TransferToServerRequest* arg, + TransferToServerResponse* result) override; + + ::grpc::Status TransferToInfeed(::grpc::ServerContext* context, + const TransferToInfeedRequest* arg, + TransferToInfeedResponse* result) override; + + ::grpc::Status TransferFromOutfeed( + ::grpc::ServerContext* context, const TransferFromOutfeedRequest* arg, + TransferFromOutfeedResponse* result) override; + + ::grpc::Status ResetDevice(::grpc::ServerContext* context, + const ResetDeviceRequest* arg, + ResetDeviceResponse* result) override; + + ::grpc::Status IsConstant(::grpc::ServerContext* context, + const IsConstantRequest* arg, + IsConstantResponse* result) override; + + ::grpc::Status ComputeConstant(::grpc::ServerContext* context, + const ComputeConstantRequest* arg, + ComputeConstantResponse* result) override; + + ::grpc::Status GetShape(::grpc::ServerContext* context, + const GetShapeRequest* arg, + GetShapeResponse* result) override; + + ::grpc::Status GetComputationShape( + ::grpc::ServerContext* context, const GetComputationShapeRequest* arg, + GetComputationShapeResponse* result) override; + + ::grpc::Status GetLocalShape(::grpc::ServerContext* context, + const GetLocalShapeRequest* arg, + GetLocalShapeResponse* result) override; + + ::grpc::Status GetComputationStats(::grpc::ServerContext* context, + const ComputationStatsRequest* arg, + ComputationStatsResponse* result) override; + + ::grpc::Status SnapshotComputation( + ::grpc::ServerContext* context, const SnapshotComputationRequest* arg, + SnapshotComputationResponse* result) override; + + ::grpc::Status LoadComputationSnapshot( + ::grpc::ServerContext* context, const LoadComputationSnapshotRequest* arg, + LoadComputationSnapshotResponse* result) override; + + private: + std::unique_ptr<::xla::Service> service_; + + GRPCService() {} + GRPCService(const GRPCService&) = delete; + void operator=(const GRPCService&) = delete; +}; +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_RPC_GRPC_SERVICE_H_ diff --git a/tensorflow/compiler/xla/rpc/grpc_service_main.cc b/tensorflow/compiler/xla/rpc/grpc_service_main.cc new file mode 100644 index 0000000000000000000000000000000000000000..e29908ccec80db76e3b5b856e57382c56430c379 --- /dev/null +++ b/tensorflow/compiler/xla/rpc/grpc_service_main.cc @@ -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. +==============================================================================*/ + +// Basic server binary that exposes a xla::Service through a GRPC interface +// on a configurable port. +#include "grpc++/security/server_credentials.h" +#include "grpc++/server.h" +#include "grpc++/server_builder.h" +#include "tensorflow/compiler/xla/rpc/grpc_service.h" +#include "tensorflow/core/lib/strings/stringprintf.h" +#include "tensorflow/core/platform/init_main.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/util/command_line_flags.h" + +namespace xla { +namespace { + +int RealMain(int argc, char** argv) { + int32 port = 1685; + std::vector flag_list = { + tensorflow::Flag("port", &port, "port to listen on"), + }; + string usage = tensorflow::Flags::Usage(argv[0], flag_list); + bool parsed_values_ok = tensorflow::Flags::Parse(&argc, argv, flag_list); + if (!parsed_values_ok) { + LOG(ERROR) << usage; + return 2; + } + tensorflow::port::InitMain(argv[0], &argc, &argv); + + std::unique_ptr service = + xla::GRPCService::NewService().ConsumeValueOrDie(); + + ::grpc::ServerBuilder builder; + string server_address(tensorflow::strings::Printf("localhost:%d", port)); + + builder.AddListeningPort(server_address, ::grpc::InsecureServerCredentials()); + builder.RegisterService(service.get()); + std::unique_ptr<::grpc::Server> server(builder.BuildAndStart()); + + LOG(INFO) << "Server listening on " << server_address; + server->Wait(); + + return 0; +} + +} // namespace +} // namespace xla + +int main(int argc, char** argv) { return xla::RealMain(argc, argv); } diff --git a/tensorflow/compiler/xla/rpc/grpc_stub.cc b/tensorflow/compiler/xla/rpc/grpc_stub.cc new file mode 100644 index 0000000000000000000000000000000000000000..e1f2b0abe39b10dd82b700941748bc4f4e8cb2f8 --- /dev/null +++ b/tensorflow/compiler/xla/rpc/grpc_stub.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/compiler/xla/rpc/grpc_stub.h" +#include "tensorflow/core/distributed_runtime/rpc/grpc_util.h" + +namespace xla { + +GRPCStub::~GRPCStub() = default; + +tensorflow::Status MakeRPC( + const std::function<::grpc::Status(::grpc::ClientContext*)>& rpc_method) { + ::grpc::ClientContext context; + ::grpc::Status s = rpc_method(&context); + return tensorflow::FromGrpcStatus(s); +} + +tensorflow::Status GRPCStub::TransferToClient( + const TransferToClientRequest* request, + TransferToClientResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->TransferToClient(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::TransferToServer( + const TransferToServerRequest* request, + TransferToServerResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->TransferToServer(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::TransferToInfeed( + const TransferToInfeedRequest* request, + TransferToInfeedResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->TransferToInfeed(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::TransferFromOutfeed( + const TransferFromOutfeedRequest* request, + TransferFromOutfeedResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->TransferFromOutfeed(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::ResetDevice(const ResetDeviceRequest* request, + ResetDeviceResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->ResetDevice(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::LoadComputationSnapshot( + const LoadComputationSnapshotRequest* request, + LoadComputationSnapshotResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->LoadComputationSnapshot(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::Execute(const ExecuteRequest* request, + ExecuteResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->Execute(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::ExecuteGraph(const ExecuteGraphRequest* request, + ExecuteResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->ExecuteGraph(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::ExecuteParallel( + const ExecuteParallelRequest* request, ExecuteParallelResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->ExecuteParallel(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::ExecuteGraphParallel( + const ExecuteGraphParallelRequest* request, + ExecuteParallelResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->ExecuteGraphParallel(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::ExecuteAsync(const ExecuteAsyncRequest* request, + ExecuteAsyncResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->ExecuteAsync(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::WaitForExecution( + const WaitForExecutionRequest* request, + WaitForExecutionResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->WaitForExecution(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::DeconstructTuple( + const DeconstructTupleRequest* request, + DeconstructTupleResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->DeconstructTuple(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::GetComputationStats( + const ComputationStatsRequest* request, + ComputationStatsResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->GetComputationStats(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::GetComputationGraphStats( + const ComputationGraphStatsRequest* request, + ComputationStatsResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->GetComputationGraphStats(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::GetComputationShape( + const GetComputationShapeRequest* request, + GetComputationShapeResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->GetComputationShape(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::GetShape(const GetShapeRequest* request, + GetShapeResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->GetShape(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::GetDeviceHandles( + const GetDeviceHandlesRequest* request, + GetDeviceHandlesResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->GetDeviceHandles(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::CreateChannelHandle( + const CreateChannelHandleRequest* request, + CreateChannelHandleResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->CreateChannelHandle(context, *request, response); + }); +} + +// Methods used by ComputationBuilder. +tensorflow::Status GRPCStub::Computation(const ComputationRequest* request, + ComputationResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->Computation(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::Op(const OpRequest* request, + OpResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->CreateOp(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::GetLocalShape(const GetLocalShapeRequest* request, + GetLocalShapeResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->GetLocalShape(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::SetReturnValue( + const SetReturnValueRequest* request, SetReturnValueResponse* responses) { + return MakeRPC([this, request, responses](::grpc::ClientContext* context) { + return grpc_stub_->SetReturnValue(context, *request, responses); + }); +} + +tensorflow::Status GRPCStub::IsConstant(const IsConstantRequest* request, + IsConstantResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->IsConstant(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::ComputeConstant( + const ComputeConstantRequest* request, ComputeConstantResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->ComputeConstant(context, *request, response); + }); +} + +tensorflow::Status GRPCStub::ComputeConstantGraph( + const ComputeConstantGraphRequest* request, + ComputeConstantResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->ComputeConstantGraph(context, *request, response); + }); +} + +// Methods used by Computation. +tensorflow::Status GRPCStub::SnapshotComputation( + const SnapshotComputationRequest* request, + SnapshotComputationResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->SnapshotComputation(context, *request, response); + }); +} + +// Methods used by GlobalData. +tensorflow::Status GRPCStub::Unregister(const UnregisterRequest* request, + UnregisterResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->Unregister(context, *request, response); + }); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/rpc/grpc_stub.h b/tensorflow/compiler/xla/rpc/grpc_stub.h new file mode 100644 index 0000000000000000000000000000000000000000..fd9810d4f1a5e084b73e83007ea7f9f8b0462c72 --- /dev/null +++ b/tensorflow/compiler/xla/rpc/grpc_stub.h @@ -0,0 +1,141 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT 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_RPC_GRPC_STUB_H_ +#define TENSORFLOW_COMPILER_XLA_RPC_GRPC_STUB_H_ + +#include "tensorflow/compiler/xla/rpc/xla_service.grpc.pb.h" +#include "tensorflow/compiler/xla/service_interface.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/platform/macros.h" + +namespace xla { + +class GRPCStub : public ServiceInterface { + public: + explicit GRPCStub(grpc::XlaService::Stub* stub) : grpc_stub_(stub) {} + ~GRPCStub() override; + + tensorflow::Status TransferToClient( + const TransferToClientRequest* arg, + TransferToClientResponse* result) override; + + tensorflow::Status TransferToServer( + const TransferToServerRequest* arg, + TransferToServerResponse* result) override; + + tensorflow::Status TransferToInfeed( + const TransferToInfeedRequest* arg, + TransferToInfeedResponse* result) override; + + tensorflow::Status TransferFromOutfeed( + const TransferFromOutfeedRequest* arg, + TransferFromOutfeedResponse* result) override; + + tensorflow::Status ResetDevice(const ResetDeviceRequest* arg, + ResetDeviceResponse* result) override; + + tensorflow::Status LoadComputationSnapshot( + const LoadComputationSnapshotRequest* request, + LoadComputationSnapshotResponse* result) override; + + tensorflow::Status Execute(const ExecuteRequest* arg, + ExecuteResponse* result) override; + + tensorflow::Status ExecuteGraph(const ExecuteGraphRequest* request, + ExecuteResponse* response) override; + + tensorflow::Status ExecuteParallel(const ExecuteParallelRequest* arg, + ExecuteParallelResponse* result) override; + + tensorflow::Status ExecuteGraphParallel( + const ExecuteGraphParallelRequest* request, + ExecuteParallelResponse* response) override; + + tensorflow::Status ExecuteAsync(const ExecuteAsyncRequest* arg, + ExecuteAsyncResponse* result) override; + + tensorflow::Status WaitForExecution( + const WaitForExecutionRequest* arg, + WaitForExecutionResponse* result) override; + + tensorflow::Status DeconstructTuple( + const DeconstructTupleRequest* arg, + DeconstructTupleResponse* result) override; + + tensorflow::Status GetComputationStats( + const ComputationStatsRequest* arg, + ComputationStatsResponse* result) override; + + tensorflow::Status GetComputationGraphStats( + const ComputationGraphStatsRequest* request, + ComputationStatsResponse* response) override; + + tensorflow::Status GetComputationShape( + const GetComputationShapeRequest* arg, + GetComputationShapeResponse* result) override; + + tensorflow::Status GetShape(const GetShapeRequest* arg, + GetShapeResponse* result) override; + + tensorflow::Status GetDeviceHandles( + const GetDeviceHandlesRequest* arg, + GetDeviceHandlesResponse* result) override; + + tensorflow::Status CreateChannelHandle( + const CreateChannelHandleRequest* arg, + CreateChannelHandleResponse* result) override; + + // Methods used by ComputationBuilder. + tensorflow::Status Computation(const ComputationRequest* arg, + ComputationResponse* result) override; + + tensorflow::Status Op(const OpRequest* arg, OpResponse* result) override; + tensorflow::Status GetLocalShape(const GetLocalShapeRequest* arg, + GetLocalShapeResponse* result) override; + + tensorflow::Status SetReturnValue(const SetReturnValueRequest* arg, + SetReturnValueResponse* results) override; + + tensorflow::Status IsConstant(const IsConstantRequest* arg, + IsConstantResponse* result) override; + + tensorflow::Status ComputeConstant(const ComputeConstantRequest* arg, + ComputeConstantResponse* result) override; + + tensorflow::Status ComputeConstantGraph( + const ComputeConstantGraphRequest* arg, + ComputeConstantResponse* result) override; + + // Methods used by Computation. + tensorflow::Status SnapshotComputation( + const SnapshotComputationRequest* ag, + SnapshotComputationResponse* result) override; + + // Methods used by GlobalData. + tensorflow::Status Unregister(const UnregisterRequest* arg, + UnregisterResponse* result) override; + + grpc::XlaService::Stub* service() { return grpc_stub_; } + + private: + grpc::XlaService::Stub* grpc_stub_; + + TF_DISALLOW_COPY_AND_ASSIGN(GRPCStub); +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_RPC_GRPC_STUB_H_ diff --git a/tensorflow/compiler/xla/rpc/xla_service.proto b/tensorflow/compiler/xla/rpc/xla_service.proto new file mode 100644 index 0000000000000000000000000000000000000000..c47164ee1b7657ae378a053f553442bee751753e --- /dev/null +++ b/tensorflow/compiler/xla/rpc/xla_service.proto @@ -0,0 +1,225 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// XLA service API. +// +// Users 1) build up computations and 2) create allocations via this API. +// Computations are composed of data flowing between arbitrarily-sized +// vector-oriented operations. +// +// Users build up computations using a ComputationHandle, and talk about +// allocations using GlobalDataHandles. +// +// There are currently no checkpointing capabilities or distribution/replication +// guarantees. The service runs on a single machine (e.g. one task) and that is +// its failure domain. +// +// Canonical example of "alpha * X + Y": +// * Make a computation. +// * Add alpha and X and Y as parameters. +// * Request the multiplication of alpha and X. +// * Request the addition of that result and Y. +// +// Then, pass the computation and appropriately shaped inputs to the XLA +// service's Execute method, which provides a result as a GlobalDataHandle. +// +// All data in XLA computations are conceptually immutable. +// +// Note: this API is subject to change / refinement over time -- use the +// provided client libraries to insulate code from changes to this service API. + +syntax = "proto3"; + +import "tensorflow/compiler/xla/xla.proto"; +import "tensorflow/compiler/xla/xla_data.proto"; + +package xla; + +service XlaService { + ///////////////////////// + // Global data requests + + // Unregisters a global allocation. + // + // If the handle given is not currently allocated, a NOT_FOUND status is + // returned. + rpc Unregister(UnregisterRequest) returns (UnregisterResponse) { + } + + // Deconstructs a tuple. Returns a newly created GlobalDataHandle for each + // element in the tuple. + rpc DeconstructTuple(DeconstructTupleRequest) + returns (DeconstructTupleResponse) { + } + + // Unpack requests that a global data handle, with a tuple shape, has global + // data handles created for each of its constituent members. This is the + // equivalent of the "destructuring assignment" present in various programming + // languages. + rpc Unpack(UnpackRequest) returns (UnpackResponse) { + } + + // Requests the shape of the referenced global data. + rpc GetShape(GetShapeRequest) returns (GetShapeResponse) { + } + + // Requests the program shape of the referenced computation. + rpc GetComputationShape(GetComputationShapeRequest) + returns (GetComputationShapeResponse) { + } + + // Requests the statistics of the given computation. + rpc GetComputationStats(ComputationStatsRequest) + returns (ComputationStatsResponse) { + } + + // Requests the statistics of the given computation. + // + // TODO(b/74197823): This is a part of a NOT YET ready refactor. + rpc GetComputationGraphStats(ComputationGraphStatsRequest) + returns (ComputationStatsResponse) { + } + + // Loads a variable number of values with a given element type from ColumnIO. + rpc LoadData(LoadDataRequest) returns (LoadDataResponse) { + } + + // Transfers the given global data to the client in the form of a Literal. + rpc TransferToClient(TransferToClientRequest) + returns (TransferToClientResponse) { + } + + // Transfers the given literal to the server to be stored in a global + // allocation, which is returned. + rpc TransferToServer(TransferToServerRequest) + returns (TransferToServerResponse) { + } + + // Transfers the given literal to the Infeed buffer of the device. + rpc TransferToInfeed(TransferToInfeedRequest) + returns (TransferToInfeedResponse) { + } + + // Transferred literal from the Outfeed buffer of the device. + rpc TransferFromOutfeed(TransferFromOutfeedRequest) + returns (TransferFromOutfeedResponse) { + } + + // Resets the device, clearing all existing state on the device. + rpc ResetDevice(ResetDeviceRequest) returns (ResetDeviceResponse) { + } + + // Tests if an expression is a compile-time constant. + rpc IsConstant(IsConstantRequest) returns (IsConstantResponse) { + } + + // Computes the value of a constant expression. + rpc ComputeConstant(ComputeConstantRequest) + returns (ComputeConstantResponse) { + } + + // Computes the value of a constant expression. The request contains the + // computation graph for the constant expression. + rpc ComputeConstantGraph(ComputeConstantGraphRequest) + returns (ComputeConstantResponse) { + } + + // Retrieves the inferred shape for a value within a computation. + rpc GetLocalShape(GetLocalShapeRequest) returns (GetLocalShapeResponse) { + } + + // Requests one or more device handles from the target. The returned device + // handles can be used to specify the device on which to execute computations + // or transfer data. + rpc GetDeviceHandles(GetDeviceHandlesRequest) + returns (GetDeviceHandlesResponse) { + } + + // Creates a channel handle that can be used to transfer data between + // two computations via a pair of Send and Recv instructions. + rpc CreateChannelHandle(CreateChannelHandleRequest) + returns (CreateChannelHandleResponse) { + } + + // Requests that the referenced computation be specialized for the provided + // arguments for subsequent execution. This permits things such as value + // specialization. + rpc Specialize(SpecializeRequest) returns (SpecializeResponse) { + } + + // Modifies the provided computation so that subsequent executions + // will compute the provided ComputationDataHandle, rather than the + // last expression enqueued on that Computation. + rpc SetReturnValue(SetReturnValueRequest) returns (SetReturnValueResponse) { + } + + // Computation creates a new computation with the given name. + // A unique ComputationHandle is returned. + rpc Computation(ComputationRequest) returns (ComputationResponse) { + } + + // Adds a new op to a computation. + rpc CreateOp(OpRequest) returns (OpResponse) { + } + + // Invokes the provided computation with the provided global data passed as + // immutable arguments. Returns global data output and execution timing. + rpc Execute(ExecuteRequest) returns (ExecuteResponse) { + } + + // Invokes the provided computation with the provided global data passed as + // immutable arguments. The request contains the whole computation graph. + // Returns global data output and execution timing. + rpc ExecuteGraph(ExecuteGraphRequest) returns (ExecuteResponse) { + } + + // Invokes the provided list of computations in parallel with the provided + // global data for each computation. Returns a list of global data output and + // execution timing. + rpc ExecuteParallel(ExecuteParallelRequest) + returns (ExecuteParallelResponse) { + } + + // Invokes the provided list of computations in parallel with the provided + // global data for each computation. Returns a list of global data output and + // execution timing. + // + // TODO(b/74197823): This is a part of a NOT YET ready refactor. + rpc ExecuteGraphParallel(ExecuteGraphParallelRequest) + returns (ExecuteParallelResponse) { + } + + // Invokes the provided computation with the provided global data passed as + // immutable arguments. Returns a handle to the execution. + rpc ExecuteAsync(ExecuteAsyncRequest) returns (ExecuteAsyncResponse) { + } + + // Waits until the given execution (aysnchronously launched) is complete, and + // returns the global data output. + rpc WaitForExecution(WaitForExecutionRequest) + returns (WaitForExecutionResponse) { + } + + // Serializes a computation to proto form, so it can be loaded via + // LoadComputationSnapshot. + rpc SnapshotComputation(SnapshotComputationRequest) + returns (SnapshotComputationResponse) { + } + + // Loads a computation from a captured snapshot. + rpc LoadComputationSnapshot(LoadComputationSnapshotRequest) + returns (LoadComputationSnapshotResponse) { + } +} diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD index db91e804072676c609d4d1fa3110bd587f5f2bc0..afb344e5ae2aa5e420bf762dcfad4fad34247bf3 100644 --- a/tensorflow/compiler/xla/service/BUILD +++ b/tensorflow/compiler/xla/service/BUILD @@ -302,6 +302,29 @@ tf_cc_test( ], ) +cc_library( + name = "pattern_matcher", + hdrs = ["pattern_matcher.h"], + deps = [ + ":hlo", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/core:lib", + ], +) + +tf_cc_test( + name = "pattern_matcher_test", + srcs = ["pattern_matcher_test.cc"], + deps = [ + ":hlo", + ":pattern_matcher", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/compiler/xla/tools/parser:hlo_parser", + "//tensorflow/core:test", + ], +) + cc_library( name = "hlo_reachability", srcs = ["hlo_reachability.cc"], @@ -732,6 +755,7 @@ cc_library( ":hlo", ":hlo_execution_profile", ":hlo_graph_dumper", + ":hlo_proto", ":pool", ":session_proto", ":shaped_buffer", @@ -1197,6 +1221,23 @@ cc_library( ], ) +tf_cc_test( + name = "hlo_creation_utils_test", + srcs = ["hlo_creation_utils_test.cc"], + deps = [ + ":hlo", + ":hlo_creation_utils", + ":hlo_evaluator", + ":hlo_matchers", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:test", + ], +) + cc_library( name = "batchnorm_expander", srcs = ["batchnorm_expander.cc"], @@ -1260,6 +1301,7 @@ cc_library( ":hlo_creation_utils", ":hlo_pass", ":hlo_query", + ":pattern_matcher", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", @@ -1912,10 +1954,12 @@ cc_library( deps = [ ":computation_layout", ":hlo", + ":hlo_dce", ":hlo_graph_dumper", ":hlo_pass", ":logical_buffer", ":tuple_points_to_analysis", + ":tuple_simplifier", "//tensorflow/compiler/xla:shape_layout", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", @@ -1991,6 +2035,7 @@ cc_library( srcs = ["hlo_verifier.cc"], hdrs = ["hlo_verifier.h"], deps = [ + ":hlo", ":hlo_pass", ":shape_inference", "//tensorflow/compiler/xla:status_macros", @@ -2391,6 +2436,7 @@ cc_library( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", ], ) @@ -2535,6 +2581,7 @@ cc_library( srcs = ["hlo_runner.cc"], hdrs = ["hlo_runner.h"], deps = [ + ":computation_placer", ":executable", ":hlo", ":transfer_manager", @@ -2551,6 +2598,7 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:stream_executor_no_cuda", "//third_party/eigen3", + "@com_google_absl//absl/memory", ], ) diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier.cc b/tensorflow/compiler/xla/service/algebraic_simplifier.cc index 6cb1bd56695772a38c377280da4ea357027519e5..8e785de68cb1fbe4ce9fd58a661bdc208725483b 100644 --- a/tensorflow/compiler/xla/service/algebraic_simplifier.cc +++ b/tensorflow/compiler/xla/service/algebraic_simplifier.cc @@ -30,6 +30,7 @@ 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/pattern_matcher.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" @@ -44,8 +45,11 @@ limitations under the License. #include "tensorflow/core/platform/types.h" namespace xla { + namespace { +namespace m = match; + // Returns whether operand is a literal with the given value. bool IsLiteralWithValue(const HloInstruction* operand, int8 value) { return operand->opcode() == HloOpcode::kConstant && @@ -105,6 +109,7 @@ HloComputation* CreateScalarBinaryComputation(HloModule* module, module->AddEmbeddedComputation(b.Build(scalar_op)); return scalar_computation; } + } // namespace // AlgebraicSimplifierVisitor traverses the HLO computation and reduces certain @@ -350,8 +355,9 @@ bool AlgebraicSimplifierVisitor::ReplaceInstructionIfSameShape( } Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) { - auto lhs = add->mutable_operand(0); - auto rhs = add->mutable_operand(1); + HloInstruction *lhs, *rhs; + CHECK(Match(add, m::Add(m::Op(&lhs), m::Op(&rhs)))); + // A + 0 => A VLOG(10) << "trying transform [A + 0 => A]: " << add->ToString(); if (IsAll(rhs, 0) && ReplaceInstructionIfSameShape(add, lhs)) { @@ -366,7 +372,7 @@ Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) { // Canonicalization: Put constants on the right. This makes the reassociation // rules below simpler. VLOG(10) << "trying transform [Const + A => A + Const]"; - if (lhs->IsConstant() && !rhs->IsConstant()) { + if (Match(add, m::Add(m::Constant(), m::NonConstant()))) { return ReplaceWithNewInstruction( add, HloInstruction::CreateBinary(add->shape(), HloOpcode::kAdd, rhs, lhs)); @@ -379,16 +385,13 @@ Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) { // (A + C1) + (B + C2) => A + B + (C1 + C2). // VLOG(10) << "trying transform [(A + C1) + C2 => A + (C1 + C2)]"; - if (rhs->IsConstant() && lhs->opcode() == HloOpcode::kAdd && - !lhs->operand(0)->IsConstant() && lhs->operand(1)->IsConstant()) { - auto* c1 = lhs->mutable_operand(1); - auto* c2 = rhs; - + HloInstruction *a, *c1, *c2; + if (Match(add, m::Add(m::Add(m::NonConstant(&a), m::Constant(&c1)), + m::Constant(&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), + add, HloInstruction::CreateBinary(add->shape(), HloOpcode::kAdd, a, sum_of_constants)); } @@ -397,11 +400,11 @@ Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) { Status AlgebraicSimplifierVisitor::HandleBitcast(HloInstruction* bitcast) { // If a bitcast feeds a bitcast, make it a single bitcast. - if (bitcast->operand(0)->opcode() == HloOpcode::kBitcast) { + HloInstruction* op; + if (Match(bitcast, m::Bitcast(m::Bitcast(m::Op(&op))))) { return ReplaceWithNewInstruction( - bitcast, HloInstruction::CreateUnary( - bitcast->shape(), HloOpcode::kBitcast, - bitcast->mutable_operand(0)->mutable_operand(0))); + bitcast, + HloInstruction::CreateUnary(bitcast->shape(), HloOpcode::kBitcast, op)); } // All bitcasts can be eliminated (assuming layout constraints are // satisified). @@ -418,11 +421,10 @@ Status AlgebraicSimplifierVisitor::HandleBitcastConvert( Status AlgebraicSimplifierVisitor::HandleCopy(HloInstruction* copy) { // If a copy feeds a copy, make it a single copy. - if (copy->operand(0)->opcode() == HloOpcode::kCopy) { + HloInstruction* op; + if (Match(copy, m::Copy(m::Copy(m::Op(&op))))) { return ReplaceWithNewInstruction( - copy, HloInstruction::CreateUnary( - copy->shape(), HloOpcode::kCopy, - copy->mutable_operand(0)->mutable_operand(0))); + copy, HloInstruction::CreateUnary(copy->shape(), HloOpcode::kCopy, op)); } // All copies can be eliminated (assuming layout constraints are satisified). ReplaceInstructionIfSameShape(copy, copy->mutable_operand(0)); @@ -462,12 +464,10 @@ Status AlgebraicSimplifierVisitor::HandleConcatenate( } else if (operands.size() == 2) { // A binary concat with a broadcasted scalar as an operand can be converted // into a pad which is simpler to fold into other operations. - bool is_effective_low_pad = - operands[0]->opcode() == HloOpcode::kBroadcast && - ShapeUtil::IsScalar(operands[0]->operand(0)->shape()); - bool is_effective_high_pad = - operands[1]->opcode() == HloOpcode::kBroadcast && - ShapeUtil::IsScalar(operands[1]->operand(0)->shape()); + bool is_effective_low_pad = Match( + operands[0], m::Broadcast(m::Op().WithShape(m::Shape().IsScalar()))); + bool is_effective_high_pad = Match( + operands[1], m::Broadcast(m::Op().WithShape(m::Shape().IsScalar()))); if (!is_effective_low_pad && !is_effective_high_pad) { return Status::OK(); } @@ -537,8 +537,8 @@ Status AlgebraicSimplifierVisitor::HandleConstant(HloInstruction* constant) { } Status AlgebraicSimplifierVisitor::HandleSubtract(HloInstruction* sub) { - auto lhs = sub->mutable_operand(0); - auto rhs = sub->mutable_operand(1); + HloInstruction *lhs, *rhs; + CHECK(Match(sub, m::Subtract(m::Op(&lhs), m::Op(&rhs)))); // A - 0 => A VLOG(10) << "trying transform [A - 0 => A]: " << sub->ToString(); if (IsAll(rhs, 0) && ReplaceInstructionIfSameShape(sub, lhs)) { @@ -547,7 +547,7 @@ Status AlgebraicSimplifierVisitor::HandleSubtract(HloInstruction* sub) { // Canonicalize subtraction of a constant to addition. VLOG(10) << "trying transform [A - Const => A + (-Const)]"; - if (rhs->IsConstant() && !lhs->IsConstant()) { + if (Match(sub, m::Subtract(m::NonConstant(&lhs), m::Constant(&rhs)))) { HloInstruction* negative_const = computation_->AddInstruction( HloInstruction::CreateUnary(rhs->shape(), HloOpcode::kNegate, rhs)); return ReplaceWithNewInstruction( @@ -559,56 +559,53 @@ Status AlgebraicSimplifierVisitor::HandleSubtract(HloInstruction* sub) { } Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) { - auto lhs = divide->mutable_operand(0); - auto rhs = divide->mutable_operand(1); + Shape* shape; + HloInstruction *a, *b, *c, *d; + CHECK(Match(divide, m::Divide(m::Op(&a), m::Op(&b)))); // A/1 => A VLOG(10) << "trying transform [A/1 => A]: " << divide->ToString(); - if (IsAll(rhs, 1) && ReplaceInstructionIfSameShape(divide, lhs)) { + if (IsAll(b, 1) && ReplaceInstructionIfSameShape(divide, a)) { return Status::OK(); } // exp(A)/exp(B) => exp(A-B) - if (lhs->opcode() == HloOpcode::kExp && rhs->opcode() == HloOpcode::kExp) { + if (Match(divide, m::Divide(m::Exp(m::Op(&a)), m::Exp(m::Op(&b))) + .WithShape(m::Shape(&shape)))) { VLOG(10) << "transform [exp(A)/exp(B) => exp(A-B)]: " << divide->ToString(); - HloInstruction* subtract = - computation_->AddInstruction(HloInstruction::CreateBinary( - divide->shape(), HloOpcode::kSubtract, lhs->mutable_operand(0), - rhs->mutable_operand(0))); + HloInstruction* subtract = computation_->AddInstruction( + HloInstruction::CreateBinary(*shape, HloOpcode::kSubtract, a, b)); return ReplaceWithNewInstruction( - divide, HloInstruction::CreateUnary(divide->shape(), HloOpcode::kExp, - subtract)); + divide, HloInstruction::CreateUnary(*shape, HloOpcode::kExp, subtract)); } // A/exp(B) => A*exp(-B) - if (rhs->opcode() == HloOpcode::kExp) { + if (Match(divide, m::Divide(m::Op(&a), m::Exp(m::Op(&b))))) { VLOG(10) << "transform [A/exp(B) => A*exp(-B)]: " << divide->ToString(); - HloInstruction* negate = - computation_->AddInstruction(HloInstruction::CreateUnary( - divide->shape(), HloOpcode::kNegate, rhs->mutable_operand(0))); + HloInstruction* negate = computation_->AddInstruction( + HloInstruction::CreateUnary(divide->shape(), HloOpcode::kNegate, b)); HloInstruction* new_exp = computation_->AddInstruction( HloInstruction::CreateUnary(divide->shape(), HloOpcode::kExp, negate)); return ReplaceWithNewInstruction( - divide, HloInstruction::CreateBinary( - divide->shape(), HloOpcode::kMultiply, lhs, new_exp)); + divide, HloInstruction::CreateBinary(divide->shape(), + HloOpcode::kMultiply, a, new_exp)); } // A/pow(B,C) => A*pow(B,-C) - if (rhs->opcode() == HloOpcode::kPower) { + if (Match(divide, m::Divide(m::Op(&a), m::Power(m::Op(&b), m::Op(&c))))) { VLOG(10) << "transform [A/pow(B,C) => A*pow(B,-C)]: " << divide->ToString(); // The output shape of the created negate operator should be the same as the // input. - const Shape& negate_shape = rhs->operand(1)->shape(); - HloInstruction* negate = - computation_->AddInstruction(HloInstruction::CreateUnary( - negate_shape, HloOpcode::kNegate, rhs->mutable_operand(1))); + const Shape& negate_shape = c->shape(); + HloInstruction* negate = computation_->AddInstruction( + HloInstruction::CreateUnary(negate_shape, HloOpcode::kNegate, c)); // And the power operator should retain the output shape of the old one. - const Shape& new_power_shape = rhs->shape(); - HloInstruction* new_power = computation_->AddInstruction( - HloInstruction::CreateBinary(new_power_shape, HloOpcode::kPower, - rhs->mutable_operand(0), negate)); + const Shape& new_power_shape = b->shape(); + HloInstruction* new_power = + computation_->AddInstruction(HloInstruction::CreateBinary( + new_power_shape, HloOpcode::kPower, b, negate)); return ReplaceWithNewInstruction( divide, HloInstruction::CreateBinary( - divide->shape(), HloOpcode::kMultiply, lhs, new_power)); + divide->shape(), HloOpcode::kMultiply, a, new_power)); } // Simplifying integral division would produce unexpected results. @@ -620,28 +617,24 @@ Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) { // // (Backends can do this transformation, but generally only if the constant is // a scalar.) - if (lhs->opcode() != HloOpcode::kConstant && - rhs->opcode() == HloOpcode::kConstant) { + if (Match(divide, m::Divide(m::NonConstant(&a), m::Constant(&b)))) { HloInstruction* one = computation_->AddInstruction(HloInstruction::CreateConstant( - Literal::One(lhs->shape().element_type()).CloneToUnique())); - HloInstruction* inverse = - computation_->AddInstruction(HloInstruction::CreateBinary( - rhs->shape(), HloOpcode::kDivide, one, rhs)); + Literal::One(a->shape().element_type()).CloneToUnique())); + HloInstruction* inverse = computation_->AddInstruction( + HloInstruction::CreateBinary(b->shape(), HloOpcode::kDivide, one, b)); return ReplaceWithNewInstruction( - divide, HloInstruction::CreateBinary( - divide->shape(), HloOpcode::kMultiply, lhs, inverse)); + divide, HloInstruction::CreateBinary(divide->shape(), + HloOpcode::kMultiply, a, inverse)); } // (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(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))); + if (Match(divide, m::Divide(m::Divide(m::Op(&a), m::Op(&b)), + m::Divide(m::Op(&c), m::Op(&d))))) { + TF_ASSIGN_OR_RETURN(auto a_times_d, + MakeBinaryHlo(HloOpcode::kMultiply, a, d)); + TF_ASSIGN_OR_RETURN(auto b_times_c, + MakeBinaryHlo(HloOpcode::kMultiply, b, c)); TF_ASSIGN_OR_RETURN(auto new_divide, MakeBinaryHlo(HloOpcode::kDivide, a_times_d, b_times_c)); @@ -649,24 +642,21 @@ Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) { } // (A / B) / C => A / (B * C) - if (lhs->opcode() == HloOpcode::kDivide) { - TF_ASSIGN_OR_RETURN( - auto b_times_c, - MakeBinaryHlo(HloOpcode::kMultiply, lhs->mutable_operand(1), rhs)); + if (Match(divide, m::Divide(m::Divide(m::Op(&a), m::Op(&b)), m::Op(&c)))) { + TF_ASSIGN_OR_RETURN(auto b_times_c, + MakeBinaryHlo(HloOpcode::kMultiply, b, c)); return ReplaceWithNewInstruction( - divide, - HloInstruction::CreateBinary(divide->shape(), HloOpcode::kDivide, - lhs->mutable_operand(0), b_times_c)); + divide, HloInstruction::CreateBinary(divide->shape(), + HloOpcode::kDivide, a, b_times_c)); } // A / (B / C) => (A*C) / B - if (rhs->opcode() == HloOpcode::kDivide) { - TF_ASSIGN_OR_RETURN(auto a_times_c, MakeBinaryHlo(HloOpcode::kMultiply, lhs, - rhs->mutable_operand(1))); + if (Match(divide, m::Divide(m::Op(&a), m::Divide(m::Op(&b), m::Op(&c))))) { + TF_ASSIGN_OR_RETURN(auto a_times_c, + MakeBinaryHlo(HloOpcode::kMultiply, a, c)); return ReplaceWithNewInstruction( - divide, - HloInstruction::CreateBinary(divide->shape(), HloOpcode::kDivide, - a_times_c, rhs->mutable_operand(0))); + divide, HloInstruction::CreateBinary(divide->shape(), + HloOpcode::kDivide, a_times_c, b)); } return Status::OK(); @@ -674,8 +664,8 @@ Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) { StatusOr AlgebraicSimplifierVisitor::HandleDotStrengthReduction( HloInstruction* dot) { - HloInstruction* lhs = dot->mutable_operand(0); - HloInstruction* rhs = dot->mutable_operand(1); + HloInstruction *lhs, *rhs; + CHECK(Match(dot, m::Dot(m::Op(&lhs), m::Op(&rhs)))); int64 lhs_collapsing_dim = dot->dot_dimension_numbers().lhs_contracting_dimensions(0); if (lhs->IsRank2Transpose()) { @@ -792,8 +782,8 @@ StatusOr AlgebraicSimplifierVisitor::OptimizeDotOfConcat( const int64 lhs_contracting_dim = dnums.lhs_contracting_dimensions(0); const int64 rhs_contracting_dim = dnums.rhs_contracting_dimensions(0); - HloInstruction* lhs = dot->mutable_operand(0); - HloInstruction* rhs = dot->mutable_operand(1); + HloInstruction *lhs, *rhs; + CHECK(Match(dot, m::Dot(m::Op(&lhs), m::Op(&rhs)))); TF_ASSIGN_OR_RETURN( HloInstruction * optimized_lhs_concat, @@ -923,8 +913,8 @@ StatusOr AlgebraicSimplifierVisitor::OptimizeDotOfConcatHelper( } Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot) { - auto lhs = dot->mutable_operand(0); - auto rhs = dot->mutable_operand(1); + HloInstruction *lhs, *rhs; + CHECK(Match(dot, m::Dot(m::Op(&lhs), m::Op(&rhs)))); // Only optimize F32 dot operations where the dot, rhs and lhs are rank 2 or // below. @@ -976,8 +966,8 @@ Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot) { } Status AlgebraicSimplifierVisitor::HandleMultiply(HloInstruction* multiply) { - auto lhs = multiply->mutable_operand(0); - auto rhs = multiply->mutable_operand(1); + HloInstruction *lhs, *rhs; + CHECK(Match(multiply, m::Multiply(m::Op(&lhs), m::Op(&rhs)))); // A*1 => A VLOG(10) << "trying transform [A*1 => A]: " << multiply->ToString(); if (IsAll(rhs, 1) && ReplaceInstructionIfSameShape(multiply, lhs)) { @@ -990,10 +980,9 @@ Status AlgebraicSimplifierVisitor::HandleMultiply(HloInstruction* multiply) { } // exp(A) * exp(B) => exp(A+B) - if (lhs->opcode() == HloOpcode::kExp && rhs->opcode() == HloOpcode::kExp) { + if (Match(multiply, m::Multiply(m::Exp(m::Op(&lhs)), m::Exp(m::Op(&rhs))))) { auto add = computation_->AddInstruction(HloInstruction::CreateBinary( - multiply->shape(), HloOpcode::kAdd, lhs->mutable_operand(0), - rhs->mutable_operand(0))); + multiply->shape(), HloOpcode::kAdd, lhs, rhs)); return ReplaceWithNewInstruction( multiply, HloInstruction::CreateUnary(multiply->shape(), HloOpcode::kExp, add)); @@ -1004,20 +993,19 @@ Status AlgebraicSimplifierVisitor::HandleMultiply(HloInstruction* multiply) { Status AlgebraicSimplifierVisitor::HandleLog(HloInstruction* log) { // ln(exp(A)) => A VLOG(10) << "trying transform [ln(exp(A)) => A]: " << log->ToString(); - auto operand = log->mutable_operand(0); - if (operand->opcode() == HloOpcode::kExp && - ReplaceInstructionIfSameShape(log, operand->mutable_operand(0))) { + HloInstruction *a, *b; + if (Match(log, m::Log(m::Exp(m::Op(&a)))) && + ReplaceInstructionIfSameShape(log, a)) { return Status::OK(); } // ln(pow(A,B)) => B*ln(A) - if (operand->opcode() == HloOpcode::kPower) { - auto new_log = computation_->AddInstruction(HloInstruction::CreateUnary( - log->shape(), HloOpcode::kLog, operand->mutable_operand(0))); + if (Match(log, m::Log(m::Power(m::Op(&a), m::Op(&b))))) { + auto new_log = computation_->AddInstruction( + HloInstruction::CreateUnary(log->shape(), HloOpcode::kLog, a)); return ReplaceWithNewInstruction( - log, - HloInstruction::CreateBinary(log->shape(), HloOpcode::kMultiply, - new_log, operand->mutable_operand(1))); + log, HloInstruction::CreateBinary(log->shape(), HloOpcode::kMultiply, + new_log, b)); } return Status::OK(); @@ -1120,7 +1108,8 @@ bool OutputIsSubsetOfOperandElements(HloInstruction* instruction, } // namespace Status AlgebraicSimplifierVisitor::HandleBroadcast(HloInstruction* broadcast) { - auto operand = broadcast->mutable_operand(0); + HloInstruction* operand; + CHECK(Match(broadcast, m::Broadcast(m::Op(&operand)))); auto dims = broadcast->dimensions(); // A degenerate broadcast of a reshape that does not change the number of // elements can be replaced by a reshape. @@ -1231,30 +1220,28 @@ Status AlgebraicSimplifierVisitor::HandleConvert(HloInstruction* convert) { // Complex(Real(c), Imag(c)) -> c Status AlgebraicSimplifierVisitor::HandleComplex(HloInstruction* complex) { - auto real = complex->mutable_operand(0); - auto imag = complex->mutable_operand(1); - if (real->opcode() == HloOpcode::kReal && - imag->opcode() == HloOpcode::kImag && - real->operand(0) == imag->operand(0)) { - return ReplaceInstruction(complex, real->mutable_operand(0)); + HloInstruction *c0, *c1; + if (Match(complex, m::Complex(m::Real(m::Op(&c0)), m::Imag(m::Op(&c1)))) && + c0 == c1) { + return ReplaceInstruction(complex, c0); } return Status::OK(); } // Real(Complex(r, i)) -> r Status AlgebraicSimplifierVisitor::HandleReal(HloInstruction* real) { - auto operand = real->mutable_operand(0); - if (operand->opcode() == HloOpcode::kComplex) { - return ReplaceInstruction(real, operand->mutable_operand(0)); + HloInstruction* op; + if (Match(real, m::Real(m::Complex(m::Op(&op), m::Op())))) { + return ReplaceInstruction(real, op); } return Status::OK(); } // Imag(Complex(r, i)) -> i Status AlgebraicSimplifierVisitor::HandleImag(HloInstruction* imag) { - auto operand = imag->mutable_operand(0); - if (operand->opcode() == HloOpcode::kComplex) { - return ReplaceInstruction(imag, operand->mutable_operand(1)); + HloInstruction* op; + if (Match(imag, m::Imag(m::Complex(m::Op(), m::Op(&op))))) { + return ReplaceInstruction(imag, op); } return Status::OK(); } @@ -1351,8 +1338,8 @@ Status AlgebraicSimplifierVisitor::HandlePad(HloInstruction* pad) { Status AlgebraicSimplifierVisitor::HandlePower(HloInstruction* power) { VLOG(10) << "trying transform [pow(A, 0) => 1]: " << power->ToString(); - auto lhs = power->mutable_operand(0); - auto rhs = power->mutable_operand(1); + HloInstruction *lhs, *rhs; + CHECK(Match(power, m::Power(m::Op(&lhs), m::Op(&rhs)))); if (IsAll(rhs, 0)) { auto one = HloInstruction::CreateConstant( Literal::One(power->shape().element_type()).CloneToUnique()); @@ -1372,9 +1359,10 @@ Status AlgebraicSimplifierVisitor::HandlePower(HloInstruction* power) { } // pow(exp(A),B) => exp(A*B) - if (lhs->opcode() == HloOpcode::kExp) { + HloInstruction *a, *b; + if (Match(power, m::Power(m::Exp(m::Op(&a)), m::Op(&b)))) { auto a_times_b = computation_->AddInstruction(HloInstruction::CreateBinary( - power->shape(), HloOpcode::kMultiply, lhs->operands()[0], rhs)); + power->shape(), HloOpcode::kMultiply, a, b)); return ReplaceWithNewInstruction( power, HloInstruction::CreateUnary(power->shape(), HloOpcode::kExp, a_times_b)); @@ -1424,7 +1412,6 @@ Status AlgebraicSimplifierVisitor::HandlePower(HloInstruction* power) { return Status::OK(); } -// TODO(b/74536353): do this simplification for BroadcastDimOne as well. StatusOr AlgebraicSimplifierVisitor:: TryToSinkReshapeOrBroadcastAfterOpWithUniqueNonScalarOperand( HloInstruction* reshape_or_broadcast) { @@ -1707,7 +1694,7 @@ Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* reduce) { HloInstruction::CreateReshape(reduce->shape(), arg)); return ReplaceWithNewInstruction( reduce, HloInstruction::CreateMap(reduce->shape(), - {reshape, init_value}, function)); + {init_value, reshape}, function)); } return Status::OK(); } diff --git a/tensorflow/compiler/xla/service/allocation_tracker.cc b/tensorflow/compiler/xla/service/allocation_tracker.cc index 4f819a743c48f30df8dde00ece72a0b4e1748802..6bf65825cd02a77e51e9fc521516164d93483084 100644 --- a/tensorflow/compiler/xla/service/allocation_tracker.cc +++ b/tensorflow/compiler/xla/service/allocation_tracker.cc @@ -31,52 +31,51 @@ limitations under the License. namespace xla { StatusOr AllocationTracker::Register( - std::unique_ptr shaped_buffer, const string& tag) { + ShapedBuffer shaped_buffer, const string& tag) { tensorflow::mutex_lock lock(mutex_); VLOG(2) << "Register"; - std::vector> replicated_buffers; + 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) { + 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::vector> replicated_buffers, - const string& tag) { + std::vector replicated_buffers, const string& tag) { VLOG(2) << "RegisterInternal(" << "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()) { + 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()); + shaped_buffer.platform()->Name().c_str()); } } int64 handle = next_handle_++; for (auto& shaped_buffer : replicated_buffers) { std::vector shape_indices; - ShapeUtil::ForEachSubshape(shaped_buffer->on_device_shape(), + 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()); + AddAllocationOrIncrementRefCount(shaped_buffer.buffer(index), + shaped_buffer.device_ordinal()); } - handle_to_shaped_buffers_[handle].emplace_back(std::move(shaped_buffer)); + handle_to_shaped_buffers_[handle].emplace_back( + MakeUnique(std::move(shaped_buffer))); } GlobalDataHandle result; @@ -146,13 +145,13 @@ StatusOr> AllocationTracker::DeconstructTuple( for (int i = 0; i < ShapeUtil::TupleElementCount(shaped_buffer->on_device_shape()); ++i) { - auto element_buffer = MakeUnique( + auto element_buffer = ShapedBuffer( ShapeUtil::GetTupleElementShape(shaped_buffer->on_host_shape(), i), ShapeUtil::GetTupleElementShape(shaped_buffer->on_device_shape(), i), shaped_buffer->platform(), shaped_buffer->device_ordinal()); - element_buffer->set_buffer(shaped_buffer->buffer(/*index=*/{i}), - /*index=*/{}); - std::vector> replicated_buffers; + 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, @@ -204,7 +203,7 @@ StatusOr> AllocationTracker::ResolveInternal( } void AllocationTracker::AddAllocationOrIncrementRefCount( - perftools::gputools::DeviceMemoryBase device_memory, int device_ordinal) { + se::DeviceMemoryBase device_memory, int device_ordinal) { AllocationMap& allocation_map = opaque_to_allocation_map_[device_ordinal]; auto it = allocation_map.find(device_memory.opaque()); if (it == allocation_map.end()) { @@ -215,8 +214,8 @@ void AllocationTracker::AddAllocationOrIncrementRefCount( } } -Status AllocationTracker::DecrementRefCount( - perftools::gputools::DeviceMemoryBase device_memory, int device_ordinal) { +Status AllocationTracker::DecrementRefCount(se::DeviceMemoryBase device_memory, + int device_ordinal) { AllocationMap& allocation_map = opaque_to_allocation_map_[device_ordinal]; auto it = allocation_map.find(device_memory.opaque()); TF_RET_CHECK(it != allocation_map.end()); diff --git a/tensorflow/compiler/xla/service/allocation_tracker.h b/tensorflow/compiler/xla/service/allocation_tracker.h index 038aee8541b297d6f91fe2b3bce7455fd9a7084e..2bfcd5371293635f6525d0c5506c606e2762b943 100644 --- a/tensorflow/compiler/xla/service/allocation_tracker.h +++ b/tensorflow/compiler/xla/service/allocation_tracker.h @@ -45,14 +45,13 @@ class AllocationTracker { // Registers a shaped buffer of device memory, and returns a corresponding // 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); + StatusOr Register(ShapedBuffer 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); + std::vector replicated_buffers, const string& tag); // Unregister the allocation for the given data handle. Status Unregister(const GlobalDataHandle& data); @@ -77,7 +76,7 @@ class AllocationTracker { // Data structure encapsulating single memory allocation on the device. struct Allocation { // The pointer to this allocation. - perftools::gputools::DeviceMemoryBase device_memory; + se::DeviceMemoryBase device_memory; // The device that the memory is allocated on. int device_ordinal; @@ -95,21 +94,21 @@ class AllocationTracker { // Internal helper which registers a vector of shaped buffers, one per // replica. StatusOr RegisterInternal( - std::vector> replicated_buffers, - 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. - void AddAllocationOrIncrementRefCount( - perftools::gputools::DeviceMemoryBase device_memory, int device_ordinal) + void AddAllocationOrIncrementRefCount(se::DeviceMemoryBase device_memory, + int device_ordinal) EXCLUSIVE_LOCKS_REQUIRED(mutex_); // Decrements the reference count of the given device memory. Then, if it is // zero, deallocate the memory. - Status DecrementRefCount(perftools::gputools::DeviceMemoryBase device_memory, + Status DecrementRefCount(se::DeviceMemoryBase device_memory, int device_ordinal) EXCLUSIVE_LOCKS_REQUIRED(mutex_); // A map from device memory opaque value to allocation. One such map is @@ -132,6 +131,9 @@ class AllocationTracker { // A map from data handle to a vector of shaped buffers that represent the // buffers for different replicas. + // + // The ShapedBuffers in this map's vectors need to be unique_ptrs, because our + // public API returns pointers to them. tensorflow::gtl::FlatMap>> handle_to_shaped_buffers_ GUARDED_BY(mutex_); diff --git a/tensorflow/compiler/xla/service/backend.cc b/tensorflow/compiler/xla/service/backend.cc index 05f2d062784147108a94ffb7bb0ca42ddfe4f010..a582dbffd688a4ca904539a462aa268a0b4d3cb6 100644 --- a/tensorflow/compiler/xla/service/backend.cc +++ b/tensorflow/compiler/xla/service/backend.cc @@ -36,19 +36,14 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" -namespace se = ::perftools::gputools; - namespace xla { -BackendOptions& BackendOptions::set_platform( - perftools::gputools::Platform* platform) { +BackendOptions& BackendOptions::set_platform(se::Platform* platform) { platform_ = platform; return *this; } -perftools::gputools::Platform* BackendOptions::platform() const { - return platform_; -} +se::Platform* BackendOptions::platform() const { return platform_; } BackendOptions& BackendOptions::set_intra_op_parallelism_threads( int num_threads) { @@ -77,7 +72,7 @@ struct Backend::EigenThreadPoolWrapper { /* static */ StatusOr> Backend::CreateBackend( const BackendOptions& options) { - perftools::gputools::Platform* platform = options.platform(); + se::Platform* platform = options.platform(); TF_ASSIGN_OR_RETURN(auto compiler, Compiler::GetForPlatform(platform)); TF_ASSIGN_OR_RETURN(auto stream_executors, PlatformUtil::GetStreamExecutors(platform)); @@ -121,7 +116,7 @@ StatusOr Backend::BorrowStream( } Backend::Backend( - perftools::gputools::Platform* platform, Compiler* compiler, + se::Platform* platform, Compiler* compiler, tensorflow::gtl::ArraySlice stream_executors, TransferManager* transfer_manager, ComputationPlacer* computation_placer, int intra_op_parallelism_threads) @@ -178,7 +173,7 @@ tensorflow::thread::ThreadPool* Backend::eigen_intra_op_thread_pool() const { return intra_op_thread_pool_wrapper_->pool.get(); } -StatusOr Backend::stream_executor( +StatusOr Backend::stream_executor( int device_ordinal) const { if (device_ordinal < 0 || device_ordinal > stream_executors_.back()->device_ordinal()) { @@ -201,9 +196,9 @@ StatusOr Backend::devices_equivalent(int device_ordinal_a, // bit crude but works for GPUs which is the important case where we compile // an executable for one GPU and want to know if it will run (well) on // another. - TF_ASSIGN_OR_RETURN(perftools::gputools::StreamExecutor * executor_a, + TF_ASSIGN_OR_RETURN(se::StreamExecutor * executor_a, stream_executor(device_ordinal_a)); - TF_ASSIGN_OR_RETURN(perftools::gputools::StreamExecutor * executor_b, + TF_ASSIGN_OR_RETURN(se::StreamExecutor * executor_b, stream_executor(device_ordinal_b)); return (executor_a->GetDeviceDescription().name() == executor_b->GetDeviceDescription().name()); diff --git a/tensorflow/compiler/xla/service/backend.h b/tensorflow/compiler/xla/service/backend.h index b5ca483b7274d20c31e932d748b6a4c9dea926f9..d32a0a400d8bd51fed6bfb6c6a581e16c8249a5c 100644 --- a/tensorflow/compiler/xla/service/backend.h +++ b/tensorflow/compiler/xla/service/backend.h @@ -44,8 +44,8 @@ namespace xla { class BackendOptions { public: // Set the platform backing the backend, or nullptr for the default platform. - BackendOptions& set_platform(perftools::gputools::Platform* platform); - perftools::gputools::Platform* platform() const; + BackendOptions& set_platform(se::Platform* platform); + se::Platform* platform() const; // Sets the thread pool size for parallel execution of an individual operator. // The default value of -1 will result in initializing the thread pool with @@ -54,7 +54,7 @@ class BackendOptions { int intra_op_parallelism_threads() const; private: - perftools::gputools::Platform* platform_ = nullptr; + se::Platform* platform_ = nullptr; int intra_op_parallelism_threads_ = -1; }; @@ -66,7 +66,7 @@ class BackendOptions { // StreamPtr stream = backend->BorrowStream().ConsumeValueOrDie(); class Backend { public: - using StreamPtr = Pool::SmartPtr; + using StreamPtr = Pool::SmartPtr; // Creates a new backend. static StatusOr> CreateBackend( @@ -79,7 +79,7 @@ class Backend { ~Backend(); // Accessors for the various objects. - perftools::gputools::Platform* platform() const { return platform_; } + se::Platform* platform() const { return platform_; } Compiler* compiler() const { return compiler_; } DeviceMemoryAllocator* memory_allocator() const { return memory_allocator_.get(); @@ -96,19 +96,17 @@ class Backend { // Returns stream executors of all supported devices for this backend. The // executors are ordered by the device ordinal. - const std::vector& stream_executors() - const { + const std::vector& stream_executors() const { return stream_executors_; } // Returns the stream executor for the given device ordinal. - StatusOr stream_executor( - int device_ordinal) const; + StatusOr stream_executor(int device_ordinal) const; // Returns the stream executor for the default device ordinal. This stream // executor can only be used when the number of computations is 1 (replication // can be > 1). - perftools::gputools::StreamExecutor* default_stream_executor() const { + se::StreamExecutor* default_stream_executor() const { CHECK(!stream_executors_.empty()); return stream_executors_[0]; } @@ -117,8 +115,7 @@ class Backend { // internal pool, or by constructing/initializating it, and returns the result // to the caller. StatusOr BorrowStream(int device_ordinal); - StatusOr BorrowStream( - perftools::gputools::StreamExecutor* executor); + StatusOr BorrowStream(se::StreamExecutor* executor); // Returns a function to borrow a stream, as `BorrowStream` above does. // Purely for convenience, the caller could rather make this anonymous @@ -157,29 +154,26 @@ class Backend { private: struct EigenThreadPoolWrapper; - Backend(perftools::gputools::Platform* platform, Compiler* compiler, - tensorflow::gtl::ArraySlice - stream_executors, + Backend(se::Platform* platform, Compiler* compiler, + tensorflow::gtl::ArraySlice stream_executors, TransferManager* transfer_manager, ComputationPlacer* computation_placer, int intra_op_parallelism_threads); Backend(const Backend&) = delete; Backend& operator=(const Backend&) = delete; - perftools::gputools::Platform* platform_; + se::Platform* platform_; Compiler* compiler_; TransferManager* transfer_manager_; ComputationPlacer* computation_placer_; // Vector of stream executors. stream_executors_[0] is the default executor. - std::vector stream_executors_; + std::vector stream_executors_; tensorflow::mutex mu_; // Mapping from stream executor to stream pools, used by `BorrowStream` above. - std::map> - stream_pools_ GUARDED_BY(mu_); + std::map> stream_pools_ GUARDED_BY(mu_); // The default memory allocator to use. std::unique_ptr memory_allocator_; diff --git a/tensorflow/compiler/xla/service/bfloat16_propagation.cc b/tensorflow/compiler/xla/service/bfloat16_propagation.cc index c26d2feef584faeff013a602409cdd58c2d44a5a..43ebe92c5ec1c945780f76ca4178a94f948a81b9 100644 --- a/tensorflow/compiler/xla/service/bfloat16_propagation.cc +++ b/tensorflow/compiler/xla/service/bfloat16_propagation.cc @@ -392,7 +392,6 @@ void BFloat16Propagation::AdjustCalledComputationRoot(HloInstruction* hlo) { 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: diff --git a/tensorflow/compiler/xla/service/bfloat16_propagation_test.cc b/tensorflow/compiler/xla/service/bfloat16_propagation_test.cc index 88f83014164ff726a11e45e762b9c082cf12720d..183db1652e498edb0b94e9c9a272e2b8a7fc53ba 100644 --- a/tensorflow/compiler/xla/service/bfloat16_propagation_test.cc +++ b/tensorflow/compiler/xla/service/bfloat16_propagation_test.cc @@ -426,8 +426,62 @@ TEST_F(BFloat16PropagationTest, SelectOverTuples) { EXPECT_TRUE(OutputsBF16(xpose)); } -// Tests that BF16 is propagated properly through while computations. -TEST_F(BFloat16PropagationTest, PropagateThroughWhile) { +// Tests that BF16 is propagated properly through a while computation with +// non-tuple input/output. +TEST_F(BFloat16PropagationTest, PropagateThroughSimpleWhile) { + 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* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param0, param1)); + + auto builder_cond = HloComputation::Builder("cond"); + auto cond_param = builder_cond.AddInstruction( + HloInstruction::CreateParameter(0, shape, "cond_param")); + auto cond_dot = builder_cond.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kDot, cond_param, cond_param)); + auto cond_root = 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, shape, "body_param")); + auto body_dot = builder_body.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kDot, body_param, body_param)); + auto body = module->AddEmbeddedComputation(builder_body.Build()); + + auto while_hlo = builder.AddInstruction( + HloInstruction::CreateWhile(shape, cond, body, add)); + + auto dot = builder.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kDot, while_hlo, while_hlo)); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(PropagatePrecision(module.get())); + + EXPECT_EQ(computation->root_instruction(), dot); + EXPECT_TRUE( + ShapeUtil::Equal(cond_root->shape(), ShapeUtil::MakeShape(PRED, {}))); + EXPECT_TRUE(OutputsBF16(add)); + EXPECT_TRUE(OutputsBF16(body_dot)); + EXPECT_TRUE(OutputsBF16(body_param)); + EXPECT_TRUE(OutputsBF16(cond_param)); + EXPECT_FALSE(OutputsBF16(dot)); +} + +// Tests that BF16 is propagated properly through while computations with +// tuple-shaped input/output. +TEST_F(BFloat16PropagationTest, PropagateThroughTupleWhile) { auto module = CreateNewModule(); auto builder = HloComputation::Builder(TestName()); Shape shape = ShapeUtil::MakeShape(F32, {4, 4}); diff --git a/tensorflow/compiler/xla/service/compile_only_service.cc b/tensorflow/compiler/xla/service/compile_only_service.cc index c83da9eddc8f8b156dd9acfc99b393bf844575da..c9f78a0f9f1c0e889cd2c761e3129ec329a7b647 100644 --- a/tensorflow/compiler/xla/service/compile_only_service.cc +++ b/tensorflow/compiler/xla/service/compile_only_service.cc @@ -37,7 +37,7 @@ limitations under the License. namespace xla { /* static */ StatusOr> -CompileOnlyService::NewService(perftools::gputools::Platform* platform) { +CompileOnlyService::NewService(se::Platform* platform) { ServiceOptions default_options; default_options.set_platform(platform); return NewService(default_options); @@ -45,7 +45,7 @@ CompileOnlyService::NewService(perftools::gputools::Platform* platform) { /* static */ StatusOr> CompileOnlyService::NewService(const ServiceOptions& options) { - perftools::gputools::Platform* platform = options.platform(); + se::Platform* platform = options.platform(); if (platform == nullptr) { TF_ASSIGN_OR_RETURN(platform, PlatformUtil::GetDefaultPlatform()); } @@ -61,6 +61,33 @@ CompileOnlyService::CompileOnlyService(const ServiceOptions& options, Compiler* compiler) : Service(options, /*execute_backend=*/nullptr), compiler_(compiler) {} +StatusOr>> +CompileOnlyService::CompileAheadOfTime( + const tensorflow::gtl::ArraySlice computations, + const AotCompilationOptions& options) { + std::vector> hlo_modules; + for (const AotXlaComputationInstance& instance : computations) { + TF_RET_CHECK(instance.computation.has_program_shape()); + + const DebugOptions& debug_options = options.debug_options(); + const auto& program_shape = instance.computation.program_shape(); + ExecutionOptions execution_options; + *execution_options.mutable_debug_options() = debug_options; + TF_ASSIGN_OR_RETURN( + std::unique_ptr module_config, + CreateModuleConfig(program_shape, instance.argument_layouts, + &execution_options)); + + TF_ASSIGN_OR_RETURN( + std::unique_ptr hlo_module, + HloModule::CreateFromProto(instance.computation, *module_config)); + TF_RETURN_IF_ERROR(MaybeDumpHloModule(*hlo_module)); + hlo_modules.push_back(std::move(hlo_module)); + } + + return compiler_->CompileAheadOfTime(std::move(hlo_modules), options); +} + StatusOr>> CompileOnlyService::CompileAheadOfTime( const tensorflow::gtl::ArraySlice computations, diff --git a/tensorflow/compiler/xla/service/compile_only_service.h b/tensorflow/compiler/xla/service/compile_only_service.h index 9859941c6c17460939e5b6817f1c7c415e63443c..c10609e67fcdec459baf25a95173bbf700994be9 100644 --- a/tensorflow/compiler/xla/service/compile_only_service.h +++ b/tensorflow/compiler/xla/service/compile_only_service.h @@ -34,7 +34,7 @@ class CompileOnlyService : public Service { // platform that the service should target. If platform is null then the // default platform is used. static StatusOr> NewService( - perftools::gputools::Platform* platform); + se::Platform* platform); static StatusOr> NewService( const ServiceOptions& options); @@ -53,6 +53,25 @@ class CompileOnlyService : public Service { const tensorflow::gtl::ArraySlice computations, const AotCompilationOptions& Options); + // A description of a xla computation to compile using CompileAheadOfTime. + // + // TODO(b/74197823): This is a part of a NOT YET ready refactor. + struct AotXlaComputationInstance { + HloModuleProto computation; + std::vector argument_layouts; + const Shape* result_layout = nullptr; + }; + + // Compiles a list of xla computations for ahead-of-time execution. This is + // intended for use in static compilation. See + // |CompileOnlyClient::CompileAheadOfTime| for additional details. + // + // TODO(b/74197823): This is a part of a NOT YET ready refactor. + StatusOr>> + CompileAheadOfTime( + const tensorflow::gtl::ArraySlice computations, + const AotCompilationOptions& options); + // Override Service methods that require or imply the existence of an // execute backend. Note that this does not include TransferToClient, as // computing constants produces global data that we may wish to transfer. diff --git a/tensorflow/compiler/xla/service/compiler.cc b/tensorflow/compiler/xla/service/compiler.cc index 0392d4af48a040c4a648f7bf9bf21a62ce03a990..8b01a6c4b5004d03e6e7d23b99b923fdcdeaff99 100644 --- a/tensorflow/compiler/xla/service/compiler.cc +++ b/tensorflow/compiler/xla/service/compiler.cc @@ -23,26 +23,21 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/macros.h" -namespace se = ::perftools::gputools; - namespace xla { /* static */ tensorflow::mutex Compiler::platform_compiler_mutex_( tensorflow::LINKER_INITIALIZED); -/* static */ std::map* +/* static */ std::map* Compiler::GetPlatformCompilerFactories() { - static auto* r = - new std::map; + static auto* r = new std::map; return r; } /* static */ -std::map>* +std::map>* Compiler::GetPlatformCompilers() { - static auto* r = new std::map>; + static auto* r = new std::map>; return r; } diff --git a/tensorflow/compiler/xla/service/compiler.h b/tensorflow/compiler/xla/service/compiler.h index b4b53ae2ed425a48de5bcb6ba5c37b5d37e1f371..5c14591d93cc995a0b75efb14da8ec98d5859ff5 100644 --- a/tensorflow/compiler/xla/service/compiler.h +++ b/tensorflow/compiler/xla/service/compiler.h @@ -70,7 +70,7 @@ class AotCompilationOptions { virtual ~AotCompilationOptions() = default; // Returns the ID of the platform to which these options apply. - virtual perftools::gputools::Platform::Id PlatformId() const = 0; + virtual se::Platform::Id PlatformId() const = 0; // Optional allocator that may be used for allocating temp space on the device // during compilation. @@ -109,7 +109,7 @@ class Compiler { virtual ~Compiler() {} // Returns the ID of the platform that this compiler targets. - virtual perftools::gputools::Platform::Id PlatformId() const = 0; + virtual se::Platform::Id PlatformId() const = 0; // Runs Hlo passes to optimize the given Hlo module, returns the optimized // module. @@ -120,8 +120,7 @@ class Compiler { // algorithm over those buffers, to see which variant is fastest. Any space // allocated should be deallocated before this function returns. virtual StatusOr> RunHloPasses( - std::unique_ptr module, - perftools::gputools::StreamExecutor* executor, + std::unique_ptr module, se::StreamExecutor* executor, DeviceMemoryAllocator* device_allocator) = 0; // Compiles the HLO module for execution on a device given by the executor, @@ -137,8 +136,7 @@ class Compiler { // // Use the overload below to compile computations that run in parallel. virtual StatusOr> RunBackend( - std::unique_ptr module, - perftools::gputools::StreamExecutor* executor, + std::unique_ptr module, se::StreamExecutor* executor, DeviceMemoryAllocator* device_allocator) = 0; // Compiles a set of HLO modules that can run in parallel, potentially @@ -151,8 +149,7 @@ class Compiler { // modules to RunHloPasses and RunBackends. virtual StatusOr>> Compile( std::vector> modules, - std::vector> - stream_exec, + std::vector> stream_exec, DeviceMemoryAllocator* device_allocator) = 0; // Compiles the HLO module for ahead-of-time execution. This is intended for @@ -171,14 +168,12 @@ class Compiler { // be a singleton, so no ownership is transferred. // // Precondition: a platform kind must not be registered more than once. - static void RegisterCompilerFactory( - perftools::gputools::Platform::Id platform_id, - CompilerFactory compiler_factory); + static void RegisterCompilerFactory(se::Platform::Id platform_id, + CompilerFactory compiler_factory); // Returns the compiler singleton pointer if it is available for the given // platform, or an error status if it is not. - static StatusOr GetForPlatform( - const perftools::gputools::Platform* platform); + static StatusOr GetForPlatform(const se::Platform* platform); // Returns a function that computes the size in bytes of the logical // buffer that contains a shape. @@ -198,12 +193,12 @@ class Compiler { static tensorflow::mutex platform_compiler_mutex_; // Map from platform kind to compiler factory. - static std::map* + static std::map* GetPlatformCompilerFactories(); // Map from platform kind to compiler instance, if we made one already (based // on the factories above). - static std::map>* + static std::map>* GetPlatformCompilers(); }; diff --git a/tensorflow/compiler/xla/service/computation_layout.cc b/tensorflow/compiler/xla/service/computation_layout.cc index d2d4f14fcec35f5b51a2670a646154ce8bb9bfc1..cb61f3da39fb8eef69fd81066d87a1da91a62935 100644 --- a/tensorflow/compiler/xla/service/computation_layout.cc +++ b/tensorflow/compiler/xla/service/computation_layout.cc @@ -23,12 +23,15 @@ limitations under the License. namespace xla { -ComputationLayout::ComputationLayout(const ProgramShape& program_shape) +ComputationLayout::ComputationLayout(const ProgramShape& program_shape, + bool ignore_layouts) : result_layout_(program_shape.result()) { for (auto& shape : program_shape.parameters()) { parameter_layouts_.emplace_back(shape); } - SetToDefaultLayout(); + if (ignore_layouts) { + SetToDefaultLayout(); + } } void ComputationLayout::SetToDefaultLayout() { diff --git a/tensorflow/compiler/xla/service/computation_layout.h b/tensorflow/compiler/xla/service/computation_layout.h index 80e102411c7885669947d89f378b1ec61e3e4e96..53c3a3f7b738687db3098acfaef1ae87860d0440 100644 --- a/tensorflow/compiler/xla/service/computation_layout.h +++ b/tensorflow/compiler/xla/service/computation_layout.h @@ -34,8 +34,9 @@ class ComputationLayout { public: // Constructs a ComputationLayout from a ProgramShape. The layouts of the // parameters and results are set to the default layout. Layouts in the - // ProgramShape are ignored. - explicit ComputationLayout(const ProgramShape& program_shape); + // ProgramShape are ignored if ignore_layouts is true. + explicit ComputationLayout(const ProgramShape& program_shape, + bool ignore_layouts = true); // Returns the layout of a particular parameter. const ShapeLayout& parameter_layout(int64 param_no) const { diff --git a/tensorflow/compiler/xla/service/computation_placer.cc b/tensorflow/compiler/xla/service/computation_placer.cc index 657fba6b6231104bf47f9dec80f7cd36a0ba3efd..7c1bacff92b231661477b9931a3066fd91110445 100644 --- a/tensorflow/compiler/xla/service/computation_placer.cc +++ b/tensorflow/compiler/xla/service/computation_placer.cc @@ -32,8 +32,6 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" -namespace se = ::perftools::gputools; - namespace xla { Status DeviceAssignment::Serialize(DeviceAssignmentProto* proto) const { @@ -132,11 +130,9 @@ StatusOr ComputationPlacer::AssignDevices( ComputationPlacer::platform_computation_placer_mutex_( tensorflow::LINKER_INITIALIZED); -/* static */ std::map* +/* static */ std::map* ComputationPlacer::GetPlatformComputationPlacers() { - static auto* r = - new std::map; + static auto* r = new std::map; return r; } @@ -147,10 +143,10 @@ static std::unique_ptr CreateComputationPlacer() { } static bool InitModule() { - xla::ComputationPlacer::RegisterComputationPlacer(se::host::kHostPlatformId, - &CreateComputationPlacer); - xla::ComputationPlacer::RegisterComputationPlacer(se::cuda::kCudaPlatformId, - &CreateComputationPlacer); + xla::ComputationPlacer::RegisterComputationPlacer( + stream_executor::host::kHostPlatformId, &CreateComputationPlacer); + xla::ComputationPlacer::RegisterComputationPlacer( + stream_executor::cuda::kCudaPlatformId, &CreateComputationPlacer); return true; } static bool module_initialized = InitModule(); diff --git a/tensorflow/compiler/xla/service/computation_placer.h b/tensorflow/compiler/xla/service/computation_placer.h index 737ccabaa7a61931b6e2787f75b02857562d4820..737d00e93ecb51a9bd544bbcbe99d93374d108fb 100644 --- a/tensorflow/compiler/xla/service/computation_placer.h +++ b/tensorflow/compiler/xla/service/computation_placer.h @@ -80,13 +80,13 @@ class ComputationPlacer { // Registers a computation placer creation function for a particular platform. static void RegisterComputationPlacer( - perftools::gputools::Platform::Id platform_id, + se::Platform::Id platform_id, ComputationPlacerCreationFunction creation_function); // Returns the computation placer singleton pointer if it is available for the // given platform, or an error status if it is not. static StatusOr GetForPlatform( - const perftools::gputools::Platform* platform); + const se::Platform* platform); private: // The mutex that guards the platform-to-computation placer map. @@ -101,10 +101,9 @@ class ComputationPlacer { }; // Map from platform kind to computation placer singleton. - static std::map* - GetPlatformComputationPlacers(); + static std::map* GetPlatformComputationPlacers(); - perftools::gputools::Platform::Id platform_id_; + se::Platform::Id platform_id_; TF_DISALLOW_COPY_AND_ASSIGN(ComputationPlacer); }; diff --git a/tensorflow/compiler/xla/service/conditional_simplifier.cc b/tensorflow/compiler/xla/service/conditional_simplifier.cc index f35de080853f7ec986565cb2df1050946ac3f244..e560abc87f84566905333181c159edd3ca297563 100644 --- a/tensorflow/compiler/xla/service/conditional_simplifier.cc +++ b/tensorflow/compiler/xla/service/conditional_simplifier.cc @@ -69,7 +69,7 @@ static StatusOr TryRemoveConditional(HloInstruction* conditional) { conditional->shape(), {conditional->mutable_operand(2)}, conditional->false_computation())); } - + conditional->SetupDerivedInstruction(call_op); TF_RETURN_IF_ERROR(computation->ReplaceInstruction(conditional, call_op)); TF_RETURN_IF_ERROR(CallInliner::Inline(call_op).status()); diff --git a/tensorflow/compiler/xla/service/cpu/compiler_functor.cc b/tensorflow/compiler/xla/service/cpu/compiler_functor.cc index 61b2da7a7dce7f6fba46a23cc8e5462a3899a18c..6a7eb85e3baec3517b8f3ddef6a8dcfae9c9e614 100644 --- a/tensorflow/compiler/xla/service/cpu/compiler_functor.cc +++ b/tensorflow/compiler/xla/service/cpu/compiler_functor.cc @@ -25,11 +25,11 @@ limitations under the License. #include "llvm/ADT/StringRef.h" #include "llvm/Analysis/TargetLibraryInfo.h" #include "llvm/Analysis/TargetTransformInfo.h" -#include "llvm/ExecutionEngine/ObjectMemoryBuffer.h" #include "llvm/IR/LegacyPassManager.h" #include "llvm/IR/Verifier.h" #include "llvm/MC/MCContext.h" #include "llvm/Object/ObjectFile.h" +#include "llvm/Support/SmallVectorMemoryBuffer.h" #include "llvm/Support/raw_ostream.h" #include "llvm/Target/TargetMachine.h" #include "llvm/Transforms/IPO.h" @@ -158,7 +158,7 @@ std::unique_ptr CompilerFunctor::operator()( // Construct ObjectFile from machine code buffer. return std::unique_ptr( - new llvm::ObjectMemoryBuffer(std::move(stream_buffer))); + new llvm::SmallVectorMemoryBuffer(std::move(stream_buffer))); } static std::vector VectorFunctionsForTargetLibraryInfoImpl() { diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc index e43777c5e5e8afcf08e1e334c8847f6b94d0d047..e8472fd36b360fcb1f73cd97128fe7975bf2cd5c 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc @@ -100,8 +100,6 @@ limitations under the License. #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/lib/strings/strcat.h" -namespace se = ::perftools::gputools; - namespace xla { namespace cpu { @@ -440,8 +438,7 @@ Status VerifyLlvmModule(const llvm::Module& llvm_module) { } // namespace StatusOr> CpuCompiler::RunHloPasses( - std::unique_ptr module, - perftools::gputools::StreamExecutor* /*stream_exec*/, + std::unique_ptr module, se::StreamExecutor* /*stream_exec*/, DeviceMemoryAllocator* /*device_allocator*/) { VLOG(2) << "Before optimization:"; XLA_VLOG_LINES(2, module->ToString()); @@ -454,8 +451,7 @@ StatusOr> CpuCompiler::RunHloPasses( } StatusOr> CpuCompiler::RunBackend( - std::unique_ptr module, - perftools::gputools::StreamExecutor* stream_exec, + std::unique_ptr module, se::StreamExecutor* stream_exec, DeviceMemoryAllocator* /*device_allocator*/) { const string timer_message = "Compiling [" + module->name() + "] for CPU using JIT"; @@ -938,9 +934,9 @@ HloCostAnalysis::ShapeSizeFunction CpuCompiler::ShapeSizeBytesFunction() const { } // namespace xla static bool InitModule() { - xla::Compiler::RegisterCompilerFactory(se::host::kHostPlatformId, []() { - return xla::MakeUnique(); - }); + xla::Compiler::RegisterCompilerFactory( + stream_executor::host::kHostPlatformId, + []() { return xla::MakeUnique(); }); return true; } static bool module_initialized = InitModule(); diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.h b/tensorflow/compiler/xla/service/cpu/cpu_compiler.h index 3498139ab95d21383c6dc008ae5614b7bfe91148..151af38438a980e40c06a1801a936cb620c6c4ba 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.h +++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.h @@ -53,7 +53,7 @@ class CpuAotCompilationOptions : public AotCompilationOptions { RelocationModel relocation_model); ~CpuAotCompilationOptions() override; - perftools::gputools::Platform::Id PlatformId() const override; + se::Platform::Id PlatformId() const override; // The triple used for compilation, similar to clang's -target flag. const string& triple() const { return triple_; } @@ -112,25 +112,23 @@ class CpuCompiler : public LLVMCompiler { // Bring in // StatusOr>> Compile( // std::vector> modules, - // std::vector> + // std::vector> // stream_execs) using LLVMCompiler::Compile; StatusOr> RunHloPasses( - std::unique_ptr module, - perftools::gputools::StreamExecutor* stream_exec, + std::unique_ptr module, se::StreamExecutor* stream_exec, DeviceMemoryAllocator* device_allocator) override; StatusOr> RunBackend( - std::unique_ptr module, - perftools::gputools::StreamExecutor* stream_exec, + std::unique_ptr module, se::StreamExecutor* stream_exec, DeviceMemoryAllocator* device_allocator) override; StatusOr>> CompileAheadOfTime(std::vector> modules, const AotCompilationOptions& options) override; - perftools::gputools::Platform::Id PlatformId() const override; + se::Platform::Id PlatformId() const override; HloCostAnalysis::ShapeSizeFunction ShapeSizeBytesFunction() const override; diff --git a/tensorflow/compiler/xla/service/cpu/cpu_executable.cc b/tensorflow/compiler/xla/service/cpu/cpu_executable.cc index c053703c3524a47ee1de9681c1b986edbf109430..97e550abe44a32f25b0aab039e95b0ce2eecde4c 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_executable.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_executable.cc @@ -45,8 +45,6 @@ limitations under the License. #include "tensorflow/core/platform/types.h" #include "tensorflow/stream_executor/host/host_stream.h" -namespace se = ::perftools::gputools; - namespace xla { namespace cpu { @@ -75,7 +73,7 @@ CpuExecutable::CpuExecutable( Status CpuExecutable::AllocateBuffers( DeviceMemoryAllocator* memory_allocator, int device_ordinal, - std::vector* buffers) { + std::vector* buffers) { CHECK_EQ(buffers->size(), assignment_->Allocations().size()); VLOG(3) << "Allocating " << assignment_->Allocations().size() << " allocations for module " << module().name(); @@ -245,19 +243,18 @@ static Status DeallocateTempBuffers( return Status::OK(); } -StatusOr> CpuExecutable::CreateResultShapedBuffer( +StatusOr CpuExecutable::CreateResultShapedBuffer( const ServiceExecutableRunOptions* run_options, - tensorflow::gtl::ArraySlice - allocated_buffers, + tensorflow::gtl::ArraySlice allocated_buffers, std::vector* buffers_in_result) { se::Stream* stream = run_options->stream(); - auto result_buffer = MakeUnique( + ShapedBuffer result_buffer( /*on_host_shape=*/result_shape(), /*on_device_shape=*/result_shape(), stream->parent()->platform(), stream->parent()->device_ordinal()); // Copy DeviceMemoryBase values which contain the array(s) of the result into // the respective location in ShapedBuffer which is returned to the caller. - TF_RETURN_IF_ERROR(result_buffer->buffers().ForEachMutableElementWithStatus( + TF_RETURN_IF_ERROR(result_buffer.buffers().ForEachMutableElementWithStatus( [&](const ShapeIndex& index, se::DeviceMemoryBase* device_memory) { const auto& sources = this->GetRootPointsToSet().element(index); // The points to set is unambiguous so the set should be a @@ -284,7 +281,7 @@ StatusOr> CpuExecutable::CreateResultShapedBuffer( return std::move(result_buffer); } -StatusOr> CpuExecutable::ExecuteOnStream( +StatusOr CpuExecutable::ExecuteOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments, HloExecutionProfile* hlo_execution_profile) { @@ -303,7 +300,7 @@ StatusOr> CpuExecutable::ExecuteOnStream( std::vector buffers_in_result(assignment_->Allocations().size(), false); TF_ASSIGN_OR_RETURN( - std::unique_ptr result_buffer, + ShapedBuffer result_buffer, CreateResultShapedBuffer(run_options, buffers, &buffers_in_result)); // Free all buffers not in the result. @@ -313,7 +310,7 @@ StatusOr> CpuExecutable::ExecuteOnStream( return std::move(result_buffer); } -StatusOr> CpuExecutable::ExecuteAsyncOnStream( +StatusOr CpuExecutable::ExecuteAsyncOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments) { if (hlo_profiling_enabled()) { @@ -322,7 +319,7 @@ StatusOr> CpuExecutable::ExecuteAsyncOnStream( "supported on CPU."); } - auto* host_stream = dynamic_cast( + auto* host_stream = dynamic_cast( run_options->stream()->implementation()); se::Stream* stream = run_options->stream(); DeviceMemoryAllocator* memory_allocator = run_options->allocator(); @@ -333,7 +330,7 @@ StatusOr> CpuExecutable::ExecuteAsyncOnStream( std::vector buffers_in_result(assignment_->Allocations().size(), false); TF_ASSIGN_OR_RETURN( - std::unique_ptr result_buffer, + ShapedBuffer result_buffer, CreateResultShapedBuffer(run_options, buffers, &buffers_in_result)); LogLiveAddresses(buffers, buffers_in_result); diff --git a/tensorflow/compiler/xla/service/cpu/cpu_executable.h b/tensorflow/compiler/xla/service/cpu/cpu_executable.h index d3502b3a03e27c8f90ed74c4d826dfab1c4e8b75..06b6943cb5a3f20f4d8aff8acfc4ba50b9067477 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_executable.h +++ b/tensorflow/compiler/xla/service/cpu/cpu_executable.h @@ -55,12 +55,12 @@ class CpuExecutable : public Executable { std::unique_ptr hlo_profile_index_map); ~CpuExecutable() override {} - StatusOr> ExecuteOnStream( + StatusOr ExecuteOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments, HloExecutionProfile* hlo_execution_profile) override; - StatusOr> ExecuteAsyncOnStream( + StatusOr ExecuteAsyncOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments) override; @@ -90,29 +90,27 @@ class CpuExecutable : public Executable { // assignment. Each vector element corresponds to a particular Index. If // a vector element already contains a non-null DeviceMemoryBase, then no // buffer is assigned for this element. - Status AllocateBuffers( - DeviceMemoryAllocator* memory_allocator, int device_ordinal, - std::vector* buffers); + Status AllocateBuffers(DeviceMemoryAllocator* memory_allocator, + int device_ordinal, + std::vector* buffers); // Calls the generated function performing the computation with the given // arguments using the supplied buffers. Status ExecuteComputeFunction( const ExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments, - tensorflow::gtl::ArraySlice - buffers, + tensorflow::gtl::ArraySlice buffers, HloExecutionProfile* hlo_execution_profile); - // Create a ShapedBuffer for holding the result of the computation. The + // Creates a ShapedBuffer for holding the result of the computation. The // addresses (DeviceMemoryBases) are set according to buffer assignment. // 'buffers_in_result' should point to a vector of the same size as // 'allocated_buffers'. An element in buffers_in_result is set to true if the // corresponding buffer is live out of the computation (and thus contained in // the returned ShapedBuffer). - StatusOr> CreateResultShapedBuffer( + StatusOr CreateResultShapedBuffer( const ServiceExecutableRunOptions* run_options, - tensorflow::gtl::ArraySlice - allocated_buffers, + tensorflow::gtl::ArraySlice allocated_buffers, std::vector* buffers_in_result); // Returns the points-to set of the root instruction of the entry diff --git a/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.cc b/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.cc index f5e61aef534da57ce13d3ee9bbeaeaec31f53d2e..9b39e7f5765ae5eb6a25c06eef4d74b1c00e5c91 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.cc @@ -34,8 +34,6 @@ limitations under the License. #include "tensorflow/core/platform/notification.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" -namespace se = ::perftools::gputools; - namespace xla { namespace { @@ -241,21 +239,20 @@ Status CpuTransferManager::TransferLiteralFromOutfeed( } StatusOr CpuTransferManager::TransferTupleBuffersFromOutfeed( - perftools::gputools::StreamExecutor* executor, + se::StreamExecutor* executor, tensorflow::gtl::ArraySlice> buffer_data) { return TransferBuffersFromOutfeedInternal(executor, buffer_data, /*is_tuple=*/true); } StatusOr CpuTransferManager::TransferArrayBufferFromOutfeed( - perftools::gputools::StreamExecutor* executor, void* destination, - int64 size_bytes) { + se::StreamExecutor* executor, void* destination, int64 size_bytes) { return TransferBuffersFromOutfeedInternal( executor, {{destination, size_bytes}}, /*is_tuple=*/false); } StatusOr CpuTransferManager::TransferBuffersFromOutfeedInternal( - perftools::gputools::StreamExecutor* executor, + se::StreamExecutor* executor, tensorflow::gtl::ArraySlice> buffer_data, bool is_tuple) { std::vector> buffers; @@ -306,8 +303,8 @@ static std::unique_ptr CreateCpuTransferManager() { } static bool InitModule() { - xla::TransferManager::RegisterTransferManager(se::host::kHostPlatformId, - &CreateCpuTransferManager); + xla::TransferManager::RegisterTransferManager( + stream_executor::host::kHostPlatformId, &CreateCpuTransferManager); return true; } static bool module_initialized = InitModule(); diff --git a/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.h b/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.h index 6c7524d94716464218ba18ad9950f702d2759f89..3ecb0d236498371f48caf63249f9cd4e8777752b 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.h +++ b/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.h @@ -37,36 +37,35 @@ class CpuTransferManager : public GenericTransferManager { CpuTransferManager(); ~CpuTransferManager() override {} - Status TransferLiteralToInfeed(perftools::gputools::StreamExecutor* executor, + Status TransferLiteralToInfeed(se::StreamExecutor* executor, const Literal& literal) override; - Status TransferBufferToInfeed(perftools::gputools::StreamExecutor* executor, - int64 size, const void* source) override; - Status TransferLiteralFromOutfeed( - perftools::gputools::StreamExecutor* executor, const Shape& literal_shape, - Literal* literal) override; + Status TransferBufferToInfeed(se::StreamExecutor* executor, int64 size, + const void* source) override; + Status TransferLiteralFromOutfeed(se::StreamExecutor* executor, + const Shape& literal_shape, + Literal* literal) override; private: // Transfers infeed data to device. InfeedBuffer->Done() must be // called to clean up the memory allocated for InfeedBuffer. StatusOr TransferBufferToInfeedInternal( - perftools::gputools::StreamExecutor* executor, int64 size, - const void* source); + se::StreamExecutor* executor, int64 size, const void* source); // Helper that transfers a tuple of element buffers from the device's outfeed. StatusOr TransferTupleBuffersFromOutfeed( - perftools::gputools::StreamExecutor* executor, + se::StreamExecutor* executor, tensorflow::gtl::ArraySlice> buffer_data); // Helper that transfers an array buffer from the device's outfeed. - StatusOr TransferArrayBufferFromOutfeed( - perftools::gputools::StreamExecutor* executor, void* destination, - int64 size_bytes); + StatusOr TransferArrayBufferFromOutfeed(se::StreamExecutor* executor, + void* destination, + int64 size_bytes); // On success, returns the shape that was transferred from the outfeed -- if // is_tuple is true, the returned shape will be a tuple of the returned shapes // for the given buffers. StatusOr TransferBuffersFromOutfeedInternal( - perftools::gputools::StreamExecutor* executor, + se::StreamExecutor* executor, tensorflow::gtl::ArraySlice> buffer_data, bool is_tuple); diff --git a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc index 29afd8ea5f9822ea9ae969ae035511a58de4888e..495fecc4aa8b3cf8fcb3ab63d82d8146546854da 100644 --- a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc @@ -1070,7 +1070,8 @@ static bool AreValidGemmShapes(const Shape& lhs_shape, const Shape& rhs_shape, // 1) be matrices with no padding, and // 2) have an allowed element type. PrimitiveType output_primitive_type = output_shape.element_type(); - return (output_primitive_type == F32 || output_primitive_type == F16) && + return (output_primitive_type == F64 || 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/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc index 99c5e16db70c6a203b4751c1ed8a106c0dc260e6..e97113dfa0f59e791d614c0093d0781e49c48ee4 100644 --- a/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc @@ -115,7 +115,7 @@ llvm_ir::ElementGenerator CpuElementalIrEmitter::MakeElementGenerator( for (int i = 0; i < hlo->operand_count(); i++) { TF_ASSIGN_OR_RETURN(llvm::Value * operand_value, operand_to_generator.at(hlo->operand(i))( - ElementwiseSourceIndex(index, *hlo, 0))); + ElementwiseSourceIndex(index, *hlo, i))); operands.push_back(operand_value); } return ir_emitter_->EmitScalarCall(hlo->shape().element_type(), diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc index 3405277d449f2d9e558f2d3f83277163655af592..f990ee278526c396d9b55ae49abdaa9bfd2f9d5f 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc @@ -2076,7 +2076,7 @@ Status IrEmitter::HandleFusion(HloInstruction* fusion) { TF_RETURN_IF_ERROR(ElementTypesSameAndSupported( /*instruction=*/*root, /*operands=*/{lhs, rhs}, - /*supported_types=*/{F16, F32})); + /*supported_types=*/{F16, F32, F64})); llvm_ir::IrArray lhs_array(GetIrArrayFor(lhs)); llvm_ir::IrArray rhs_array(GetIrArrayFor(rhs)); diff --git a/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.cc b/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.cc index 07a9f0efcb64db4b2ff0c6518d4b48eee9a505e0..a2bd4fa195bd93211914da172e7aed59b713ba8b 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.cc +++ b/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.cc @@ -49,8 +49,6 @@ limitations under the License. #include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/types.h" -namespace se = ::perftools::gputools; - namespace xla { namespace cpu { @@ -325,7 +323,7 @@ const void** Executor::GetOperandBuffers(HloInstruction* instruction) { Status ParallelCpuExecutable::AllocateBuffers( DeviceMemoryAllocator* memory_allocator, int device_ordinal, - std::vector* buffers) { + std::vector* buffers) { CHECK_EQ(buffers->size(), assignment_->Allocations().size()); VLOG(3) << "Allocating " << assignment_->Allocations().size() << " allocations for module " << module().name(); @@ -449,7 +447,7 @@ Status ParallelCpuExecutable::ExecuteComputeFunctions( return Status::OK(); } -StatusOr> ParallelCpuExecutable::ExecuteOnStream( +StatusOr ParallelCpuExecutable::ExecuteOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments, HloExecutionProfile* hlo_execution_profile) { @@ -461,7 +459,7 @@ StatusOr> ParallelCpuExecutable::ExecuteOnStream( DeviceMemoryAllocator* memory_allocator = run_options->allocator(); std::vector buffers(assignment_->Allocations().size()); - auto result_buffer = MakeUnique( + ShapedBuffer result_buffer( /*on_host_shape=*/result_shape(), /*on_device_shape=*/result_shape(), stream->parent()->platform(), stream->parent()->device_ordinal()); @@ -474,7 +472,7 @@ StatusOr> ParallelCpuExecutable::ExecuteOnStream( // Copy DeviceMemoryBase values which into the respective location in // ShapedBuffer which is returned to the caller. std::vector buffers_in_result(assignment_->Allocations().size(), false); - TF_RETURN_IF_ERROR(result_buffer->buffers().ForEachMutableElementWithStatus( + TF_RETURN_IF_ERROR(result_buffer.buffers().ForEachMutableElementWithStatus( [&](const ShapeIndex& index, se::DeviceMemoryBase* device_memory) { const auto& sources = this->GetRootPointsToSet().element(index); @@ -513,8 +511,7 @@ StatusOr> ParallelCpuExecutable::ExecuteOnStream( return std::move(result_buffer); } -StatusOr> -ParallelCpuExecutable::ExecuteAsyncOnStream( +StatusOr ParallelCpuExecutable::ExecuteAsyncOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments) { // TODO(b/30671675): Implement asynchronous execution mode. diff --git a/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.h b/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.h index 87c0a3df458eb4b3f217192597e0de1576304367..5ce84fa99640c782d85e8bcedd19820788737a38 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.h +++ b/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.h @@ -59,12 +59,12 @@ class ParallelCpuExecutable : public Executable { std::unique_ptr hlo_profile_index_map); ~ParallelCpuExecutable() override {} - StatusOr> ExecuteOnStream( + StatusOr ExecuteOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments, HloExecutionProfile* hlo_execution_profile) override; - StatusOr> ExecuteAsyncOnStream( + StatusOr ExecuteAsyncOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments) override; @@ -89,17 +89,16 @@ class ParallelCpuExecutable : public Executable { // assignment. Each vector element corresponds to a particular Index. If // a vector element already contains a non-null DeviceMemoryBase, then no // buffer is assigned for this element. - Status AllocateBuffers( - DeviceMemoryAllocator* memory_allocator, int device_ordinal, - std::vector* buffers); + Status AllocateBuffers(DeviceMemoryAllocator* memory_allocator, + int device_ordinal, + std::vector* buffers); // Calls the generated functions in 'function_names_', performing the // computation with the given arguments using the supplied buffers. Status ExecuteComputeFunctions( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments, - tensorflow::gtl::ArraySlice - buffers, + tensorflow::gtl::ArraySlice buffers, HloExecutionProfile* hlo_execution_profile); // Returns the points-to set of the root instruction of the entry diff --git a/tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.cc b/tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.cc index 1e439cde11cf74272101b80c867a308e51ab26a6..54af40506dab48b3c2a3a44eb0b5f5fb213a32ec 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.cc @@ -29,7 +29,8 @@ ParallelLoopEmitter::ParallelLoopEmitter( : LoopEmitter(target_element_generator, target_array, ir_builder), dynamic_loop_bounds_(dynamic_loop_bounds) {} -llvm_ir::IrArray::Index ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock( +std::vector +ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock( tensorflow::StringPiece loop_name) { CHECK(!ShapeUtil::IsTuple(shape_)); CHECK(!ShapeUtil::IsScalar(shape_)); @@ -69,7 +70,7 @@ llvm_ir::IrArray::Index ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock( exit_bb_ = loop_nest.GetOuterLoopExitBasicBlock(); CHECK(exit_bb_ != nullptr); - return array_index; + return {array_index}; } } // namespace cpu diff --git a/tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.h b/tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.h index ce92e36a944de33b991d97460f0b2e859ad56081..755715634aa70a822b21d25dcae20a8fe053477a 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.h +++ b/tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.h @@ -60,7 +60,7 @@ class ParallelLoopEmitter : public llvm_ir::LoopEmitter { ParallelLoopEmitter& operator=(const ParallelLoopEmitter&) = delete; ~ParallelLoopEmitter() override = default; - llvm_ir::IrArray::Index EmitIndexAndSetExitBasicBlock( + std::vector EmitIndexAndSetExitBasicBlock( tensorflow::StringPiece loop_name) override; private: diff --git a/tensorflow/compiler/xla/service/device_memory_allocator.cc b/tensorflow/compiler/xla/service/device_memory_allocator.cc index 78e7aa48accdbb51a8477455f5f9c004828c068f..35db4fd2a22cc1615ade77a801cb28c504db09a6 100644 --- a/tensorflow/compiler/xla/service/device_memory_allocator.cc +++ b/tensorflow/compiler/xla/service/device_memory_allocator.cc @@ -24,19 +24,16 @@ limitations under the License. namespace xla { StreamExecutorMemoryAllocator::StreamExecutorMemoryAllocator( - const perftools::gputools::Platform* platform, - tensorflow::gtl::ArraySlice - stream_executors) + const se::Platform* platform, + tensorflow::gtl::ArraySlice stream_executors) : DeviceMemoryAllocator(platform), stream_executors_(stream_executors.begin(), stream_executors.end()) {} -StatusOr -StreamExecutorMemoryAllocator::Allocate(int device_ordinal, uint64 size, - bool retry_on_failure) { - TF_ASSIGN_OR_RETURN(perftools::gputools::StreamExecutor * stream_executor, +StatusOr StreamExecutorMemoryAllocator::Allocate( + int device_ordinal, uint64 size, bool retry_on_failure) { + TF_ASSIGN_OR_RETURN(se::StreamExecutor * stream_executor, GetStreamExecutor(device_ordinal)); - perftools::gputools::DeviceMemoryBase result = - stream_executor->AllocateArray(size); + se::DeviceMemoryBase result = stream_executor->AllocateArray(size); if (size > 0 && result == nullptr) { return ResourceExhausted( "Failed to allocate request for %s (%lluB) on device ordinal %d", @@ -47,22 +44,22 @@ StreamExecutorMemoryAllocator::Allocate(int device_ordinal, uint64 size, } tensorflow::Status StreamExecutorMemoryAllocator::Deallocate( - int device_ordinal, perftools::gputools::DeviceMemoryBase* mem) { + int device_ordinal, se::DeviceMemoryBase* mem) { if (!mem->is_null()) { - TF_ASSIGN_OR_RETURN(perftools::gputools::StreamExecutor * stream_executor, + TF_ASSIGN_OR_RETURN(se::StreamExecutor * stream_executor, GetStreamExecutor(device_ordinal)); // We make a local copy of 'mem' so the original is not zeroed out by the // Deallocate() call below. This gives us a better chance of // catching double-free bugs, since Deallocate silently succeeds for null // values. - perftools::gputools::DeviceMemoryBase mem_copy(*mem); + se::DeviceMemoryBase mem_copy(*mem); stream_executor->Deallocate(&mem_copy); } return tensorflow::Status::OK(); } -StatusOr -StreamExecutorMemoryAllocator::GetStreamExecutor(int device_ordinal) { +StatusOr StreamExecutorMemoryAllocator::GetStreamExecutor( + int device_ordinal) { if (device_ordinal < 0) { return InvalidArgument("device ordinal value (%d) must be non-negative", device_ordinal); diff --git a/tensorflow/compiler/xla/service/device_memory_allocator.h b/tensorflow/compiler/xla/service/device_memory_allocator.h index 39dfad84c1c1c1c461c24de555ecd919cea47d83..da45c4d45a1c56fd39b1e3e17ff131de59ceeced 100644 --- a/tensorflow/compiler/xla/service/device_memory_allocator.h +++ b/tensorflow/compiler/xla/service/device_memory_allocator.h @@ -33,30 +33,42 @@ class DeviceMemoryAllocator { public: // Parameter platform indicates which platform the allocator allocates memory // on. Must be non-null. - explicit DeviceMemoryAllocator(const perftools::gputools::Platform* platform) + explicit DeviceMemoryAllocator(const se::Platform* platform) : platform_(platform) {} virtual ~DeviceMemoryAllocator() {} // 'retry_on_failure': If false, and the first attempt to allocate the memory - // fails, the allocation should return immediately without retrying. - // An example use case is optional scratch spaces where a failure - // has only performance impact. + // fails, the allocation should return immediately without retrying. An + // example use case is optional scratch spaces where a failure has only + // performance impact. + // // Allocate() should return a null pointer for a size-0 allocation. // Deallocate() must be a no-op for null pointers. - virtual StatusOr Allocate( - int device_ordinal, uint64 size, bool retry_on_failure = true) = 0; - virtual tensorflow::Status Deallocate( - int device_ordinal, perftools::gputools::DeviceMemoryBase* mem) = 0; + virtual StatusOr Allocate(int device_ordinal, + uint64 size, + bool retry_on_failure) = 0; + + // Two-arg version of Allocate(), which sets retry-on-failure to true. + // + // (We don't simply use a default argument on the virtual Allocate function + // because default args on virtual functions are disallowed by the Google + // style guide.) + StatusOr Allocate(int device_ordinal, uint64 size) { + return Allocate(device_ordinal, size, /*retry_on_failure=*/true); + } + + virtual tensorflow::Status Deallocate(int device_ordinal, + se::DeviceMemoryBase* mem) = 0; // Return the platform that the allocator allocates memory on. - const perftools::gputools::Platform* platform() const { return platform_; } + const se::Platform* platform() const { return platform_; } // Can we call Deallocate() as soon as a computation has been scheduled on // a stream, or do we have to wait for the computation to complete first? virtual bool AllowsAsynchronousDeallocation() const = 0; protected: - const perftools::gputools::Platform* platform_; + const se::Platform* platform_; }; // Default memory allocator for a platform which uses @@ -64,25 +76,27 @@ class DeviceMemoryAllocator { class StreamExecutorMemoryAllocator : public DeviceMemoryAllocator { public: StreamExecutorMemoryAllocator( - const perftools::gputools::Platform* platform, - tensorflow::gtl::ArraySlice - stream_executors); + const se::Platform* platform, + tensorflow::gtl::ArraySlice stream_executors); - StatusOr Allocate( - int device_ordinal, uint64 size, bool retry_on_failure = true) override; - tensorflow::Status Deallocate( - int device_ordinal, perftools::gputools::DeviceMemoryBase* mem) override; + StatusOr Allocate(int device_ordinal, uint64 size, + bool retry_on_failure) override; + + // Pull in two-arg overload that sets retry_on_failure to true. + using DeviceMemoryAllocator::Allocate; + + tensorflow::Status Deallocate(int device_ordinal, + se::DeviceMemoryBase* mem) override; bool AllowsAsynchronousDeallocation() const override; private: - StatusOr GetStreamExecutor( - int device_ordinal); + StatusOr GetStreamExecutor(int device_ordinal); // A vector indexed by device ordinal of StreamExecutors for each device of // the allocator's platform type. If an element is nullptr, then the device // with the respective device ordinal is not supported by XLA. - std::vector stream_executors_; + std::vector stream_executors_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h index 3f7089d6ca1e1a3b9bb42028327ba54ba4b93974..0528b076027603796a445d8b0e9cbcebd1b513a7 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h @@ -147,6 +147,9 @@ class DfsHloVisitorBase { virtual Status HandleLog(HloInstructionPtr hlo) { return HandleElementwiseUnary(hlo); } + virtual Status HandleClz(HloInstructionPtr hlo) { + return HandleElementwiseUnary(hlo); + } virtual Status HandleCos(HloInstructionPtr hlo) { return HandleElementwiseUnary(hlo); } @@ -199,7 +202,6 @@ class DfsHloVisitorBase { virtual Status HandleReduce(HloInstructionPtr hlo) = 0; virtual Status HandleBitcast(HloInstructionPtr hlo) = 0; virtual Status HandleBroadcast(HloInstructionPtr hlo) = 0; - virtual Status HandleBroadcastDimOne(HloInstructionPtr hlo) = 0; virtual Status HandleReshape(HloInstructionPtr hlo) = 0; virtual Status HandleTranspose(HloInstructionPtr hlo) = 0; virtual Status HandleParameter(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 e6680ee9b87e1a01782204047c3b2104995c11ed..240faebe62f5cee4f61b3c36b5e8f653cfd6db8e 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h @@ -158,9 +158,6 @@ class DfsHloVisitorWithDefaultBase Status HandleBroadcast(HloInstructionPtr broadcast) override { return DefaultAction(broadcast); } - Status HandleBroadcastDimOne(HloInstructionPtr broadcastDimOne) override { - return DefaultAction(broadcastDimOne); - } Status HandlePad(HloInstructionPtr pad) override { return DefaultAction(pad); } diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc index b6a0903b0eeaa04d8bc1488378c148b2016c5d48..56e35e260469e383848db7bdbdbaacd9c9f9e56c 100644 --- a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc @@ -293,6 +293,12 @@ StatusOr ElementalIrEmitter::EmitIntegerUnaryOp( return operand_value; } } + case HloOpcode::kClz: { + auto is_zero_undef = ir_builder_->getFalse(); + return llvm_ir::EmitCallToIntrinsic( + llvm::Intrinsic::ctlz, {operand_value, is_zero_undef}, + {operand_value->getType()}, ir_builder_); + } case HloOpcode::kSign: { bool is_signed = primitive_util::IsSignedIntegralType(op->shape().element_type()); @@ -1334,6 +1340,7 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator( case HloOpcode::kAbs: case HloOpcode::kRoundNearestAfz: case HloOpcode::kCeil: + case HloOpcode::kClz: case HloOpcode::kConvert: case HloOpcode::kBitcastConvert: case HloOpcode::kCopy: diff --git a/tensorflow/compiler/xla/service/executable.cc b/tensorflow/compiler/xla/service/executable.cc index 471d2fd6cebcd7a00dfea4aca08da08af534b05f..8218b5f7c87a26e6eb4563e72fdb598eb2e87a99 100644 --- a/tensorflow/compiler/xla/service/executable.cc +++ b/tensorflow/compiler/xla/service/executable.cc @@ -29,18 +29,19 @@ using tensorflow::gtl::ArraySlice; namespace xla { -StatusOr>> -Executable::ExecuteOnStreams( +StatusOr> Executable::ExecuteOnStreams( ArraySlice run_options, ArraySlice> arguments) { TF_RET_CHECK(run_options.size() == arguments.size()); - std::vector> return_values(run_options.size()); + std::vector return_values; + return_values.reserve(run_options.size()); if (run_options.size() == 1) { - TF_ASSIGN_OR_RETURN(return_values[0], + TF_ASSIGN_OR_RETURN(auto rv, ExecuteOnStream(&run_options[0], arguments[0], /*hlo_execution_profile=*/nullptr)); + return_values.push_back(std::move(rv)); return std::move(return_values); } @@ -48,8 +49,9 @@ Executable::ExecuteOnStreams( // We cannot BlockHostUntilDone() on the already-launched executions in case // of error, since if the executions communicate, the initially launched // executions may never complete if not all executions are running. - TF_ASSIGN_OR_RETURN(return_values[i], + TF_ASSIGN_OR_RETURN(auto rv, ExecuteAsyncOnStream(&run_options[i], arguments[i])); + return_values.push_back(std::move(rv)); } for (const auto& options : run_options) { TF_RET_CHECK(options.stream() != nullptr); @@ -58,13 +60,13 @@ Executable::ExecuteOnStreams( return std::move(return_values); } -StatusOr> Executable::ExecuteOnStreamWrapper( +StatusOr Executable::ExecuteOnStreamWrapper( const ServiceExecutableRunOptions* run_options, ExecutionProfile* profile, ArraySlice arguments) { - perftools::gputools::Stream* stream = run_options->stream(); - std::unique_ptr timer; + se::Stream* stream = run_options->stream(); + std::unique_ptr timer; if (profile != nullptr) { - timer.reset(new perftools::gputools::Timer(stream->parent())); + timer.reset(new se::Timer(stream->parent())); stream->InitTimer(timer.get()).ThenStartTimer(timer.get()); } @@ -78,7 +80,7 @@ StatusOr> Executable::ExecuteOnStreamWrapper( &hlo_profile_index_map()) : nullptr; - StatusOr> return_value = + StatusOr return_value = ExecuteOnStream(run_options, arguments, profile_ptr.get()); TF_RETURN_IF_ERROR(return_value.status()); @@ -161,4 +163,24 @@ Status Executable::DumpSessionModule() { result); } +/* static */ Status Executable::DumpToDirectory(const string& directory_path, + string filename, + const HloSession& hlo_session) { + tensorflow::Env* env = tensorflow::Env::Default(); + if (!env->IsDirectory(directory_path).ok()) { + // NB! CreateDir does not work reliably with multiple XLA threads -- two + // threads can race to observe the absence of the dump directory and + // simultaneously try to create it, causing the "losing" thread to get a + // "directory already exists" error. + TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(directory_path)); + } + filename = SanitizeFileName(std::move(filename)); + string file_path = tensorflow::io::JoinPath(directory_path, filename); + string result; + TF_RET_CHECK( + tensorflow::SerializeToStringDeterministic(hlo_session, &result)); + return tensorflow::WriteStringToFile(tensorflow::Env::Default(), file_path, + result); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/executable.h b/tensorflow/compiler/xla/service/executable.h index a157235f8af6ea64a488510e427bbae502c46ca6..bdbe119120f233c8ee7e1c23de27c0021838c8f7 100644 --- a/tensorflow/compiler/xla/service/executable.h +++ b/tensorflow/compiler/xla/service/executable.h @@ -22,6 +22,7 @@ limitations under the License. #include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" #include "tensorflow/compiler/xla/service/computation_layout.h" #include "tensorflow/compiler/xla/service/device_memory_allocator.h" +#include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_execution_profile.h" #include "tensorflow/compiler/xla/service/hlo_graph_dumper.h" #include "tensorflow/compiler/xla/service/hlo_module.h" @@ -62,14 +63,14 @@ class Executable { // enabled. // // Returns a shaped buffer containing the result of the computation. - virtual StatusOr> ExecuteOnStream( + virtual StatusOr ExecuteOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments, HloExecutionProfile* hlo_execution_profile) = 0; // Same as ExecuteOnStream(), but this call is non-blocking and returns as // soon as all of the operations are enqueued for launch on the stream. - virtual StatusOr> ExecuteAsyncOnStream( + virtual StatusOr ExecuteAsyncOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments) = 0; @@ -77,7 +78,7 @@ class Executable { // streams. arguments[i] contains the arguments to the execution on // run_options[i]->stream() and the returned value is at index i of the // returned vector. - virtual StatusOr>> ExecuteOnStreams( + virtual StatusOr> ExecuteOnStreams( tensorflow::gtl::ArraySlice run_options, tensorflow::gtl::ArraySlice< @@ -90,14 +91,14 @@ class Executable { // has completed. virtual Status PopulateExecutionProfile( HloExecutionProfile* hlo_execution_profile, - perftools::gputools::StreamExecutor* executor) { + se::StreamExecutor* executor) { return Status::OK(); } // Convenience wrapper for calling Executable::ExecuteOnStream. Sets up a // timer for the execution, sets up HLO profiling if enabled, and fills in the // given ExecutionProfile if non-null. - StatusOr> ExecuteOnStreamWrapper( + StatusOr ExecuteOnStreamWrapper( const ServiceExecutableRunOptions* run_options, ExecutionProfile* profile, tensorflow::gtl::ArraySlice arguments); @@ -155,6 +156,10 @@ class Executable { static Status DumpToDirectory(const string& directory_path, string filename, const SessionModule& session_module); + // Dump hlo_session to directory_path/filename. + static Status DumpToDirectory(const string& directory_path, string filename, + const HloSession& hlo_session); + protected: mutable tensorflow::mutex mutex_; diff --git a/tensorflow/compiler/xla/service/gather_expander.cc b/tensorflow/compiler/xla/service/gather_expander.cc index 221ff7900f398166c193c495848a2afcfd4edc81..2d3e4b1fcdf6675955714cab262a8b2ca8ff4297 100644 --- a/tensorflow/compiler/xla/service/gather_expander.cc +++ b/tensorflow/compiler/xla/service/gather_expander.cc @@ -28,9 +28,15 @@ using tensorflow::gtl::ArraySlice; static StatusOr TransposeIndexVectorDimToLast( HloInstruction* gather_indices, int64 index_vector_dim) { const Shape& gather_indices_shape = gather_indices->shape(); + + if (gather_indices_shape.dimensions_size() == index_vector_dim) { + return gather_indices; + } + 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++) { @@ -42,55 +48,35 @@ static StatusOr TransposeIndexVectorDimToLast( 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; - } - - DCHECK_EQ(index_vector_dim, gather_indices_shape.dimensions_size()); - - std::vector result_shape_dims; - c_copy(gather_indices_shape.dimensions(), - std::back_inserter(result_shape_dims)); - result_shape_dims.push_back(1); - - return MakeReshapeHlo(result_shape_dims, gather_indices); -} - // 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. + // Transpose the non-index-vector dimensions to the front. TF_ASSIGN_OR_RETURN( - HloInstruction * descalarized_gather_indices, - DeScalarizeGatherIndices(gather_indices, index_vector_dim)); + HloInstruction * transposed_gather_indices, + TransposeIndexVectorDimToLast(gather_indices, index_vector_dim)); + bool indices_are_scalar = + index_vector_dim == gather_indices->shape().dimensions_size(); - // Transpose the non-index-vector dimensions to the front. - TF_ASSIGN_OR_RETURN(HloInstruction * transposed_gather_indices, - TransposeIndexVectorDimToLast(descalarized_gather_indices, - index_vector_dim)); + // The number of dimensions in gather_indices that are index dimensions. + const int64 index_dims_in_gather_indices = indices_are_scalar ? 0 : 1; // 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(transposed_gather_indices, - {1, shape.dimensions(0)}); + if (shape.dimensions_size() == index_dims_in_gather_indices) { + return PrependDegenerateDims(transposed_gather_indices, 1); } else { - return CollapseFirstNDims(transposed_gather_indices, - shape.dimensions_size() - 1); + // Collapse all but the dimensions (0 or 1) in gather_indices containing the + // index vectors. + return CollapseFirstNDims( + transposed_gather_indices, + shape.dimensions_size() - index_dims_in_gather_indices); } } @@ -112,11 +98,7 @@ static StatusOr AdjustGatherDimsInAccumulator( // dynamic-slice. In that case, there is a leading degenerate gather // dimension that we added to make this special case play well with the // general while loop which we need to remove now. - CHECK_EQ(accumulator->shape().dimensions(0), 1); - ArraySlice reshaped_dim_sizes = - AsInt64Slice(accumulator->shape().dimensions()); - reshaped_dim_sizes.remove_prefix(1); - return MakeReshapeHlo(reshaped_dim_sizes, accumulator); + return ElideDegenerateDims(accumulator, {0}); } return ExpandFirstDimIntoNDims(accumulator, output_gather_dim_bounds); @@ -161,50 +143,73 @@ static StatusOr ExpandIndexVectorIntoOperandSpace( static StatusOr> GatherLoopBody( const HloInstruction& gather, HloInstruction* induction_var, const std::vector& incoming_loop_state) { + const GatherDimensionNumbers& dim_numbers = gather.gather_dimension_numbers(); 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); + bool has_scalar_indices = gather_indices->shape().dimensions_size() == 1; + CHECK_EQ(has_scalar_indices, + dim_numbers.index_vector_dim() == + gather.operand(1)->shape().dimensions_size()); 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})); + HloInstruction* index_vector; - TF_ASSIGN_OR_RETURN(HloInstruction * index_vector, - ElideDegenerateDims(index_vector_2d, {0})); + if (has_scalar_indices) { + // In this case gather_indices has rank 1 and induction_var_as_vector (of + // shape {1}) is an index into this rank 1 tensor. + TF_ASSIGN_OR_RETURN( + index_vector, + MakeDynamicSliceHlo(gather_indices, induction_var_as_vector, {1})); + } else { + // In this case gather_indices has rank 2 and induction_var_as_vector (of + // shape {1}) is an index into just the first dimension of this rank 2 + // tensor. + TF_ASSIGN_OR_RETURN( + HloInstruction * index_into_gather_indices, + PadVectorWithZeros(induction_var_as_vector, + /*zeros_to_prepend=*/0, /*zeros_to_append=*/1)); + + int64 index_vector_size = gather_indices->shape().dimensions(1); + TF_ASSIGN_OR_RETURN( + HloInstruction * index_vector_2d, + MakeDynamicSliceHlo(gather_indices, index_into_gather_indices, + {1, index_vector_size})); + + TF_ASSIGN_OR_RETURN(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_start, + ExpandIndexVectorIntoOperandSpace(index_vector, dim_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_with_dims_elided, + ElideDegenerateDims(gathered_slice, + AsInt64Slice(dim_numbers.elided_window_dims()))); + TF_ASSIGN_OR_RETURN( HloInstruction * gathered_slice_for_update, - ExpandFirstDimIntoNDims(gathered_slice, - {1, gathered_slice->shape().dimensions(0)})); + PrependDegenerateDims(gathered_slice_with_dims_elided, 1)); 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())); + /*zeros_to_append=*/ + gathered_slice_with_dims_elided->shape().dimensions_size())); TF_ASSIGN_OR_RETURN( HloInstruction * updated_accumulator, @@ -220,26 +225,20 @@ static StatusOr> GatherLoopBody( static StatusOr CreateGatherLoopAccumulatorInitValue( HloComputation* computation, PrimitiveType element_type, - ArraySlice window_bounds, int64 gather_loop_trip_count) { + ArraySlice window_bounds, int64 gather_loop_trip_count, + const GatherDimensionNumbers& dim_numbers) { 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)); + for (int64 i = 0; i < window_bounds.size(); i++) { + if (!c_binary_search(dim_numbers.elided_window_dims(), i)) { + accumulator_state_shape_dims.push_back(window_bounds[i]); + } + } 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. @@ -338,7 +337,8 @@ StatusOr GatherExpander::ExpandGather( HloInstruction * accumulator_init, CreateGatherLoopAccumulatorInitValue( computation, output_shape.element_type(), - gather_instr->gather_window_bounds(), gather_loop_trip_count)); + gather_instr->gather_window_bounds(), gather_loop_trip_count, + gather_instr->gather_dimension_numbers())); StatusOr> gather_loop_result_or_error = WhileUtil::MakeCountedLoop( @@ -353,14 +353,10 @@ StatusOr GatherExpander::ExpandGather( 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, - AdjustGatherDimsInAccumulator(gather_indices->shape(), - accumulator_with_window_dims_elided, + AdjustGatherDimsInAccumulator(gather_indices->shape(), accumulator_result, dim_numbers.index_vector_dim())); return PermuteGatherAndWindowDims( diff --git a/tensorflow/compiler/xla/service/gather_expander_test.cc b/tensorflow/compiler/xla/service/gather_expander_test.cc index ba41ee8428cbe7132103df24d552565a8dc2f9f6..1c72ca066502eb549bf8638cdf0b7827b06f92d7 100644 --- a/tensorflow/compiler/xla/service/gather_expander_test.cc +++ b/tensorflow/compiler/xla/service/gather_expander_test.cc @@ -47,5 +47,62 @@ ENTRY main { "indices are not supported.")); } +TEST(GatherExpanderTest, AvoidDegenerateDims) { + 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} +} +)"; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + tools::Parse(hlo_text)); + TF_ASSERT_OK_AND_ASSIGN(bool changed, GatherExpander{}.Run(module.get())); + ASSERT_TRUE(changed); + + HloInstruction* while_instr = nullptr; + for (auto* instr : module->entry_computation()->instructions()) { + if (instr->opcode() == HloOpcode::kWhile) { + ASSERT_EQ(while_instr, nullptr) + << "Expected exactly one while instruction in the entry computation " + "after gather expansion"; + while_instr = instr; + } + } + + ASSERT_NE(while_instr, nullptr) + << "Expected exactly one while instruction in the entry computation " + "after gather expansion"; + + // We want to avoid create while loop with shapes that have degenerate + // dimensions for TF gather. In this case we expect the loop state to be of + // the shape (sNN[], s32[3,3]{1,0}, s32[2]{0}, s32[2,3]{1,0}). The leading + // sNN is an implementation detail from WhileUtil::MakeCountedLoop so we don't + // check it here (though in theory the form of the while loop state is itself + // an implementation detail from WhileUtil::MakeCountedLoop). + + const Shape& while_shape = while_instr->shape(); + ASSERT_TRUE(ShapeUtil::IsTuple(while_shape)); + ASSERT_EQ(ShapeUtil::TupleElementCount(while_shape), 4); + + EXPECT_TRUE(ShapeUtil::SameDimensions( + ShapeUtil::MakeShape(S32, {3, 3}), + ShapeUtil::GetTupleElementShape(while_shape, 1))); + + EXPECT_TRUE(ShapeUtil::SameDimensions( + ShapeUtil::MakeShape(S32, {2}), + ShapeUtil::GetTupleElementShape(while_shape, 2))); + + EXPECT_TRUE(ShapeUtil::SameDimensions( + ShapeUtil::MakeShape(S32, {2, 3}), + ShapeUtil::GetTupleElementShape(while_shape, 3))); +} } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/generic_transfer_manager.cc b/tensorflow/compiler/xla/service/generic_transfer_manager.cc index a99e2b7794a399047fb5a77a140bd333214e3f23..ddb687314ee8221ba9282f230db498b3a5d23499 100644 --- a/tensorflow/compiler/xla/service/generic_transfer_manager.cc +++ b/tensorflow/compiler/xla/service/generic_transfer_manager.cc @@ -32,8 +32,6 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" -namespace se = ::perftools::gputools; - namespace xla { GenericTransferManager::GenericTransferManager(se::Platform::Id platform_id, @@ -45,9 +43,9 @@ se::Platform::Id GenericTransferManager::PlatformId() const { } Status GenericTransferManager::WriteSingleTupleIndexTable( - perftools::gputools::StreamExecutor* executor, + se::StreamExecutor* executor, tensorflow::gtl::ArraySlice elements, - const Shape& shape, perftools::gputools::DeviceMemoryBase* region) { + const Shape& shape, se::DeviceMemoryBase* region) { TF_RET_CHECK(elements.size() == ShapeUtil::TupleElementCount(shape)); std::vector element_pointers; @@ -144,20 +142,19 @@ Status GenericTransferManager::TransferLiteralToInfeed( } Status GenericTransferManager::TransferBufferToInfeed( - perftools::gputools::StreamExecutor* executor, int64 size, - const void* source) { + se::StreamExecutor* executor, int64 size, const void* source) { return Unimplemented("Generic transfer to Infeed"); } Status GenericTransferManager::TransferLiteralFromOutfeed( - perftools::gputools::StreamExecutor* executor, const Shape& literal_shape, + se::StreamExecutor* executor, const Shape& literal_shape, Literal* literal) { return Unimplemented( "Outfeed is not supported on this platform (b/30467474)"); } Status GenericTransferManager::ResetDevices( - tensorflow::gtl::ArraySlice + tensorflow::gtl::ArraySlice /*executors*/) { return Unimplemented( "Device reset is not yet supported on this platform (b/30481585)"); diff --git a/tensorflow/compiler/xla/service/generic_transfer_manager.h b/tensorflow/compiler/xla/service/generic_transfer_manager.h index 63a7c820cf4e5fbbdf870086a4fb5316ac50d10b..0579099de40ba3e43f69e4e6474b56691064c692 100644 --- a/tensorflow/compiler/xla/service/generic_transfer_manager.h +++ b/tensorflow/compiler/xla/service/generic_transfer_manager.h @@ -36,46 +36,41 @@ namespace xla { // infeed. class GenericTransferManager : public TransferManager { public: - GenericTransferManager(perftools::gputools::Platform::Id platform_id, - size_t pointer_size); + GenericTransferManager(se::Platform::Id platform_id, size_t pointer_size); ~GenericTransferManager() override {} - perftools::gputools::Platform::Id PlatformId() const override; + se::Platform::Id PlatformId() const override; StatusOr> TransferLiteralFromDevice( - perftools::gputools::StreamExecutor* executor, - const ShapedBuffer& device_buffer) override; + se::StreamExecutor* executor, const ShapedBuffer& device_buffer) override; - Status TransferLiteralToDevice(perftools::gputools::StreamExecutor* executor, + Status TransferLiteralToDevice(se::StreamExecutor* executor, const Literal& literal, const ShapedBuffer& device_buffer) override; - Status TransferLiteralToInfeed(perftools::gputools::StreamExecutor* executor, + Status TransferLiteralToInfeed(se::StreamExecutor* executor, const Literal& literal) override; - Status TransferLiteralFromOutfeed( - perftools::gputools::StreamExecutor* executor, const Shape& literal_shape, - Literal* literal) override; + Status TransferLiteralFromOutfeed(se::StreamExecutor* executor, + const Shape& literal_shape, + Literal* literal) override; Status ResetDevices( - tensorflow::gtl::ArraySlice - executors) override; + tensorflow::gtl::ArraySlice executors) override; int64 GetByteSizeRequirement(const Shape& shape) const override; protected: - Status TransferBufferToInfeed(perftools::gputools::StreamExecutor* executor, - int64 size, const void* source) override; + Status TransferBufferToInfeed(se::StreamExecutor* executor, int64 size, + const void* source) override; Status WriteSingleTupleIndexTable( - perftools::gputools::StreamExecutor* executor, - tensorflow::gtl::ArraySlice - elements, - const Shape& shape, - perftools::gputools::DeviceMemoryBase* region) override; + se::StreamExecutor* executor, + tensorflow::gtl::ArraySlice elements, + const Shape& shape, se::DeviceMemoryBase* region) override; private: // The platform this transfer manager targets. - const perftools::gputools::Platform::Id platform_id_; + const se::Platform::Id platform_id_; // The size in bytes of pointers on this platform. const size_t pointer_size_; diff --git a/tensorflow/compiler/xla/service/gpu/buffer_allocations.cc b/tensorflow/compiler/xla/service/gpu/buffer_allocations.cc index 2029c303d47e9a62135b003c3bd9be6f8b3438d4..837f05244f7a8c71483cc30eeac9e1c219e6bbd2 100644 --- a/tensorflow/compiler/xla/service/gpu/buffer_allocations.cc +++ b/tensorflow/compiler/xla/service/gpu/buffer_allocations.cc @@ -28,8 +28,6 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" -namespace se = ::perftools::gputools; - namespace xla { namespace gpu { diff --git a/tensorflow/compiler/xla/service/gpu/buffer_allocations.h b/tensorflow/compiler/xla/service/gpu/buffer_allocations.h index ea7f0eb3745f2e0e0bfd88c3dca79d6ad25884ed..c2fc35be4ca4bc6df85ee21fb6b564bfd6de3ec0 100644 --- a/tensorflow/compiler/xla/service/gpu/buffer_allocations.h +++ b/tensorflow/compiler/xla/service/gpu/buffer_allocations.h @@ -41,7 +41,7 @@ class BufferAllocations { // user-specified result buffers) to the given buffer index. The builder // will skip allocating buffers for registered buffer indices. void RegisterBuffer(BufferAllocation::Index index, - perftools::gputools::DeviceMemoryBase address); + se::DeviceMemoryBase address); // Builds a BufferAllocations object from the given buffer assignment. // `memory_allocator` is what this function uses to allocate device memory. @@ -52,8 +52,7 @@ class BufferAllocations { DeviceMemoryAllocator* memory_allocator); private: - std::map - registered_buffers_; + std::map registered_buffers_; }; BufferAllocations(const BufferAllocations&) = delete; @@ -65,22 +64,20 @@ class BufferAllocations { // Returns the device address of buffer `buffer_index`. `buffer_index` must be // a valid index, i.e., in [0, buffer_count). This function returns null if // `buffer_index` is not assigned to a buffer address. - perftools::gputools::DeviceMemoryBase GetDeviceAddress( + se::DeviceMemoryBase GetDeviceAddress( BufferAllocation::Index buffer_index) const; // Same as above, but also adjusts the returned address for the offset and // size contained in the given slice. - perftools::gputools::DeviceMemoryBase GetDeviceAddress( + se::DeviceMemoryBase GetDeviceAddress( const BufferAllocation::Slice& buffer_slice) const; - perftools::gputools::DeviceMemoryBase GetTempBufferBase() const { - return temp_buffer_base_; - } + se::DeviceMemoryBase GetTempBufferBase() const { return temp_buffer_base_; } // Tears down all buffers allocated by this object that are not in // `live_addresses`. tensorflow::Status TearDown( - const std::set& live_addresses, + const std::set& live_addresses, const BufferAssignment& buffer_assignment); private: @@ -92,15 +89,15 @@ class BufferAllocations { // Sets the device address of buffer `buffer_index`. void SetBuffer(BufferAllocation::Index buffer_index, - perftools::gputools::DeviceMemoryBase buffer); + se::DeviceMemoryBase buffer); // An array of device pointers that stores the address of each buffer // indexed by Index. Each element can point to a temporary buffer, an // input buffer, or nullptr if no buffer is needed for that Index. - std::vector buffers_; + std::vector buffers_; // The base address of the memory block that contains all temporary buffers. - perftools::gputools::DeviceMemoryBase temp_buffer_base_; + se::DeviceMemoryBase temp_buffer_base_; int device_ordinal_; diff --git a/tensorflow/compiler/xla/service/gpu/conditional_thunk.cc b/tensorflow/compiler/xla/service/gpu/conditional_thunk.cc index 790ca535b11ee47724ef6227de40726d940d6153..dce8de2e301ecfaa4674b8be48b8c02bfabf3f4b 100644 --- a/tensorflow/compiler/xla/service/gpu/conditional_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/conditional_thunk.cc @@ -42,11 +42,10 @@ Status ConditionalThunk::Initialize(const GpuExecutable& executable) { } Status ConditionalThunk::ExecuteOnStream( - const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) { + const BufferAllocations& buffer_allocations, se::Stream* stream) { // Copy the predicate value from device. bool predicate; - perftools::gputools::DeviceMemoryBase predicate_address = + se::DeviceMemoryBase predicate_address = buffer_allocations.GetDeviceAddress(predicate_buffer_index_); stream->ThenMemcpy(&predicate, predicate_address, sizeof(bool)); diff --git a/tensorflow/compiler/xla/service/gpu/conditional_thunk.h b/tensorflow/compiler/xla/service/gpu/conditional_thunk.h index 7725c46a3b4b51af34a4dd977885353ff32c21f6..e40872688fdad24d24db5f7cebb3206c77652dce 100644 --- a/tensorflow/compiler/xla/service/gpu/conditional_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/conditional_thunk.h @@ -49,7 +49,7 @@ class ConditionalThunk : public Thunk { Status Initialize(const GpuExecutable& executable) override; Status ExecuteOnStream(const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) override; + se::Stream* stream) override; private: BufferAllocation::Slice predicate_buffer_index_; diff --git a/tensorflow/compiler/xla/service/gpu/convolution_thunk.cc b/tensorflow/compiler/xla/service/gpu/convolution_thunk.cc index 461747b699b542ae0c8735aea34cc9e57c1fb387..64d3b84b8c73d82800270aebcebf7f7a8fec5fe4 100644 --- a/tensorflow/compiler/xla/service/gpu/convolution_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/convolution_thunk.cc @@ -25,8 +25,6 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" -namespace se = ::perftools::gputools; - namespace xla { namespace gpu { diff --git a/tensorflow/compiler/xla/service/gpu/convolution_thunk.h b/tensorflow/compiler/xla/service/gpu/convolution_thunk.h index 900d9cb6243088b56a1825fb3ab8c06cf8d74726..6d845025b1aef2b0a5f147401b6db0598ba94d6d 100644 --- a/tensorflow/compiler/xla/service/gpu/convolution_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/convolution_thunk.h @@ -66,23 +66,21 @@ class ConvolutionThunk : public Thunk { // Does the convolution for the thunk on "stream". Status ExecuteOnStream(const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) override; + se::Stream* stream) override; private: class ScratchAllocator; - Status Convolve( - 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 perftools::gputools::dnn::AlgorithmConfig& algorithm_config, - perftools::gputools::Stream* stream, ScratchAllocator* scratch_allocator, - perftools::gputools::dnn::ProfileResult* profile_result); + Status Convolve(const se::dnn::BatchDescriptor& input_descriptor, + se::DeviceMemory input_data, + const se::dnn::FilterDescriptor& filter_descriptor, + se::DeviceMemory filter_data, + const se::dnn::BatchDescriptor& output_descriptor, + se::DeviceMemory output_data, + const se::dnn::ConvolutionDescriptor& convolution_descriptor, + const se::dnn::AlgorithmConfig& algorithm_config, + se::Stream* stream, ScratchAllocator* scratch_allocator, + se::dnn::ProfileResult* profile_result); const CudnnConvKind convolution_kind_; diff --git a/tensorflow/compiler/xla/service/gpu/copy_thunk.cc b/tensorflow/compiler/xla/service/gpu/copy_thunk.cc index f4498663b1c039b3175376baf8f27c4ecec678ec..bf912fbd14de5874062a79db28186ab233575233 100644 --- a/tensorflow/compiler/xla/service/gpu/copy_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/copy_thunk.cc @@ -30,9 +30,8 @@ HostToDeviceCopyThunk::HostToDeviceCopyThunk( mem_size_(mem_size) {} tensorflow::Status HostToDeviceCopyThunk::ExecuteOnStream( - const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) { - perftools::gputools::DeviceMemoryBase destination_data = + const BufferAllocations& buffer_allocations, se::Stream* stream) { + se::DeviceMemoryBase destination_data = buffer_allocations.GetDeviceAddress(destination_buffer_); stream->ThenMemcpy(&destination_data, source_address_, mem_size_); return tensorflow::Status::OK(); @@ -48,11 +47,10 @@ DeviceToDeviceCopyThunk::DeviceToDeviceCopyThunk( mem_size_(mem_size) {} tensorflow::Status DeviceToDeviceCopyThunk::ExecuteOnStream( - const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) { - perftools::gputools::DeviceMemoryBase destination_data = + const BufferAllocations& buffer_allocations, se::Stream* stream) { + se::DeviceMemoryBase destination_data = buffer_allocations.GetDeviceAddress(destination_buffer_); - perftools::gputools::DeviceMemoryBase source_data = + se::DeviceMemoryBase source_data = buffer_allocations.GetDeviceAddress(source_buffer_); stream->ThenMemcpy(&destination_data, source_data, mem_size_); return tensorflow::Status::OK(); diff --git a/tensorflow/compiler/xla/service/gpu/copy_thunk.h b/tensorflow/compiler/xla/service/gpu/copy_thunk.h index e2783fd255239d31edc89701ea208f33ebb8d3fb..2e7eb5f3445bc9294fa455ef31fb816cdba4726c 100644 --- a/tensorflow/compiler/xla/service/gpu/copy_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/copy_thunk.h @@ -40,8 +40,7 @@ class HostToDeviceCopyThunk : public Thunk { HostToDeviceCopyThunk& operator=(const HostToDeviceCopyThunk&) = delete; tensorflow::Status ExecuteOnStream( - const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) override; + const BufferAllocations& buffer_allocations, se::Stream* stream) override; private: const void* source_address_; @@ -64,8 +63,7 @@ class DeviceToDeviceCopyThunk : public Thunk { DeviceToDeviceCopyThunk& operator=(const DeviceToDeviceCopyThunk&) = delete; tensorflow::Status ExecuteOnStream( - const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) override; + const BufferAllocations& buffer_allocations, se::Stream* stream) override; private: const BufferAllocation::Slice source_buffer_; diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_batchnorm_thunk.cc b/tensorflow/compiler/xla/service/gpu/cudnn_batchnorm_thunk.cc index 58d9c8caff31e878487fbef01afce566e6187fd9..68099fd63847ef9993f9bc7ac0e28b2939631b35 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_batchnorm_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/cudnn_batchnorm_thunk.cc @@ -28,7 +28,6 @@ limitations under the License. namespace xla { namespace gpu { -namespace se = ::perftools::gputools; namespace dnn = se::dnn; static std::pair> version = + auto version = stream_exec->AsDnn()->GetVersion(); - if (version.ok() && std::get<0>(version.ValueOrDie()) >= 7) { + if (version.ok() && version.ValueOrDie().major_version() >= 7) { return true; } diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.h b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.h index 516210ec2e500cf03774d27408300ac3346e7b4f..bc5d1ce94afd2075a006899f0f6bcf64352e5e99 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.h +++ b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.h @@ -33,9 +33,8 @@ class CudnnConvolutionAlgorithmPicker : public HloPassInterface { // 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) + CudnnConvolutionAlgorithmPicker(se::StreamExecutor* stream_exec, + DeviceMemoryAllocator* allocator) : stream_exec_(stream_exec), allocator_(allocator) {} tensorflow::StringPiece name() const override { @@ -52,7 +51,7 @@ class CudnnConvolutionAlgorithmPicker : public HloPassInterface { const Shape& output_shape, const Window& window, const ConvolutionDimensionNumbers& dnums, HloInstruction* instr); - perftools::gputools::StreamExecutor* stream_exec_; // never null + se::StreamExecutor* stream_exec_; // never null DeviceMemoryAllocator* allocator_; // may be null }; diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.cc b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.cc index e4ae839e1dd4cb3a744a3f6a3329cabdaeb3f38d..10b4c3de89989c52cfea5273c3d5b0beef76abd2 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.cc +++ b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.cc @@ -22,8 +22,6 @@ namespace xla { namespace gpu { namespace { -namespace se = ::perftools::gputools; - using se::DeviceMemory; using se::DeviceMemoryBase; using se::Stream; @@ -215,14 +213,12 @@ string CudnnConvKindToString(CudnnConvKind kind) { 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 Shape& output_shape, se::DeviceMemoryBase input_buf, + se::DeviceMemoryBase filter_buf, se::DeviceMemoryBase output_buf, + se::DeviceMemoryBase scratch_buf, const Window& window, const ConvolutionDimensionNumbers& dnums, - perftools::gputools::dnn::AlgorithmConfig algorithm, - perftools::gputools::Stream* stream, - perftools::gputools::dnn::ProfileResult* profile_result) { + se::dnn::AlgorithmConfig algorithm, se::Stream* stream, + se::dnn::ProfileResult* profile_result) { ScratchBufAllocator scratch_allocator(scratch_buf); return RunCudnnConvolution(kind, input_shape, filter_shape, output_shape, input_buf, filter_buf, output_buf, @@ -232,14 +228,12 @@ Status RunCudnnConvolution( 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) { + const Shape& output_shape, se::DeviceMemoryBase input_buf, + se::DeviceMemoryBase filter_buf, se::DeviceMemoryBase output_buf, + se::ScratchAllocator* scratch_allocator, const Window& window, + const ConvolutionDimensionNumbers& dnums, + se::dnn::AlgorithmConfig algorithm, se::Stream* stream, + se::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); diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.h b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.h index 3dbfa2730da359d3c7937140508017c4a7b02d6c..944e4ac686d45408b08ff1faa321510c1c8920ba 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.h +++ b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_runner.h @@ -72,25 +72,21 @@ string CudnnConvKindToString(CudnnConvKind kind); // 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 Shape& output_shape, se::DeviceMemoryBase input_buf, + se::DeviceMemoryBase filter_buf, se::DeviceMemoryBase output_buf, + se::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); + se::dnn::AlgorithmConfig algorithm, se::Stream* stream, + se::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); + const Shape& output_shape, se::DeviceMemoryBase input_buf, + se::DeviceMemoryBase filter_buf, se::DeviceMemoryBase output_buf, + se::ScratchAllocator* scratch_allocator, const Window& window, + const ConvolutionDimensionNumbers& dnums, + se::dnn::AlgorithmConfig algorithm, se::Stream* stream, + se::dnn::ProfileResult* profile_result = nullptr); } // namespace gpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/fft_thunk.cc b/tensorflow/compiler/xla/service/gpu/fft_thunk.cc index 66931bdc8b1030b2b2e7731ce6327c1e908d4ee6..cc747addbd152eb82b0b2ef92b8653fc861f97be 100644 --- a/tensorflow/compiler/xla/service/gpu/fft_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/fft_thunk.cc @@ -24,8 +24,6 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" -namespace se = ::perftools::gputools; - namespace xla { namespace gpu { diff --git a/tensorflow/compiler/xla/service/gpu/fft_thunk.h b/tensorflow/compiler/xla/service/gpu/fft_thunk.h index 52fb8c376d7acea0f15aaa865c23fa2382717338..24b1dca99865fe21d0ca3af91e0d169f7b74a78a 100644 --- a/tensorflow/compiler/xla/service/gpu/fft_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/fft_thunk.h @@ -34,24 +34,24 @@ namespace gpu { // released on destruction. // // Not thread-safe in that AllocateBytes, destructor are not locked. -class FftScratchAllocator : public perftools::gputools::ScratchAllocator { +class FftScratchAllocator : public se::ScratchAllocator { public: FftScratchAllocator(int device_ordinal, DeviceMemoryAllocator* memory_allocator); ~FftScratchAllocator() override; - int64 GetMemoryLimitInBytes(perftools::gputools::Stream* stream) override; + int64 GetMemoryLimitInBytes(se::Stream* stream) override; int64 TotalAllocatedBytes() { return total_allocated_bytes_; } - perftools::gputools::port::StatusOr> - AllocateBytes(perftools::gputools::Stream* stream, int64 byte_size) override; + se::port::StatusOr> AllocateBytes( + se::Stream* stream, int64 byte_size) override; private: const int device_ordinal_; DeviceMemoryAllocator* memory_allocator_; - std::vector allocated_buffers_; + std::vector allocated_buffers_; int64 total_allocated_bytes_ = 0; }; @@ -74,16 +74,15 @@ class FftThunk : public Thunk { // Does the FFT for the thunk on "stream". tensorflow::Status ExecuteOnStream( - const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) override; + const BufferAllocations& buffer_allocations, se::Stream* stream) override; private: - const perftools::gputools::fft::Type fft_type_; + const se::fft::Type fft_type_; const std::vector fft_length_; float scale_factor_; - std::unique_ptr fft_plan_; + std::unique_ptr fft_plan_; const BufferAllocation::Slice input_buffer_; const BufferAllocation::Slice output_buffer_; diff --git a/tensorflow/compiler/xla/service/gpu/for_thunk.cc b/tensorflow/compiler/xla/service/gpu/for_thunk.cc index 283d21ca222a236a69e4bab1b6504665d4d1cdd3..6e6966df3987eef29b2122c3ef8f11b7cd0bfe14 100644 --- a/tensorflow/compiler/xla/service/gpu/for_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/for_thunk.cc @@ -36,8 +36,7 @@ tensorflow::Status ForThunk::Initialize(const GpuExecutable& executable) { } tensorflow::Status ForThunk::ExecuteOnStream( - const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) { + const BufferAllocations& buffer_allocations, se::Stream* stream) { for (int64 i = 0; i < loop_limit_; ++i) { // Invoke loop body thunk sequence. TF_RETURN_IF_ERROR( diff --git a/tensorflow/compiler/xla/service/gpu/for_thunk.h b/tensorflow/compiler/xla/service/gpu/for_thunk.h index 832494d17e9c4e1d9e92e18ef331df1cf3689024..c78d1c50686297aea8235af928aba562697f49bc 100644 --- a/tensorflow/compiler/xla/service/gpu/for_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/for_thunk.h @@ -38,8 +38,7 @@ class ForThunk : public Thunk { tensorflow::Status Initialize(const GpuExecutable& executable) override; tensorflow::Status ExecuteOnStream( - const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) override; + const BufferAllocations& buffer_allocations, se::Stream* stream) override; private: const int64 loop_limit_; diff --git a/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc b/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc index 38668ff455a44c7ef99b57b750f1a3b18a90bd2c..0ec12f52d8b398218ec370fc74bfdf6f97f43893 100644 --- a/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc @@ -22,8 +22,6 @@ limitations under the License. #include "tensorflow/core/platform/stream_executor_no_cuda.h" #include "tensorflow/core/platform/types.h" -namespace se = ::perftools::gputools; - namespace xla { namespace gpu { diff --git a/tensorflow/compiler/xla/service/gpu/gemm_thunk.h b/tensorflow/compiler/xla/service/gpu/gemm_thunk.h index df3edcefef898d465cd5ddc53e5d06a966a31f88..a18f425bc38fd3fbbb345901514c4ac16dbe97ec 100644 --- a/tensorflow/compiler/xla/service/gpu/gemm_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/gemm_thunk.h @@ -50,14 +50,12 @@ class GemmThunk : public Thunk { // Does the gemm operation for the thunk on "stream", which must be non-null. tensorflow::Status ExecuteOnStream( - const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) override; + const BufferAllocations& buffer_allocations, se::Stream* stream) override; // Returns true if we'll perform autotuning if run on the given stream. 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* stream) override { + bool ShouldHaltAllActivityBeforeRunning(se::Stream* stream) override { return autotune_results_.count( stream->parent()->GetDeviceDescription().name()) != 0; } @@ -79,8 +77,7 @@ class GemmThunk : public Thunk { // results. The map's value is the best algorithm we've found for this thunk // on this device, or an error if none of the algorithms worked and we should // use the regular gemm without an algorithm. - std::unordered_map> + std::unordered_map> autotune_results_; }; diff --git a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc index 07be2a0cf90c326af6e41764e79950db546e43e4..30bfc9351a5273b4cf854e269c5f576f6dd1bef7 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc @@ -91,8 +91,6 @@ limitations under the License. #include "tensorflow/core/platform/tracing.h" #include "tensorflow/stream_executor/cuda/cuda_diagnostics.h" -namespace se = ::perftools::gputools; - namespace xla { namespace gpu { @@ -779,9 +777,9 @@ se::Platform::Id GpuCompiler::PlatformId() const { } // namespace xla static bool InitModule() { - xla::Compiler::RegisterCompilerFactory(se::cuda::kCudaPlatformId, []() { - return xla::MakeUnique(); - }); + xla::Compiler::RegisterCompilerFactory( + stream_executor::cuda::kCudaPlatformId, + []() { return xla::MakeUnique(); }); return true; } static bool module_initialized = InitModule(); diff --git a/tensorflow/compiler/xla/service/gpu/gpu_compiler.h b/tensorflow/compiler/xla/service/gpu/gpu_compiler.h index c352d4d8462fadb266c55ad437de998e86a6528e..f3b02ae5d8867bdf1d970e809bff95a15d9f54d2 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_compiler.h +++ b/tensorflow/compiler/xla/service/gpu/gpu_compiler.h @@ -45,25 +45,23 @@ class GpuCompiler : public LLVMCompiler { // Bring in // StatusOr>> Compile( // std::vector> modules, - // std::vector> + // std::vector> // stream_execs) using LLVMCompiler::Compile; StatusOr> RunHloPasses( - std::unique_ptr module, - perftools::gputools::StreamExecutor* stream_exec, + std::unique_ptr module, se::StreamExecutor* stream_exec, DeviceMemoryAllocator* device_allocator) override; StatusOr> RunBackend( - std::unique_ptr module, - perftools::gputools::StreamExecutor* stream_exec, + std::unique_ptr module, se::StreamExecutor* stream_exec, DeviceMemoryAllocator* device_allocator) override; StatusOr>> CompileAheadOfTime(std::vector> module, AotCompilationOptions const& options) override; - perftools::gputools::Platform::Id PlatformId() const override; + se::Platform::Id PlatformId() const override; HloCostAnalysis::ShapeSizeFunction ShapeSizeBytesFunction() const override { // Capture just the pointer size, not the entire GpuCompiler object. diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc index 28f93447953b90d8a7fa4386e2355066c0405aec..62ce15bc59d34e3bd1768313f2894871bf97ef21 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc @@ -34,8 +34,6 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" -namespace se = ::perftools::gputools; - namespace xla { namespace gpu { namespace { @@ -252,7 +250,7 @@ Status GpuExecutable::ExecuteThunks( return Status::OK(); } -StatusOr> GpuExecutable::ExecuteOnStream( +StatusOr GpuExecutable::ExecuteOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments, HloExecutionProfile* hlo_execution_profile) { @@ -299,13 +297,13 @@ StatusOr> GpuExecutable::ExecuteOnStream( HloInstruction* root = hlo_module_->entry_computation()->root_instruction(); auto device_ordinal = executor->device_ordinal(); - auto shaped_buffer = MakeUnique( - root->shape(), root->shape(), executor->platform(), device_ordinal); + auto shaped_buffer = ShapedBuffer(root->shape(), root->shape(), + executor->platform(), device_ordinal); // Copy DeviceMemoryBase values which contain the array(s) of the result into // the respective location in ShapedBuffer. std::set buffers_in_result; - TF_RETURN_IF_ERROR(shaped_buffer->buffers().ForEachMutableElementWithStatus( + TF_RETURN_IF_ERROR(shaped_buffer.buffers().ForEachMutableElementWithStatus( [&buffer_allocations, &buffers_in_result, &shaped_buffer, this]( const ShapeIndex& index, se::DeviceMemoryBase* device_memory) { const auto& sources = this->GetRootPointsToSet().element(index); @@ -324,7 +322,7 @@ StatusOr> GpuExecutable::ExecuteOnStream( this->assignment_->GetUniqueSlice(src_hlo, sources[0]->index())); CHECK(!slice.allocation()->is_entry_computation_parameter()); - perftools::gputools::DeviceMemoryBase src_base = + se::DeviceMemoryBase src_base = buffer_allocations->GetDeviceAddress(slice.index()); CHECK(!src_base.is_null() || src_base.size() == 0); *device_memory = src_base; @@ -337,7 +335,7 @@ StatusOr> GpuExecutable::ExecuteOnStream( return std::move(shaped_buffer); } -StatusOr> GpuExecutable::ExecuteAsyncOnStream( +StatusOr GpuExecutable::ExecuteAsyncOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments) { // TODO(b/30671675): Implement asynchronous execution mode. diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.h b/tensorflow/compiler/xla/service/gpu/gpu_executable.h index dcb3991f41a31db84d8e9e555ae7d13c3ac84b97..361bc30b2f367c23051e06a1169f9c15c19f4a6a 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_executable.h +++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.h @@ -74,12 +74,12 @@ class GpuExecutable : public Executable { // ExecuteOnStream will fail if the compute capability of the stream doesn't // match the compute capability passed to this object's constructor. - StatusOr> ExecuteOnStream( + StatusOr ExecuteOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments, HloExecutionProfile* hlo_execution_profile) override; - StatusOr> ExecuteAsyncOnStream( + StatusOr ExecuteAsyncOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments) override; diff --git a/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.cc b/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.cc index af9897769fda371e47af06c19abce9a06015e094..f13727ca9b6954f6be9b9277018fcc64ee326954 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.cc @@ -33,8 +33,6 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" -namespace se = ::perftools::gputools; - namespace xla { // TODO(b/30467474) Once GPU infeed implementation settles, consider @@ -153,8 +151,8 @@ static std::unique_ptr CreateGpuTransferManager() { } static bool InitModule() { - xla::TransferManager::RegisterTransferManager(se::cuda::kCudaPlatformId, - &CreateGpuTransferManager); + xla::TransferManager::RegisterTransferManager( + stream_executor::cuda::kCudaPlatformId, &CreateGpuTransferManager); return true; } static bool module_initialized = InitModule(); diff --git a/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.h b/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.h index 9aa369c668364079504ead3491903e2590a142cc..d040a99975230578c270deabdfe60c61649e778c 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.h +++ b/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.h @@ -36,21 +36,20 @@ class GpuTransferManager : public GenericTransferManager { GpuTransferManager(); ~GpuTransferManager() override {} - Status TransferLiteralToInfeed(perftools::gputools::StreamExecutor* executor, + Status TransferLiteralToInfeed(se::StreamExecutor* executor, const Literal& literal) override; - Status TransferBufferToInfeed(perftools::gputools::StreamExecutor* executor, - int64 size, const void* source) override; + Status TransferBufferToInfeed(se::StreamExecutor* executor, int64 size, + const void* source) override; private: // Initiates the infeed data transfers. InfeedBuffer->Done() must be // called to clean up the memory allocated for InfeedBuffer. StatusOr TransferBufferToInfeedInternal( - perftools::gputools::StreamExecutor* executor, int64 size, - const void* source); + se::StreamExecutor* executor, int64 size, const void* source); // Enqueues infeed data buffers with the infeed manager after their // transfer completes. - Status EnqueueBuffersToInfeed(perftools::gputools::StreamExecutor* executor, + Status EnqueueBuffersToInfeed(se::StreamExecutor* executor, std::vector buffers); TF_DISALLOW_COPY_AND_ASSIGN(GpuTransferManager); diff --git a/tensorflow/compiler/xla/service/gpu/infeed_manager.cc b/tensorflow/compiler/xla/service/gpu/infeed_manager.cc index ee5b447c9cd0b1fde4d3a0943d5d4cb8cc5b3376..3ddc1c0789d746bf021256638342364aac63e0e3 100644 --- a/tensorflow/compiler/xla/service/gpu/infeed_manager.cc +++ b/tensorflow/compiler/xla/service/gpu/infeed_manager.cc @@ -19,8 +19,6 @@ limitations under the License. #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/core/platform/logging.h" -namespace se = ::perftools::gputools; - namespace xla { namespace gpu { diff --git a/tensorflow/compiler/xla/service/gpu/infeed_manager.h b/tensorflow/compiler/xla/service/gpu/infeed_manager.h index 73d5a5ce35497f156a181371bfb97fc37a8eb09e..d5f2216d460a45085536b15f9bf6e3bd3579f9c8 100644 --- a/tensorflow/compiler/xla/service/gpu/infeed_manager.h +++ b/tensorflow/compiler/xla/service/gpu/infeed_manager.h @@ -46,7 +46,7 @@ namespace gpu { // the client. The client manages the memory of the buffer. class InfeedBuffer { public: - InfeedBuffer(perftools::gputools::StreamExecutor* executor, int64 length) + InfeedBuffer(se::StreamExecutor* executor, int64 length) : executor_(executor), length_(length) { device_memory_ = executor_->AllocateArray(length); CHECK(!device_memory_.is_null()); @@ -60,14 +60,12 @@ class InfeedBuffer { // client to manage memory for the infeed buffers. void Done() { delete this; } - perftools::gputools::DeviceMemoryBase* device_memory() { - return &device_memory_; - } + se::DeviceMemoryBase* device_memory() { return &device_memory_; } private: - perftools::gputools::StreamExecutor* executor_; // Not owned. + se::StreamExecutor* executor_; // Not owned. const int64 length_; - perftools::gputools::DeviceMemoryBase device_memory_; + se::DeviceMemoryBase device_memory_; }; // Client-side class used to enqueue infeed buffers. @@ -100,8 +98,7 @@ class InfeedManager { // new stream on the first invocation. On subsequent invocations, if // the cached executor is not the same as the requested executor, // returns null. - perftools::gputools::Stream* GetStream( - perftools::gputools::StreamExecutor* executor); + se::Stream* GetStream(se::StreamExecutor* executor); private: // TODO(b/30467474): Revisit if this mutex becomes a point of @@ -121,10 +118,10 @@ class InfeedManager { tensorflow::gtl::FlatSet dequeued_buffer_; // Cached host to device stream for queuing infeed data. - std::unique_ptr host_to_device_stream_; + std::unique_ptr host_to_device_stream_; // Executor that the host_to_device_stream belongs to. Not owned. - perftools::gputools::StreamExecutor* host_to_device_executor_; + se::StreamExecutor* host_to_device_executor_; }; // Singleton creator-or-accessor: Returns the GPU infeed manager. diff --git a/tensorflow/compiler/xla/service/gpu/infeed_thunk.cc b/tensorflow/compiler/xla/service/gpu/infeed_thunk.cc index 2ac95ceb692447c7ac6dbbcd8b9a38876f7a77b6..ea34d5b30c91e8b809e3e17a904e27e589fd6b5f 100644 --- a/tensorflow/compiler/xla/service/gpu/infeed_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/infeed_thunk.cc @@ -31,10 +31,10 @@ InfeedThunk::InfeedThunk( destination_buffer_(destination_buffer) {} Status InfeedThunk::ExecuteOnStream(const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) { + se::Stream* stream) { VLOG(2) << "Infeeding to GPU "; - perftools::gputools::DeviceMemoryBase destination_address = + se::DeviceMemoryBase destination_address = buffer_allocations.GetDeviceAddress(destination_buffer_); InfeedManager* infeed_manager = GetOrCreateInfeedManager(); @@ -45,7 +45,7 @@ Status InfeedThunk::ExecuteOnStream(const BufferAllocations& buffer_allocations, std::vector tuple_element_addresses; for (BufferAllocation::Slice tuple_element_buffer : tuple_element_buffers_) { - perftools::gputools::DeviceMemoryBase tuple_element_address = + se::DeviceMemoryBase tuple_element_address = buffer_allocations.GetDeviceAddress(tuple_element_buffer); InfeedBuffer* buffer = infeed_manager->BlockingDequeueBuffer(); diff --git a/tensorflow/compiler/xla/service/gpu/infeed_thunk.h b/tensorflow/compiler/xla/service/gpu/infeed_thunk.h index 86918705fa0305217f11753e383200c7bd71474b..93713cb12defd95bdd69cb0aa7ad7b4e37fc8fae 100644 --- a/tensorflow/compiler/xla/service/gpu/infeed_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/infeed_thunk.h @@ -44,7 +44,7 @@ class InfeedThunk : public Thunk { InfeedThunk& operator=(const InfeedThunk&) = delete; Status ExecuteOnStream(const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) override; + se::Stream* stream) override; private: const std::vector tuple_element_buffers_; diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc index 32413f975a40c1abc334b16e81097bb44f56a44a..532d436ee82b985a4efe300f90223e1298e85765 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc @@ -160,14 +160,19 @@ static HloInstruction* CreateCudnnConv( 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. + // Our CustomCall takes four arguments: The conv lhs and rhs, the cudnn + // algorithm to use, and a boolean indicating whether to use tensor cores. + // + // It's up to a later pass to choose the algorithm and decide whether to use + // tensor cores, so to indicate that we haven't yet made a choice, we speicfy + // -1 and false for those args. HloInstruction* negative_one = computation->AddInstruction( HloInstruction::CreateConstant(Literal::CreateR0(-1))); + HloInstruction* false_constant = computation->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); HloInstruction* custom_call = computation->AddInstruction(HloInstruction::CreateCustomCall( - call_shape, {lhs, rhs, negative_one}, call_target)); + call_shape, {lhs, rhs, negative_one, false_constant}, call_target)); custom_call->set_window(window); custom_call->set_convolution_dimension_numbers(dnums); return custom_call; diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_context.h b/tensorflow/compiler/xla/service/gpu/ir_emitter_context.h index 3790ed313b9d0e167185a8b12c812132ee78811f..a78b4ff83075fd7ef330bb97ce217a198d450cf8 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter_context.h +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_context.h @@ -32,7 +32,7 @@ class IrEmitterContext { public: IrEmitterContext(const HloModule* hlo_module, const BufferAssignment* buffer_assignment, - const perftools::gputools::DeviceDescription* device_desc, + const se::DeviceDescription* device_desc, llvm::Module* llvm_module) : hlo_module_(hlo_module), buffer_assignment_(buffer_assignment), @@ -47,7 +47,7 @@ class IrEmitterContext { const BufferAssignment& buffer_assignment() const { return *buffer_assignment_; } - const perftools::gputools::DeviceDescription& device_description() const { + const se::DeviceDescription& device_description() const { return *device_desc_; } llvm::Module* llvm_module() { return llvm_module_; } @@ -56,7 +56,7 @@ class IrEmitterContext { private: const HloModule* hlo_module_; const BufferAssignment* buffer_assignment_; - const perftools::gputools::DeviceDescription* device_desc_; + const se::DeviceDescription* device_desc_; llvm::Module* llvm_module_; NameUniquer name_uniquer_; }; diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc index d29cc21ab1c697f8481ed1e94846d4df5ec5c1dc..26e497762f2a6f23767c5b98f339eefdef0b7468 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc @@ -536,7 +536,27 @@ Status IrEmitterUnnested::HandleFusion(HloInstruction* fusion) { thunk_sequence_->emplace_back(BuildGemmThunk(fusion)); return Status::OK(); } - thunk_sequence_->emplace_back(BuildKernelThunk(fusion)); + + int max_unroll_factor = fusion->GetModule() + ->config() + .debug_options() + .xla_gpu_max_kernel_unroll_factor(); + + // Find the largest possible power of two to unroll by. + // TODO(kramerb): Make this smarter. + int unroll_factor = 1; + if (!fusion->IsMultiOutputFusion()) { + CHECK(fusion->fusion_kind() == HloInstruction::FusionKind::kLoop); + int64 num_elements = ShapeUtil::ElementsIn(fusion->shape()); + for (int i = max_unroll_factor; i > 1; i /= 2) { + if (num_elements % i == 0) { + unroll_factor = i; + break; + } + } + } + + thunk_sequence_->emplace_back(BuildKernelThunk(fusion, unroll_factor)); return IrEmitter::HandleFusion(fusion); } @@ -2021,7 +2041,7 @@ Status IrEmitterUnnested::HandleGather(HloInstruction* gather) { } std::unique_ptr IrEmitterUnnested::BuildKernelThunk( - const HloInstruction* inst) { + const HloInstruction* inst, int unroll_factor) { const BufferAssignment& buffer_assn = ir_emitter_context_->buffer_assignment(); @@ -2113,7 +2133,7 @@ std::unique_ptr IrEmitterUnnested::BuildKernelThunk( } return MakeUnique(buffers, llvm_ir::AsString(kernel->getName()), - inst); + inst, unroll_factor); } std::unique_ptr IrEmitterUnnested::BuildHostToDeviceCopyThunk( @@ -2485,21 +2505,28 @@ std::unique_ptr IrEmitterUnnested::BuildConditionalThunk( Status IrEmitterUnnested::EmitTargetElementLoopInThunk( const HloInstruction& hlo, const llvm_ir::ElementGenerator& element_generator, KernelThunk* thunk) { + int unroll_factor = thunk->unroll_factor(); VLOG(3) << bindings_.ToString(); const Shape& element_shape = hlo.IsMultiOutputFusion() ? ShapeUtil::GetSubshape(hlo.shape(), {0}) : hlo.shape(); + VLOG(3) << "EmitTargetElementLoopInThunk " + << ShapeUtil::HumanStringWithLayout(hlo.shape()) + << " for unroll_factor " << unroll_factor; LaunchDimensions launch_dimensions = CalculateLaunchDimensions( - element_shape, ir_emitter_context_->device_description()); + element_shape, ir_emitter_context_->device_description(), unroll_factor); UpdateLaunchDimensions(launch_dimensions, thunk, ir_emitter_context_->llvm_module()); if (!hlo.IsMultiOutputFusion()) { return ParallelLoopEmitter(element_generator, GetIrArray(hlo, hlo), - launch_dimensions, &ir_builder_) + launch_dimensions, &ir_builder_, unroll_factor) .EmitLoop(IrName(&hlo)); } + CHECK_EQ(unroll_factor, 1) + << "multi-output fusion does not support unrolling"; + // For multiple outputs fusion, we need to emit each operand and the root. std::vector output_arrays; for (int64 i = 0; i < ShapeUtil::TupleElementCount(hlo.shape()); ++i) { diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h index 66c62e2d2de3ed1668271a21943dc73ed3d77651..b842f480c6257c1a8bee8cdac55e29c5db6801a0 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h @@ -150,8 +150,10 @@ class IrEmitterUnnested : public IrEmitter { // 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); + // Thunk object. The kernel implementation will be unrolled if unroll_factor + // is greater than one. + std::unique_ptr BuildKernelThunk(const HloInstruction* inst, + int unroll_factor = 1); // Returns a FftThunk that calls cuFFT to implement `inst`. std::unique_ptr BuildFftThunk(const HloInstruction* inst); diff --git a/tensorflow/compiler/xla/service/gpu/kernel_thunk.cc b/tensorflow/compiler/xla/service/gpu/kernel_thunk.cc index c20a781a33fe89af4740ed31dd5bfb1a64473057..d376ef7a245eb9ed86939f44c611b6dde5606b23 100644 --- a/tensorflow/compiler/xla/service/gpu/kernel_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/kernel_thunk.cc @@ -23,17 +23,17 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" -namespace se = ::perftools::gputools; - namespace xla { namespace gpu { KernelThunk::KernelThunk( tensorflow::gtl::ArraySlice args, - const string& kernel_name, const HloInstruction* hlo_instruction) + const string& kernel_name, const HloInstruction* hlo_instruction, + int unroll_factor) : Thunk(Kind::kKernel, hlo_instruction), args_(args.begin(), args.end()), - kernel_name_(kernel_name) {} + kernel_name_(kernel_name), + unroll_factor_(unroll_factor) {} tensorflow::Status KernelThunk::Initialize(const GpuExecutable& executable) { tensorflow::mutex_lock lock(mutex_); diff --git a/tensorflow/compiler/xla/service/gpu/kernel_thunk.h b/tensorflow/compiler/xla/service/gpu/kernel_thunk.h index 9ae455e2fcc253a7a08ff95764721048a16b0bf7..b556befe66b6bebba1a958f553f0a9b2c4eebbe4 100644 --- a/tensorflow/compiler/xla/service/gpu/kernel_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/kernel_thunk.h @@ -47,20 +47,21 @@ class KernelThunk : public Thunk { // // `hlo_instruction` is as in Thunk. Other arguments are as the class members. KernelThunk(tensorflow::gtl::ArraySlice args, - const string& kernel_name, const HloInstruction* hlo_instruction); + const string& kernel_name, const HloInstruction* hlo_instruction, + int unroll_factor); KernelThunk(const KernelThunk&) = delete; KernelThunk& operator=(const KernelThunk&) = delete; ~KernelThunk() override = default; const string& kernel_name() const { return kernel_name_; } + int unroll_factor() const { return unroll_factor_; } void SetLaunchDimensions(const LaunchDimensions& launch_dims); tensorflow::Status Initialize(const GpuExecutable& executable) override; // Executes the kernel for the thunk on "stream", which must be non-null. tensorflow::Status ExecuteOnStream( - const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) override; + const BufferAllocations& buffer_allocations, se::Stream* stream) override; private: // Buffers passed to the kernel as arguments. @@ -69,6 +70,10 @@ class KernelThunk : public Thunk { // Entry kernel name for the computation. const string kernel_name_; + // The number of times this kernel should be unrolled. This works as a + // multiplier on the number of elements produced by a GPU thread. + const int unroll_factor_; + // The thread and block dimension used to launch the kernel. // Will be set by IrEmitterUnnested. LaunchDimensions launch_dimensions_; @@ -76,13 +81,11 @@ class KernelThunk : public Thunk { // Describes how to load this kernel. ExecuteOnStream reuses this loader // specification for all executions. mutable tensorflow::mutex mutex_; - std::unique_ptr loader_spec_ - GUARDED_BY(mutex_); + std::unique_ptr loader_spec_ GUARDED_BY(mutex_); // Loaded kernels for each `StreamExecutor` - std::unordered_map - kernel_cache_ GUARDED_BY(mutex_); + std::unordered_map kernel_cache_ + GUARDED_BY(mutex_); }; } // namespace gpu 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 defd281d74bd38f7da3f268e0f55970fc1af8263..df9d9be889ce839ee665cd4820b169c124d9fcde 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 @@ -34,7 +34,7 @@ limitations under the License. #include "llvm/Analysis/TargetTransformInfo.h" #include "llvm/Bitcode/BitcodeReader.h" #include "llvm/Bitcode/BitcodeWriter.h" -#include "llvm/CodeGen/CommandFlags.def" +#include "llvm/CodeGen/CommandFlags.inc" #include "llvm/IR/LLVMContext.h" #include "llvm/IR/LegacyPassManager.h" #include "llvm/IR/Module.h" diff --git a/tensorflow/compiler/xla/service/gpu/memset_thunk.cc b/tensorflow/compiler/xla/service/gpu/memset_thunk.cc index 18e673542c5b47cb90d31a8eff62a5e4adb78d1d..d4100a898b5bb9eec382c34932c2db104c9e985b 100644 --- a/tensorflow/compiler/xla/service/gpu/memset_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/memset_thunk.cc @@ -19,8 +19,6 @@ limitations under the License. namespace xla { namespace gpu { -namespace se = ::perftools::gputools; - Status MemzeroThunk::ExecuteOnStream( const BufferAllocations& buffer_allocations, se::Stream* stream) { se::DeviceMemoryBase dest_data = buffer_allocations.GetDeviceAddress(dest_); diff --git a/tensorflow/compiler/xla/service/gpu/memset_thunk.h b/tensorflow/compiler/xla/service/gpu/memset_thunk.h index b4bb74d1dd6dc9d09c5e4d439d57dfe8b57c2ed9..51c332d287d139335b356fc66411b5ffaa448b5a 100644 --- a/tensorflow/compiler/xla/service/gpu/memset_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/memset_thunk.h @@ -36,7 +36,7 @@ class MemzeroThunk : public Thunk { : Thunk(Kind::kMemzero, hlo), dest_(dest) {} Status ExecuteOnStream(const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) override; + se::Stream* stream) override; private: const BufferAllocation::Slice dest_; @@ -52,7 +52,7 @@ class Memset32BitValueThunk : public Thunk { : Thunk(Kind::kMemset32BitValue, hlo), value_(value), dest_(dest) {} Status ExecuteOnStream(const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) override; + se::Stream* stream) override; private: uint32 value_; diff --git a/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc b/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc index 388dcc008b07a76ff9ed07df04181e49a8734f51..d8c07dc3119fb81a3ef22822acb11b7c4d5bbca5 100644 --- a/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc +++ b/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc @@ -32,25 +32,32 @@ namespace gpu { ParallelLoopEmitter::ParallelLoopEmitter( BodyEmitter body_emitter, const Shape& shape, - const LaunchDimensions& launch_dimensions, llvm::IRBuilder<>* ir_builder) + const LaunchDimensions& launch_dimensions, llvm::IRBuilder<>* ir_builder, + int unroll_factor) : LoopEmitter(body_emitter, shape, ir_builder), - launch_dimensions_(launch_dimensions) {} + launch_dimensions_(launch_dimensions), + unroll_factor_(unroll_factor) {} ParallelLoopEmitter::ParallelLoopEmitter( const llvm_ir::ElementGenerator& target_element_generator, tensorflow::gtl::ArraySlice target_arrays, - const LaunchDimensions& launch_dimensions, llvm::IRBuilder<>* ir_builder) + const LaunchDimensions& launch_dimensions, llvm::IRBuilder<>* ir_builder, + int unroll_factor) : LoopEmitter(target_element_generator, target_arrays, ir_builder), - launch_dimensions_(launch_dimensions) {} + launch_dimensions_(launch_dimensions), + unroll_factor_(unroll_factor) {} ParallelLoopEmitter::ParallelLoopEmitter( const llvm_ir::ElementGenerator& target_element_generator, const llvm_ir::IrArray& target_array, - const LaunchDimensions& launch_dimensions, llvm::IRBuilder<>* ir_builder) + const LaunchDimensions& launch_dimensions, llvm::IRBuilder<>* ir_builder, + int unroll_factor) : LoopEmitter(target_element_generator, target_array, ir_builder), - launch_dimensions_(launch_dimensions) {} + launch_dimensions_(launch_dimensions), + unroll_factor_(unroll_factor) {} -llvm_ir::IrArray::Index ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock( +std::vector +ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock( tensorflow::StringPiece loop_name) { // Emit the following code in LLVM IR: // linear_index = blockIdx.x * blockDim.x + threadIdx.x; @@ -63,6 +70,9 @@ llvm_ir::IrArray::Index ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock( // "It is guaranteed that [...] 0 <= %ctaid.x < %nctaid.x" // // %nctaid.x is currently specified as 2147483647. + VLOG(3) << "EmitIndexAndSetExitBasicBlock unroll_factor " << unroll_factor_; + std::vector array_indices; + llvm::Value* block_id = llvm_ir::EmitCallToIntrinsic( llvm::Intrinsic::nvvm_read_ptx_sreg_ctaid_x, {}, {}, ir_builder_); llvm_ir::AddRangeMetadata(0, launch_dimensions_.block_count(), @@ -81,7 +91,7 @@ llvm_ir::IrArray::Index ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock( thread_id = ir_builder_->CreateZExt(thread_id, ir_builder_->getInt64Ty(), "thread_id"); - llvm::Value* linear_index = ir_builder_->CreateAdd( + llvm::Value* linear_index_base = ir_builder_->CreateAdd( ir_builder_->CreateMul( block_id, ir_builder_->getInt64(launch_dimensions_.threads_per_block()), "", @@ -99,15 +109,30 @@ llvm_ir::IrArray::Index ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock( llvm_ir::EmitCallToIntrinsic( llvm::Intrinsic::assume, {ir_builder_->CreateICmpULT( - linear_index, + linear_index_base, ir_builder_->getInt64(launch_dimensions_.threads_per_block() * launch_dimensions_.block_count()), "linear_index_in_range")}, {}, ir_builder_); + if (unroll_factor_ > 1) { + linear_index_base = ir_builder_->CreateMul( + linear_index_base, ir_builder_->getInt64(unroll_factor_), + "linear_index_base", /*HasNUW=*/true, /*HasNSW=*/true); + } + + array_indices.emplace_back(linear_index_base, shape_, ir_builder_); + for (int i = 1; i < unroll_factor_; ++i) { + llvm::Value* linear_index = ir_builder_->CreateAdd( + linear_index_base, ir_builder_->getInt64(i), "linear_index", + /*HasNUW=*/true, /*HasNSW=*/true); + array_indices.emplace_back(linear_index, shape_, ir_builder_); + } + auto if_in_bounds = llvm_ir::EmitIfThenElse( ir_builder_->CreateICmpULT( - linear_index, ir_builder_->getInt64(ShapeUtil::ElementsIn(shape_))), + linear_index_base, + ir_builder_->getInt64(ShapeUtil::ElementsIn(shape_))), llvm_ir::IrName(loop_name, "in_bounds"), ir_builder_, false); // Set exit_bb_ to the exit block of the if structure. @@ -116,7 +141,8 @@ llvm_ir::IrArray::Index ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock( // Set IR builder insertion point to the body of the if structure. llvm_ir::SetToFirstInsertPoint(if_in_bounds.true_block, ir_builder_); - return llvm_ir::IrArray::Index(linear_index, shape_, ir_builder_); + + return array_indices; } } // namespace gpu diff --git a/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.h b/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.h index 8ed63a854a74fc06c3c389f40fe1f5970885deac..25318b3bed8bf4a2dfe3a4a974269d0405c3bfec 100644 --- a/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.h +++ b/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.h @@ -34,13 +34,13 @@ class ParallelLoopEmitter : public llvm_ir::LoopEmitter { // The meanings of other parameters are the same as LoopEmitter. ParallelLoopEmitter(BodyEmitter body_emitter, const Shape& shape, const LaunchDimensions& launch_dimensions, - llvm::IRBuilder<>* ir_builder); + llvm::IRBuilder<>* ir_builder, int unroll_factor = 1); // Constructs a ParallelLoopEmitter from an element generator that generates // each element of the given target array. ParallelLoopEmitter(const llvm_ir::ElementGenerator& target_element_generator, const llvm_ir::IrArray& target_array, const LaunchDimensions& launch_dimensions, - llvm::IRBuilder<>* ir_builder); + llvm::IRBuilder<>* ir_builder, int unroll_factor = 1); // Constructs a loop emitter for a loop that generates on element of each of N // arrays on each iteration. @@ -50,18 +50,20 @@ class ParallelLoopEmitter : public llvm_ir::LoopEmitter { ParallelLoopEmitter( const llvm_ir::ElementGenerator& target_element_generator, tensorflow::gtl::ArraySlice target_arrays, - const LaunchDimensions& launch_dimensions, llvm::IRBuilder<>* ir_builder); + const LaunchDimensions& launch_dimensions, llvm::IRBuilder<>* ir_builder, + int unroll_factor = 1); ParallelLoopEmitter(const ParallelLoopEmitter&) = delete; ParallelLoopEmitter& operator=(const ParallelLoopEmitter&) = delete; ~ParallelLoopEmitter() override = default; - llvm_ir::IrArray::Index EmitIndexAndSetExitBasicBlock( + std::vector EmitIndexAndSetExitBasicBlock( tensorflow::StringPiece loop_name) override; private: // The thread and block dimension to parallelize the loop on. const LaunchDimensions launch_dimensions_; + const int unroll_factor_; }; } // namespace gpu diff --git a/tensorflow/compiler/xla/service/gpu/partition_assignment.cc b/tensorflow/compiler/xla/service/gpu/partition_assignment.cc index 6cf280df05496716a0780d61ded92efd9982734c..d3fd0544fb68809125e9b9f7a5e5b7eff8c6ef43 100644 --- a/tensorflow/compiler/xla/service/gpu/partition_assignment.cc +++ b/tensorflow/compiler/xla/service/gpu/partition_assignment.cc @@ -29,8 +29,6 @@ limitations under the License. #include "tensorflow/core/lib/strings/stringprintf.h" #include "tensorflow/core/platform/logging.h" -namespace se = ::perftools::gputools; - namespace xla { namespace gpu { @@ -44,12 +42,16 @@ std::ostream& operator<<(std::ostream& out, // Calculates the launch dimensions used to invoke `hlo`. LaunchDimensions CalculateLaunchDimensions( - const Shape& shape, const se::DeviceDescription& device_desc) { + const Shape& shape, const se::DeviceDescription& device_desc, + int unroll_factor) { int64 num_elements = ShapeUtil::ElementsIn(shape); if (num_elements <= 1) { return LaunchDimensions(); } + CHECK_EQ(num_elements % unroll_factor, 0); + num_elements = num_elements / unroll_factor; + // Since we don't do any inter-warp communication, we're free to choose any // block size we want, subject to hardware constraints. We choose the // smallest block size that allows the GPU to reach full occupancy (assuming diff --git a/tensorflow/compiler/xla/service/gpu/partition_assignment.h b/tensorflow/compiler/xla/service/gpu/partition_assignment.h index 0bf463a6ef95d5a32784838c08ad239752fd1acf..c125474edb1036090a926020f2b1e7fcf64c751a 100644 --- a/tensorflow/compiler/xla/service/gpu/partition_assignment.h +++ b/tensorflow/compiler/xla/service/gpu/partition_assignment.h @@ -57,8 +57,8 @@ std::ostream& operator<<(std::ostream& out, const LaunchDimensions& launch_dims); LaunchDimensions CalculateLaunchDimensions( - const Shape& shape, - const perftools::gputools::DeviceDescription& device_desc); + const Shape& shape, const se::DeviceDescription& device_desc, + int unroll_factor = 1); } // namespace gpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/sequential_thunk.cc b/tensorflow/compiler/xla/service/gpu/sequential_thunk.cc index d8a43091d4037a0edd125a4a1b6cb5ad7c7065f0..c8510808f10a731af90154447bd3e1e037db6348 100644 --- a/tensorflow/compiler/xla/service/gpu/sequential_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/sequential_thunk.cc @@ -33,8 +33,7 @@ tensorflow::Status SequentialThunk::Initialize( } tensorflow::Status SequentialThunk::ExecuteOnStream( - const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) { + const BufferAllocations& buffer_allocations, se::Stream* stream) { for (const auto& thunk : thunks_) { TF_RETURN_IF_ERROR(thunk->ExecuteOnStream(buffer_allocations, stream)); } diff --git a/tensorflow/compiler/xla/service/gpu/sequential_thunk.h b/tensorflow/compiler/xla/service/gpu/sequential_thunk.h index 32c5b748aba14239d6795d14e442c1c3b43d010e..df17b8d67b80321c7088243eae46e7a723b4ede9 100644 --- a/tensorflow/compiler/xla/service/gpu/sequential_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/sequential_thunk.h @@ -40,8 +40,7 @@ class SequentialThunk : public Thunk { tensorflow::Status Initialize(const GpuExecutable& executable) override; tensorflow::Status ExecuteOnStream( - const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) override; + const BufferAllocations& buffer_allocations, se::Stream* stream) override; private: // The list of sub-thunks. diff --git a/tensorflow/compiler/xla/service/gpu/thunk.h b/tensorflow/compiler/xla/service/gpu/thunk.h index 9eea958d1214b131d49cb4e28f1944860408d3a8..a0c785ed913109e987d058124c8ef49019c98500 100644 --- a/tensorflow/compiler/xla/service/gpu/thunk.h +++ b/tensorflow/compiler/xla/service/gpu/thunk.h @@ -85,8 +85,7 @@ class Thunk { // This value is not required to be constant for a given Thunk. For example, // a Thunk that performs autotuning may return true for its first run and // false thereafter. - virtual bool ShouldHaltAllActivityBeforeRunning( - perftools::gputools::Stream* /*stream*/) { + virtual bool ShouldHaltAllActivityBeforeRunning(se::Stream* /*stream*/) { return false; } @@ -104,8 +103,7 @@ class Thunk { // called after Initialize and can be called multiple times over Thunk's // lifetime. Stream argument must be non-null. virtual tensorflow::Status ExecuteOnStream( - const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) = 0; + const BufferAllocations& buffer_allocations, se::Stream* stream) = 0; private: Kind kind_; diff --git a/tensorflow/compiler/xla/service/gpu/tuple_thunk.cc b/tensorflow/compiler/xla/service/gpu/tuple_thunk.cc index bd65e72393a59e72671ff0cc32c37eaa48856255..ecb54857ccc40ead21e5a18d79a37b545680021d 100644 --- a/tensorflow/compiler/xla/service/gpu/tuple_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/tuple_thunk.cc @@ -17,8 +17,6 @@ limitations under the License. #include "tensorflow/compiler/xla/util.h" -namespace se = ::perftools::gputools; - namespace xla { namespace gpu { diff --git a/tensorflow/compiler/xla/service/gpu/tuple_thunk.h b/tensorflow/compiler/xla/service/gpu/tuple_thunk.h index 3b1a496328540ae69a449e7080903d31284885d1..8b459c29a136a6e7853e68a1bead7d12c0d08ad0 100644 --- a/tensorflow/compiler/xla/service/gpu/tuple_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/tuple_thunk.h @@ -46,8 +46,7 @@ class TupleThunk : public Thunk { TupleThunk& operator=(const TupleThunk&) = delete; tensorflow::Status ExecuteOnStream( - const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) override; + const BufferAllocations& buffer_allocations, se::Stream* stream) override; private: const std::vector tuple_element_buffers_; diff --git a/tensorflow/compiler/xla/service/gpu/while_thunk.cc b/tensorflow/compiler/xla/service/gpu/while_thunk.cc index c21559af6d2e5dfb5aaf62afcdcaed514e0914c9..a9f3d619a3ffd6d849572355e2902375e43508fa 100644 --- a/tensorflow/compiler/xla/service/gpu/while_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/while_thunk.cc @@ -41,8 +41,8 @@ Status WhileThunk::Initialize(const GpuExecutable& executable) { } Status WhileThunk::ExecuteOnStream(const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) { - perftools::gputools::DeviceMemoryBase condition_result_data = + se::Stream* stream) { + se::DeviceMemoryBase condition_result_data = buffer_allocations.GetDeviceAddress(condition_result_buffer_index_); while (true) { diff --git a/tensorflow/compiler/xla/service/gpu/while_thunk.h b/tensorflow/compiler/xla/service/gpu/while_thunk.h index 4c9f45de9e42494df58706d0a4a3eb0c4220b8b8..e589ca78a7ea00e7592d6e09ead9c270f902702f 100644 --- a/tensorflow/compiler/xla/service/gpu/while_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/while_thunk.h @@ -47,7 +47,7 @@ class WhileThunk : public Thunk { Status Initialize(const GpuExecutable& executable) override; Status ExecuteOnStream(const BufferAllocations& buffer_allocations, - perftools::gputools::Stream* stream) override; + se::Stream* stream) override; private: const BufferAllocation::Slice condition_result_buffer_index_; diff --git a/tensorflow/compiler/xla/service/hlo.proto b/tensorflow/compiler/xla/service/hlo.proto index 0b446c654779db410ebbd91ef9a5bab14d08a278..0c3eb7dcb44f1203ae10e238f194e41e9ecaea74 100644 --- a/tensorflow/compiler/xla/service/hlo.proto +++ b/tensorflow/compiler/xla/service/hlo.proto @@ -135,6 +135,10 @@ message HloInstructionProto { xla.GatherDimensionNumbers gather_dimension_numbers = 33; repeated int64 gather_window_bounds = 34; + // Compute Host. + string channel_name = 41; + int64 cost_estimate_ns = 42; + // The id of this instruction. int64 id = 35; @@ -292,3 +296,20 @@ message HloProto { HloOrderingProto hlo_ordering = 2; BufferAssignmentProto buffer_assignment = 3; } + +// Encapsulates HloProto together with the arguments, result, and +// execution_platform. This message is used for purposes such as +// analysis/replay/file-storage. +message HloSession { + // The hlo graph. + HloProto hlo = 1; + + // The arguments passed to the graph. + repeated LiteralProto arguments = 2; + + // The result of the graph. + LiteralProto result = 3; + + // The name of the platform used to run the graph. + string execution_platform = 4; +} diff --git a/tensorflow/compiler/xla/service/hlo_constant_folding.cc b/tensorflow/compiler/xla/service/hlo_constant_folding.cc index 7aa38c6b79ed904bb4a518c4b7aaa1e079c27ea8..35ecd4428d0dfde2de445ea34472d2c78148c6c9 100644 --- a/tensorflow/compiler/xla/service/hlo_constant_folding.cc +++ b/tensorflow/compiler/xla/service/hlo_constant_folding.cc @@ -69,8 +69,7 @@ StatusOr HloConstantFolding::Run(HloModule* module) { // Broadcasts dramatically increase the size of constants, which is often // detrimental to performance and memory capacity, so do not fold // broadcasts. - if (instruction->opcode() == HloOpcode::kBroadcast || - instruction->opcode() == HloOpcode::kBroadcastDimOne) { + if (instruction->opcode() == HloOpcode::kBroadcast) { continue; } diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis.cc b/tensorflow/compiler/xla/service/hlo_cost_analysis.cc index ea4dd62fdb5bb3be40987d1a6ea96b3a58b0053b..44e4f75f75b275653e1a07111943843fc6f78b33 100644 --- a/tensorflow/compiler/xla/service/hlo_cost_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_cost_analysis.cc @@ -336,11 +336,6 @@ Status HloCostAnalysis::HandleBroadcast(const HloInstruction*) { return Status::OK(); } -Status HloCostAnalysis::HandleBroadcastDimOne( - const HloInstruction* broadcastDimOne) { - return Status::OK(); -} - Status HloCostAnalysis::HandlePad(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 a9f6845747aa2081df936d388551bbc0b75b787b..d17678d20f2a23fd98d18b77d5fb25853901a789 100644 --- a/tensorflow/compiler/xla/service/hlo_cost_analysis.h +++ b/tensorflow/compiler/xla/service/hlo_cost_analysis.h @@ -95,7 +95,6 @@ class HloCostAnalysis : public ConstDfsHloVisitor { Status HandleSelectAndScatter(const HloInstruction* instruction) override; Status HandleBitcast(const HloInstruction* bitcast) override; Status HandleBroadcast(const HloInstruction* broadcast) override; - Status HandleBroadcastDimOne(const HloInstruction* broadcastDimOne) override; Status HandlePad(const HloInstruction* pad) override; Status HandleReshape(const HloInstruction* reshape) override; Status HandleTranspose(const HloInstruction* transpose) override; diff --git a/tensorflow/compiler/xla/service/hlo_creation_utils.cc b/tensorflow/compiler/xla/service/hlo_creation_utils.cc index b186767ce792cd89ae77fe9a03b3a2ecf296b804..9a89888480b8c79dfb1f79a50e9686bf45aa49b3 100644 --- a/tensorflow/compiler/xla/service/hlo_creation_utils.cc +++ b/tensorflow/compiler/xla/service/hlo_creation_utils.cc @@ -163,6 +163,8 @@ StatusOr MakeConcatHlo(ArraySlice operands, } StatusOr CollapseFirstNDims(HloInstruction* operand, int64 n) { + CHECK_GT(n, 0); + const Shape& operand_shape = operand->shape(); CHECK_GE(operand_shape.dimensions_size(), n); int64 new_shape_leading_bound = 1; @@ -184,6 +186,17 @@ StatusOr CollapseFirstNDims(HloInstruction* operand, int64 n) { return MakeReshapeHlo(output_shape, operand); } +StatusOr PrependDegenerateDims(HloInstruction* operand, + int64 n) { + CHECK_GT(n, 0); + std::vector new_shape_dims; + const Shape& operand_shape = operand->shape(); + new_shape_dims.reserve(n + operand_shape.dimensions_size()); + new_shape_dims.insert(new_shape_dims.begin(), n, 1); + c_copy(operand_shape.dimensions(), std::back_inserter(new_shape_dims)); + return MakeReshapeHlo(new_shape_dims, operand); +} + StatusOr ExpandFirstDimIntoNDims( HloInstruction* operand, ArraySlice expanded_dims) { CHECK_GT(operand->shape().dimensions_size(), 0); diff --git a/tensorflow/compiler/xla/service/hlo_creation_utils.h b/tensorflow/compiler/xla/service/hlo_creation_utils.h index d99e32a737e6aaa2ff746cf6c00d4300cf62f4e1..c9a7361a6af0c2a0839c59a0ea695ec1b9a98bd4 100644 --- a/tensorflow/compiler/xla/service/hlo_creation_utils.h +++ b/tensorflow/compiler/xla/service/hlo_creation_utils.h @@ -103,12 +103,22 @@ StatusOr MakeConcatHlo( // their operand(s). // Collapses (via reshape) the first N (logical) dimensions of `operand` into a -// single leading dimension. `operand` must have rank > n. +// single leading dimension. `operand` must have rank > `n` and `n` must not be +// 0. // // 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); +// Prepends `n` degenerate dimensions (dimensions with bound = 1) to `operand` +// using a reshape. +// +// For instance if operand has shape f32[3,4,5] then this returns the operand +// reshaped to f32[1,3,4,5]. If the operand is a f32 scalar (i.e. has shape +// f32[]) then this returns the operand reshaped to f32[1]. +StatusOr PrependDegenerateDims(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 diff --git a/tensorflow/compiler/xla/service/hlo_creation_utils_test.cc b/tensorflow/compiler/xla/service/hlo_creation_utils_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..6b681a5bf6f34b724bed52d60df59c2ac1068b52 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_creation_utils_test.cc @@ -0,0 +1,234 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT 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/ptr_util.h" +#include "tensorflow/compiler/xla/service/hlo_evaluator.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/core/platform/test.h" + +namespace xla { +namespace { +using tensorflow::gtl::ArraySlice; + +std::unique_ptr CreateModuleWithProgramShape( + PrimitiveType primitive_type, ArraySlice input_shape_dims, + ArraySlice output_shape_dims, HloInstruction** param, + HloComputation** entry_computation) { + Shape input_shape = ShapeUtil::MakeShape(primitive_type, input_shape_dims); + Shape output_shape = ShapeUtil::MakeShape(primitive_type, output_shape_dims); + std::unique_ptr module = MakeUnique("test"); + *entry_computation = module->AddEntryComputation( + CreateComputationWithSignature({&input_shape}, output_shape, "entry") + .ValueOrDie()); + *param = (*entry_computation)->parameter_instruction(0); + return module; +} + +TEST(HloCreationUtilsTest, CollapseFirst1Dim) { + HloInstruction* param; + HloComputation* entry_computation; + + std::unique_ptr module = CreateModuleWithProgramShape( + S32, + /*input_shape_dims=*/{2}, /*output_shape_dims=*/{2}, ¶m, + &entry_computation); + + TF_ASSERT_OK_AND_ASSIGN(HloInstruction * first_1_dims_collapsed, + CollapseFirstNDims(param, 1)); + entry_computation->set_root_instruction(first_1_dims_collapsed); + + HloEvaluator evaluator; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result_literal, + evaluator.Evaluate>( + *module, {Literal::CreateR1({3, 4})})); + CHECK_EQ(*result_literal, *Literal::CreateR1({3, 4})); +} + +TEST(HloCreationUtilsTest, CollapseFirst2Dims) { + HloInstruction* param; + HloComputation* entry_computation; + + std::unique_ptr module = CreateModuleWithProgramShape( + S32, + /*input_shape_dims=*/{2, 3, 2}, /*output_shape_dims=*/{6, 2}, ¶m, + &entry_computation); + + TF_ASSERT_OK_AND_ASSIGN(HloInstruction * first_2_dims_collapsed, + CollapseFirstNDims(param, 2)); + entry_computation->set_root_instruction(first_2_dims_collapsed); + + HloEvaluator evaluator; + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr result_literal, + evaluator.Evaluate>( + *module, + {Literal::CreateR3( + {{{1, 2}, {3, 4}, {5, 6}}, {{-1, -2}, {-3, -4}, {-5, -6}}})})); + CHECK_EQ(*result_literal, + *Literal::CreateR2( + {{1, 2}, {3, 4}, {5, 6}, {-1, -2}, {-3, -4}, {-5, -6}})); +} + +TEST(HloCreationUtilsTest, Prepend1DegenerateDim) { + HloInstruction* param; + HloComputation* entry_computation; + + std::unique_ptr module = CreateModuleWithProgramShape( + S32, + /*input_shape_dims=*/{2}, /*output_shape_dims=*/{1, 2}, ¶m, + &entry_computation); + + TF_ASSERT_OK_AND_ASSIGN(HloInstruction * with_1_degenerate_dim_prepended, + PrependDegenerateDims(param, 1)); + entry_computation->set_root_instruction(with_1_degenerate_dim_prepended); + + HloEvaluator evaluator; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result_literal, + evaluator.Evaluate>( + *module, {Literal::CreateR1({9, 10})})); + CHECK_EQ(*result_literal, *Literal::CreateR2({{9, 10}})); +} + +TEST(HloCreationUtilsTest, Prepend2DegenerateDims) { + HloInstruction* param; + HloComputation* entry_computation; + + std::unique_ptr module = CreateModuleWithProgramShape( + S32, + /*input_shape_dims=*/{2}, /*output_shape_dims=*/{1, 1, 2}, ¶m, + &entry_computation); + + TF_ASSERT_OK_AND_ASSIGN(HloInstruction * with_2_degenerate_dims_prepended, + PrependDegenerateDims(param, 2)); + entry_computation->set_root_instruction(with_2_degenerate_dims_prepended); + + HloEvaluator evaluator; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result_literal, + evaluator.Evaluate>( + *module, {Literal::CreateR1({9, 10})})); + CHECK_EQ(*result_literal, *Literal::CreateR3({{{9, 10}}})); +} + +TEST(HloCreationUtilsTest, Prepend2DegenerateDimsToScalar) { + HloInstruction* param; + HloComputation* entry_computation; + + std::unique_ptr module = CreateModuleWithProgramShape( + S32, + /*input_shape_dims=*/{}, /*output_shape_dims=*/{1, 1}, ¶m, + &entry_computation); + + TF_ASSERT_OK_AND_ASSIGN(HloInstruction * with_2_degenerate_dims_prepended, + PrependDegenerateDims(param, 2)); + entry_computation->set_root_instruction(with_2_degenerate_dims_prepended); + + HloEvaluator evaluator; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result_literal, + evaluator.Evaluate>( + *module, {Literal::CreateR0(9)})); + CHECK_EQ(*result_literal, *Literal::CreateR2({{9}})); +} + +TEST(HloCreationUtilsTest, ExpandFirstDimInto3Dims) { + HloInstruction* param; + HloComputation* entry_computation; + + std::unique_ptr module = CreateModuleWithProgramShape( + S32, + /*input_shape_dims=*/{6}, /*output_shape_dims=*/{3, 1, 2}, ¶m, + &entry_computation); + + TF_ASSERT_OK_AND_ASSIGN(HloInstruction * first_dim_expanded, + ExpandFirstDimIntoNDims(param, {3, 1, 2})); + entry_computation->set_root_instruction(first_dim_expanded); + + HloEvaluator evaluator; + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr result_literal, + evaluator.Evaluate>( + *module, {Literal::CreateR1({1, 2, 3, 4, 5, 6})})); + CHECK_EQ(*result_literal, + *Literal::CreateR3({{{1, 2}}, {{3, 4}}, {{5, 6}}})); +} + +TEST(HloCreationUtilsTest, PadVectorWithZeros) { + HloInstruction* param; + HloComputation* entry_computation; + + std::unique_ptr module = CreateModuleWithProgramShape( + S32, + /*input_shape_dims=*/{2}, /*output_shape_dims=*/{6}, ¶m, + &entry_computation); + + TF_ASSERT_OK_AND_ASSIGN( + HloInstruction * zero_padded_param, + PadVectorWithZeros(param, /*zeros_to_prepend=*/3, /*zeros_to_append=*/1)); + entry_computation->set_root_instruction(zero_padded_param); + + HloEvaluator evaluator; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result_literal, + evaluator.Evaluate>( + *module, {Literal::CreateR1({3, 4})})); + CHECK_EQ(*result_literal, *Literal::CreateR1({0, 0, 0, 3, 4, 0})); +} + +TEST(HloCreationUtilsTest, BroadcastZeros_S32) { + HloInstruction* param; + HloComputation* entry_computation; + + std::unique_ptr module = CreateModuleWithProgramShape( + S32, + /*input_shape_dims=*/{}, /*output_shape_dims=*/{2, 2}, ¶m, + &entry_computation); + + TF_ASSERT_OK_AND_ASSIGN( + HloInstruction * zeros, + BroadcastZeros(module->entry_computation(), S32, {2, 2})); + entry_computation->set_root_instruction(zeros); + + HloEvaluator evaluator; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result_literal, + evaluator.Evaluate>( + *module, {Literal::CreateR0(0)})); + CHECK_EQ(*result_literal, *Literal::CreateR2({{0, 0}, {0, 0}})); +} + +TEST(HloCreationUtilsTest, BroadcastZeros_F32) { + HloInstruction* param; + HloComputation* entry_computation; + + std::unique_ptr module = CreateModuleWithProgramShape( + F32, + /*input_shape_dims=*/{}, /*output_shape_dims=*/{2, 2}, ¶m, + &entry_computation); + + TF_ASSERT_OK_AND_ASSIGN( + HloInstruction * zeros, + BroadcastZeros(module->entry_computation(), F32, {2, 2})); + entry_computation->set_root_instruction(zeros); + + HloEvaluator evaluator; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result_literal, + evaluator.Evaluate>( + *module, {Literal::CreateR0(0.0f)})); + CHECK_EQ(*result_literal, + *Literal::CreateR2({{0.0f, 0.0f}, {0.0f, 0.0f}})); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_element_type_converter.cc b/tensorflow/compiler/xla/service/hlo_element_type_converter.cc index c782d1b0add17c70e0f54826917df251d5a613e2..d236f83aeb9254b9c6e6d04629758ac2c8fd0da3 100644 --- a/tensorflow/compiler/xla/service/hlo_element_type_converter.cc +++ b/tensorflow/compiler/xla/service/hlo_element_type_converter.cc @@ -178,24 +178,37 @@ StatusOr HloElementTypeConverter::Run(HloModule* module) { if (hlo->shape().element_type() == eliminate_type_) { Shape shape = ShapeUtil::ChangeElementType(hlo->shape(), replace_with_type_); + new_hlo = computation->AddInstruction( hlo->CloneWithNewOperands(shape, new_operands, hlo->GetModule())); + TF_RETURN_IF_ERROR(new_hlo->CopyAllControlDepsFrom(hlo)); + new_hlo = ToElementType(new_hlo, eliminate_type_); } else if (ShapeUtil::IsTuple(hlo->shape())) { Shape old_shape = hlo->shape(); Shape new_shape = GetConvertedTupleShape(hlo->shape(), eliminate_type_, replace_with_type_); + new_hlo = computation->AddInstruction(hlo->CloneWithNewOperands( new_shape, new_operands, hlo->GetModule())); + TF_RETURN_IF_ERROR(new_hlo->CopyAllControlDepsFrom(hlo)); + // Convert the elements of the result of `new_hlo` to produce a new // tuple with shape `old_shape`. new_hlo = ConvertTupleElements(new_hlo, old_shape); } else { new_hlo = computation->AddInstruction(hlo->CloneWithNewOperands( hlo->shape(), new_operands, hlo->GetModule())); + TF_RETURN_IF_ERROR(new_hlo->CopyAllControlDepsFrom(hlo)); } - TF_RETURN_IF_ERROR(computation->ReplaceInstruction(hlo, new_hlo)); + TF_RETURN_IF_ERROR(hlo->ReplaceAllUsesWith(new_hlo)); + TF_RETURN_IF_ERROR(hlo->DropAllControlDeps()); + + // NB! We want to replace and remove side effecting instructions like Rng + // as well so we can't rely HloComputation::ReplaceInstruction to reliably + // remove the replaced instruction. + TF_RETURN_IF_ERROR(computation->RemoveInstruction(hlo)); changed = true; } } diff --git a/tensorflow/compiler/xla/service/hlo_element_type_converter_test.cc b/tensorflow/compiler/xla/service/hlo_element_type_converter_test.cc index cb94d9f19b825d1321263a4737b66a6bf198a772..5c5a059e0fd895f03bc26a975609b57333237faf 100644 --- a/tensorflow/compiler/xla/service/hlo_element_type_converter_test.cc +++ b/tensorflow/compiler/xla/service/hlo_element_type_converter_test.cc @@ -22,6 +22,12 @@ namespace { namespace op = xla::testing::opcode_matchers; +using ::testing::Contains; +using ::testing::ElementsAre; +using ::testing::Eq; +using ::testing::Not; +using ::testing::ResultOf; + class HloElementTypeConverterTest : public HloTestBase { public: std::unique_ptr CreateModuleFromHloString( @@ -117,5 +123,65 @@ TEST_F(HloElementTypeConverterTest, BatchNormGradBF16Converted) { op::Convert(op::GetTupleElement(batch_norm, 2)))); } +TEST_F(HloElementTypeConverterTest, RngIsRemoved) { + const string& hlo_string = R"( +HloModule RngIsRemoved + +ENTRY main { + constant.3 = bf16[] constant(0) + constant.4 = bf16[] constant(1) + ROOT rng = bf16[1,1000,20]{2,1,0} rng(constant.3, constant.4), distribution=rng_uniform +} + )"; + auto module = CreateModuleFromHloString(hlo_string); + HloElementTypeConverter type_converter(BF16, F32); + TF_ASSERT_OK_AND_ASSIGN(bool converted, type_converter.Run(module.get())); + EXPECT_TRUE(converted); + + std::function is_bf16_rng = + [](const HloInstruction* inst) { + return inst->shape().element_type() == BF16 && + inst->opcode() == HloOpcode::kRng; + }; + + EXPECT_THAT(module->entry_computation()->instructions(), + Not(Contains(ResultOf(is_bf16_rng, Eq(true))))); +} + +TEST_F(HloElementTypeConverterTest, RngCtrlDep) { + const string& hlo_string = R"( +HloModule RngIsRemoved + +ENTRY main { + constant.3 = bf16[] constant(0) + constant.4 = bf16[] constant(1) + rng0 = bf16[1,2000,20]{2,1,0} rng(constant.3, constant.4), distribution=rng_uniform + ROOT rng1 = bf16[1,1000,20]{2,1,0} rng(constant.3, constant.4), control-predecessors={%rng0}, distribution=rng_uniform +} + )"; + auto module = CreateModuleFromHloString(hlo_string); + + HloElementTypeConverter type_converter(BF16, F32); + TF_ASSERT_OK_AND_ASSIGN(bool converted, type_converter.Run(module.get())); + EXPECT_TRUE(converted); + + HloInstruction *rng0, *rng1; + for (auto* inst : module->entry_computation()->instructions()) { + if (inst->opcode() == HloOpcode::kRng) { + const Shape& shape = inst->shape(); + ASSERT_EQ(shape.dimensions_size(), 3); + ASSERT_TRUE(shape.dimensions(1) == 2000 || shape.dimensions(1) == 1000); + if (shape.dimensions(1) == 2000) { + rng0 = inst; + } else { + rng1 = inst; + } + } + } + + EXPECT_THAT(rng0->control_successors(), ElementsAre(rng1)); + EXPECT_THAT(rng1->control_predecessors(), ElementsAre(rng0)); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.cc b/tensorflow/compiler/xla/service/hlo_evaluator.cc index b4f9a9db9cbcae56fbf60ad9d2ef4b3e0ffe2a90..c5e30148345fec2029bf533fcfa9deb89662ec83 100644 --- a/tensorflow/compiler/xla/service/hlo_evaluator.cc +++ b/tensorflow/compiler/xla/service/hlo_evaluator.cc @@ -1604,8 +1604,8 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { // Evaluate computation with specified literal operands. auto curr_val_literal = Literal::CreateR0(curr_val); auto result_val_literal = Literal::CreateR0(result_val); - std::vector args = {curr_val_literal.get(), - result_val_literal.get()}; + std::vector args = {result_val_literal.get(), + curr_val_literal.get()}; std::unique_ptr computed_result = embedded_evaluator.Evaluate(*function, args) @@ -1804,7 +1804,7 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { const auto result_val_literal = Literal::CreateR0(result_val); const std::vector args = { - curr_val_literal.get(), result_val_literal.get()}; + result_val_literal.get(), curr_val_literal.get()}; std::unique_ptr computed_result = embedded_evaluator.Evaluate(*function, args) .ConsumeValueOrDie(); @@ -1853,6 +1853,34 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { return Status::OK(); } + // Enable CLZ only for int32 and uint32. + template < + typename NativeT, + typename std::enable_if< + (std::is_floating_point::value || + std::is_integral::value || is_complex_t::value) && + !(std::is_same::value || + std::is_same::value)>::type* = nullptr> + Status HandleClz(HloInstruction* clz) { + return InvalidArgument("Unsupported type for Clz"); + } + + template ::value || + std::is_same::value>::type* = nullptr> + Status HandleClz(HloInstruction* clz) { + TF_ASSIGN_OR_RETURN(parent_->evaluated_[clz], + ElementWiseUnaryOp(clz, [](ElementwiseT elem_operand) { + return 31 - tensorflow::Log2Floor(elem_operand); + })); + return Status::OK(); + } + + Status HandleClz(HloInstruction* clz) override { + return HandleClz(clz); + } + template ::value>::type* = nullptr> Status HandleSin(HloInstruction* sin) { diff --git a/tensorflow/compiler/xla/service/hlo_execution_profile.h b/tensorflow/compiler/xla/service/hlo_execution_profile.h index 6fb91b9bef9d1df82b8806ce79cc147823edeb3d..be989846ef5cd2645da88ac9bbfea9534dd47821 100644 --- a/tensorflow/compiler/xla/service/hlo_execution_profile.h +++ b/tensorflow/compiler/xla/service/hlo_execution_profile.h @@ -88,7 +88,7 @@ std::unique_ptr CreateHloProfilePrinterData( // down how much time each HLO took. class HloExecutionProfile { public: - using DeviceDescription = perftools::gputools::DeviceDescription; + using DeviceDescription = se::DeviceDescription; HloExecutionProfile(const HloProfilePrinterData* hlo_profile_printer_data, const HloProfileIndexMap* hlo_profile_index_map); diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc index c35783c456c63b9a651d1221cf9a3d70af38ba66..516e14b4642ae6665a2d15c91715dc9b057ab41a 100644 --- a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc +++ b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc @@ -909,6 +909,7 @@ ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { case HloOpcode::kBitcastConvert: case HloOpcode::kCeil: case HloOpcode::kClamp: + case HloOpcode::kClz: case HloOpcode::kComplex: case HloOpcode::kConvert: case HloOpcode::kCos: @@ -956,7 +957,6 @@ ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { case HloOpcode::kTuple: return kWhite; case HloOpcode::kBroadcast: - case HloOpcode::kBroadcastDimOne: // De-emphasize nodes which broadcast a scalar within a fusion node -- // these are essentially free. if (instr->IsFused() && diff --git a/tensorflow/compiler/xla/service/hlo_instruction.cc b/tensorflow/compiler/xla/service/hlo_instruction.cc index 8149e47cb532b832d8b77a2422b12cd09fa8a973..a714d0e114245021c28da26beae444dbd3d99bb5 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.cc +++ b/tensorflow/compiler/xla/service/hlo_instruction.cc @@ -159,6 +159,23 @@ StatusOr> HloInstruction::CreateFromProto( instruction->fft_length_.push_back(fft_len); } + if (proto.has_sharding()) { + TF_ASSIGN_OR_RETURN(const auto& sharding, + HloSharding::FromProto(proto.sharding())); + instruction->set_sharding(sharding); + } + + if (proto.has_gather_dimension_numbers()) { + instruction->gather_dimension_numbers_ = + MakeUnique(proto.gather_dimension_numbers()); + } + for (int64 bound : proto.gather_window_bounds()) { + instruction->gather_window_bounds_.push_back(bound); + } + + instruction->channel_name_ = proto.channel_name(); + instruction->cost_estimate_ns_ = proto.cost_estimate_ns(); + return std::move(instruction); } @@ -237,6 +254,7 @@ HloInstruction::CreateGetTupleElement(const Shape& shape, case HloOpcode::kCeil: case HloOpcode::kCopy: case HloOpcode::kCos: + case HloOpcode::kClz: case HloOpcode::kExp: case HloOpcode::kFloor: case HloOpcode::kImag: @@ -683,15 +701,6 @@ HloInstruction::CreateSelectAndScatter( return instruction; } -/* static */ std::unique_ptr -HloInstruction::CreateBroadcastDimOne(const Shape& shape, - HloInstruction* operand) { - auto instruction = - WrapUnique(new HloInstruction(HloOpcode::kBroadcastDimOne, shape)); - instruction->AppendOperand(operand); - return instruction; -} - /* static */ std::unique_ptr HloInstruction::CreateBroadcastSequence( const Shape& output_shape, HloInstruction* operand, @@ -821,6 +830,16 @@ static string FusionNodeName(HloInstruction::FusionKind fusion_kind) { return instruction; } +void HloInstruction::SetupDerivedInstruction( + HloInstruction* derived_instruction) const { + if (sharding_ != nullptr) { + derived_instruction->set_sharding(*sharding_); + } else { + derived_instruction->clear_sharding(); + } + derived_instruction->set_metadata(metadata_); +} + HloInstruction* HloInstruction::AddFusionOperand(HloInstruction* new_operand) { CHECK_EQ(opcode(), HloOpcode::kFusion); CHECK_EQ(operand_count(), @@ -1230,6 +1249,7 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( case HloOpcode::kRoundNearestAfz: case HloOpcode::kBitcast: case HloOpcode::kCeil: + case HloOpcode::kClz: case HloOpcode::kCopy: case HloOpcode::kCos: case HloOpcode::kExp: @@ -1284,10 +1304,6 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( CHECK_EQ(new_operands.size(), 1); clone = CreateBroadcast(shape, new_operands[0], dimensions_); break; - case HloOpcode::kBroadcastDimOne: - CHECK_EQ(new_operands.size(), 1); - clone = CreateBroadcastDimOne(shape, new_operands[0]); - break; case HloOpcode::kCall: clone = CreateCall(shape, new_operands, to_apply()); break; @@ -1463,10 +1479,7 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( case HloOpcode::kTrace: LOG(FATAL) << "Not yet implemented, clone: " << HloOpcodeString(opcode_); } - clone->set_metadata(metadata_); - if (has_sharding()) { - clone->set_sharding(sharding()); - } + SetupDerivedInstruction(clone.get()); clone->set_parent(parent_); return clone; } @@ -1665,14 +1678,35 @@ Status HloInstruction::AddControlDependencyTo(HloInstruction* instruction) { } Status HloInstruction::RemoveControlDependencyTo(HloInstruction* instruction) { - auto succ_it = std::find(control_successors_.begin(), - control_successors_.end(), instruction); - TF_RET_CHECK(succ_it != control_successors_.end()); - control_successors_.erase(succ_it); - auto pred_it = std::find(instruction->control_predecessors_.begin(), - instruction->control_predecessors_.end(), this); - TF_RET_CHECK(pred_it != instruction->control_predecessors_.end()); - instruction->control_predecessors_.erase(pred_it); + TF_RET_CHECK(instruction->parent() == parent()); + TF_RETURN_IF_ERROR(EraseElementFromVector(&control_successors_, instruction)); + TF_RETURN_IF_ERROR( + EraseElementFromVector(&instruction->control_predecessors_, this)); + return Status::OK(); +} + +Status HloInstruction::DropAllControlDeps() { + for (auto* ctrl_succ : control_successors_) { + TF_RETURN_IF_ERROR( + EraseElementFromVector(&ctrl_succ->control_predecessors_, this)); + } + for (auto* ctrl_pred : control_predecessors_) { + TF_RETURN_IF_ERROR( + EraseElementFromVector(&ctrl_pred->control_successors_, this)); + } + control_successors_.clear(); + control_predecessors_.clear(); + return Status::OK(); +} + +Status HloInstruction::CopyAllControlDepsFrom(const HloInstruction* inst) { + for (auto* ctrl_pred : inst->control_predecessors()) { + TF_RETURN_IF_ERROR(ctrl_pred->AddControlDependencyTo(this)); + } + + for (auto* ctrl_succ : inst->control_successors()) { + TF_RETURN_IF_ERROR(this->AddControlDependencyTo(ctrl_succ)); + } return Status::OK(); } @@ -1717,6 +1751,7 @@ bool HloInstruction::IdenticalSlowPath( case HloOpcode::kAdd: case HloOpcode::kCeil: case HloOpcode::kClamp: + case HloOpcode::kClz: case HloOpcode::kComplex: case HloOpcode::kCopy: case HloOpcode::kCos: @@ -1839,8 +1874,6 @@ bool HloInstruction::IdenticalSlowPath( // Remaining instructions with special values. case HloOpcode::kBitcast: - case HloOpcode::kBroadcastDimOne: - case HloOpcode::kDynamicUpdateSlice: return eq_shapes(shape(), other.shape()); case HloOpcode::kBroadcast: return eq_shapes(shape(), other.shape()) && @@ -1859,6 +1892,8 @@ bool HloInstruction::IdenticalSlowPath( case HloOpcode::kDynamicSlice: return eq_shapes(shape(), other.shape()) && dynamic_slice_sizes_ == other.dynamic_slice_sizes_; + case HloOpcode::kDynamicUpdateSlice: + return eq_shapes(shape(), other.shape()); case HloOpcode::kCall: case HloOpcode::kMap: return eq_computations(to_apply(), other.to_apply()); @@ -2416,6 +2451,9 @@ HloInstructionProto HloInstruction::ToProto() const { proto.add_fft_length(fft_len); } + proto.set_channel_name(channel_name_); + proto.set_cost_estimate_ns(cost_estimate_ns_); + return proto; } @@ -2639,6 +2677,8 @@ Status HloInstruction::Visit(DfsHloVisitorBase* visitor) { return visitor->HandleFloor(this); case HloOpcode::kCeil: return visitor->HandleCeil(this); + case HloOpcode::kClz: + return visitor->HandleClz(this); case HloOpcode::kLog: return visitor->HandleLog(this); case HloOpcode::kTanh: @@ -2659,8 +2699,6 @@ Status HloInstruction::Visit(DfsHloVisitorBase* visitor) { return visitor->HandleBitcast(this); case HloOpcode::kBroadcast: return visitor->HandleBroadcast(this); - case HloOpcode::kBroadcastDimOne: - return visitor->HandleBroadcastDimOne(this); case HloOpcode::kPad: return visitor->HandlePad(this); case HloOpcode::kReshape: @@ -2982,6 +3020,7 @@ bool HloInstruction::IsElementwise() const { case HloOpcode::kAbs: case HloOpcode::kRoundNearestAfz: case HloOpcode::kCeil: + case HloOpcode::kClz: case HloOpcode::kConvert: case HloOpcode::kBitcastConvert: case HloOpcode::kCopy: diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h index a6cb19f331695da996769e8bb5ad16cb5cfd9bf1..f3da3fc256efe8dc9cc85bcbc6cc0c064c770d5c 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.h +++ b/tensorflow/compiler/xla/service/hlo_instruction.h @@ -401,10 +401,6 @@ class HloInstruction { const Shape& shape, HloInstruction* operand, tensorflow::gtl::ArraySlice broadcast_dimensions); - // Creates a broadcast-size-one-dimensions instruction. - static std::unique_ptr CreateBroadcastDimOne( - const Shape& shape, HloInstruction* operand); - // Creates a sequence of instructions that performs an explicit broadcast of // the operand to the target shape. // @@ -561,6 +557,18 @@ class HloInstruction { // 'instruction'. Status RemoveControlDependencyTo(HloInstruction* instruction); + // Drops all control predecessors and successors from this HLO instruction. + Status DropAllControlDeps(); + + // Copies the control predecessors and successors on this HLO instruction to + // `inst`. Does not do a deep copy so this makes sense only if `inst` and + // this HLO are in the same module. + // + // Depending on the use cases we see in practice, in the future we may + // consider folding the logic here into Clone, CloneWithNewOperands and + // ReplaceAllUsesWith by treating control dependencies like data dependencies. + Status CopyAllControlDepsFrom(const HloInstruction* inst); + // Returns the set of control predecessors (successors) of this // instruction. Control predecessors (successors) must execute before (after) // the current instruction. @@ -948,6 +956,21 @@ class HloInstruction { void clear_sharding() { sharding_ = nullptr; } // Return true if this operator has a sharding assigned. bool has_sharding() const { return sharding_ != nullptr; } + // Checks whether the instruction has compatible sharding with the other + // instruction. + bool has_compatible_sharding(const HloInstruction* other) const { + if (!has_sharding()) { + return !other->has_sharding(); + } + return other->has_sharding() ? sharding() == other->sharding() : false; + } + + // When creating a new instruction which either replaces, or shifts up (kCopy + // insertion case), another instruction, we need to make sure the certain + // properties of the new instruction are copied into the derived one. As of + // today, the metadata and sharding will be propagated to the derived + // instruction. + void SetupDerivedInstruction(HloInstruction* derived_instruction) const; // Adds a new operand the fusion instruction. HloInstruction* AddFusionOperand(HloInstruction* new_operand); @@ -1145,17 +1168,17 @@ class HloInstruction { // Clones the HLO instruction. The clone will have the same opcode, shape, and // operands. After creation the clone has no uses. "this" (the instruction // cloned from) is not changed. Suffix is the string to append to the name of - // the instruction to form the name of the cloned instruction. - // If the module pointer is not nullptr, it will be the module where - // the cloned computations will be added to (in order to support deep - // cloning). + // the instruction to form the name of the cloned instruction. If the module + // pointer is not nullptr, it will be the module where the cloned computations + // will be added to (in order to support deep cloning). Ignores the control + // predecessors and successors of this HLO instruction. std::unique_ptr Clone(const string& suffix = "clone", HloModule* module = nullptr) const; - // Clones the HLO instruction as above but with new shape and operands. - // If the module pointer is not nullptr, it will be the module where - // the cloned computations will be added to (in order to support deep - // cloning). + // Clones the HLO instruction as above but with new shape and operands. If + // the module pointer is not nullptr, it will be the module where the cloned + // computations will be added to (in order to support deep cloning). Ignores + // the control predecessors and successors of this HLO instruction. std::unique_ptr CloneWithNewOperands( const Shape& shape, tensorflow::gtl::ArraySlice operands, HloModule* module = nullptr) const; @@ -1446,7 +1469,7 @@ class HloInstruction { string channel_name_; // Estimate of the duration of a host computation in nanoseconds. - int64 cost_estimate_ns_; + int64 cost_estimate_ns_ = 0; // Computations called by this instruction. std::vector called_computations_; diff --git a/tensorflow/compiler/xla/service/hlo_opcode.h b/tensorflow/compiler/xla/service/hlo_opcode.h index dddc72480f93c4c3cc29f41db99fa773dc8d6b68..ca763076a16af1150a8623fb7dbf22c46a5ca263 100644 --- a/tensorflow/compiler/xla/service/hlo_opcode.h +++ b/tensorflow/compiler/xla/service/hlo_opcode.h @@ -54,10 +54,10 @@ namespace xla { V(kBitcast, "bitcast") \ V(kBitcastConvert, "bitcast-convert") \ V(kBroadcast, "broadcast") \ - V(kBroadcastDimOne, "broadcast-dim-one") \ V(kCall, "call", kHloOpcodeIsVariadic) \ V(kCeil, "ceil") \ V(kClamp, "clamp") \ + V(kClz, "count-leading-zeros") \ V(kComplex, "complex") \ V(kConcatenate, "concatenate", kHloOpcodeIsVariadic) \ V(kConditional, "conditional") \ diff --git a/tensorflow/compiler/xla/service/hlo_runner.cc b/tensorflow/compiler/xla/service/hlo_runner.cc index ec7d8210a70ad7498f77fe807abd53544d4b0487..df5ffd0b7d6f0f4575c5f3d8aaa7bff4a1b1e771 100644 --- a/tensorflow/compiler/xla/service/hlo_runner.cc +++ b/tensorflow/compiler/xla/service/hlo_runner.cc @@ -16,27 +16,20 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_runner.h" -#include #include #include +#include "absl/memory/memory.h" #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/ptr_util.h" -#include "tensorflow/compiler/xla/service/backend.h" -#include "tensorflow/compiler/xla/service/executable.h" -#include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/transfer_manager.h" #include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/tools/parser/hlo_parser.h" -#include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/common_runtime/eigen_thread_pool.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" -namespace se = ::perftools::gputools; - namespace xla { /*static*/ StatusOr> @@ -91,15 +84,6 @@ HloRunner::ReadModuleFromHloTextFile(const std::string& filename, return tools::Parse(hlo_string, config); } -// Define this in .cc file to avoid having to include eigen or forward declare -// these types in the header. -struct HloRunner::EigenThreadPoolWrapper { - std::unique_ptr pool; - std::unique_ptr device; -}; - -HloRunner::HloRunner() {} - HloRunner::HloRunner(se::Platform* platform) { BackendOptions backend_options; backend_options.set_platform(platform); @@ -113,61 +97,45 @@ StatusOr> HloRunner::Execute( std::unique_ptr module, const tensorflow::gtl::ArraySlice arguments, bool run_hlo_passes) { - if (run_hlo_passes) { - TF_ASSIGN_OR_RETURN( - module, backend().compiler()->RunHloPasses( - std::move(module), backend().default_stream_executor(), - /*device_allocator=*/nullptr)); - } - TF_ASSIGN_OR_RETURN( - std::unique_ptr executable, - backend().compiler()->RunBackend(std::move(module), - backend().default_stream_executor(), - /*device_allocator=*/nullptr)); - + TF_ASSIGN_OR_RETURN(std::unique_ptr executable, + CreateExecutable(std::move(module), run_hlo_passes)); se::Stream stream(backend().default_stream_executor()); stream.Init(); - ExecutableRunOptions run_options; - run_options.set_device_ordinal(backend().default_device_ordinal()); - run_options.set_stream(&stream); - run_options.set_allocator(backend().memory_allocator()); - run_options.set_inter_op_thread_pool(backend().inter_op_thread_pool()); - run_options.set_intra_op_thread_pool( - backend().eigen_intra_op_thread_pool_device()); - - ServiceExecutableRunOptions service_run_options( - run_options, backend().StreamBorrower(), - backend().inter_op_thread_pool()); + ServiceExecutableRunOptions service_run_options(GetServiceRunOptionsForDevice( + backend().default_device_ordinal(), &stream, nullptr)); + const ExecutableRunOptions& run_options = service_run_options.run_options(); // Copy arguments to device. - std::vector> argument_buffers; - std::vector argument_buffer_ptrs; + std::vector argument_buffers; for (Literal* argument : arguments) { TF_ASSIGN_OR_RETURN( - std::unique_ptr argument_buffer, + ScopedShapedBuffer argument_buffer, backend().transfer_manager()->AllocateScopedShapedBuffer( argument->shape(), run_options.allocator(), run_options.device_ordinal())); TF_RETURN_IF_ERROR(backend().transfer_manager()->TransferLiteralToDevice( - stream.parent(), *argument, *argument_buffer)); + stream.parent(), *argument, argument_buffer)); argument_buffers.push_back(std::move(argument_buffer)); - argument_buffer_ptrs.push_back(argument_buffers.back().get()); + } + + std::vector argument_buffer_ptrs; + argument_buffer_ptrs.reserve(argument_buffers.size()); + for (const auto& buf : argument_buffers) { + argument_buffer_ptrs.push_back(&buf); } TF_ASSIGN_OR_RETURN( - std::unique_ptr result, + ShapedBuffer result, 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. - TF_ASSIGN_OR_RETURN( - std::unique_ptr scoped_result, - ScopedShapedBuffer::MakeScoped(result.get(), run_options.allocator())); + ScopedShapedBuffer scoped_result(std::move(result), run_options.allocator()); auto result_literal = backend().transfer_manager()->TransferLiteralFromDevice( - stream.parent(), *scoped_result); + stream.parent(), scoped_result); if (result_literal.ok()) { VLOG(4) << "Executed binary and got result: " << result_literal.ValueOrDie()->ToString(); @@ -178,10 +146,158 @@ StatusOr> HloRunner::Execute( return result_literal; } +StatusOr>> HloRunner::ExecuteReplicated( + std::unique_ptr module, + const ReplicatedExecuteOptions& options) { + TF_ASSIGN_OR_RETURN( + std::unique_ptr executable, + CreateExecutable(std::move(module), options.run_hlo_passes)); + TF_ASSIGN_OR_RETURN( + DeviceAssignment device_assignment, + backend().computation_placer()->AssignDevices(options.num_replicas, 1)); + std::vector> streams; + std::vector service_run_options; + + std::vector argument_buffers; + // This reserve() call is necessary for correctness, because + // argument_buffer_ptrs contains pointers into the elements of + // argument_buffers. + argument_buffers.reserve(options.num_replicas * options.arguments.size()); + + // Plus one so we can safely get &argument_buffer_ptrs[0] in case there are + // no arguments. + std::vector argument_buffer_ptrs( + options.num_replicas * options.arguments.size() + 1); + std::vector> + argument_buffer_slices; + int64 index = 0; + for (int64 i = 0; i < options.num_replicas; ++i) { + int64 device = device_assignment(i, 0); + TF_ASSIGN_OR_RETURN(se::StreamExecutor * executor, + backend().stream_executor(device)); + streams.push_back(absl::make_unique(executor)); + streams.back()->Init(); + service_run_options.emplace_back(GetServiceRunOptionsForDevice( + device, streams.back().get(), &device_assignment)); + + // Copy arguments to device. + for (const Literal* argument : options.arguments) { + TF_ASSIGN_OR_RETURN( + ScopedShapedBuffer argument_buffer, + backend().transfer_manager()->AllocateScopedShapedBuffer( + argument->shape(), backend().memory_allocator(), device)); + TF_RETURN_IF_ERROR(backend().transfer_manager()->TransferLiteralToDevice( + executor, *argument, argument_buffer)); + argument_buffers.push_back(std::move(argument_buffer)); + argument_buffer_ptrs[index++] = &argument_buffers.back(); + } + argument_buffer_slices.emplace_back( + &argument_buffer_ptrs[index - options.arguments.size()], + options.arguments.size()); + } + + std::unique_ptr pool; + int64 num_threads = (options.infeed != nullptr) ? options.num_replicas : 0; + if (ShapeUtil::IsInitialized(options.outfeed_shape)) { + num_threads += options.num_replicas; + } + if (num_threads > 0) { + pool = absl::make_unique( + tensorflow::Env::Default(), "infeed_outfeed", + /*num_threads=*/num_threads); + } + if (options.infeed != nullptr) { + for (int64 i = 0; i < options.num_replicas; ++i) { + int64 device = device_assignment(i, 0); + pool->Schedule([this, device, &options]() { + se::StreamExecutor* executor = + backend().stream_executor(device).ValueOrDie(); + VLOG(1) << "Starting infeed on device " << device; + for (int64 step = 1; + options.infeed_steps < 0 || step <= options.infeed_steps; ++step) { + TF_CHECK_OK(backend().transfer_manager()->TransferLiteralToInfeed( + executor, *options.infeed)); + if (step % 100 == 0) { + VLOG(1) << "Infeed step " << step; + } + } + }); + } + } + if (ShapeUtil::IsInitialized(options.outfeed_shape)) { + for (int64 i = 0; i < options.num_replicas; ++i) { + int64 device = device_assignment(i, 0); + pool->Schedule([this, device, &options]() { + se::StreamExecutor* executor = + backend().stream_executor(device).ValueOrDie(); + VLOG(1) << "Starting outfeed on device " << device; + for (int64 step = 1; + options.infeed_steps < 0 || step <= options.infeed_steps; ++step) { + auto literal = absl::make_unique(); + TF_CHECK_OK(backend().transfer_manager()->TransferLiteralFromOutfeed( + executor, options.outfeed_shape, literal.get())); + if (options.outfeed_values != nullptr) { + options.outfeed_values->push_back(std::move(literal)); + } + if (step % 100 == 0) { + VLOG(1) << "Outfeed step " << step; + } + } + }); + } + } + + LOG(INFO) << "Replicated execution started"; + TF_ASSIGN_OR_RETURN(std::vector results, + executable->ExecuteOnStreams(service_run_options, + argument_buffer_slices)); + LOG(INFO) << "Replicated execution terminated"; + + std::vector> exec_results; + for (int64 i = 0; i < options.num_replicas; ++i) { + ScopedShapedBuffer result(std::move(results[i]), + backend().memory_allocator()); + TF_ASSIGN_OR_RETURN(std::unique_ptr literal, + backend().transfer_manager()->TransferLiteralFromDevice( + streams[i]->parent(), result)); + exec_results.push_back(std::move(literal)); + } + return std::move(exec_results); +} + +StatusOr> HloRunner::CreateExecutable( + std::unique_ptr module, bool run_hlo_passes) { + if (run_hlo_passes) { + TF_ASSIGN_OR_RETURN( + module, backend().compiler()->RunHloPasses( + std::move(module), backend().default_stream_executor(), + backend().memory_allocator())); + } + return backend().compiler()->RunBackend(std::move(module), + backend().default_stream_executor(), + backend().memory_allocator()); +} + +ServiceExecutableRunOptions HloRunner::GetServiceRunOptionsForDevice( + int64 device, se::Stream* stream, DeviceAssignment* device_assignment) { + ExecutableRunOptions run_options; + run_options.set_device_ordinal(device); + run_options.set_stream(stream); + run_options.set_allocator(backend().memory_allocator()); + run_options.set_inter_op_thread_pool(backend().inter_op_thread_pool()); + run_options.set_intra_op_thread_pool( + backend().eigen_intra_op_thread_pool_device()); + if (device_assignment != nullptr) { + run_options.set_device_assignment(device_assignment); + } + return ServiceExecutableRunOptions(run_options, backend().StreamBorrower(), + backend().inter_op_thread_pool()); +} + Backend& HloRunner::backend() { if (!backend_) { backend_ = Backend::CreateDefaultBackend().ConsumeValueOrDie(); - VLOG(1) << "executing on platform " << backend().platform()->Name(); + VLOG(1) << "Executing on platform " << backend().platform()->Name(); } return *backend_; } diff --git a/tensorflow/compiler/xla/service/hlo_runner.h b/tensorflow/compiler/xla/service/hlo_runner.h index 06ce22a5b9fc7b3d6c10857c84196094c0eed303..53f7c6fe4a09111c5ee24f2290f0f4aeed0a4401 100644 --- a/tensorflow/compiler/xla/service/hlo_runner.h +++ b/tensorflow/compiler/xla/service/hlo_runner.h @@ -16,12 +16,16 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_RUNNER_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_RUNNER_H_ +#include #include +#include #include #include #include "tensorflow/compiler/xla/service/backend.h" #include "tensorflow/compiler/xla/service/compiler.h" +#include "tensorflow/compiler/xla/service/computation_placer.h" +#include "tensorflow/compiler/xla/service/executable.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -40,9 +44,43 @@ namespace xla { // file), or parsed from a hlo textual IR string. class HloRunner { public: - HloRunner(); - - HloRunner(::perftools::gputools::Platform* platform); + // The options used to configure a ExecuteReplicated() call. + struct ReplicatedExecuteOptions { + // The number of devices the HLO module should be replicated onto. + int64 num_replicas = 1; + + // The arguments to be fed to each replica. Since this is used for a + // replicated execution, all the arguments are the same for all replicas. + std::vector arguments; + + // If the HLO module being run has an infeed instruction, this will be the + // data which will be fed to it, for as many as infeed_steps steps. + const Literal* infeed = nullptr; + + // The number of times the infeed literal should be fed to the HLO module. + // For a clean exit, this should match the iterations-per-loop parameter + // used when generating the HLO module proto (that is usually the main + // while bounary counter). A value higher then iterations-per-loop would + // lead to infeed threads feeding to a gone computation, while a lower + // value would trigger a stuck ExecuteReplicated() call (the computation + // will be trying to infeed data which will never come). + int64 infeed_steps = -1; + + // The shape of the outfeed operation. If empty, the HLO module does not + // generate any outfeed. + Shape outfeed_shape; + + // A pointer to a vector where the outfeed values will be stored. If + // nullptr, the values will be read and discarded. + std::vector>* outfeed_values = nullptr; + + // Whether the HLO passes should be run on the input module. Usually + // saved modules are coming from after the HLO pass pipeline, so triggering + // another run will likely cause errors. + bool run_hlo_passes = false; + }; + + explicit HloRunner(se::Platform* platform); ~HloRunner(); @@ -86,6 +124,13 @@ class HloRunner { return Execute(std::move(module), argument_pointers, run_hlo_passes); } + // Executes a given HLO module into a set of replicas, and returns a map + // with the replica number as key, and the corresponding returned literal as + // value. + StatusOr>> ExecuteReplicated( + std::unique_ptr module, + const ReplicatedExecuteOptions& options); + // If backend is not created in the constructor, creates and returns the // default backend. If creation fails, crashes the program. // @@ -94,9 +139,17 @@ class HloRunner { Backend& backend(); private: - struct EigenThreadPoolWrapper; - - std::unique_ptr thread_pool_wrapper_; + // Creates an executable object given an HLO module. If run_hlo_passes is + // true, the HLO passes will be run before. + StatusOr> CreateExecutable( + std::unique_ptr module, bool run_hlo_passes); + + // Creates a ServiceExecutableRunOptions object to configure a run on device, + // using the provided stream object. If device_assignment is not nullptr, it + // will be used to configure the replication parameters. Replicated executions + // should pass the device_assignment parameter. + ServiceExecutableRunOptions GetServiceRunOptionsForDevice( + int64 device, se::Stream* stream, DeviceAssignment* device_assignment); std::unique_ptr backend_; }; diff --git a/tensorflow/compiler/xla/service/hlo_sharding.cc b/tensorflow/compiler/xla/service/hlo_sharding.cc index e8e45f1ee968992901988e8b85d4e9ae28f2abe9..1b42349b0b3ad9634bb910b3843affed6a0ca334 100644 --- a/tensorflow/compiler/xla/service/hlo_sharding.cc +++ b/tensorflow/compiler/xla/service/hlo_sharding.cc @@ -376,6 +376,16 @@ HloSharding HloSharding::TransformShardedTileShape( return HloSharding::Tile(new_tile_shape, tile_assignment()); } +HloSharding HloSharding::GetSubSharding(const Shape& shape, + const ShapeIndex& index) const { + CHECK(IsTuple()); + + ShapeTree sub_shape_tree(ShapeUtil::GetSubshape(shape, index), + Replicate()); + sub_shape_tree.CopySubtreeFrom(GetAsShapeTree(shape), index, {}); + return Tuple(sub_shape_tree); +} + std::ostream& operator<<(std::ostream& out, const HloSharding& sharding) { out << sharding.ToString(); return out; diff --git a/tensorflow/compiler/xla/service/hlo_sharding.h b/tensorflow/compiler/xla/service/hlo_sharding.h index 06204acbca30648e73382cb4641139e852664b77..2b8e757f42991f697df37d3d34bfdff6a36bc509 100644 --- a/tensorflow/compiler/xla/service/hlo_sharding.h +++ b/tensorflow/compiler/xla/service/hlo_sharding.h @@ -175,6 +175,10 @@ class HloSharding { } } + // Retrieves the sub sharding at a given index, out of a tuple sharding. + // REQUIRES: IsTuple() + HloSharding GetSubSharding(const Shape& shape, const ShapeIndex& index) const; + bool operator==(const HloSharding& other) const { return replicated_ == other.replicated_ && maximal_ == other.maximal_ && ShapeUtil::Compatible(tile_shape_, other.tile_shape_) && diff --git a/tensorflow/compiler/xla/service/hlo_verifier.cc b/tensorflow/compiler/xla/service/hlo_verifier.cc index 63ec5964eb935239e86233c1ae94e2bcce3b0461..8a30cbf9cd622ffb64d345ddaf0dc88f34850bfc 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.cc +++ b/tensorflow/compiler/xla/service/hlo_verifier.cc @@ -15,6 +15,7 @@ limitations under the License. #include +#include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/hlo_verifier.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/core/lib/core/errors.h" @@ -174,34 +175,17 @@ Status ShapeVerifier::HandleBroadcast(HloInstruction* broadcast) { TF_RETURN_IF_ERROR(CheckShape(broadcast, broadcast->shape())); TF_RET_CHECK(ShapeUtil::Rank(operand_shape) == broadcast->dimensions().size()); - for (int64 i = 0; i < ShapeUtil::Rank(operand_shape); ++i) { - int64 output_dimension = broadcast->dimensions()[i]; + for (int64 operand_dimension = 0; + operand_dimension < ShapeUtil::Rank(operand_shape); + ++operand_dimension) { + int64 output_dimension = broadcast->dimensions()[operand_dimension]; TF_RET_CHECK(broadcast->shape().dimensions(output_dimension) == - operand_shape.dimensions(i)) + operand_shape.dimensions(operand_dimension)) << broadcast->ToString() << " operand shape " << operand_shape; } return tensorflow::Status::OK(); } -Status ShapeVerifier::HandleBroadcastDimOne(HloInstruction* broadcastDimOne) { - const Shape& operand_shape = broadcastDimOne->operand(0)->shape(); - int64 operand_rank = ShapeUtil::Rank(operand_shape); - const Shape& output_shape = broadcastDimOne->shape(); - // Check for mixed precision. - TF_RETURN_IF_ERROR(CheckShape(broadcastDimOne, output_shape)); - TF_RET_CHECK(operand_rank == ShapeUtil::Rank(output_shape)); - for (int64 i = 0; i < operand_rank; ++i) { - int64 operand_dimension = operand_shape.dimensions(i); - int64 output_dimension = output_shape.dimensions(i); - TF_RET_CHECK(operand_dimension == 1 || - operand_dimension == output_dimension) - << "Dimension " << i << " of broadcastDimOne " - << broadcastDimOne->ToString() << " is " << operand_dimension - << ", expected 1 or " << output_dimension; - } - return tensorflow::Status::OK(); -} - Status ShapeVerifier::HandleReshape(HloInstruction* reshape) { // Check for mixed precision. TF_RETURN_IF_ERROR(CheckShape(reshape, reshape->shape())); @@ -748,6 +732,73 @@ Status HloVerifier::CheckFusionInstruction(HloInstruction* fusion) const { return tensorflow::Status::OK(); } +Status HloVerifier::CheckWhileInstruction(HloInstruction* instruction) { + auto* while_cond = instruction->while_condition(); + auto* while_body = instruction->while_body(); + if (while_cond->num_parameters() != 1) { + return FailedPrecondition( + "While condition must have exactly 1 parameter; had %lld : %s", + while_cond->num_parameters(), while_cond->ToString().c_str()); + } + if (while_body->num_parameters() != 1) { + return FailedPrecondition( + "While body must have exactly 1 parameter; had %lld : %s", + while_body->num_parameters(), while_body->ToString().c_str()); + } + if (instruction->operand_count() != 1) { + return FailedPrecondition( + "While loop must have exactly one operand; had %lld : %s", + instruction->operand_count(), instruction->ToString().c_str()); + } + auto* init = instruction->operand(0); + auto* cond_param = while_cond->parameter_instruction(0); + if (!ShapeUtil::Compatible(init->shape(), cond_param->shape())) { + return FailedPrecondition( + "While condition's parameter must have the same shape as the " + "loop's 'init'. init: %s, param: %s", + init->ToString().c_str(), cond_param->ToString().c_str()); + } + auto* cond_root = while_cond->root_instruction(); + if (!ShapeUtil::Compatible(cond_root->shape(), + ShapeUtil::MakeShape(PRED, {}))) { + return FailedPrecondition("While condition should have shape PRED: %s", + cond_root->ToString().c_str()); + } + auto* body_param = while_body->parameter_instruction(0); + if (!ShapeUtil::Compatible(init->shape(), body_param->shape())) { + return FailedPrecondition( + "While body's parameter must have the same shape as the loop's" + " 'init'. init: %s, param: %s", + init->ToString().c_str(), body_param->ToString().c_str()); + } + auto* body_root = while_body->root_instruction(); + if (!ShapeUtil::Compatible(init->shape(), body_root->shape())) { + return FailedPrecondition( + "While body should have same shape as the loop's 'init'." + "init: %s, body: %s", + init->ToString().c_str(), body_root->ToString().c_str()); + } + return tensorflow::Status::OK(); +} + +Status HloVerifier::CheckElementwiseInstruction(HloInstruction* instruction) { + const Shape& out_shape = instruction->shape(); + for (HloInstruction* operand : instruction->operands()) { + const Shape& operand_shape = operand->shape(); + if (!ShapeUtil::IsScalar(operand_shape) && + !ShapeUtil::CompatibleIgnoringElementType(operand_shape, out_shape)) { + return FailedPrecondition( + "Implicit broadcast is not allowed in HLO." + "Found non-compatible shapes for instruction %s.\n" + "output: %s\noperand: %s\n", + HloOpcodeString(instruction->opcode()).c_str(), + ShapeUtil::HumanString(out_shape).c_str(), + ShapeUtil::HumanString(operand_shape).c_str()); + } + } + return tensorflow::Status::OK(); +} + StatusOr HloVerifier::Run(HloModule* module) { TF_RETURN_IF_ERROR(VerifyHloStructure(module)); @@ -788,39 +839,9 @@ StatusOr HloVerifier::Run(HloModule* module) { << 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(); - TF_RET_CHECK(while_cond->num_parameters() == 1) - << "While condition must have exactly 1 parameter; had " - << while_cond->num_parameters() << ": " << while_cond->ToString(); - TF_RET_CHECK(while_body->num_parameters() == 1) - << "While body must have exactly 1 parameter; had " - << while_body->num_parameters() << ": " << while_body->ToString(); - TF_RET_CHECK(instruction->operand_count() == 1) - << "While loop must have exactly one operand; had " - << instruction->operand_count() << ": " << instruction->ToString(); - - auto* init = instruction->operand(0); - auto* cond_param = while_cond->parameter_instruction(0); - TF_RET_CHECK(ShapeUtil::Compatible(init->shape(), cond_param->shape())) - << "While condition's parameter must have the same shape as the " - "loop's 'init'. init: " - << init->ToString() << ", param: " << cond_param->ToString(); - auto* cond_root = while_cond->root_instruction(); - TF_RET_CHECK(ShapeUtil::Compatible(cond_root->shape(), - ShapeUtil::MakeShape(PRED, {}))) - << "While condition should have shape PRED: " - << cond_root->ToString(); - - auto* body_param = while_body->parameter_instruction(0); - TF_RET_CHECK(ShapeUtil::Compatible(init->shape(), body_param->shape())) - << "While body's parameter must have the same shape as the loop's " - "'init'. init: " - << init->ToString() << ", param: " << body_param->ToString(); - auto* body_root = while_body->root_instruction(); - TF_RET_CHECK(ShapeUtil::Compatible(init->shape(), body_root->shape())) - << "While body should have same shape as the loop's 'init'. init: " - << init->ToString() << ", body: " << body_root->ToString(); + TF_RETURN_IF_ERROR(CheckWhileInstruction(instruction)); + } else if (instruction->IsElementwise()) { + TF_RETURN_IF_ERROR(CheckElementwiseInstruction(instruction)); } auto previous = instructions.find(instruction->name()); diff --git a/tensorflow/compiler/xla/service/hlo_verifier.h b/tensorflow/compiler/xla/service/hlo_verifier.h index a4dff977ba268137d8ab94c576b4b511e911806f..6208887547a14d22b512ef308dd2668af2f4468d 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.h +++ b/tensorflow/compiler/xla/service/hlo_verifier.h @@ -54,7 +54,6 @@ class ShapeVerifier : public DfsHloVisitor { Status HandleReduce(HloInstruction* reduce) override; Status HandleBitcast(HloInstruction* bitcast) override; Status HandleBroadcast(HloInstruction* broadcast) override; - Status HandleBroadcastDimOne(HloInstruction* broadcastDimOne) override; Status HandleReshape(HloInstruction* reshape) override; Status HandleTranspose(HloInstruction* transpose) override; Status HandleParameter(HloInstruction*) override; @@ -103,7 +102,7 @@ class ShapeVerifier : public DfsHloVisitor { Status CheckTernaryShape(const HloInstruction* instruction); Status CheckVariadicShape(const HloInstruction* instruction); - // Checks if the given two instructions shares the same channel id. + // Checks if the given two instructions share the same channel id. Status CheckSameChannel(const HloInstruction* instr1, const HloInstruction* instr2); @@ -145,9 +144,15 @@ class HloVerifier : public HloPassInterface { // CHECKs various invariants of a fusion instruction. Status CheckFusionInstruction(HloInstruction* fusion) const; + Status CheckWhileInstruction(HloInstruction* instruction); + + // Checks that the non-scalar operand shapes are compatible to the output + // shape, i.e., that there are no implicit broadcasts of size-one dimensions. + Status CheckElementwiseInstruction(HloInstruction* instruction); + // Creates a ShapeVerifier that checks that shapes match inferred - // expectations. This is a factory function because ShapeVerifier, Note that - // ShapeVerifier, being a DfsHloVisitor, is stateful. We want a clean object + // expectations. This is a factory function because ShapeVerifier, + // being a DfsHloVisitor, is stateful. We want a clean object // for each run of the verifier. ShapeVerifierFactory shape_verifier_factory_; }; diff --git a/tensorflow/compiler/xla/service/instruction_fusion.cc b/tensorflow/compiler/xla/service/instruction_fusion.cc index 3f4dbf897df7e1fd62f4229ed90c949c59da9d46..b9ccfeddb565b7b44f5b38281a49df6cd0fdc766 100644 --- a/tensorflow/compiler/xla/service/instruction_fusion.cc +++ b/tensorflow/compiler/xla/service/instruction_fusion.cc @@ -37,9 +37,9 @@ namespace xla { case HloOpcode::kBitcast: case HloOpcode::kBitcastConvert: case HloOpcode::kBroadcast: - case HloOpcode::kBroadcastDimOne: case HloOpcode::kCeil: case HloOpcode::kClamp: + case HloOpcode::kClz: case HloOpcode::kComplex: case HloOpcode::kConcatenate: case HloOpcode::kConstant: @@ -143,8 +143,7 @@ bool InstructionFusion::EffectivelyUnary(HloInstruction* hlo) { }); return std::count_if(hlo->operands().begin(), hlo->operands().end(), [output_rank](HloInstruction* operand) { - if (operand->opcode() == HloOpcode::kBroadcast || - operand->opcode() == HloOpcode::kBroadcastDimOne) { + if (operand->opcode() == HloOpcode::kBroadcast) { return false; } if (operand->opcode() == HloOpcode::kConstant && @@ -249,8 +248,7 @@ StatusOr InstructionFusion::Run(HloModule* module) { auto reachability = computation->ComputeReachability(); auto cheap_to_duplicate = [this](HloInstruction* producer) { - if (producer->opcode() == HloOpcode::kBroadcast || - producer->opcode() == HloOpcode::kBroadcastDimOne) { + if (producer->opcode() == HloOpcode::kBroadcast) { return true; } if (producer->opcode() == HloOpcode::kConstant && diff --git a/tensorflow/compiler/xla/service/interpreter/compiler.cc b/tensorflow/compiler/xla/service/interpreter/compiler.cc index 9171e859c6f84ceef9664aa1eb90a07c87dfab40..76b3ecad26fe92e910fd3fe0e405c726da7e14b7 100644 --- a/tensorflow/compiler/xla/service/interpreter/compiler.cc +++ b/tensorflow/compiler/xla/service/interpreter/compiler.cc @@ -41,9 +41,6 @@ limitations under the License. namespace xla { namespace interpreter { -namespace se = ::perftools::gputools; -namespace sep = ::perftools::gputools::interpreter; - Status InterpreterCompiler::RunHloOptimization(HloModule* hlo_module) { HloPassPipeline pipeline("Interpreter"); @@ -96,7 +93,7 @@ InterpreterCompiler::CompileAheadOfTime( } se::Platform::Id InterpreterCompiler::PlatformId() const { - return sep::kInterpreterPlatformId; + return se::interpreter::kXlaInterpreterPlatformId; } HloCostAnalysis::ShapeSizeFunction InterpreterCompiler::ShapeSizeBytesFunction() @@ -109,11 +106,12 @@ static std::unique_ptr CreateComputationPlacer() { } static bool InitModule() { - xla::Compiler::RegisterCompilerFactory(sep::kInterpreterPlatformId, []() { - return xla::MakeUnique(); - }); - xla::ComputationPlacer::RegisterComputationPlacer(sep::kInterpreterPlatformId, - &CreateComputationPlacer); + xla::Compiler::RegisterCompilerFactory( + se::interpreter::kXlaInterpreterPlatformId, []() { + return xla::MakeUnique(); + }); + xla::ComputationPlacer::RegisterComputationPlacer( + se::interpreter::kXlaInterpreterPlatformId, &CreateComputationPlacer); return true; } diff --git a/tensorflow/compiler/xla/service/interpreter/compiler.h b/tensorflow/compiler/xla/service/interpreter/compiler.h index c8660c04d86a82e7dfcfd1658310c2a0e4fa0083..e90ae3e818522e6e4fd9d9f5acb846800bc899ca 100644 --- a/tensorflow/compiler/xla/service/interpreter/compiler.h +++ b/tensorflow/compiler/xla/service/interpreter/compiler.h @@ -44,19 +44,16 @@ class InterpreterCompiler : public Compiler { ~InterpreterCompiler() override {} StatusOr> RunHloPasses( - std::unique_ptr hlo_module, - perftools::gputools::StreamExecutor* stream_exec, + std::unique_ptr hlo_module, se::StreamExecutor* stream_exec, DeviceMemoryAllocator* device_allocator) override; StatusOr> RunBackend( - std::unique_ptr hlo_module, - perftools::gputools::StreamExecutor* stream_exec, + std::unique_ptr hlo_module, se::StreamExecutor* stream_exec, DeviceMemoryAllocator* device_allocator) override; StatusOr>> Compile( std::vector> hlo_modules, - std::vector> - stream_exec, + std::vector> stream_exec, DeviceMemoryAllocator* device_allocator) override; StatusOr>> @@ -65,7 +62,7 @@ class InterpreterCompiler : public Compiler { HloCostAnalysis::ShapeSizeFunction ShapeSizeBytesFunction() const override; - perftools::gputools::Platform::Id PlatformId() const override; + se::Platform::Id PlatformId() const override; private: Status RunHloOptimization(HloModule* hlo_module); diff --git a/tensorflow/compiler/xla/service/interpreter/executable.cc b/tensorflow/compiler/xla/service/interpreter/executable.cc index 883063d0f075f5b0d79edc01bcd27a7c579272f4..6553000336b072d56ea01cad1f2f66c7bd339e1a 100644 --- a/tensorflow/compiler/xla/service/interpreter/executable.cc +++ b/tensorflow/compiler/xla/service/interpreter/executable.cc @@ -38,8 +38,6 @@ limitations under the License. namespace xla { namespace interpreter { -namespace se = ::perftools::gputools; - InterpreterExecutable::InterpreterExecutable( std::unique_ptr hlo_module) : Executable(std::move(hlo_module), /*hlo_profile_printer=*/nullptr, @@ -47,7 +45,7 @@ InterpreterExecutable::InterpreterExecutable( InterpreterExecutable::~InterpreterExecutable() {} -StatusOr> InterpreterExecutable::ExecuteOnStream( +StatusOr InterpreterExecutable::ExecuteOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments, HloExecutionProfile* hlo_execution_profile) { @@ -90,12 +88,12 @@ StatusOr> InterpreterExecutable::ExecuteOnStream( evaluator.Evaluate>(*computation, arg_literals)); // Transform the result literal back into a ShapedBuffer. - TF_ASSIGN_OR_RETURN(std::unique_ptr result, + TF_ASSIGN_OR_RETURN(ShapedBuffer result, transfer_manager->AllocateShapedBuffer( result_literal->shape(), run_options->allocator(), executor->device_ordinal())); TF_RETURN_IF_ERROR(transfer_manager->TransferLiteralToDevice( - executor, *result_literal, *result)); + executor, *result_literal, result)); uint64 end_micros = tensorflow::Env::Default()->NowMicros(); @@ -108,8 +106,7 @@ StatusOr> InterpreterExecutable::ExecuteOnStream( return std::move(result); } -StatusOr> -InterpreterExecutable::ExecuteAsyncOnStream( +StatusOr InterpreterExecutable::ExecuteAsyncOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments) { return tensorflow::errors::Unimplemented( diff --git a/tensorflow/compiler/xla/service/interpreter/executable.h b/tensorflow/compiler/xla/service/interpreter/executable.h index 410110a1adf04c83001c38ed03f5d60dd203dc7e..c825a9a368d43b5ce1cdc4701e3bbf57374d81f3 100644 --- a/tensorflow/compiler/xla/service/interpreter/executable.h +++ b/tensorflow/compiler/xla/service/interpreter/executable.h @@ -43,12 +43,12 @@ class InterpreterExecutable : public Executable { InterpreterExecutable(std::unique_ptr hlo_module); ~InterpreterExecutable() override; - StatusOr> ExecuteOnStream( + StatusOr ExecuteOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments, HloExecutionProfile* hlo_execution_profile) override; - StatusOr> ExecuteAsyncOnStream( + StatusOr ExecuteAsyncOnStream( const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments) override; diff --git a/tensorflow/compiler/xla/service/interpreter/executor.cc b/tensorflow/compiler/xla/service/interpreter/executor.cc index 68371910d76f42c0b6d4b1adad9d6a83bdb858e6..97e9fa2c8e8ecd918ffe3df2fd4e731f3b91e6db 100644 --- a/tensorflow/compiler/xla/service/interpreter/executor.cc +++ b/tensorflow/compiler/xla/service/interpreter/executor.cc @@ -19,8 +19,7 @@ limitations under the License. #include "tensorflow/compiler/xla/status_macros.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace interpreter { host::HostStream *AsExecutorStream(Stream *stream) { @@ -28,84 +27,85 @@ host::HostStream *AsExecutorStream(Stream *stream) { return dynamic_cast(stream->implementation()); } -InterpreterExecutor::InterpreterExecutor(const PluginConfig &plugin_config) +XlaInterpreterExecutor::XlaInterpreterExecutor( + const PluginConfig &plugin_config) : plugin_config_(plugin_config) {} -InterpreterExecutor::~InterpreterExecutor() {} +XlaInterpreterExecutor::~XlaInterpreterExecutor() {} -void *InterpreterExecutor::Allocate(uint64 size) { return new char[size]; } +void *XlaInterpreterExecutor::Allocate(uint64 size) { return new char[size]; } -void *InterpreterExecutor::AllocateSubBuffer(DeviceMemoryBase *parent, - uint64 offset_bytes, - uint64 /*size_bytes*/) { +void *XlaInterpreterExecutor::AllocateSubBuffer(DeviceMemoryBase *parent, + uint64 offset_bytes, + uint64 /*size_bytes*/) { return parent + offset_bytes; } -void InterpreterExecutor::Deallocate(DeviceMemoryBase *mem) { +void XlaInterpreterExecutor::Deallocate(DeviceMemoryBase *mem) { if (!mem->is_sub_buffer()) { delete[] static_cast(mem->opaque()); } } -bool InterpreterExecutor::Memcpy(Stream *stream, void *host_dst, - const DeviceMemoryBase &dev_src, uint64 size) { +bool XlaInterpreterExecutor::Memcpy(Stream *stream, void *host_dst, + const DeviceMemoryBase &dev_src, + uint64 size) { AsExecutorStream(stream)->EnqueueTask([this, host_dst, dev_src, size]() { port::Status ok = SynchronousMemcpy(host_dst, dev_src, size); }); return true; } -bool InterpreterExecutor::Memcpy(Stream *stream, DeviceMemoryBase *dev_dst, - const void *host_src, uint64 size) { +bool XlaInterpreterExecutor::Memcpy(Stream *stream, DeviceMemoryBase *dev_dst, + const void *host_src, uint64 size) { AsExecutorStream(stream)->EnqueueTask([this, dev_dst, host_src, size]() { port::Status ok = SynchronousMemcpy(dev_dst, host_src, size); }); return true; } -port::Status InterpreterExecutor::SynchronousMemcpy(DeviceMemoryBase *dev_dst, - const void *host_src, - uint64 size) { +port::Status XlaInterpreterExecutor::SynchronousMemcpy( + DeviceMemoryBase *dev_dst, const void *host_src, uint64 size) { memcpy(dev_dst->opaque(), host_src, size); return port::Status::OK(); } -port::Status InterpreterExecutor::SynchronousMemcpy( +port::Status XlaInterpreterExecutor::SynchronousMemcpy( void *host_dst, const DeviceMemoryBase &dev_src, uint64 size) { memcpy(host_dst, dev_src.opaque(), size); return port::Status::OK(); } -bool InterpreterExecutor::HostCallback(Stream *stream, - std::function callback) { +bool XlaInterpreterExecutor::HostCallback(Stream *stream, + std::function callback) { AsExecutorStream(stream)->EnqueueTask(callback); return true; } -bool InterpreterExecutor::CreateStreamDependency(Stream *dependent, - Stream *other) { +bool XlaInterpreterExecutor::CreateStreamDependency(Stream *dependent, + Stream *other) { AsExecutorStream(dependent)->EnqueueTask( [other]() { SE_CHECK_OK(other->BlockHostUntilDone()); }); AsExecutorStream(dependent)->BlockUntilDone(); return true; } -bool InterpreterExecutor::StartTimer(Stream *stream, Timer *timer) { +bool XlaInterpreterExecutor::StartTimer(Stream *stream, Timer *timer) { dynamic_cast(timer->implementation())->Start(stream); return true; } -bool InterpreterExecutor::StopTimer(Stream *stream, Timer *timer) { +bool XlaInterpreterExecutor::StopTimer(Stream *stream, Timer *timer) { dynamic_cast(timer->implementation())->Stop(stream); return true; } -port::Status InterpreterExecutor::BlockHostUntilDone(Stream *stream) { +port::Status XlaInterpreterExecutor::BlockHostUntilDone(Stream *stream) { AsExecutorStream(stream)->BlockUntilDone(); return port::Status::OK(); } -DeviceDescription *InterpreterExecutor::PopulateDeviceDescription() const { +DeviceDescription *XlaInterpreterExecutor::PopulateDeviceDescription() const { internal::DeviceDescriptionBuilder builder; builder.set_device_address_bits(64); @@ -118,5 +118,4 @@ DeviceDescription *InterpreterExecutor::PopulateDeviceDescription() const { } } // namespace interpreter -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/compiler/xla/service/interpreter/executor.h b/tensorflow/compiler/xla/service/interpreter/executor.h index c5d07e906dafb033905c50c604069e80e1ce80cd..9b109022fbfc698f7dadc678ef837da270a5e74a 100644 --- a/tensorflow/compiler/xla/service/interpreter/executor.h +++ b/tensorflow/compiler/xla/service/interpreter/executor.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -// Declares the InterpreterExecutor class, which is a CPU-only implementation of -// the StreamExecutor interface. For now, this is used for testing and to +// Declares the XlaInterpreterExecutor class, which is a CPU-only implementation +// of the StreamExecutor interface. For now, this is used for testing and to // examine the performance of host-based StreamExecutor code. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_INTERPRETER_EXECUTOR_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_INTERPRETER_EXECUTOR_H_ @@ -44,16 +44,15 @@ limitations under the License. #include "tensorflow/stream_executor/stream_executor_internal.h" #include "tensorflow/stream_executor/timer.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace interpreter { using Args = tensorflow::gtl::ArraySlice; -class InterpreterExecutor : public internal::StreamExecutorInterface { +class XlaInterpreterExecutor : public internal::StreamExecutorInterface { public: - explicit InterpreterExecutor(const PluginConfig &plugin_config); - ~InterpreterExecutor() override; + explicit XlaInterpreterExecutor(const PluginConfig &plugin_config); + ~XlaInterpreterExecutor() override; port::Status Init(int device_ordinal, DeviceOptions device_options) override { return port::Status::OK(); @@ -213,7 +212,6 @@ class InterpreterExecutor : public internal::StreamExecutorInterface { }; } // namespace interpreter -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_COMPILER_XLA_SERVICE_INTERPRETER_EXECUTOR_H_ diff --git a/tensorflow/compiler/xla/service/interpreter/interpreter_transfer_manager.cc b/tensorflow/compiler/xla/service/interpreter/interpreter_transfer_manager.cc index cf98ecd7749d61261bf072cdb1882c7603f39556..d27cd7502f10a1f615fc5b0d610acafdf55e3e43 100644 --- a/tensorflow/compiler/xla/service/interpreter/interpreter_transfer_manager.cc +++ b/tensorflow/compiler/xla/service/interpreter/interpreter_transfer_manager.cc @@ -21,12 +21,10 @@ limitations under the License. #include "tensorflow/compiler/xla/service/interpreter/platform_id.h" #include "tensorflow/compiler/xla/service/transfer_manager.h" -namespace sei = ::perftools::gputools::interpreter; - namespace xla { InterpreterTransferManager::InterpreterTransferManager() - : GenericTransferManager(sei::kInterpreterPlatformId, + : GenericTransferManager(se::interpreter::kXlaInterpreterPlatformId, /*pointer_size=*/sizeof(void*)) {} } // namespace xla @@ -38,7 +36,8 @@ CreateInterpreterTransferManager() { static bool InitModule() { xla::TransferManager::RegisterTransferManager( - sei::kInterpreterPlatformId, &CreateInterpreterTransferManager); + stream_executor::interpreter::kXlaInterpreterPlatformId, + &CreateInterpreterTransferManager); return true; } diff --git a/tensorflow/compiler/xla/service/interpreter/platform.cc b/tensorflow/compiler/xla/service/interpreter/platform.cc index a60e7fc59f7c5f0b1b24e026b34e195ca0fe5ebb..ce2f4d378c02592da01be12c0647da47d00b01ba 100644 --- a/tensorflow/compiler/xla/service/interpreter/platform.cc +++ b/tensorflow/compiler/xla/service/interpreter/platform.cc @@ -28,24 +28,22 @@ limitations under the License. #include "tensorflow/stream_executor/multi_platform_manager.h" #include "tensorflow/stream_executor/platform.h" -namespace se = ::perftools::gputools; -namespace sep = ::perftools::gputools::interpreter; - -namespace perftools { -namespace gputools { +namespace stream_executor { namespace interpreter { -InterpreterPlatform::InterpreterPlatform() : name_("Interpreter") {} +XlaInterpreterPlatform::XlaInterpreterPlatform() : name_("Interpreter") {} -InterpreterPlatform::~InterpreterPlatform() {} +XlaInterpreterPlatform::~XlaInterpreterPlatform() {} -Platform::Id InterpreterPlatform::id() const { return kInterpreterPlatformId; } +Platform::Id XlaInterpreterPlatform::id() const { + return kXlaInterpreterPlatformId; +} -int InterpreterPlatform::VisibleDeviceCount() const { return 1; } +int XlaInterpreterPlatform::VisibleDeviceCount() const { return 1; } -const string& InterpreterPlatform::Name() const { return name_; } +const string& XlaInterpreterPlatform::Name() const { return name_; } -port::StatusOr InterpreterPlatform::ExecutorForDevice( +port::StatusOr XlaInterpreterPlatform::ExecutorForDevice( int ordinal) { StreamExecutorConfig config; config.ordinal = ordinal; @@ -55,7 +53,7 @@ port::StatusOr InterpreterPlatform::ExecutorForDevice( } port::StatusOr -InterpreterPlatform::ExecutorForDeviceWithPluginConfig( +XlaInterpreterPlatform::ExecutorForDeviceWithPluginConfig( int device_ordinal, const PluginConfig& plugin_config) { StreamExecutorConfig config; config.ordinal = device_ordinal; @@ -64,16 +62,17 @@ InterpreterPlatform::ExecutorForDeviceWithPluginConfig( return GetExecutor(config); } -port::StatusOr InterpreterPlatform::GetExecutor( +port::StatusOr XlaInterpreterPlatform::GetExecutor( const StreamExecutorConfig& config) { return executor_cache_.GetOrCreate( config, [&]() { return GetUncachedExecutor(config); }); } port::StatusOr> -InterpreterPlatform::GetUncachedExecutor(const StreamExecutorConfig& config) { +XlaInterpreterPlatform::GetUncachedExecutor( + const StreamExecutorConfig& config) { auto executor = port::MakeUnique( - this, port::MakeUnique(config.plugin_config)); + this, port::MakeUnique(config.plugin_config)); auto init_status = executor->Init(config.ordinal, config.device_options); if (!init_status.ok()) { return port::Status{ @@ -86,26 +85,26 @@ InterpreterPlatform::GetUncachedExecutor(const StreamExecutorConfig& config) { return std::move(executor); } -void InterpreterPlatform::RegisterTraceListener( +void XlaInterpreterPlatform::RegisterTraceListener( std::unique_ptr listener) { LOG(FATAL) << "not yet implemented: register executor trace listener"; } -void InterpreterPlatform::UnregisterTraceListener(TraceListener* listener) { +void XlaInterpreterPlatform::UnregisterTraceListener(TraceListener* listener) { LOG(FATAL) << "not yet implemented: unregister executor trace listener"; } -static void InitializeInterpreterPlatform() { - std::unique_ptr platform(new sep::InterpreterPlatform); - SE_CHECK_OK(se::MultiPlatformManager::RegisterPlatform(std::move(platform))); +static void InitializeXlaInterpreterPlatform() { + std::unique_ptr platform(new XlaInterpreterPlatform); + SE_CHECK_OK(MultiPlatformManager::RegisterPlatform(std::move(platform))); } } // namespace interpreter -} // namespace gputools -} // namespace perftools +} // namespace stream_executor -REGISTER_MODULE_INITIALIZER(interpreter_platform, - sep::InitializeInterpreterPlatform()); +REGISTER_MODULE_INITIALIZER( + interpreter_platform, + stream_executor::interpreter::InitializeXlaInterpreterPlatform()); DECLARE_MODULE_INITIALIZER(multi_platform_manager); diff --git a/tensorflow/compiler/xla/service/interpreter/platform.h b/tensorflow/compiler/xla/service/interpreter/platform.h index c66ddb907d1c5a8e99d3178a202a77a72a646ce5..d68c5aa20dda7ac246ed4aa667851e385a604c04 100644 --- a/tensorflow/compiler/xla/service/interpreter/platform.h +++ b/tensorflow/compiler/xla/service/interpreter/platform.h @@ -23,14 +23,13 @@ limitations under the License. #include "tensorflow/stream_executor/stream_executor.h" #include "tensorflow/stream_executor/trace_listener.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace interpreter { -class InterpreterPlatform : public Platform { +class XlaInterpreterPlatform : public Platform { public: - InterpreterPlatform(); - ~InterpreterPlatform() override; + XlaInterpreterPlatform(); + ~XlaInterpreterPlatform() override; Platform::Id id() const override; @@ -60,11 +59,10 @@ class InterpreterPlatform : public Platform { // Cache of created StreamExecutors. ExecutorCache executor_cache_; - SE_DISALLOW_COPY_AND_ASSIGN(InterpreterPlatform); + SE_DISALLOW_COPY_AND_ASSIGN(XlaInterpreterPlatform); }; } // namespace interpreter -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_COMPILER_XLA_SERVICE_INTERPRETER_PLATFORM_H_ diff --git a/tensorflow/compiler/xla/service/interpreter/platform_id.cc b/tensorflow/compiler/xla/service/interpreter/platform_id.cc index 1a0373cf86e26b564e0e732e8de1a0a5d868bfa6..3272396ce5045129a7689a160ec859d11fbbe9fa 100644 --- a/tensorflow/compiler/xla/service/interpreter/platform_id.cc +++ b/tensorflow/compiler/xla/service/interpreter/platform_id.cc @@ -14,12 +14,10 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/service/interpreter/platform_id.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace interpreter { -PLATFORM_DEFINE_ID(kInterpreterPlatformId); +PLATFORM_DEFINE_ID(kXlaInterpreterPlatformId); } // namespace interpreter -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/compiler/xla/service/interpreter/platform_id.h b/tensorflow/compiler/xla/service/interpreter/platform_id.h index 905efef1690d3bd32461353fe32dd394eb6bca9e..a6cc10bcc1eb756a3146d4a834efa4cd3ceb2d27 100644 --- a/tensorflow/compiler/xla/service/interpreter/platform_id.h +++ b/tensorflow/compiler/xla/service/interpreter/platform_id.h @@ -18,14 +18,12 @@ limitations under the License. #include "tensorflow/stream_executor/platform.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace interpreter { -extern const Platform::Id kInterpreterPlatformId; +extern const Platform::Id kXlaInterpreterPlatformId; } // namespace interpreter -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_COMPILER_XLA_SERVICE_INTERPRETER_PLATFORM_ID_H_ diff --git a/tensorflow/compiler/xla/service/layout_assignment.cc b/tensorflow/compiler/xla/service/layout_assignment.cc index 39f9120e552f014dd2759bff2892157402d9c47a..7067b6f86a0fb24fb946ad236bca9bbd48d53722 100644 --- a/tensorflow/compiler/xla/service/layout_assignment.cc +++ b/tensorflow/compiler/xla/service/layout_assignment.cc @@ -31,10 +31,12 @@ limitations under the License. #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/computation_layout.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_dce.h" #include "tensorflow/compiler/xla/service/hlo_graph_dumper.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/logical_buffer.h" +#include "tensorflow/compiler/xla/service/tuple_simplifier.h" #include "tensorflow/compiler/xla/shape_layout.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -57,76 +59,6 @@ namespace xla { // 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 @@ -470,9 +402,9 @@ string LayoutConstraints::ToString() const { } Status LayoutAssignment::AddMandatoryConstraints( - const ComputationLayout& computation_layout, - const ChannelLayoutConstraints* channel_constraints, - HloComputation* computation, LayoutConstraints* constraints) { + const ComputationLayout* computation_layout, + ChannelLayoutConstraints* channel_constraints, HloComputation* computation, + LayoutConstraints* constraints) { VLOG(3) << "Adding mandatory layout constraints to computation " << computation->name(); @@ -494,11 +426,16 @@ Status LayoutAssignment::AddMandatoryConstraints( TF_RETURN_IF_ERROR(constraints->SetOperandLayout( instruction->outfeed_shape(), instruction, 0)); } else if (instruction->opcode() == HloOpcode::kParameter) { - // Parameter layouts must match the respective layout in - // ComputationLayout. - shape_with_layout = - &computation_layout.parameter_layout(instruction->parameter_number()) - .shape(); + if (computation_layout != nullptr) { + const ShapeLayout& parameter_layout = + computation_layout->parameter_layout( + instruction->parameter_number()); + if (parameter_layout.LayoutIsSet()) { + // Parameter layouts must match the respective layout in + // ComputationLayout, if there is one. + shape_with_layout = ¶meter_layout.shape(); + } + } } if (shape_with_layout != nullptr) { TF_RETURN_IF_ERROR( @@ -563,9 +500,8 @@ Status LayoutAssignment::AddMandatoryConstraints( HloComputation* body = instruction->while_body(); HloComputation* condition = instruction->while_condition(); const HloInstruction* init = instruction->operand(0); - const ComputationLayout& body_layout = - FindOrDie(computation_layouts_, body); - const ComputationLayout& condition_layout = + ComputationLayout& body_layout = FindOrDie(computation_layouts_, body); + ComputationLayout& condition_layout = FindOrDie(computation_layouts_, condition); // Check a few invariants irrespective of layout. @@ -578,26 +514,19 @@ Status LayoutAssignment::AddMandatoryConstraints( condition_layout.parameter_shape(0))); DCHECK(ShapeUtil::Compatible(body_layout.result_shape(), init->shape())); - // Return error if earlier layout assignment of the embedded computations - // has produced conflicting layouts. - if (!ShapeUtil::Equal(body_layout.result_shape(), - body_layout.parameter_shape(0))) { - return InternalError( - "Parameter and result of body computation %s of while instruction " - "%s have different layouts: %s vs %s", - body->name().c_str(), instruction->name().c_str(), - ShapeUtil::HumanString(body_layout.result_shape()).c_str(), - ShapeUtil::HumanString(body_layout.parameter_shape(0)).c_str()); + if (body_layout.result_layout() != body_layout.parameter_layout(0)) { + VLOG(2) << "Reset %while body parameter layout: body=" << body->name() + << " while=" << instruction->name() + << " shape=" << body_layout.result_layout().ToString(); + *body_layout.mutable_parameter_layout(0) = body_layout.result_layout(); } - if (!ShapeUtil::Equal(body->root_instruction()->shape(), - condition->parameter_instruction(0)->shape())) { - return InternalError( - "Parameter of condition computation %s of while instruction " - "%s does not match body computation %s result: %s vs %s", - condition->name().c_str(), instruction->name().c_str(), - body->name().c_str(), - ShapeUtil::HumanString(condition_layout.parameter_shape(0)).c_str(), - ShapeUtil::HumanString(body_layout.result_shape()).c_str()); + if (condition_layout.parameter_layout(0) != + body_layout.parameter_layout(0)) { + VLOG(2) << "Reset %while condition parameter layout: cond=" + << condition->name() << " while=" << instruction->name() + << " shape=" << body_layout.parameter_layout(0).ToString(); + *condition_layout.mutable_parameter_layout(0) = + body_layout.parameter_layout(0); } // Constrain the output and the operand of the while instruction to match @@ -627,7 +556,20 @@ Status LayoutAssignment::AddMandatoryConstraints( true_computation_layout.parameter_shape(0))); DCHECK(ShapeUtil::Compatible( false_operand->shape(), false_computation_layout.parameter_shape(0))); - + if (true_computation_layout.result_layout() != + false_computation_layout.result_layout()) { + // We assign layouts in DFS fashion, so the true and false computations + // might have negotiated a different layout. But for the conditional + // instruction POV the layout must match, so we run again on the false + // computation, this time with proper computation layout. + VLOG(2) << "Reset %conditional false computation result layout: " + "false_computation=" + << false_computation->name() + << " conditional=" << instruction->name() << " shape=" + << true_computation_layout.result_layout().ToString(); + *false_computation_layout.mutable_result_layout() = + true_computation_layout.result_layout(); + } TF_RETURN_IF_ERROR(constraints->SetInstructionLayout( true_computation_layout.result_shape(), instruction)); TF_RETURN_IF_ERROR(constraints->SetOperandLayout( @@ -663,10 +605,14 @@ Status LayoutAssignment::AddMandatoryConstraints( } } } - - // Finally set the result layout to match ComputationLayout. - return constraints->SetResultLayout( - computation_layout.result_layout().shape()); + // Finally set the result layout to match ComputationLayout, if there is one. + if (computation_layout != nullptr) { + const ShapeLayout& result_layout = computation_layout->result_layout(); + if (result_layout.LayoutIsSet()) { + TF_RETURN_IF_ERROR(constraints->SetResultLayout(result_layout.shape())); + } + } + return Status::OK(); } namespace { @@ -793,6 +739,106 @@ Status CheckConstantLayout(HloInstruction* constant) { } // namespace +StatusOr LayoutAssignment::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)); + SetupCopiedInstruction(*instruction, gte, {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)); + SetupCopiedInstruction(*instruction, tuple_copy, {}); + 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)); + RegisterAddedCopy(copy); + SetupCopiedInstruction(*instruction, copy, {}); + 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 LayoutAssignment::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())) { + VLOG(5) << "Operand " << operand->ToString() << " layout matches in " + << instruction->ToString(); + // Operand layout already matches our constraint. Nothing to do. + return Status::OK(); + } + VLOG(4) << "Operand " << operand->ToString() << " layout does not match " + << operand_layout.ToString() << " in " << instruction->ToString(); + + TF_ASSIGN_OR_RETURN(HloInstruction * operand_copy, + CreateCopyWithNewLayout(operand_layout.shape(), operand)); + + VLOG(4) << "New copy of " << operand->ToString() << " is " + << operand_copy->ToString(); + return instruction->ReplaceOperandWith(operand_no, operand_copy); +} + +void LayoutAssignment::SetupCopiedInstruction(const HloInstruction& instruction, + HloInstruction* copy, + const ShapeIndex& index) { + if (instruction.has_sharding()) { + // If the index is empty, we want to copy the whole sharding, in case the + // sharding is a tuple sharding. + HloSharding sharding = + !index.empty() && instruction.sharding().IsTuple() + ? instruction.sharding().GetSubSharding(instruction.shape(), index) + : instruction.sharding(); + // We propagate the sharding to the copied instruction only if it is a + // special sharding, like tiled ones, or special devices like the + // HostCompute module. + // Otherwise it is preferable to leave the new instruction without device, + // and let the automatic device placer to choose the best location. + if (!sharding.HasUniqueDevice() || + HloSharding::IsReservedDevice(sharding.UniqueDevice().ValueOrDie())) { + copy->set_sharding(sharding); + } + } + copy->set_metadata(instruction.metadata()); +} + Status LayoutAssignment::CheckLayouts(HloModule* module) { TF_ASSIGN_OR_RETURN(auto points_to_analysis, TuplePointsToAnalysis::Run(module)); @@ -873,15 +919,16 @@ Status LayoutAssignment::CheckLayouts(HloModule* module) { } } } - - // Finally verify the result layout matches the layout of the entry + // Finally verify the result layout, if set, matches the layout of the entry // computation root. - TF_RET_CHECK(ShapeUtil::Equal( - module->entry_computation()->root_instruction()->shape(), + const ShapeLayout& result_layout = FindOrDie(computation_layouts_, module->entry_computation()) - .result_layout() - .shape())); - + .result_layout(); + if (result_layout.LayoutIsSet()) { + TF_RET_CHECK(ShapeUtil::Equal( + module->entry_computation()->root_instruction()->shape(), + result_layout.shape())); + } return Status::OK(); } @@ -890,18 +937,13 @@ LayoutAssignment::LayoutAssignment( ChannelLayoutConstraints* channel_constraints) : entry_computation_layout_(entry_computation_layout), channel_layout_constraints_(channel_constraints) { - VLOG(1) << "entry computation layout given to layout assignment: " + VLOG(1) << "Entry computation layout given to layout assignment: " << entry_computation_layout_->ToString(); // Layouts of all parameter instructions must be set. for (const ShapeLayout& parameter_layout : entry_computation_layout_->parameter_layouts()) { CHECK(parameter_layout.LayoutIsSet()); } - // If the result layout is not set, then choose the default. - // TODO(b/29118294): Choose a better layout in this case. - if (!entry_computation_layout_->result_layout().LayoutIsSet()) { - entry_computation_layout_->mutable_result_layout()->SetToDefaultLayout(); - } } std::unique_ptr LayoutAssignment::ChooseOperandLayoutFromOutputLayout( @@ -1461,16 +1503,60 @@ Status LayoutAssignment::AssignLayouts(const LayoutConstraints& constraints, return Status::OK(); } +Status LayoutAssignment::CalculateComputationLayout( + HloComputation* computation) { + ComputationLayout computation_layout(computation->ComputeProgramShape(), + /*ignore_layouts=*/false); + InsertOrDie(&computation_layouts_, computation, computation_layout); + VLOG(2) << " Calculated ComputationLayout = " + << computation_layout.ToString(); + return Status::OK(); +} + +Status LayoutAssignment::ClearComputationLayouts(HloComputation* computation) { + // Clear existing layouts of the instructions. All layouts must be assigned + // by the LayoutAssignment pass, except for those on infeeds, parameters, + // and the computation result. The latter two are specified in + // computation_layout, so we only need to keep the existing layouts for + // infeeds. Clearing the layouts here avoids hiding potential bugs in the + // layout assignment pass that may accidently use the existing layout. + for (HloInstruction* instruction : computation->instructions()) { + if (instruction->opcode() == HloOpcode::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()); + } + } + return Status::OK(); +} + Status LayoutAssignment::RunOnComputation( - const ComputationLayout& computation_layout, + ComputationLayout* computation_layout, const TuplePointsToAnalysis& points_to_analysis, HloComputation* computation, ChannelLayoutConstraints* channel_constraints) { - DCHECK(computation_layout.LayoutIsSet()); - InsertOrDie(&computation_layouts_, computation, computation_layout); VLOG(2) << "LayoutAssignment::RunOnComputation(" << computation->name() << ")"; - VLOG(2) << " ComputationLayout = " << computation_layout.ToString(); + TF_RETURN_IF_ERROR(ClearComputationLayouts(computation)); + if (computation_layout != nullptr) { + auto it = computation_layouts_.find(computation); + if (it == computation_layouts_.end()) { + VLOG(2) << " New ComputationLayout = " << computation_layout->ToString(); + computation_layouts_.emplace(computation, *computation_layout); + } else { + TF_RET_CHECK(computation_layout == &it->second || + computation_layout == entry_computation_layout_); + VLOG(2) << " Existing ComputationLayout = " + << computation_layout->ToString(); + } + } else { + VLOG(2) << " No ComputationLayout specified (will be calculated)"; + } // Construct LayoutConstraints with all layout constraints of the computation. LayoutConstraints constraints(points_to_analysis, computation); @@ -1513,12 +1599,19 @@ Status LayoutAssignment::RunOnComputation( CHECK_LT(constraints.unconstrained_buffer_ids().size(), unconstrained_count); } - // All logical buffers should have constraints at this point. All that // remains is assign the constraints to the buffers and infer layouts for // aliased buffers. TF_RETURN_IF_ERROR(AssignLayouts(constraints, computation)); + // If the computation layout wasn't specified, now it is the time to compute + // it according to the parameters and root instruction layouts. + // This allows the first pass through this API to record the best flowing + // layout to parameters and root instruction. + if (computation_layout == nullptr) { + TF_RETURN_IF_ERROR(CalculateComputationLayout(computation)); + } + // Record the layouts assigned for any communication ops in // channel_constraints so that they are constrained for future modules. for (HloInstruction* instruction : computation->instructions()) { @@ -1533,6 +1626,34 @@ Status LayoutAssignment::RunOnComputation( return Status::OK(); } +Status LayoutAssignment::PropagateComputationLayouts( + HloComputation* computation, ComputationLayout* computation_layout) { + ComputationLayout computed_computation_layout( + computation->ComputeProgramShape(), + /*ignore_layouts=*/false); + for (int64 i = 0; i < computed_computation_layout.parameter_count(); ++i) { + ShapeLayout* param_layout = computation_layout->mutable_parameter_layout(i); + if (!param_layout->LayoutIsSet()) { + VLOG(4) << "Assigning layout to parameter " << i << " of computation " + << computation->name() << ": " + << computed_computation_layout.parameter_layout(i).ToString(); + *param_layout = computed_computation_layout.parameter_layout(i); + } else { + TF_RET_CHECK(computed_computation_layout.parameter_layout(i) == + *param_layout); + } + } + ShapeLayout* result_layout = computation_layout->mutable_result_layout(); + if (!result_layout->LayoutIsSet()) { + VLOG(4) << "Assigning result layout of computation " << computation->name() + << ": " << computed_computation_layout.result_layout().ToString(); + *result_layout = computed_computation_layout.result_layout(); + } else { + TF_RET_CHECK(computed_computation_layout.result_layout() == *result_layout); + } + return Status::OK(); +} + StatusOr LayoutAssignment::Run(HloModule* module) { VLOG(2) << "Running layout assignment on module " << module->name(); XLA_VLOG_LINES(3, module->ToString()); @@ -1541,52 +1662,45 @@ StatusOr LayoutAssignment::Run(HloModule* module) { "before layout assignment", module->config().debug_options()); } - - TF_ASSIGN_OR_RETURN(auto points_to_analysis, - TuplePointsToAnalysis::Run(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 - // computation_layout, so we only need to keep the existing layouts for - // infeeds. Clearing the layouts here avoids hiding potential bugs in the - // layout assignment pass that may accidently use the existing layout. - for (HloInstruction* instruction : computation->instructions()) { - if (instruction->opcode() == HloOpcode::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()); + TF_RETURN_IF_ERROR(Init()); + + // We do two passes. The first one we pass a nullptr ComputationLayout to + // the RunOnComputation() calls (for non entry computations), and we register + // the ComputationLayout which are naturally flowing in DFS fashion to the + // parameters and root instruction. + // Walking in DFS mode though, means that we can end up with incorrect layouts + // when seen from an outer instruction, which has across-computation + // constraints to impose. + // For example, the kWhile instruction needs to enforce the same layouts for + // the parameters and root of the bosy, as well as the condition parameters. + // Similarly, the kConditional instruction needs to enforce the same layouts + // for the root of the true and false computations. + // So in the first pass, while allowing the layouts to flow to parameters and + // root, we also fix up the eventually inconsistent ComputationLayout, which + // will be then made mandatory by the second pass. + for (int64 i = 0; i < 2; ++i) { + TF_RETURN_IF_ERROR(ClearPreviousPassSideEffects(module)); + TF_ASSIGN_OR_RETURN(auto points_to_analysis, + TuplePointsToAnalysis::Run(module)); + for (auto* computation : module->MakeComputationPostOrder()) { + if (computation->IsFusionComputation()) { + continue; } - if (instruction->opcode() != HloOpcode::kInfeed) { - LayoutUtil::ClearLayout(instruction->mutable_shape()); + if (computation == module->entry_computation()) { + TF_RETURN_IF_ERROR(RunOnComputation( + entry_computation_layout_, *points_to_analysis, + module->entry_computation(), channel_layout_constraints_)); + } else { + ComputationLayout* computation_layout = + (i == 0) ? nullptr : &FindOrDie(computation_layouts_, computation); + TF_RETURN_IF_ERROR(RunOnComputation(computation_layout, + *points_to_analysis, computation, + channel_layout_constraints_)); } } - if (computation == module->entry_computation()) { - TF_RETURN_IF_ERROR(RunOnComputation( - *entry_computation_layout_, *points_to_analysis, - module->entry_computation(), channel_layout_constraints_)); - } else { - ComputationLayout computation_layout(computation->ComputeProgramShape()); - // Setting all embedded computations to the default layout is potentially - // suboptimal. - computation_layout.SetToDefaultLayout(); - TF_RETURN_IF_ERROR(RunOnComputation(computation_layout, - *points_to_analysis, computation, - channel_layout_constraints_)); - } } - + TF_RETURN_IF_ERROR(PropagateComputationLayouts(module->entry_computation(), + entry_computation_layout_)); TF_RETURN_IF_ERROR(CheckLayouts(module)); VLOG(3) << "After layout assignment:"; @@ -1596,9 +1710,54 @@ StatusOr LayoutAssignment::Run(HloModule* module) { "after layout assignment", module->config().debug_options()); } - // All layouts are reset then reassigned by this pass. return true; } +Status LayoutAssignment::Init() { + computation_layouts_.clear(); + return Status::OK(); +} + +Status LayoutAssignment::ClearPreviousPassSideEffects(HloModule* module) { + // Clear all the copies which have been added, and all the related + // instructions (like GTE and tuples). + int64 removed_copies = 0; + for (HloComputation* computation : module->computations()) { + for (HloInstruction* instruction : + computation->MakeInstructionPostOrder()) { + if (instruction->opcode() == HloOpcode::kCopy && + added_copies_.count(instruction) > 0) { + VLOG(5) << "Removing added copy: " << instruction->ToString(); + TF_RETURN_IF_ERROR( + instruction->ReplaceAllUsesWith(instruction->mutable_operand(0))); + TF_RETURN_IF_ERROR(computation->RemoveInstruction(instruction)); + ++removed_copies; + } + } + } + added_copies_.clear(); + if (removed_copies > 0) { + TupleSimplifier tuple_simplifier; + HloDCE dce; + TF_RETURN_IF_ERROR(tuple_simplifier.Run(module).status()); + TF_RETURN_IF_ERROR(dce.Run(module).status()); + } + return Status::OK(); +} + +Status LayoutAssignment::AddCopyForOperand(HloInstruction* instruction, + int64 operand_number) { + HloInstruction* operand = instruction->mutable_operand(operand_number); + if (operand->opcode() != HloOpcode::kCopy || operand->user_count() > 1) { + HloInstruction* copy = + instruction->parent()->AddInstruction(HloInstruction::CreateUnary( + operand->shape(), HloOpcode::kCopy, operand)); + SetupCopiedInstruction(*operand, copy, {}); + LayoutUtil::ClearLayout(copy->mutable_shape()); + TF_RETURN_IF_ERROR(instruction->ReplaceOperandWith(operand_number, copy)); + } + return Status::OK(); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/layout_assignment.h b/tensorflow/compiler/xla/service/layout_assignment.h index 680f88048a1f0cd5ede7991640003ef407d4facf..8b4e07995afffa196f689ffc45985b00bbeeb77d 100644 --- a/tensorflow/compiler/xla/service/layout_assignment.h +++ b/tensorflow/compiler/xla/service/layout_assignment.h @@ -39,6 +39,7 @@ limitations under the License. #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/lib/gtl/flatset.h" #include "tensorflow/core/platform/types.h" namespace xla { @@ -362,12 +363,15 @@ class LayoutAssignment : public HloPassInterface { int64 operand_no); private: + // Initializes the layout assignment object for a new Run() call. + Status Init(); + // Adds constraints which must be satisfied for correctness on all // backends. Called once prior to propagating constraints. - Status AddMandatoryConstraints( - const ComputationLayout& computation_layout, - const ChannelLayoutConstraints* channel_constraints, - HloComputation* computation, LayoutConstraints* constraints); + Status AddMandatoryConstraints(const ComputationLayout* computation_layout, + ChannelLayoutConstraints* channel_constraints, + HloComputation* computation, + LayoutConstraints* constraints); // This method can be overridden to add backend-specific constraints to the // layout of the instructions of a computation. This method is called after @@ -378,10 +382,12 @@ class LayoutAssignment : public HloPassInterface { } // Construct contraints and assign layouts to all instructions in the - // computation satisfying the given ComputationLayout. Layouts constraints are - // added, then propagated until all LogicalBuffers in the computation are - // constrained. - Status RunOnComputation(const ComputationLayout& computation_layout, + // computation satisfying the given ComputationLayout, if not nullptr. + // Otherwise the ComputationLayout will be calculated by propagating the + // computation instruction contraints. + // Layouts constraints are added, then propagated until all LogicalBuffers in + // the computation are constrained. + Status RunOnComputation(ComputationLayout* computation_layout, const TuplePointsToAnalysis& points_to_analysis, HloComputation* computation, ChannelLayoutConstraints* channel_constraints); @@ -402,14 +408,72 @@ class LayoutAssignment : public HloPassInterface { // necessary conditions. Status CheckLayouts(HloModule* module); + // Computes the ComputationLayout of the given computation based of the + // layouts assigned to parameters and root instruction, and inserts it to the + // computation_layouts_ map. + Status CalculateComputationLayout(HloComputation* computation); + + // Clears all the layouts which can be cleared within a computation. + Status ClearComputationLayouts(HloComputation* computation); + + // Clears the side effects of a previous pass, like added copy instructions. + Status ClearPreviousPassSideEffects(HloModule* module); + + // Propagates the layouts computed by the layout assignment pass on the given + // computation, to the computation layout passed in to this API. + // This API propagates missing layout, and also checks that the caller + // specified have been respected, by comparing those with the parameters and + // root computation instruction. + Status PropagateComputationLayouts(HloComputation* computation, + ComputationLayout* computation_layout); + ComputationLayout* entry_computation_layout_; protected: + // Sets up the copy instruction according to the characteristic (sharding, + // metadata, ...) of the reference instruction. The index argument is used + // when the instruction is a tuple, and in such case the index represents + // the location from where the copy instruction was created from. + // If the index is empty, the whole sharding will be propagated, even in case + // the intruction has a tuple sharding. + static void SetupCopiedInstruction(const HloInstruction& instruction, + HloInstruction* copy, + const ShapeIndex& index); + + // 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); + + // 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); + + // Registers a copy instruction added by the layout assignment pass. + void RegisterAddedCopy(HloInstruction* copy) { + CHECK_EQ(copy->opcode(), HloOpcode::kCopy); + added_copies_.insert(copy); + } + + // Adds a copy for the operand of an instruction, unless such operand is + // already a copy, and has a single user (which is forcibly the instruction + // itself). + Status AddCopyForOperand(HloInstruction* instruction, int64 operand_number); + // Map containing the layouts of all computations assigned so // far. Computations are handled in a topological sort where computations are // handled before their caller instructions so the layouts of caller // instructions can be set to match the computation. std::map computation_layouts_; + + // Every copy added to the module by the layout assignment pass is registered + // here. + tensorflow::gtl::FlatSet added_copies_; + ChannelLayoutConstraints* channel_layout_constraints_; }; diff --git a/tensorflow/compiler/xla/service/llvm_compiler.cc b/tensorflow/compiler/xla/service/llvm_compiler.cc index 911b243fe28a5baf8a4b8ed752b892265f5388ac..b17c9d504501a907e27d5152e0082799e87443c7 100644 --- a/tensorflow/compiler/xla/service/llvm_compiler.cc +++ b/tensorflow/compiler/xla/service/llvm_compiler.cc @@ -23,7 +23,7 @@ limitations under the License. namespace xla { StatusOr>> LLVMCompiler::Compile( std::vector> modules, - std::vector> stream_execs, + std::vector> stream_execs, DeviceMemoryAllocator* device_allocator) { // Tensorflow tries to enable the following behaviors in all its threads: // diff --git a/tensorflow/compiler/xla/service/llvm_compiler.h b/tensorflow/compiler/xla/service/llvm_compiler.h index d74e81bb7f622ac5e89203a3d02ca5ad839da07e..f1c623508c5307f2b1c036d3ec6823b75c7eda13 100644 --- a/tensorflow/compiler/xla/service/llvm_compiler.h +++ b/tensorflow/compiler/xla/service/llvm_compiler.h @@ -60,19 +60,18 @@ class LLVMCompiler : public Compiler { // Bring in // StatusOr> RunBackend( // std::unique_ptr module, - // perftools::gputools::StreamExecutor* stream_exec, + // se::StreamExecutor* stream_exec, // DeviceMemoryAllocator* device_allocator) // StatusOr> RunHloPasses( // std::unique_ptr module, - // perftools::gputools::StreamExecutor* stream_exec, + // se::StreamExecutor* stream_exec, // DeviceMemoryAllocator* device_allocator) using Compiler::RunBackend; using Compiler::RunHloPasses; StatusOr>> Compile( std::vector> modules, - std::vector> - stream_execs, + std::vector> stream_execs, DeviceMemoryAllocator* device_allocator) override; protected: diff --git a/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc b/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc index b6b918ec78a27b90325f72eea14b97f9aee43c54..3978acc132f34b8b195d3772ccf71d0d467984db 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc @@ -88,12 +88,12 @@ LoopEmitter::LoopEmitter(const ElementGenerator& target_element_generator, } } -IrArray::Index LoopEmitter::EmitIndexAndSetExitBasicBlock( +std::vector LoopEmitter::EmitIndexAndSetExitBasicBlock( tensorflow::StringPiece loop_name) { if (ShapeUtil::IsScalar(shape_)) { // No loop needed, so set exit_bb_ to nullptr. exit_bb_ = nullptr; - return IrArray::Index(); + return {IrArray::Index()}; } // Create loop nest with one for-loop for each dimension of the target shape. @@ -121,12 +121,14 @@ IrArray::Index LoopEmitter::EmitIndexAndSetExitBasicBlock( exit_bb_ = loop_nest.GetOuterLoopExitBasicBlock(); CHECK_NOTNULL(exit_bb_); - return array_index; + return {array_index}; } tensorflow::Status LoopEmitter::EmitLoop(tensorflow::StringPiece loop_name) { - IrArray::Index array_index = EmitIndexAndSetExitBasicBlock(loop_name); - TF_RETURN_IF_ERROR(body_emitter_(array_index)); + for (const IrArray::Index& array_index : + EmitIndexAndSetExitBasicBlock(loop_name)) { + TF_RETURN_IF_ERROR(body_emitter_(array_index)); + } // Set the insertion point of ir_builder_ to the loop exit, so that // code emitted for later instructions will be correctly placed. diff --git a/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.h b/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.h index 0fc528439a0d5bf8382dfcf2d8b3051f8900bf1d..9ff497aecd0bc964c929205c7fd410cca87d9b77 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.h +++ b/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.h @@ -63,11 +63,12 @@ class LoopEmitter { // Emits a loop nest (with a yet-to-be-filled loop body) that iterates through // every element in the given shape. Returns the multi-dimensional index that - // specifies the element. - IrArray::Index EmitIndexAndSetExitBasicBlock() { + // specifies the element, will return multiple indices if the loop is + // unrolled. + std::vector EmitIndexAndSetExitBasicBlock() { return EmitIndexAndSetExitBasicBlock(/*loop_name=*/""); } - virtual IrArray::Index EmitIndexAndSetExitBasicBlock( + virtual std::vector EmitIndexAndSetExitBasicBlock( tensorflow::StringPiece loop_name); // Emits a complete loop nest for every element in the given shape. diff --git a/tensorflow/compiler/xla/service/local_service.cc b/tensorflow/compiler/xla/service/local_service.cc index 499f280211aacd00e79b3ca0ddb3413f933b02da..0fa4061738612df76c72a18a9353f16bf6a42677 100644 --- a/tensorflow/compiler/xla/service/local_service.cc +++ b/tensorflow/compiler/xla/service/local_service.cc @@ -43,13 +43,11 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" -namespace se = ::perftools::gputools; - namespace xla { /* static */ StatusOr> LocalService::NewService( const ServiceOptions& options) { - perftools::gputools::Platform* platform = options.platform(); + se::Platform* platform = options.platform(); if (platform == nullptr) { TF_ASSIGN_OR_RETURN(platform, PlatformUtil::GetDefaultPlatform()); } diff --git a/tensorflow/compiler/xla/service/name_uniquer.cc b/tensorflow/compiler/xla/service/name_uniquer.cc index 7d8c05fffa4ab11d7dbf9956d2cb7ebd5bcdd3c4..f74bcb0b79355c8e69890487266cbc5f2a4500be 100644 --- a/tensorflow/compiler/xla/service/name_uniquer.cc +++ b/tensorflow/compiler/xla/service/name_uniquer.cc @@ -53,17 +53,18 @@ NameUniquer::NameUniquer(const string& separator) { } string NameUniquer::GetUniqueName(tensorflow::StringPiece prefix) { - string root = prefix.empty() ? "name" : prefix.ToString(); - root = GetSanitizedName(root); + string root = GetSanitizedName(prefix.empty() ? "name" : prefix.ToString()); // Strip away numeric suffix (if any). Only recognize separator if it is in // the middle of the name. + bool has_numeric_suffix = false; + int64 numeric_suffix = 0; size_t separator_index = root.rfind(separator_); if (separator_index != string::npos && (separator_index > 0) && (separator_index < root.size() - 1)) { string after_suffix = root.substr(separator_index + 1); - int64 numeric_suffix; if (tensorflow::strings::safe_strto64(after_suffix, &numeric_suffix)) { + has_numeric_suffix = true; // Remove numeric suffix from root. root = root.substr(0, separator_index); // Update count to at least the numeric suffix value to avoid future @@ -71,11 +72,11 @@ string NameUniquer::GetUniqueName(tensorflow::StringPiece prefix) { generated_names_[root] = std::max(generated_names_[root], numeric_suffix); } } - int64* count = &(generated_names_[root]); if (*count == 0) { *count = 1; - return root; + return has_numeric_suffix ? tensorflow::strings::StrCat(root, separator_, 0) + : root; } else { tensorflow::strings::StrAppend(&root, separator_, *count); // Increment lookup under old 'root' name. diff --git a/tensorflow/compiler/xla/service/name_uniquer_test.cc b/tensorflow/compiler/xla/service/name_uniquer_test.cc index 4258cf16876ab46dce6df062ab701b1b1a4a7580..2ec255558c4ed3695ec6c824458cbedac44dc297 100644 --- a/tensorflow/compiler/xla/service/name_uniquer_test.cc +++ b/tensorflow/compiler/xla/service/name_uniquer_test.cc @@ -57,11 +57,18 @@ TEST_F(NameUniquerTest, NumericSuffixes) { EXPECT_EQ("foo.55", uniquer.GetUniqueName("foo")); EXPECT_EQ("foo.55.1", uniquer.GetUniqueName("foo.55.1")); EXPECT_EQ("foo.55.2", uniquer.GetUniqueName("foo.55.1")); - EXPECT_EQ("bar", uniquer.GetUniqueName("bar.-1000")); + EXPECT_EQ("bar.0", uniquer.GetUniqueName("bar.-1000")); EXPECT_EQ("bar.1", uniquer.GetUniqueName("bar.-2000")); EXPECT_EQ("bar.2", uniquer.GetUniqueName("bar.1")); } +TEST_F(NameUniquerTest, PrefixHasSuffix) { + NameUniquer uniquer("."); + + EXPECT_EQ("foo.11.0", uniquer.GetUniqueName("foo.11.0")); + EXPECT_EQ("foo.11", uniquer.GetUniqueName("foo.11")); +} + TEST_F(NameUniquerTest, Sanitize) { NameUniquer uniquer("_"); @@ -73,7 +80,7 @@ TEST_F(NameUniquerTest, Sanitize) { EXPECT_EQ("foo_55", uniquer.GetUniqueName("foo")); // Invalid characters will be replaced with '_'. - EXPECT_EQ("bar", uniquer.GetUniqueName("bar<-1000")); + EXPECT_EQ("bar_0", uniquer.GetUniqueName("bar<-1000")); EXPECT_EQ("bar_1", uniquer.GetUniqueName("bar<-2000")); EXPECT_EQ("bar_2", uniquer.GetUniqueName("bar_1")); diff --git a/tensorflow/compiler/xla/service/pattern_matcher.h b/tensorflow/compiler/xla/service/pattern_matcher.h new file mode 100644 index 0000000000000000000000000000000000000000..586f6ef7a9c4f17f69340e77be17aec2f677a791 --- /dev/null +++ b/tensorflow/compiler/xla/service/pattern_matcher.h @@ -0,0 +1,1013 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT 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_PATTERN_MATCHER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_PATTERN_MATCHER_H_ + +#include "tensorflow/compiler/xla/layout_util.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/core/lib/core/stringpiece.h" + +namespace xla { + +// A pattern matcher for HloInstructions, Shapes, and Layouts. +// +// The Match function's first argument must be HloInstruction*, Shape*, or +// Layout*. The second argument is a pattern that will be matched against the +// first argument, as described below. +// +// Patterns are constructed using the match::Op, match::Shape, or match::Layout +// functions. By default, the returned patterns will match any HloInstruction, +// Shape, or Layout, respectively. However the match can be made more specific +// by using the pattern's modifier methods, for example: +// +// match::Op().WithOpcode(HloOpcode::kAdd).WithOperand( +// 0, match::Op().WithOpcode(HloOpcode::kConstant)) +// +// This pattern will match Add instructions whose first operand is a constant. +// +// Each pattern type has the following modifiers: +// +// Op(): +// - WithName: match operations with the given name +// - WithOpcode: match operations with the given opcode +// - WithShape: match operations whose shape matches the given pattern +// - WithOperand: match operations whose operand matches the given pattern +// +// Shape(): +// - EqualTo: matches shapes that are equal to the argument +// - CompatibleTo: matches shapes that are compatible to the argument +// - IsScalar/IsArray/IsTuple: matches scalar/array/tuple shapes +// - IsDenseArray/IsSparseArray: matches arrays with dense/sparse format +// - WithLayout: match shapes whose layout matches the given pattern +// - WithLayoutEqualTo: matches shapes whose layouts equal the argument +// - WithSubshape: matches tuple shapes whose subshape matches the given +// pattern +// - WithSubshapeEqualTo: matches shapes with a subshape equal the argument +// - WithElementType: matches array/scalar shapes with the given element +// type +// - WithRank: matches array/scalar types with the given rank +// +// Layout(): +// - EqualTo: matches layouts that are equal to the argument +// - WithDenseFormat/WithSparseFormat: matches layouts with dense/sparse +// format +// +// Op(), Shape(), and Layout() may be passed an argument of type +// HloInstruction**, Shape**, or Layout**, respectively, or const versions of +// these pointers. If the pattern is matched, the address of the matched value +// will be "captured" and stored at this location. +// +// For example: +// HloInstruction* foo = ...; +// HloInstruction* matched_operand; +// CHECK(Match(foo, +// match::Op().WithOperand(0, match::Op(&matched_operand)))); +// +// Helpers are provided for common nullary, unary, binary, and ternary +// instructions. These helpers can be called with no arguments, in which case +// they will match any instruction matching the opcode. They may also be called +// with matches for the operands and with an optional capture. (The capture must +// be the first argument.) Some examples of these helpers and their equivalents +// are provided below. +// +// Example nullary instruction: +// Recv() == Op().WithOpcode(HloOpcode::kRecv) +// Recv(&a) == Op(&a).WithOpcode(HloOpcode::kRecv) +// +// Example unary instruction: +// Abs() == Op().WithOpcode(HloOpcode::kAbs) +// Abs(Op(&a)) == Op().WithOpcode(HloOpcode::kAbs) +// .WithOperand(0, Op(&a))) +// Abs(&a, Op(&b)) == Op(&a).WithOpcode(HloOpcode::kAbs) +// .WithOperand(0, Op(&b)) +// +// Example binary instruction: +// Add() == Op().WithOpcode(HloOpcode::kAdd) +// Add(Op(&a), Op(&b)) == Op().WithOpcode(HloOpcode::kAdd) +// .WithOperand(0, Op(&a)) +// .WithOperand(1, Op(&b)) +// Add(&a, Op(&b), Op(&c)) == Op(&a).WithOpcode(HloOpcode::kAdd) +// .WithOperand(0, Op(&b)) +// .WithOperand(1, Op(&c)) +// +// Example ternary instruction: +// Clamp() == Op().WithOpcode(HloOpcode::kClamp) +// Clamp(Op(&a), Op(&b), Op(&c)) == Op().WithOpcode(HloOpcode::kClamp) +// .WithOperand(0, Op(&a)) +// .WithOperand(1, Op(&b)) +// .WithOperand(2, Op(&c)) +// Clamp(&a, Op(&b), Op(&c), Op(&d)) == Op(&a).WithOpcode(HloOpcode::kClamp) +// .WithOperand(0, Op(&b)) +// .WithOperand(1, Op(&c)) +// .WithOperand(2, Op(&d)) +// +template +bool Match(Value* value, const Pattern& pattern) { + return pattern.Match(value); +} + +namespace match { + +namespace detail { + +template +class LayoutPattern; + +// The base LayoutPattern implementation. Matches only if the layout is not +// nullptr. +class LayoutPatternBaseImpl { + public: + bool Match(const ::xla::Layout* layout) const { return layout != nullptr; } +}; + +// A LayoutPattern implementation that matches only if the layout equals a +// Layout proto. +template +class LayoutPatternEqualImpl { + public: + explicit constexpr LayoutPatternEqualImpl(const Previous& previous, + const ::xla::Layout* layout) + : previous_(previous), layout_(layout) {} + + bool Match(const ::xla::Layout* layout) const { + return previous_.Match(layout) && LayoutUtil::Equal(*layout_, *layout); + } + + private: + Previous previous_; + const ::xla::Layout* layout_; +}; + +// A LayoutPattern implementation that matches only if the layout has a given +// format. +template +class LayoutPatternFormatImpl { + public: + explicit constexpr LayoutPatternFormatImpl(const Previous& previous, + Format format) + : previous_(previous), format_(format) {} + + bool Match(const ::xla::Layout* layout) const { + return previous_.Match(layout) && layout->format() == format_; + } + + private: + Previous previous_; + Format format_; +}; + +// A pattern that matches Layouts. +template +class LayoutPattern { + public: + explicit constexpr LayoutPattern(const Impl& impl, + LayoutType** matched_layout) + : impl_(impl), matched_layout_(matched_layout) {} + + // Returns true and captures the layout iff it matches the pattern. + bool Match(const ::xla::Layout* layout) const { + if (impl_.Match(layout)) { + if (matched_layout_) { + *matched_layout_ = layout; + } + return true; + } + return false; + } + + // Returns true and captures the layout iff it matches the pattern. + bool Match(::xla::Layout* layout) const { + if (impl_.Match(layout)) { + if (matched_layout_) { + *matched_layout_ = layout; + } + return true; + } + return false; + } + + // Modifies the pattern to match only if the layout equals the given proto. + // The layout must outlive the returned pattern. + constexpr LayoutPattern> EqualTo( + const Layout* layout) const { + return LayoutPattern>( + LayoutPatternEqualImpl(impl_, layout), matched_layout_); + } + + // Modifies the pattern to match only if the layout has a dense format. + constexpr LayoutPattern> + WithDenseFormat() const { + return LayoutPattern>( + LayoutPatternFormatImpl(impl_, DENSE), matched_layout_); + } + + // Modifies the pattern to match only if the layout has a sparse format. + constexpr LayoutPattern> + WithSparseFormat() const { + return LayoutPattern>( + LayoutPatternFormatImpl(impl_, SPARSE), matched_layout_); + } + + private: + Impl impl_; + LayoutType** matched_layout_; +}; + +} // namespace detail + +// Creates a layout pattern that will capture the matched layout in the +// argument. +inline constexpr detail::LayoutPattern +Layout(const ::xla::Layout** matched_layout = nullptr) { + return detail::LayoutPattern( + detail::LayoutPatternBaseImpl(), matched_layout); +} + +// Creates a layout pattern that will capture the matched layout in the +// argument. +inline constexpr detail::LayoutPattern<::xla::Layout, + detail::LayoutPatternBaseImpl> +Layout(::xla::Layout** matched_layout) { + return detail::LayoutPattern<::xla::Layout, detail::LayoutPatternBaseImpl>( + detail::LayoutPatternBaseImpl(), matched_layout); +} + +namespace detail { + +template +class ShapePattern; + +// The base ShapePattern implementation. Matches only if the shape is not +// nullptr. +class ShapePatternBaseImpl { + public: + bool Match(const ::xla::Shape* shape) const { return shape != nullptr; } +}; + +// A ShapePattern implementation that matches only if the shape equals a Shape +// proto. +template +class ShapePatternEqualImpl { + public: + explicit constexpr ShapePatternEqualImpl(const Previous& previous, + const ::xla::Shape* shape) + : previous_(previous), shape_(shape) {} + + bool Match(const ::xla::Shape* shape) const { + return previous_.Match(shape) && ShapeUtil::Equal(*shape_, *shape); + } + + private: + Previous previous_; + const ::xla::Shape* shape_; +}; + +// A ShapePattern implementation that matches only if the shape is compatible to +// a Shape proto. +template +class ShapePatternCompatibleImpl { + public: + explicit constexpr ShapePatternCompatibleImpl(const Previous& previous, + const ::xla::Shape* shape) + : previous_(previous), shape_(shape) {} + + bool Match(const ::xla::Shape* shape) const { + return previous_.Match(shape) && ShapeUtil::Compatible(*shape_, *shape); + } + + private: + Previous previous_; + const ::xla::Shape* shape_; +}; + +// A ShapePattern implementation that matches only if the shape has a given +// element type. +template +class ShapePatternElementTypeImpl { + public: + explicit constexpr ShapePatternElementTypeImpl(const Previous& previous, + PrimitiveType element_type) + : previous_(previous), element_type_(element_type) {} + + bool Match(const ::xla::Shape* shape) const { + return previous_.Match(shape) && shape->element_type() == element_type_; + } + + private: + Previous previous_; + PrimitiveType element_type_; +}; + +// A ShapePattern implementation that matches only if the shape is scalar. +template +class ShapePatternIsScalarImpl { + public: + explicit constexpr ShapePatternIsScalarImpl(const Previous& previous) + : previous_(previous) {} + + bool Match(const ::xla::Shape* shape) const { + return previous_.Match(shape) && ShapeUtil::IsScalar(*shape); + } + + private: + Previous previous_; +}; + +// A ShapePattern implementation that matches only if the shape is an array +template +class ShapePatternIsArrayImpl { + public: + explicit constexpr ShapePatternIsArrayImpl(const Previous& previous) + : previous_(previous) {} + + bool Match(const ::xla::Shape* shape) const { + return previous_.Match(shape) && ShapeUtil::IsArray(*shape); + } + + private: + Previous previous_; +}; + +// A ShapePattern implementation that matches only if the shape is a tuple. +template +class ShapePatternIsTupleImpl { + public: + explicit constexpr ShapePatternIsTupleImpl(const Previous& previous) + : previous_(previous) {} + + bool Match(const ::xla::Shape* shape) const { + return previous_.Match(shape) && ShapeUtil::IsTuple(*shape); + } + + private: + Previous previous_; +}; + +// A ShapePattern implementation that matches only if the shape has a given +// rank. +template +class ShapePatternRankImpl { + public: + explicit constexpr ShapePatternRankImpl(const Previous& previous, int64 rank) + : previous_(previous), rank_(rank) {} + + bool Match(const ::xla::Shape* shape) const { + return previous_.Match(shape) && ShapeUtil::Rank(*shape) == rank_; + } + + private: + Previous previous_; + int64 rank_; +}; + +// A ShapePattern implementation that matches only if the shape has a layout +// that matches a given pattern. +template +class ShapePatternLayoutImpl { + public: + explicit constexpr ShapePatternLayoutImpl( + const Previous& previous, + const LayoutPattern& layout) + : previous_(previous), layout_(layout) {} + + bool Match(const ::xla::Shape* shape) const { + return previous_.Match(shape) && LayoutUtil::HasLayout(*shape) && + layout_.Match(&shape->layout()); + } + + bool Match(Shape* shape) const { + return previous_.Match(shape) && LayoutUtil::HasLayout(*shape) && + layout_.Match(shape->mutable_layout()); + } + + private: + Previous previous_; + LayoutPattern layout_; +}; + +// A ShapePattern implementation that matches only if the shape has a subshape +// that matches a given pattern. +template +class ShapePatternSubshapeImpl { + public: + explicit ShapePatternSubshapeImpl( + const Previous& previous, ShapeIndexView index, + const ShapePattern& subshape) + : previous_(previous), index_(index), subshape_(subshape) {} + + bool Match(const ::xla::Shape* shape) const { + return previous_.Match(shape) && ShapeUtil::IndexIsValid(*shape, index_) && + subshape_.Match(&ShapeUtil::GetSubshape(*shape, index_)); + } + + bool Match(::xla::Shape* shape) const { + return previous_.Match(shape) && ShapeUtil::IndexIsValid(*shape, index_) && + subshape_.Match(ShapeUtil::GetMutableSubshape(shape, index_)); + } + + private: + Previous previous_; + ShapeIndexView index_; + ShapePattern subshape_; +}; + +// A pattern that matches Shapes. +template +class ShapePattern { + public: + explicit constexpr ShapePattern(const Impl& impl, ShapeType** matched_shape) + : impl_(impl), matched_shape_(matched_shape) {} + + // Returns true and captures the shape iff it matches the pattern. + bool Match(const ::xla::Shape* shape) const { + if (impl_.Match(shape)) { + if (matched_shape_) { + *matched_shape_ = shape; + } + return true; + } + return false; + } + + // Returns true and captures the shape iff it matches the pattern. + bool Match(::xla::Shape* shape) const { + if (impl_.Match(shape)) { + if (matched_shape_) { + *matched_shape_ = shape; + } + return true; + } + return false; + } + + // Modifies the pattern to match only if the shape equals the given proto. + // The layout must outlive the returned pattern. + constexpr ShapePattern> EqualTo( + const ::xla::Shape* shape) const { + return ShapePattern>( + ShapePatternEqualImpl(impl_, shape), matched_shape_); + } + + // Modifies the pattern to match only if the shape is compatible to the given + // proto. The layout must outlive the returned pattern. + constexpr ShapePattern> + CompatibleTo(const ::xla::Shape* shape) const { + return ShapePattern>( + ShapePatternCompatibleImpl(impl_, shape), matched_shape_); + } + + // Modifies the pattern to match only if the shape has the given element type. + constexpr ShapePattern> + WithElementType(PrimitiveType element_type) const { + return ShapePattern>( + ShapePatternElementTypeImpl(impl_, element_type), matched_shape_); + } + + // Modifies the pattern to match only if the shape is scalar. + constexpr ShapePattern> IsScalar() + const { + return ShapePattern>( + ShapePatternIsScalarImpl(impl_), matched_shape_); + } + + // Modifies the pattern to match only if the shape is an array. + constexpr ShapePattern> IsArray() + const { + return ShapePattern>( + ShapePatternIsArrayImpl(impl_), matched_shape_); + } + + // Modifies the pattern to match only if the shape is a tuple. + constexpr ShapePattern> IsTuple() + const { + return ShapePattern>( + ShapePatternIsTupleImpl(impl_), matched_shape_); + } + + // Modifies the pattern to match only if the shape has the given rank. + constexpr ShapePattern> WithRank( + int64 rank) const { + return ShapePattern>( + ShapePatternRankImpl(impl_, rank), matched_shape_); + } + + // Modifies the pattern to match only if the shape has a layout that matches + // the given pattern. + template + constexpr ShapePattern> + WithLayout(const LayoutPattern& layout) const { + return ShapePattern>( + ShapePatternLayoutImpl(impl_, layout), + matched_shape_); + } + + constexpr ShapePattern< + ShapeType, + ShapePatternLayoutImpl>> + WithLayoutEqualTo(const ::xla::Layout* layout) const { + return WithLayout(Layout().EqualTo(layout)); + } + + constexpr ShapePattern< + ShapeType, + ShapePatternLayoutImpl>> + IsDenseArray() const { + return WithLayout(Layout().WithDenseFormat()); + } + + constexpr ShapePattern< + ShapeType, + ShapePatternLayoutImpl>> + IsSparseArray() const { + return WithLayout(Layout().WithSparseFormat()); + } + + // Modifies the pattern to match only if the shape has a subshape that matches + // the given pattern. + template + ShapePattern> + WithSubshape(ShapeIndexView index, + const ShapePattern& subshape) const { + return ShapePattern< + ShapeType, ShapePatternSubshapeImpl>( + ShapePatternSubshapeImpl(impl_, index, + subshape), + matched_shape_); + } + + ShapePattern>> + WithSubshapeEqualTo(ShapeIndexView index, const ::xla::Shape* shape) const { + return WithSubshape(index, + ShapePattern( + ShapePatternBaseImpl(), nullptr) + .EqualTo(shape)); + } + + ShapePattern>> + WithSubshapeCompatibleTo(ShapeIndexView index, + const ::xla::Shape* shape) const { + return WithSubshape(index, + ShapePattern( + ShapePatternBaseImpl(), nullptr) + .CompatibleTo(shape)); + } + + private: + Impl impl_; + ShapeType** matched_shape_; +}; + +} // namespace detail + +// Creates a shape pattern that will capture the matched layout in the argument. +inline constexpr detail::ShapePattern +Shape(const ::xla::Shape** matched_shape = nullptr) { + return detail::ShapePattern( + detail::ShapePatternBaseImpl(), matched_shape); +} + +// Creates a shape pattern that will capture the matched layout in the argument. +inline constexpr detail::ShapePattern<::xla::Shape, + detail::ShapePatternBaseImpl> +Shape(::xla::Shape** matched_shape) { + return detail::ShapePattern<::xla::Shape, detail::ShapePatternBaseImpl>( + detail::ShapePatternBaseImpl(), matched_shape); +} + +namespace detail { + +template +class HloInstructionPattern; + +// The base HloInstructionPattern implementation. Matches only if the +// instruction is not nullptr. +class HloInstructionPatternBaseImpl { + public: + bool Match(const ::xla::HloInstruction* inst) const { + return inst != nullptr; + } +}; + +// An HloInstructionPattern implementation that matches only if the instruction +// has a given name. +template +class HloInstructionPatternNameImpl { + public: + explicit HloInstructionPatternNameImpl(const Previous& previous, + tensorflow::StringPiece name) + : previous_(previous), name_(name) {} + + bool Match(const ::xla::HloInstruction* inst) const { + return previous_.Match(inst) && inst->name() == name_; + } + + private: + Previous previous_; + tensorflow::StringPiece name_; +}; + +// An HloInstructionPattern implementation that matches only if the instruction +// has a given opcode. +template +class HloInstructionPatternOpcodeImpl { + public: + explicit constexpr HloInstructionPatternOpcodeImpl(const Previous& previous, + HloOpcode opcode, + bool invert) + : previous_(previous), opcode_(opcode), invert_(invert) {} + + bool Match(const ::xla::HloInstruction* inst) const { + return previous_.Match(inst) && (invert_ ^ (inst->opcode() == opcode_)); + } + + private: + Previous previous_; + HloOpcode opcode_; + bool invert_; +}; + +// An HloInstructionPattern implementation that matches only if the instruction +// has a shape that matches a given pattern. +template +class HloInstructionPatternShapeImpl { + public: + explicit constexpr HloInstructionPatternShapeImpl( + const Previous& previous, const ShapePattern& shape) + : previous_(previous), shape_(shape) {} + + bool Match(const ::xla::HloInstruction* inst) const { + return previous_.Match(inst) && shape_.Match(&inst->shape()); + } + + bool Match(::xla::HloInstruction* inst) const { + return previous_.Match(inst) && shape_.Match(inst->mutable_shape()); + } + + private: + Previous previous_; + ShapePattern shape_; +}; + +// An HloInstructionPattern implementation that matches only if the instruction +// has an operand that matches a given pattern. +template +class HloInstructionPatternOperandImpl { + public: + explicit constexpr HloInstructionPatternOperandImpl( + const Previous& previous, int64 operand_index, + const HloInstructionPattern& operand) + : previous_(previous), operand_index_(operand_index), operand_(operand) {} + + bool Match(const ::xla::HloInstruction* inst) const { + return previous_.Match(inst) && operand_index_ < inst->operand_count() && + operand_.Match(inst->operand(operand_index_)); + } + + bool Match(::xla::HloInstruction* inst) const { + return previous_.Match(inst) && operand_index_ < inst->operand_count() && + operand_.Match(inst->mutable_operand(operand_index_)); + } + + private: + Previous previous_; + int64 operand_index_; + HloInstructionPattern operand_; +}; + +// A pattern that matches HloInstructions. +template +class HloInstructionPattern { + public: + explicit constexpr HloInstructionPattern(const Impl& impl, + HloInstructionType** matched_inst) + : impl_(impl), matched_inst_(matched_inst) {} + + // Returns true and captures the instruction iff it matches the pattern. + bool Match(const ::xla::HloInstruction* inst) const { + if (impl_.Match(inst)) { + if (matched_inst_) { + *matched_inst_ = inst; + } + return true; + } + return false; + } + + // Returns true and captures the instruction iff it matches the pattern. + bool Match(::xla::HloInstruction* inst) const { + if (impl_.Match(inst)) { + if (matched_inst_) { + *matched_inst_ = inst; + } + return true; + } + return false; + } + + // Modifies the pattern to match only if the instruction has the given name. + HloInstructionPattern> + WithName(tensorflow::StringPiece name) const { + return HloInstructionPattern>( + HloInstructionPatternNameImpl(impl_, name), matched_inst_); + } + + // Modifies the pattern to match only if the instruction has the given opcode. + constexpr HloInstructionPattern> + WithOpcode(HloOpcode opcode) const { + return HloInstructionPattern>( + HloInstructionPatternOpcodeImpl(impl_, opcode, false), + matched_inst_); + } + + // Modifies the pattern to match only if the instruction does not have the + // given opcode. + constexpr HloInstructionPattern> + WithoutOpcode(HloOpcode opcode) const { + return HloInstructionPattern>( + HloInstructionPatternOpcodeImpl(impl_, opcode, true), + matched_inst_); + } + + // Modifies the pattern to match only if the instruction is a constant. + constexpr HloInstructionPattern> + IsConstant() const { + return WithOpcode(HloOpcode::kConstant); + } + + // Modifies the pattern to match only if the instruction is not a constant. + constexpr HloInstructionPattern> + IsNonConstant() const { + return WithoutOpcode(HloOpcode::kConstant); + } + + // Modifies the pattern to match only if the instruction has a shape that + // matches the given pattern. + template + constexpr HloInstructionPattern< + HloInstructionType, + HloInstructionPatternShapeImpl> + WithShape(const ShapePattern& shape) const { + return HloInstructionPattern< + HloInstructionType, + HloInstructionPatternShapeImpl>( + HloInstructionPatternShapeImpl(impl_, + shape), + matched_inst_); + } + + // Modifies the pattern to match only if the instruction has an operand that + // matches the given pattern. + template + constexpr HloInstructionPattern< + HloInstructionType, + HloInstructionPatternOperandImpl> + WithOperand( + int64 operand_index, + const HloInstructionPattern& operand) const { + return HloInstructionPattern< + HloInstructionType, + HloInstructionPatternOperandImpl>( + HloInstructionPatternOperandImpl( + impl_, operand_index, operand), + matched_inst_); + } + + private: + Impl impl_; + HloInstructionType** matched_inst_; +}; + +} // namespace detail + +// Creates an instruction pattern that will capture the matched instruction in +// the argument. +inline constexpr detail::HloInstructionPattern< + const ::xla::HloInstruction, detail::HloInstructionPatternBaseImpl> +Op(const ::xla::HloInstruction** matched_inst = nullptr) { + return detail::HloInstructionPattern( + detail::HloInstructionPatternBaseImpl(), matched_inst); +} + +// Creates an instruction pattern that will capture the matched instruction in +// the argument. +inline constexpr detail::HloInstructionPattern< + ::xla::HloInstruction, detail::HloInstructionPatternBaseImpl> +Op(::xla::HloInstruction** matched_inst) { + return detail::HloInstructionPattern<::xla::HloInstruction, + detail::HloInstructionPatternBaseImpl>( + detail::HloInstructionPatternBaseImpl(), matched_inst); +} + +// Helpers for nullary instructions. +#define XLA_NULLOP_PATTERN(NAME) \ + inline auto NAME()->decltype(Op().WithOpcode(HloOpcode::k##NAME)) { \ + return Op().WithOpcode(HloOpcode::k##NAME); \ + } \ + \ + template \ + inline auto NAME(HloInstructionType** matched_inst) \ + ->decltype(Op(matched_inst).WithOpcode(HloOpcode::k##NAME)) { \ + return Op(matched_inst).WithOpcode(HloOpcode::k##NAME); \ + } +XLA_NULLOP_PATTERN(Constant) +XLA_NULLOP_PATTERN(Infeed) +XLA_NULLOP_PATTERN(Parameter) +XLA_NULLOP_PATTERN(Recv) +#undef XLA_NULLOP_PATTERN + +// Helpers for unary instructions. +#define XLA_UNOP_PATTERN(NAME) \ + inline auto NAME()->decltype(Op().WithOpcode(HloOpcode::k##NAME)) { \ + return Op().WithOpcode(HloOpcode::k##NAME); \ + } \ + \ + template \ + inline auto NAME(Arg&& arg)->decltype( \ + Op().WithOpcode(HloOpcode::k##NAME) \ + .WithOperand(0, std::forward(arg))) { \ + return Op() \ + .WithOpcode(HloOpcode::k##NAME) \ + .WithOperand(0, std::forward(arg)); \ + } \ + \ + template \ + inline auto NAME(HloInstructionType** matched_inst, Arg&& arg) \ + ->decltype(Op(matched_inst) \ + .WithOpcode(HloOpcode::k##NAME) \ + .WithOperand(0, std::forward(arg))) { \ + return Op(matched_inst) \ + .WithOpcode(HloOpcode::k##NAME) \ + .WithOperand(0, std::forward(arg)); \ + } +XLA_UNOP_PATTERN(Abs) +XLA_UNOP_PATTERN(RoundNearestAfz) +XLA_UNOP_PATTERN(Bitcast) +XLA_UNOP_PATTERN(Broadcast) +XLA_UNOP_PATTERN(Ceil) +XLA_UNOP_PATTERN(Copy) +XLA_UNOP_PATTERN(Cos) +XLA_UNOP_PATTERN(Exp) +XLA_UNOP_PATTERN(Fft) +XLA_UNOP_PATTERN(Floor) +XLA_UNOP_PATTERN(Imag) +XLA_UNOP_PATTERN(IsFinite) +XLA_UNOP_PATTERN(Log) +XLA_UNOP_PATTERN(Not) +XLA_UNOP_PATTERN(Negate) +XLA_UNOP_PATTERN(Outfeed) +XLA_UNOP_PATTERN(Real) +XLA_UNOP_PATTERN(Reduce) +XLA_UNOP_PATTERN(ReducePrecision) +XLA_UNOP_PATTERN(Reshape) +XLA_UNOP_PATTERN(Reverse) +XLA_UNOP_PATTERN(Send) +XLA_UNOP_PATTERN(Sign) +XLA_UNOP_PATTERN(Sin) +XLA_UNOP_PATTERN(Sort) +XLA_UNOP_PATTERN(Tanh) +XLA_UNOP_PATTERN(Transpose) +#undef XLA_UNOP_PATTERN + +// Helpers for binary instructions. +#define XLA_BINOP_PATTERN(NAME) \ + inline auto NAME()->decltype(Op().WithOpcode(HloOpcode::k##NAME)) { \ + return Op().WithOpcode(HloOpcode::k##NAME); \ + } \ + \ + template \ + inline auto NAME(Lhs&& lhs, Rhs&& rhs) \ + ->decltype(Op().WithOpcode(HloOpcode::k##NAME) \ + .WithOperand(0, std::forward(lhs)) \ + .WithOperand(1, std::forward(rhs))) { \ + return Op() \ + .WithOpcode(HloOpcode::k##NAME) \ + .WithOperand(0, std::forward(lhs)) \ + .WithOperand(1, std::forward(rhs)); \ + } \ + \ + template \ + inline auto NAME(HloInstructionType** matched_inst, Lhs&& lhs, Rhs&& rhs) \ + ->decltype(Op(matched_inst) \ + .WithOpcode(HloOpcode::k##NAME) \ + .WithOperand(0, std::forward(lhs)) \ + .WithOperand(1, std::forward(rhs))) { \ + return Op(matched_inst) \ + .WithOpcode(HloOpcode::k##NAME) \ + .WithOperand(0, std::forward(lhs)) \ + .WithOperand(1, std::forward(rhs)); \ + } +XLA_BINOP_PATTERN(Add) +XLA_BINOP_PATTERN(Atan2) +XLA_BINOP_PATTERN(Divide) +XLA_BINOP_PATTERN(Complex) +XLA_BINOP_PATTERN(Dot) +XLA_BINOP_PATTERN(Eq) +XLA_BINOP_PATTERN(Gather) +XLA_BINOP_PATTERN(Ge) +XLA_BINOP_PATTERN(Gt) +XLA_BINOP_PATTERN(Le) +XLA_BINOP_PATTERN(Lt) +XLA_BINOP_PATTERN(Maximum) +XLA_BINOP_PATTERN(Minimum) +XLA_BINOP_PATTERN(Multiply) +XLA_BINOP_PATTERN(Ne) +XLA_BINOP_PATTERN(Power) +XLA_BINOP_PATTERN(Remainder) +XLA_BINOP_PATTERN(Subtract) +XLA_BINOP_PATTERN(And) +XLA_BINOP_PATTERN(Or) +XLA_BINOP_PATTERN(ShiftLeft) +XLA_BINOP_PATTERN(ShiftRightArithmetic) +XLA_BINOP_PATTERN(ShiftRightLogical) +#undef XLA_BINOP_PATTERN + +// Helpers for ternary instructions. +#define XLA_TERNOP_PATTERN(NAME) \ + inline auto NAME()->decltype(Op().WithOpcode(HloOpcode::k##NAME)) { \ + return Op().WithOpcode(HloOpcode::k##NAME); \ + } \ + \ + template \ + inline auto NAME(Arg0&& arg0, Arg1&& arg1, Arg2&& arg2) \ + ->decltype(Op().WithOpcode(HloOpcode::k##NAME) \ + .WithOperand(0, std::forward(arg0)) \ + .WithOperand(1, std::forward(arg1)) \ + .WithOperand(2, std::forward(arg2))) { \ + return Op() \ + .WithOpcode(HloOpcode::k##NAME) \ + .WithOperand(0, std::forward(arg0)) \ + .WithOperand(1, std::forward(arg1)) \ + .WithOperand(2, std::forward(arg2)); \ + } \ + \ + template \ + inline auto NAME(HloInstructionType** matched_inst, Arg0&& arg0, \ + Arg1&& arg1, Arg2&& arg2) \ + ->decltype(Op(matched_inst) \ + .WithOpcode(HloOpcode::k##NAME) \ + .WithOperand(0, std::forward(arg0)) \ + .WithOperand(1, std::forward(arg1)) \ + .WithOperand(2, std::forward(arg2))) { \ + return Op(matched_inst) \ + .WithOpcode(HloOpcode::k##NAME) \ + .WithOperand(0, std::forward(arg0)) \ + .WithOperand(1, std::forward(arg1)) \ + .WithOperand(2, std::forward(arg2)); \ + } +XLA_TERNOP_PATTERN(Clamp); +XLA_TERNOP_PATTERN(Select); +#undef XLA_TERNOP_PATTERN + +// Helpers for matching non-constant instructions. +inline auto NonConstant() -> decltype(Op().IsNonConstant()) { + return Op().IsNonConstant(); +} + +template +inline auto NonConstant(HloInstructionType** matched_inst) + -> decltype(Op(matched_inst).IsNonConstant()) { + return Op(matched_inst).IsNonConstant(); +} + +} // namespace match + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_PATTERN_MATCHER_H_ diff --git a/tensorflow/compiler/xla/service/pattern_matcher_test.cc b/tensorflow/compiler/xla/service/pattern_matcher_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..c88157c312524fb273e6df368d2ef61d679d1d8b --- /dev/null +++ b/tensorflow/compiler/xla/service/pattern_matcher_test.cc @@ -0,0 +1,174 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT 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/pattern_matcher.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h" +#include "tensorflow/core/platform/test.h" + +namespace xla { +namespace { + +TEST(PatternMatcherTest, AddOp) { + constexpr char kModuleStr[] = R"(HloModule two_plus_two_module + ENTRY %two_plus_two_computation () -> f32[] { + %two = f32[] constant(2) + ROOT %two_plus_two = f32[] add(f32[] %two, f32[] %two) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto hlo_module, tools::Parse(kModuleStr)); + + const HloInstruction* matched_inst; + HloInstruction* matched_operand; + Shape* matched_shape; + Layout* matched_layout; + + ASSERT_TRUE(Match( + hlo_module->entry_computation()->root_instruction(), + match::Op(&matched_inst) + .WithName("two_plus_two") + .WithOpcode(HloOpcode::kAdd) + .WithShape( + match::Shape(&matched_shape) + .WithLayout(match::Layout(&matched_layout).WithDenseFormat())) + .WithOperand( + 0, + match::Op(&matched_operand).WithOpcode(HloOpcode::kConstant)))); + ASSERT_NE(matched_inst, nullptr); + EXPECT_EQ(matched_inst->name(), "two_plus_two"); + EXPECT_EQ(matched_inst->opcode(), HloOpcode::kAdd); + + EXPECT_TRUE(Match(hlo_module->entry_computation()->root_instruction(), + match::Add(match::Constant(), match::Constant()))); + + EXPECT_FALSE(Match(hlo_module->entry_computation()->root_instruction(), + match::Op().WithName("bad_name"))); + matched_inst = nullptr; + EXPECT_FALSE(Match(hlo_module->entry_computation()->root_instruction(), + match::Multiply(&matched_inst, match::Op(), match::Op()))); +} + +TEST(PatternMatcherTest, ScalarShape) { + auto scalar_shape = ShapeUtil::MakeShape(F32, {}); + Shape* matched_shape; + EXPECT_TRUE(Match(&scalar_shape, match::Shape(&matched_shape).IsScalar())); + EXPECT_EQ(matched_shape, &scalar_shape); + EXPECT_TRUE(Match(&scalar_shape, match::Shape().IsArray())); + EXPECT_TRUE(Match(&scalar_shape, match::Shape().IsDenseArray())); + EXPECT_FALSE(Match(&scalar_shape, match::Shape().IsTuple())); + EXPECT_TRUE(Match(&scalar_shape, match::Shape().WithElementType(F32))); + EXPECT_TRUE(Match(&scalar_shape, match::Shape().WithRank(0))); + EXPECT_FALSE(Match( + &scalar_shape, + match::Shape().WithSubshape({0}, match::Shape()).WithElementType(F32))); +} + +TEST(PatternMatcherTest, DenseArrayShape) { + auto array_shape = ShapeUtil::MakeShape(F32, {2, 3, 4}); + Shape* matched_shape; + EXPECT_TRUE(Match(&array_shape, match::Shape(&matched_shape).IsArray())); + EXPECT_EQ(matched_shape, &array_shape); + EXPECT_TRUE(Match(&array_shape, match::Shape().IsDenseArray())); + EXPECT_FALSE(Match(&array_shape, match::Shape().IsSparseArray())); + EXPECT_FALSE(Match(&array_shape, match::Shape().IsScalar())); + EXPECT_FALSE(Match(&array_shape, match::Shape().IsTuple())); + EXPECT_TRUE(Match(&array_shape, match::Shape().WithElementType(F32))); + EXPECT_TRUE(Match(&array_shape, match::Shape().WithRank(3))); + EXPECT_FALSE( + Match(&array_shape, match::Shape().WithSubshape({0}, match::Shape()))); + Layout* matched_layout; + EXPECT_FALSE(Match(&array_shape, + match::Shape().WithLayout( + match::Layout(&matched_layout).WithSparseFormat()))); + EXPECT_TRUE(Match(&array_shape, + match::Shape().WithLayout( + match::Layout(&matched_layout).WithDenseFormat()))); + EXPECT_EQ(matched_layout, &array_shape.layout()); +} + +TEST(PatternMatcherTest, SparseArrayShape) { + auto array_shape = ShapeUtil::MakeShapeWithSparseLayout(F32, {2, 3, 4}, 10); + Shape* matched_shape; + EXPECT_TRUE(Match(&array_shape, match::Shape(&matched_shape).IsArray())); + EXPECT_EQ(matched_shape, &array_shape); + EXPECT_FALSE(Match(&array_shape, match::Shape().IsDenseArray())); + EXPECT_TRUE(Match(&array_shape, match::Shape().IsSparseArray())); + EXPECT_FALSE(Match(&array_shape, match::Shape().IsScalar())); + EXPECT_FALSE(Match(&array_shape, match::Shape().IsTuple())); + EXPECT_TRUE(Match(&array_shape, match::Shape().WithElementType(F32))); + EXPECT_TRUE(Match(&array_shape, match::Shape().WithRank(3))); + EXPECT_FALSE( + Match(&array_shape, match::Shape().WithSubshape({0}, match::Shape()))); + Layout* matched_layout; + EXPECT_FALSE(Match(&array_shape, + match::Shape().WithLayout( + match::Layout(&matched_layout).WithDenseFormat()))); + EXPECT_TRUE(Match(&array_shape, + match::Shape().WithLayout( + match::Layout(&matched_layout).WithSparseFormat()))); + EXPECT_EQ(matched_layout, &array_shape.layout()); +} + +TEST(PatternMatcherTest, TupleShape) { + auto tuple_shape = ShapeUtil::MakeTupleShape({ + ShapeUtil::MakeShape(F32, {1, 2, 3}), + ShapeUtil::MakeShape(S32, {4, 5}), + }); + EXPECT_TRUE(Match(&tuple_shape, match::Shape().IsTuple())); + EXPECT_FALSE(Match(&tuple_shape, match::Shape().IsArray())); + EXPECT_FALSE(Match(&tuple_shape, match::Shape().IsScalar())); + + Shape* subshape; + ASSERT_TRUE(Match( + &tuple_shape, + match::Shape().WithSubshape( + {0}, match::Shape(&subshape).WithElementType(F32).WithRank(3)))); + ASSERT_NE(subshape, nullptr); + EXPECT_TRUE( + ShapeUtil::Equal(*subshape, ShapeUtil::GetSubshape(tuple_shape, {0}))); + EXPECT_TRUE(Match(&tuple_shape, + match::Shape().WithSubshape( + {0}, match::Shape().EqualTo( + &ShapeUtil::GetSubshape(tuple_shape, {0}))))); + EXPECT_FALSE(Match(&tuple_shape, + match::Shape().WithSubshape( + {0}, match::Shape().EqualTo( + &ShapeUtil::GetSubshape(tuple_shape, {1}))))); + + ASSERT_TRUE(Match( + &tuple_shape, + match::Shape().WithSubshape( + {1}, match::Shape(&subshape).WithElementType(S32).WithRank(2)))); + ASSERT_NE(subshape, nullptr); + EXPECT_TRUE( + ShapeUtil::Equal(*subshape, ShapeUtil::GetSubshape(tuple_shape, {1}))); + EXPECT_TRUE(Match(&tuple_shape, + match::Shape().WithSubshape( + {1}, match::Shape().EqualTo( + &ShapeUtil::GetSubshape(tuple_shape, {1}))))); + EXPECT_FALSE(Match(&tuple_shape, + match::Shape().WithSubshape( + {1}, match::Shape().EqualTo( + &ShapeUtil::GetSubshape(tuple_shape, {0}))))); + + EXPECT_FALSE( + Match(&tuple_shape, match::Shape().WithSubshape({2}, match::Shape()))); + EXPECT_FALSE( + Match(&tuple_shape, match::Shape().WithSubshape({0, 0}, match::Shape()))); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/platform_util.cc b/tensorflow/compiler/xla/service/platform_util.cc index aa974ee61a27de9c19e97d8a6eb48f9261ce4bd9..7c63c0acc7764d558b2151190f0fa79fac355cbf 100644 --- a/tensorflow/compiler/xla/service/platform_util.cc +++ b/tensorflow/compiler/xla/service/platform_util.cc @@ -29,8 +29,6 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" -namespace se = ::perftools::gputools; - namespace xla { using tensorflow::str_util::Lowercase; diff --git a/tensorflow/compiler/xla/service/platform_util.h b/tensorflow/compiler/xla/service/platform_util.h index 69188820a70707d9c9be10b20fb7de92ad4d9873..571451ba43a81d19b70e4954e45d3447f15dcedc 100644 --- a/tensorflow/compiler/xla/service/platform_util.h +++ b/tensorflow/compiler/xla/service/platform_util.h @@ -34,29 +34,27 @@ class PlatformUtil { // // Note that, even if a platform is present with zero devices, if we *do* have // compilation support for it, it will be returned in this sequence. - static StatusOr> - GetSupportedPlatforms(); + static StatusOr> GetSupportedPlatforms(); // Convenience function which returns the default supported platform for // tests. If exactly one supported platform is present, then this platform is // the default platform. If exactly two platforms are present and one of them // is the interpreter platform, then the other platform is the default // platform. Otherwise returns an error. - static StatusOr GetDefaultPlatform(); + static StatusOr GetDefaultPlatform(); // Convenience function which returns the sole supported platform. If // exactly one supported platform is present, then this platform is the // default platform. Otherwise returns an error. - static StatusOr GetSolePlatform(); + static StatusOr GetSolePlatform(); // Returns the platform according to the given name. Returns error if there is // no such platform. - static StatusOr GetPlatform( - const string& platform_name); + static StatusOr GetPlatform(const string& platform_name); // Returns exactly one platform that does not have given name. Returns error // if there is no such platform, or there are multiple such platforms. - static StatusOr GetPlatformExceptFor( + static StatusOr GetPlatformExceptFor( const string& platform_name); // Returns a vector of StreamExecutors for the given platform. The vector is @@ -64,8 +62,8 @@ class PlatformUtil { // element is nullptr, then the device is present by not supported by XLA. // // If the platform has no visible devices, a not-found error is returned. - static StatusOr> - GetStreamExecutors(perftools::gputools::Platform* platform); + static StatusOr> GetStreamExecutors( + se::Platform* platform); private: TF_DISALLOW_COPY_AND_ASSIGN(PlatformUtil); diff --git a/tensorflow/compiler/xla/service/reshape_mover.cc b/tensorflow/compiler/xla/service/reshape_mover.cc index 49ec38eb62c7b51c7a2d301d882cef032b288036..0f26a025bf125f70199637894741540f89eae7e5 100644 --- a/tensorflow/compiler/xla/service/reshape_mover.cc +++ b/tensorflow/compiler/xla/service/reshape_mover.cc @@ -155,15 +155,20 @@ HloInstruction* UpdateOperand(const HloInstruction* first_reshape_operand, case HloOpcode::kConstant: { if (first_reshape_operand->opcode() == HloOpcode::kReshape) { VLOG(5) << "Adding reshape to kConstant operand"; - return computation->AddInstruction( + HloInstruction* reshape = computation->AddInstruction( HloInstruction::CreateReshape(new_shape, operand)); + operand->SetupDerivedInstruction(reshape); + return reshape; } else { CHECK(first_reshape_operand->opcode() == HloOpcode::kTranspose); VLOG(5) << "Adding transpose to kConstant operand"; std::vector inverse_permutation = InversePermutation(first_reshape_operand->dimensions()); - return computation->AddInstruction(HloInstruction::CreateTranspose( - new_shape, operand, inverse_permutation)); + HloInstruction* transpose = + computation->AddInstruction(HloInstruction::CreateTranspose( + new_shape, operand, inverse_permutation)); + operand->SetupDerivedInstruction(transpose); + return transpose; } } case HloOpcode::kRng: { diff --git a/tensorflow/compiler/xla/service/reshape_mover_test.cc b/tensorflow/compiler/xla/service/reshape_mover_test.cc index 094f7319f462a71f4bfe972771a1de4aedbb8ee3..13e2d3258e3b92f52320201c382594962c0e3b2b 100644 --- a/tensorflow/compiler/xla/service/reshape_mover_test.cc +++ b/tensorflow/compiler/xla/service/reshape_mover_test.cc @@ -458,57 +458,6 @@ TEST_F(ReshapeMoverTest, ScalarReshapeNotMovedAcrossSelect) { EXPECT_EQ(select, computation->root_instruction()); } -// Tree looks like: -// -// param0 [1,128,1] -// | -// reshape [128,1] constant [128,1024] -// \ / -// multiply w/implicit broadcast [128,1024] -// -// The reshape mover would like to sink the reshape below the multiply. -// -// Previously we would attempt to insert a reshape of the constant to [1,128,1] -// (which is unsound, because it has a different number of elements) as -// preparation for sinking the reshape. -// -// To eliminate the unsoundness, we outlaw reshape sinking when one of the -// operands is implicitly broadcast in the elementwise consumer. -// -// TODO(b/37799338) However, it would be possible in this case to do a more -// in-depth analysis to get reshape movement to occur: -// -// 1. Note that the broadcast dimension (logical dimension 1) in the operands -// would map back to logical dimension 2 in the param0 node. -// 2. Match rank of the constant to the param0 node (by prepending a trivial 1 -// dimension). -// 3. Reshape to [128,1024] at the root. -// -// But this is not currently done. -TEST_F(ReshapeMoverTest, ImplicitlyBroadcastReshapeIsNotMovedBug37787999) { - HloComputation::Builder builder(TestName()); - auto param0 = builder.AddInstruction(HloInstruction::CreateParameter( - 0, ShapeUtil::MakeShape(F32, {1, 128, 1}), "param0")); - auto reshape = builder.AddInstruction(HloInstruction::CreateReshape( - ShapeUtil::MakeShape(F32, {128, 1}), param0)); - Array2D a(128, 1024); - auto literal = Literal::CreateR2FromArray2D(a); - auto constant = builder.AddInstruction( - HloInstruction::CreateConstant(std::move(literal))); - auto multiply = builder.AddInstruction(HloInstruction::CreateBinary( - constant->shape(), HloOpcode::kMultiply, constant, reshape)); - - auto computation = module().AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), - op::Multiply(op::Constant(), op::Reshape(param0))); - - EXPECT_FALSE(ReshapeMover().Run(&module()).ValueOrDie()); - - EXPECT_THAT(computation->root_instruction(), - op::Multiply(op::Constant(), op::Reshape(param0))); - EXPECT_EQ(multiply, computation->root_instruction()); -} - // Tree looks like this: // // add1 diff --git a/tensorflow/compiler/xla/service/service.cc b/tensorflow/compiler/xla/service/service.cc index ec883a6cf3ce9546ac54f5c2524a8eda53bad33f..a73118c68a7f7d4169c4e592fc250fec8108c0f3 100644 --- a/tensorflow/compiler/xla/service/service.cc +++ b/tensorflow/compiler/xla/service/service.cc @@ -54,8 +54,6 @@ limitations under the License. #include "tensorflow/core/platform/stream_executor_no_cuda.h" #include "tensorflow/core/platform/types.h" -namespace se = ::perftools::gputools; - using ::tensorflow::strings::Printf; using ::tensorflow::strings::StrCat; using ::xla::source_map_util::InvalidParameterArgument; @@ -95,15 +93,12 @@ tensorflow::Status RecordResult(const ShapedBuffer& result, } // namespace -ServiceOptions& ServiceOptions::set_platform( - perftools::gputools::Platform* platform) { +ServiceOptions& ServiceOptions::set_platform(se::Platform* platform) { platform_ = platform; return *this; } -perftools::gputools::Platform* ServiceOptions::platform() const { - return platform_; -} +se::Platform* ServiceOptions::platform() const { return platform_; } ServiceOptions& ServiceOptions::set_number_of_replicas(int number_of_replicas) { number_of_replicas_ = number_of_replicas; @@ -123,7 +118,7 @@ int ServiceOptions::intra_op_parallelism_threads() const { } /* static */ StatusOr> Service::NewService( - perftools::gputools::Platform* platform) { + se::Platform* platform) { ServiceOptions default_options; default_options.set_platform(platform); return NewService(default_options); @@ -131,7 +126,7 @@ int ServiceOptions::intra_op_parallelism_threads() const { /* static */ StatusOr> Service::NewService( const ServiceOptions& options) { - perftools::gputools::Platform* platform = options.platform(); + se::Platform* platform = options.platform(); std::unique_ptr execute_backend; if (platform == nullptr) { TF_ASSIGN_OR_RETURN(platform, PlatformUtil::GetDefaultPlatform()); @@ -235,8 +230,7 @@ tensorflow::Status Service::ValidateResultShapeWithLayout( StatusOr>> Service::ResolveAndValidateArguments( tensorflow::gtl::ArraySlice arguments, - tensorflow::gtl::ArraySlice - stream_executors) { + 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()); @@ -314,7 +308,10 @@ StatusOr> Service::CreateModuleConfig( computation_layout->mutable_result_layout()->CopyLayoutFromShape( shape_with_output_layout)); } else { - computation_layout->mutable_result_layout()->Clear(); + // TODO(b/78356948): We are forcing the default layout here. We should fix + // clients which expect a default layout, to be explicit about it, by + // passing the proper ExecutionOptions with shape_with_output_layout set. + computation_layout->mutable_result_layout()->SetToDefaultLayout(); } config->set_replica_count(options_.number_of_replicas()); @@ -349,8 +346,7 @@ StatusOr> Service::CreateModuleConfig( StatusOr>> Service::BuildExecutables( std::vector versioned_handles, std::vector> module_configs, - Backend* backend, - std::vector> executors, + Backend* backend, std::vector> executors, DeviceMemoryAllocator* device_allocator) { VLOG(1) << Printf("BuildExecutable on service %p", this); @@ -412,8 +408,7 @@ StatusOr>> Service::BuildExecutables( StatusOr>> Service::BuildExecutables( const std::vector& module_protos, std::vector> module_configs, - Backend* backend, - std::vector> executors, + Backend* backend, std::vector> executors, DeviceMemoryAllocator* device_allocator) { VLOG(1) << Printf("BuildExecutable on service %p", this); @@ -493,7 +488,7 @@ StatusOr> Service::BuildExecutable( StatusOr> Service::BuildAndCacheExecutable( const VersionedComputationHandle& versioned_handle, std::unique_ptr module_config, Backend* backend, - perftools::gputools::StreamExecutor* executor, ExecutionProfile* profile, + se::StreamExecutor* executor, ExecutionProfile* profile, DeviceMemoryAllocator* device_allocator) { std::shared_ptr executable = compilation_cache_.LookUp(versioned_handle, *module_config); @@ -541,7 +536,7 @@ Service::ExecuteParallelAndRegisterResult( // Streams where the computation are launched, so we can wait on the streams // to complete. std::vector::SmartPtr> streams; - std::vector> timers; + std::vector> timers; // Global data handles for the computation results, one for each computation. std::vector result_handles; @@ -558,15 +553,14 @@ Service::ExecuteParallelAndRegisterResult( // 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; + std::vector result_buffers; for (int64 replica = 0; replica < replicas.size(); ++replica) { TF_ASSIGN_OR_RETURN(Pool::SmartPtr stream, backend->BorrowStream(replicas[replica])); streams.push_back(std::move(stream)); if (replica == 0 && profile != nullptr) { - timers.emplace_back( - new perftools::gputools::Timer(streams.back()->parent())); + timers.emplace_back(new se::Timer(streams.back()->parent())); streams.back() ->InitTimer(timers.back().get()) .ThenStartTimer(timers.back().get()); @@ -591,7 +585,7 @@ Service::ExecuteParallelAndRegisterResult( backend->StreamBorrower()); // Asynchronously launch the computation. - TF_ASSIGN_OR_RETURN(std::unique_ptr result, + TF_ASSIGN_OR_RETURN(ShapedBuffer result, executables[i]->ExecuteAsyncOnStream( &run_options, arguments[i][replica])); @@ -734,9 +728,9 @@ tensorflow::Status Service::SetReturnValue(const SetReturnValueRequest* arg, return computation->SetReturnValue(arg->operand()); } -StatusOr> -Service::GetExecutors(const ExecutionOptions& execution_options, - int64 requests_size, int64 request_index) const { +StatusOr> Service::GetExecutors( + const ExecutionOptions& execution_options, int64 requests_size, + int64 request_index) const { if (execution_options.device_handles().empty()) { return FailedPrecondition( "device handles must be given to execute parallel computations"); @@ -748,7 +742,7 @@ Service::GetExecutors(const ExecutionOptions& execution_options, "handles.", requests_size, request_index, execution_options.device_handles_size()); } - std::vector executors; + std::vector executors; for (const auto& device_handle : execution_options.device_handles()) { TF_ASSIGN_OR_RETURN(auto replicas, Replicas(*execute_backend_, device_handle)); @@ -780,7 +774,7 @@ tensorflow::Status Service::ExecuteParallel(const ExecuteParallelRequest* arg, VLOG(1) << "running execute-parallel request: " << arg->ShortDebugString(); std::vector>> all_arguments; - std::vector> all_executors; + std::vector> all_executors; std::vector versioned_handles; std::vector> module_configs; std::vector computation_names; @@ -891,7 +885,7 @@ tensorflow::Status Service::ExecuteGraphParallel( VLOG(1) << "running execute-graph-parallel request"; std::vector>> all_arguments; - std::vector> all_executors; + std::vector> all_executors; std::vector module_protos; std::vector> module_configs; std::vector computation_names; @@ -1243,7 +1237,7 @@ tensorflow::Status Service::ExecuteAsync(const ExecuteAsyncRequest* arg, streams.push_back(std::move(stream)); } - std::vector> result_buffers; + std::vector result_buffers; for (size_t i = 0; i < streams.size(); ++i) { const auto& stream = streams[i]; ExecutableRunOptions options; @@ -1256,7 +1250,7 @@ tensorflow::Status Service::ExecuteAsync(const ExecuteAsyncRequest* arg, ServiceExecutableRunOptions service_options( options, execute_backend_->StreamBorrower()); - TF_ASSIGN_OR_RETURN(std::unique_ptr this_result_buffer, + TF_ASSIGN_OR_RETURN(ShapedBuffer this_result_buffer, executable->ExecuteAsyncOnStream( &service_options, replicated_arguments[i])); @@ -1356,16 +1350,16 @@ tensorflow::Status Service::TransferToServer(const TransferToServerRequest* arg, } // Allocate memory in each replica and transfer the data to all replicas. - std::vector> replicated_buffers; + std::vector replicated_buffers; for (se::StreamExecutor* executor : replicas) { TF_ASSIGN_OR_RETURN( - std::unique_ptr shaped_buffer, + ShapedBuffer 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)); + executor, *literal, shaped_buffer)); replicated_buffers.emplace_back(std::move(shaped_buffer)); } TF_ASSIGN_OR_RETURN(*result->mutable_data(), @@ -1544,6 +1538,50 @@ tensorflow::Status Service::ComputeConstant(const ComputeConstantRequest* arg, // Since the shape_with_output_layout option in ExecutionOption is // non-effective to the Evaluator results, explicit relayout here. + // + // TODO(b/77824332): Make HloEvaluator take care of the re-layout. + if (arg->has_output_layout()) { + result_literal = result_literal->Relayout(arg->output_layout()); + } + *result->mutable_literal() = result_literal->ToProto(); + + return tensorflow::Status::OK(); +} + +tensorflow::Status Service::ComputeConstantGraph( + const ComputeConstantGraphRequest* arg, ComputeConstantResponse* result) { + if (!arg->has_computation()) { + return InvalidArgument("computations may not be empty"); + } + if (!arg->computation().has_program_shape()) { + return InvalidArgument("program shape may not be empty"); + } + if (arg->computation().program_shape().parameters_size() != 0) { + return InvalidArgument( + "constant computation may not depend on any parameters."); + } + + ProgramShape program_shape = arg->computation().program_shape(); + TF_DCHECK_OK(ShapeUtil::ValidateShape(program_shape.result())); + if (arg->has_output_layout()) { + TF_RETURN_IF_ERROR(LayoutUtil::ValidateLayoutForShape( + arg->output_layout(), program_shape.result())); + } + + HloModuleConfig config(program_shape); + + TF_ASSIGN_OR_RETURN(std::unique_ptr module, + HloModule::CreateFromProto(arg->computation(), config)); + + HloEvaluator evaluator; + TF_ASSIGN_OR_RETURN(auto result_literal, + evaluator.Evaluate>( + *module, /*arg_literals=*/{})); + + // Since the result layout is non-effective to the Evaluator results, explicit + // relayout here. + // + // TODO(b/77824332): Make HloEvaluator take care of the re-layout. if (arg->has_output_layout()) { result_literal = result_literal->Relayout(arg->output_layout()); } @@ -1617,7 +1655,14 @@ tensorflow::Status Service::GetComputationStats( tensorflow::Status Service::GetComputationGraphStats( const ComputationGraphStatsRequest* arg, ComputationStatsResponse* result) { - HloModuleConfig config; + if (!arg->has_computation()) { + return InvalidArgument("Computations may not be empty."); + } + if (!arg->computation().has_program_shape()) { + return InvalidArgument("Program shape may not be empty."); + } + + HloModuleConfig config(arg->computation().program_shape()); config.set_debug_options(arg->debug_options()); TF_ASSIGN_OR_RETURN(std::unique_ptr module, HloModule::CreateFromProto(arg->computation(), config)); @@ -1902,9 +1947,9 @@ DeviceHandle Service::SingleComputationDeviceHandle() const { return device_handle; } -StatusOr> Service::Replicas( +StatusOr> Service::Replicas( const Backend& backend, const DeviceHandle& device_handle) const { - std::vector replicas; + std::vector replicas; for (int replica = 0; replica < options_.number_of_replicas(); ++replica) { // From the computation placer, find out the device ids of the replicas for // the given device handle. diff --git a/tensorflow/compiler/xla/service/service.h b/tensorflow/compiler/xla/service/service.h index 9fa72c1b8c859b2b4a1b79abb32a63560dae8ec4..476bd0597de735a9f777be78f5ab01dac1188525 100644 --- a/tensorflow/compiler/xla/service/service.h +++ b/tensorflow/compiler/xla/service/service.h @@ -53,8 +53,8 @@ namespace xla { class ServiceOptions { public: // Set the platform backing the service, or nullptr for the default platform. - ServiceOptions& set_platform(perftools::gputools::Platform* platform); - perftools::gputools::Platform* platform() const; + ServiceOptions& set_platform(se::Platform* platform); + se::Platform* platform() const; // Set the number of replicas to use when compiling replicated // programs. @@ -66,7 +66,7 @@ class ServiceOptions { int intra_op_parallelism_threads() const; private: - perftools::gputools::Platform* platform_ = nullptr; + se::Platform* platform_ = nullptr; int number_of_replicas_ = 1; int intra_op_parallelism_threads_ = -1; }; @@ -79,7 +79,7 @@ class Service : public ServiceInterface { public: // Factory method for creating a new Service. static StatusOr> NewService( - perftools::gputools::Platform* platform = nullptr); + se::Platform* platform = nullptr); static StatusOr> NewService( const ServiceOptions& options); @@ -206,6 +206,9 @@ class Service : public ServiceInterface { // Computes the value of a constant expression. tensorflow::Status ComputeConstant(const ComputeConstantRequest* arg, ComputeConstantResponse* result) override; + tensorflow::Status ComputeConstantGraph( + const ComputeConstantGraphRequest* arg, + ComputeConstantResponse* result) override; // Returns the shape (with layout) of an array associated with a given data // handle. @@ -283,7 +286,7 @@ class Service : public ServiceInterface { ExecuteResponse* result); // Prepare the executors for executing parallel. - StatusOr> GetExecutors( + StatusOr> GetExecutors( const ExecutionOptions& execution_options, int64 requests_size, int64 request_index) const; @@ -307,8 +310,7 @@ class Service : public ServiceInterface { StatusOr>> ResolveAndValidateArguments( tensorflow::gtl::ArraySlice arguments, - tensorflow::gtl::ArraySlice - stream_executors); + 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. @@ -326,7 +328,7 @@ class Service : public ServiceInterface { StatusOr> BuildExecutable( const VersionedComputationHandle& versioned_handle, std::unique_ptr module_config, Backend* backend, - perftools::gputools::StreamExecutor* executor, + se::StreamExecutor* executor, DeviceMemoryAllocator* device_allocator = nullptr); // Builds an Executable for the given HLO module proto. @@ -335,7 +337,7 @@ class Service : public ServiceInterface { StatusOr> BuildExecutable( const HloModuleProto& module_proto, std::unique_ptr module_config, Backend* backend, - perftools::gputools::StreamExecutor* executor, + se::StreamExecutor* executor, DeviceMemoryAllocator* device_allocator = nullptr); // Same as BuildExecutable() above, but builds a list of Executables for the @@ -343,14 +345,12 @@ class Service : public ServiceInterface { StatusOr>> BuildExecutables( std::vector versioned_handles, std::vector> module_configs, - Backend* backend, - std::vector> executors, + Backend* backend, std::vector> executors, DeviceMemoryAllocator* device_allocator); StatusOr>> BuildExecutables( const std::vector& module_protos, std::vector> module_configs, - Backend* backend, - std::vector> executors, + Backend* backend, std::vector> executors, DeviceMemoryAllocator* device_allocator); // Similar to BuildExecutable, but look in the compilation cache for the @@ -359,7 +359,7 @@ class Service : public ServiceInterface { StatusOr> BuildAndCacheExecutable( const VersionedComputationHandle& versioned_handle, std::unique_ptr module_config, Backend* backend, - perftools::gputools::StreamExecutor* executor, ExecutionProfile* profile, + se::StreamExecutor* executor, ExecutionProfile* profile, DeviceMemoryAllocator* device_allocator = nullptr); // Runs the given executable with the given arguments and register the result @@ -408,7 +408,7 @@ class Service : public ServiceInterface { // Returns the stream executors assigned to the replicas represented by the // given device handle. Each device_handle is a virtual replicated device that // represents a set of physical devices for the replicas. - StatusOr> Replicas( + StatusOr> Replicas( const Backend& backend, const DeviceHandle& device_handle) const; Status MaybeDumpHloModule(const HloModule& module) const; diff --git a/tensorflow/compiler/xla/service/service_executable_run_options.h b/tensorflow/compiler/xla/service/service_executable_run_options.h index 6c1f8feac7ed4423051cf2737be57dcfab508671..7f3910cdb0366078b97fb5f6a2dc498b37570926 100644 --- a/tensorflow/compiler/xla/service/service_executable_run_options.h +++ b/tensorflow/compiler/xla/service/service_executable_run_options.h @@ -28,7 +28,7 @@ namespace xla { class ServiceExecutableRunOptions { public: using StreamBorrower = - std::function::SmartPtr>(int)>; + std::function::SmartPtr>(int)>; ServiceExecutableRunOptions() : ServiceExecutableRunOptions(ExecutableRunOptions()) {} @@ -45,14 +45,13 @@ class ServiceExecutableRunOptions { ExecutableRunOptions* mutable_run_options() { return &run_options_; } // Delegate to `ExecutableRunOptions` member. - perftools::gputools::Stream* stream() const { return run_options_.stream(); } + se::Stream* stream() const { return run_options_.stream(); } DeviceMemoryAllocator* allocator() const { return run_options_.allocator(); } int device_ordinal() const { return run_options_.device_ordinal(); } // Borrows a stream and returns a smart pointer which returns the stream on // destruction. - StatusOr::SmartPtr> BorrowStream( - int device_ordinal) const { + StatusOr::SmartPtr> BorrowStream(int device_ordinal) const { return borrow_stream_ ? borrow_stream_(device_ordinal) : Status(tensorflow::error::UNIMPLEMENTED, "No stream cache"); diff --git a/tensorflow/compiler/xla/service/shape_inference.cc b/tensorflow/compiler/xla/service/shape_inference.cc index 77e12d36024dae56003ad4e59b54f9934dfc2c58..48b2922e77b78719e5d3469cbaa4fc15969de91b 100644 --- a/tensorflow/compiler/xla/service/shape_inference.cc +++ b/tensorflow/compiler/xla/service/shape_inference.cc @@ -52,6 +52,8 @@ UnaryOperation OpcodeToUnaryOperation(HloOpcode opcode) { return UNOP_ABS; case HloOpcode::kCeil: return UNOP_CEIL; + case HloOpcode::kClz: + return UNOP_CLZ; case HloOpcode::kCos: return UNOP_COS; case HloOpcode::kExp: @@ -360,6 +362,7 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, arg, primitive_util::ComplexComponentType(arg.element_type())); } return arg; + case UNOP_CLZ: case UNOP_NEGATE: case UNOP_ROUND_NEAREST_AFZ: case UNOP_SIGN: diff --git a/tensorflow/compiler/xla/service/shaped_buffer.cc b/tensorflow/compiler/xla/service/shaped_buffer.cc index 6e9986165f7eaf71a964b42b734a5ae5db5e45d7..0b5a383f6fe5b01615c8d24ccf00fb667e0fe0c2 100644 --- a/tensorflow/compiler/xla/service/shaped_buffer.cc +++ b/tensorflow/compiler/xla/service/shaped_buffer.cc @@ -28,8 +28,6 @@ limitations under the License. #include "tensorflow/core/lib/strings/stringprintf.h" #include "tensorflow/core/platform/logging.h" -namespace se = ::perftools::gputools; - namespace xla { using ::tensorflow::strings::Appendf; @@ -68,6 +66,8 @@ ShapedBuffer& ShapedBuffer::operator=(ShapedBuffer&& s) { return *this; } +ShapedBuffer::~ShapedBuffer() {} + void ShapedBuffer::clear() { for (auto& pair : buffers_) { // A default constructed DeviceMemoryBase is a null pointer. @@ -104,18 +104,6 @@ std::ostream& operator<<(std::ostream& out, const ShapedBuffer& buffer) { return out; } -/* static */ -StatusOr> ScopedShapedBuffer::MakeScoped( - ShapedBuffer* shaped_buffer, DeviceMemoryAllocator* allocator) { - auto scoped_buffer = WrapUnique(new ScopedShapedBuffer( - shaped_buffer->on_host_shape(), shaped_buffer->on_device_shape(), - allocator, shaped_buffer->device_ordinal())); - scoped_buffer->buffers_ = shaped_buffer->buffers(); - shaped_buffer->clear(); - - return std::move(scoped_buffer); -} - ScopedShapedBuffer::ScopedShapedBuffer(const Shape& on_host_shape, const Shape& on_device_shape, DeviceMemoryAllocator* allocator, @@ -128,7 +116,25 @@ ScopedShapedBuffer::ScopedShapedBuffer(ShapedBuffer shaped_buffer, DeviceMemoryAllocator* allocator) : ShapedBuffer(std::move(shaped_buffer)), allocator_(allocator) {} +ScopedShapedBuffer::ScopedShapedBuffer(ScopedShapedBuffer&& s) + : ShapedBuffer(std::move(s)), allocator_(s.allocator_) { + // Null out s.allocator_ so it doesn't try to free anything in its destructor. + s.allocator_ = nullptr; +} + +ScopedShapedBuffer& ScopedShapedBuffer::operator=(ScopedShapedBuffer&& s) { + *static_cast(this) = std::move(static_cast(s)); + allocator_ = s.allocator_; + // Null out s.allocator_ so it doesn't try to free anything in its destructor. + s.allocator_ = nullptr; + return *this; +} + ScopedShapedBuffer::~ScopedShapedBuffer() { + // allocator_ will be null if we were moved-from. + if (allocator_ == nullptr) { + return; + } // 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 // has been deallocated. @@ -144,9 +150,9 @@ ScopedShapedBuffer::~ScopedShapedBuffer() { } } -std::unique_ptr ScopedShapedBuffer::release() { - auto shaped_buffer = MakeUnique(std::move(*this)); - buffers_ = ShapeTree(); +ShapedBuffer ScopedShapedBuffer::release() { + ShapedBuffer shaped_buffer(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 b816df8385ef65b0b69ede1d6e65a1991b4bd7c6..f1b0527474cfe28fa7108fda829556d24c12b4d9 100644 --- a/tensorflow/compiler/xla/service/shaped_buffer.h +++ b/tensorflow/compiler/xla/service/shaped_buffer.h @@ -41,8 +41,15 @@ class ShapedBuffer { // determines the number of device allocations (DeviceMemoryBase) held by the // ShapedBuffer. ShapedBuffer(const Shape& on_host_shape, const Shape& on_device_shape, - const perftools::gputools::Platform* platform, - int device_ordinal); + const se::Platform* platform, int device_ordinal); + + // Movable, but not copyable. + ShapedBuffer(ShapedBuffer&& s); + ShapedBuffer& operator=(ShapedBuffer&&); + ShapedBuffer(const ShapedBuffer&) = delete; + ShapedBuffer& operator=(const ShapedBuffer&) = delete; + + virtual ~ShapedBuffer(); // Returns the shape of the on-host representation of the data held by this // ShapedBuffer. @@ -52,48 +59,36 @@ class ShapedBuffer { // ShapedBuffer. const Shape& on_device_shape() const { return on_device_shape_; } - const perftools::gputools::Platform* platform() const { return platform_; } + const se::Platform* platform() const { return platform_; } int device_ordinal() const { return device_ordinal_; } // Return the root buffer of the shape (shape index {}). - const perftools::gputools::DeviceMemoryBase& root_buffer() const { + const se::DeviceMemoryBase& root_buffer() const { return buffer(/*index=*/{}); } // Returns the buffer at the given shape index where index is defined as in // ShapeUtil::GetSubshape. - const perftools::gputools::DeviceMemoryBase& buffer( - const ShapeIndex& index) const { + const se::DeviceMemoryBase& buffer(const ShapeIndex& index) const { return buffers_.element(index); } // Sets the device memory buffer at the given index. - void set_buffer(const perftools::gputools::DeviceMemoryBase& buffer, - const ShapeIndex& index) { + void set_buffer(const se::DeviceMemoryBase& buffer, const ShapeIndex& index) { *buffers_.mutable_element(index) = buffer; } // Returns the underlying ShapeTree containing all the device addresses in the // ShapedBuffer. - const ShapeTree& buffers() const { - return buffers_; - } - ShapeTree& buffers() { - return buffers_; - } + const ShapeTree& buffers() const { return buffers_; } + ShapeTree& buffers() { return buffers_; } // Set all device memory pointers in the object to null. void clear(); 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. Shape on_host_shape_; @@ -101,13 +96,13 @@ class ShapedBuffer { Shape on_device_shape_; // The platform the memory is allocated on. - const perftools::gputools::Platform* platform_; + const se::Platform* platform_; // The device the memory is allocated on. int device_ordinal_; // The tree of device buffers. Its shape is on_device_shape(). - ShapeTree buffers_; + ShapeTree buffers_; }; std::ostream& operator<<(std::ostream& out, const ShapedBuffer& buffer); @@ -115,41 +110,45 @@ std::ostream& operator<<(std::ostream& out, const ShapedBuffer& buffer); // ShapedBuffer derived class which allocates all internal buffers on // construction and deallocates the memory when the object is // destructed. +// +// TODO(timshen): Remove inheritance between ScopedShapedBuffer and +// ShapedBuffer. There should never be a need to consider a ScopedShapedBuffer +// as a ShapedBuffer, because in that case we should just be able to pass around +// our ShapeTree. Inheritance only adds complexity. See +// discussion in cl/192849370. class ScopedShapedBuffer : public ShapedBuffer { public: - // Takes a ShapedBuffer and returns a ScopedShapedBuffer which manages the - // deallocation of the device memory held in the shaped buffer. All device - // memory pointers in the given ShapedBuffer are set to null. - static StatusOr> MakeScoped( - ShapedBuffer* shaped_buffer, DeviceMemoryAllocator* allocator); - - // Create a ScopedShapedBuffer with null DeviceMemoryBases at each index. - ScopedShapedBuffer(const Shape& on_host_shape, const Shape& on_device_shape, - DeviceMemoryAllocator* allocator, int device_ordinal); + // Creates a ScopedShapedBuffer with null DeviceMemoryBases at each index. + explicit 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); + explicit ScopedShapedBuffer(ShapedBuffer shaped_buffer, + DeviceMemoryAllocator* allocator); + + // Movable, but not copyable. + ScopedShapedBuffer(ScopedShapedBuffer&& s); + ScopedShapedBuffer& operator=(ScopedShapedBuffer&&); + ScopedShapedBuffer(const ScopedShapedBuffer&) = delete; + ScopedShapedBuffer& operator=(const ScopedShapedBuffer&) = delete; + + // All buffers in the shape are deallocated on destruction. + ~ScopedShapedBuffer() override; // 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. 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. - virtual ~ScopedShapedBuffer(); + // Releases all device memory owned by this ScopedShapedBuffer and returns 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. + ShapedBuffer release(); protected: - ScopedShapedBuffer(const ScopedShapedBuffer&) = delete; - void operator=(const ScopedShapedBuffer&) = delete; - DeviceMemoryAllocator* allocator_; }; diff --git a/tensorflow/compiler/xla/service/transfer_manager.cc b/tensorflow/compiler/xla/service/transfer_manager.cc index 2f36e2b16e0f2eed10aef811dd3cceeba6a5b8a9..98d0111d04d981dd3e26733e3311fd846b55d8ae 100644 --- a/tensorflow/compiler/xla/service/transfer_manager.cc +++ b/tensorflow/compiler/xla/service/transfer_manager.cc @@ -25,24 +25,20 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/macros.h" -namespace se = ::perftools::gputools; - namespace xla { /* static */ tensorflow::mutex TransferManager::platform_transfer_manager_mutex_( tensorflow::LINKER_INITIALIZED); -/* static */ std::map* +/* static */ std::map* TransferManager::GetPlatformTransferManagers() { - static auto* r = - new std::map; + static auto* r = new std::map; return r; } Status TransferManager::TransferArrayToDevice( - perftools::gputools::StreamExecutor* executor, const Literal& literal, - const perftools::gputools::DeviceMemoryBase& dest) { + se::StreamExecutor* executor, const Literal& literal, + const se::DeviceMemoryBase& dest) { const Shape on_device_shape = HostShapeToDeviceShape(literal.shape()); TF_RET_CHECK(ShapeUtil::IsArray(on_device_shape)) << "On-device representation of " @@ -61,8 +57,8 @@ Status TransferManager::TransferArrayToDevice( } StatusOr> TransferManager::TransferArrayFromDevice( - perftools::gputools::StreamExecutor* executor, const Shape& shape, - const perftools::gputools::DeviceMemoryBase& source) { + se::StreamExecutor* executor, const Shape& shape, + const se::DeviceMemoryBase& source) { TF_RET_CHECK(ShapeUtil::Equal(HostShapeToDeviceShape(shape), shape)) << "Shape " << ShapeUtil::HumanString(shape) << " has a differently shaped representation on-device: " @@ -112,8 +108,7 @@ StatusOr> TransferManager::TransferArrayFromDevice( } Status TransferManager::WriteTupleIndexTables( - perftools::gputools::StreamExecutor* executor, - const ShapedBuffer& device_buffer) { + se::StreamExecutor* executor, const ShapedBuffer& device_buffer) { VLOG(2) << "Writing tuple index tables for " << device_buffer; TF_RET_CHECK(executor->device_ordinal() == device_buffer.device_ordinal()); @@ -180,7 +175,7 @@ Status TransferManager::TransferBufferToDevice( return Status::OK(); } -StatusOr> TransferManager::AllocateShapedBuffer( +StatusOr TransferManager::AllocateShapedBuffer( const Shape& on_host_shape, DeviceMemoryAllocator* allocator, int device_ordinal) { if (!LayoutUtil::HasLayout(on_host_shape)) { @@ -192,31 +187,30 @@ StatusOr> TransferManager::AllocateShapedBuffer( const Shape on_device_shape = HostShapeToDeviceShape(on_host_shape); TF_RET_CHECK(LayoutUtil::HasLayout(on_device_shape)); - auto shaped_buffer = WrapUnique(new ShapedBuffer( - on_host_shape, on_device_shape, allocator->platform(), device_ordinal)); + ShapedBuffer shaped_buffer(on_host_shape, on_device_shape, + allocator->platform(), device_ordinal); // Allocate an appropriate sized buffer for each element in the shape // including the tuple pointer arrays. - for (auto& pair : shaped_buffer->buffers()) { + for (auto& pair : shaped_buffer.buffers()) { const ShapeIndex& index = pair.first; se::DeviceMemoryBase& memory_base = pair.second; const Shape& subshape = ShapeUtil::GetSubshape(on_device_shape, index); TF_ASSIGN_OR_RETURN(memory_base, - allocator->Allocate(shaped_buffer->device_ordinal(), + allocator->Allocate(shaped_buffer.device_ordinal(), GetByteSizeRequirement(subshape))); } return std::move(shaped_buffer); } -StatusOr> -TransferManager::AllocateScopedShapedBuffer(const Shape& on_host_shape, - DeviceMemoryAllocator* allocator, - int device_ordinal) { +StatusOr TransferManager::AllocateScopedShapedBuffer( + const Shape& on_host_shape, DeviceMemoryAllocator* allocator, + int device_ordinal) { TF_ASSIGN_OR_RETURN( - std::unique_ptr unscoped_buffer, + ShapedBuffer unscoped_buffer, AllocateShapedBuffer(on_host_shape, allocator, device_ordinal)); - return ScopedShapedBuffer::MakeScoped(unscoped_buffer.get(), allocator); + return ScopedShapedBuffer(std::move(unscoped_buffer), allocator); } } // namespace xla diff --git a/tensorflow/compiler/xla/service/transfer_manager.h b/tensorflow/compiler/xla/service/transfer_manager.h index 9f2b5c4aecf0b52f610171e0c2755de577b2bd9e..a6451c4bb1104cc865c010463bbfd9107f6c4f9b 100644 --- a/tensorflow/compiler/xla/service/transfer_manager.h +++ b/tensorflow/compiler/xla/service/transfer_manager.h @@ -42,7 +42,7 @@ class TransferManager { virtual ~TransferManager() {} // Returns the ID of the platform that this transfer manager acts on. - virtual perftools::gputools::Platform::Id PlatformId() const = 0; + virtual se::Platform::Id PlatformId() const = 0; // Returns the shape of the on-device representation for the given shape on // the host. This is intended for use with ShapedBuffer where buffers are @@ -58,48 +58,45 @@ class TransferManager { // DeviceShape(literal_shape) must be compatible, but need not have the same // layout. virtual StatusOr> TransferLiteralFromDevice( - perftools::gputools::StreamExecutor* executor, - const ShapedBuffer& device_buffer) = 0; + se::StreamExecutor* executor, const ShapedBuffer& device_buffer) = 0; // Transfers the given literal into the previously allocated device memory // represented by the given ShapedBuffer using the given executor. The shape // of the ShapedBuffer and DeviceShape(literal.shape()) must be compatible, // but need not have the same layout - virtual Status TransferLiteralToDevice( - perftools::gputools::StreamExecutor* executor, const Literal& literal, - const ShapedBuffer& device_buffer) = 0; + virtual Status TransferLiteralToDevice(se::StreamExecutor* executor, + const Literal& literal, + const ShapedBuffer& device_buffer) = 0; // Convenience methods for transferring an array to or from the device at a // known address. This avoids having to construct a ShapedBuffer just to // transfer an array at a known address. - Status TransferArrayToDevice( - perftools::gputools::StreamExecutor* executor, const Literal& literal, - const perftools::gputools::DeviceMemoryBase& dest); + Status TransferArrayToDevice(se::StreamExecutor* executor, + const Literal& literal, + const se::DeviceMemoryBase& dest); StatusOr> TransferArrayFromDevice( - perftools::gputools::StreamExecutor* executor, const Shape& shape, - const perftools::gputools::DeviceMemoryBase& source); + se::StreamExecutor* executor, const Shape& shape, + const se::DeviceMemoryBase& source); // Transfers the given literal into the Infeed interface of the device, // using the given executor. - virtual Status TransferLiteralToInfeed( - perftools::gputools::StreamExecutor* executor, - const Literal& literal) = 0; + virtual Status TransferLiteralToInfeed(se::StreamExecutor* executor, + const Literal& literal) = 0; // Transfers the given literal from the Outfeed interface of the device, // using the given executor. - virtual Status TransferLiteralFromOutfeed( - perftools::gputools::StreamExecutor* executor, const Shape& literal_shape, - Literal* literal) = 0; + virtual Status TransferLiteralFromOutfeed(se::StreamExecutor* executor, + const Shape& literal_shape, + Literal* literal) = 0; // Resets the devices associated with this transfer manager. virtual Status ResetDevices( - tensorflow::gtl::ArraySlice - executor) = 0; + tensorflow::gtl::ArraySlice executor) = 0; // Given an allocated ShapedBuffer, constructs the tuple index table(s) in // each buffer of the given ShapedBuffer corresponding to tuple shapes. If the // ShapedBuffer is array-shaped this method does nothing. - Status WriteTupleIndexTables(perftools::gputools::StreamExecutor* executor, + Status WriteTupleIndexTables(se::StreamExecutor* executor, const ShapedBuffer& device_buffer); // Determines the byte size requirement for the given shape on the underlying @@ -110,10 +107,10 @@ class TransferManager { // Allocate a ShapedBuffer which can hold data with the given on-host // shape. The on-device shape may be different as indicated by // HostShapeToDeviceShape. - StatusOr> AllocateShapedBuffer( - const Shape& on_host_shape, DeviceMemoryAllocator* allocator, - int device_ordinal); - StatusOr> AllocateScopedShapedBuffer( + StatusOr AllocateShapedBuffer(const Shape& on_host_shape, + DeviceMemoryAllocator* allocator, + int device_ordinal); + StatusOr AllocateScopedShapedBuffer( const Shape& on_host_shape, DeviceMemoryAllocator* allocator, int device_ordinal); @@ -127,13 +124,13 @@ class TransferManager { // Precondition: a platform kind must not be registered more than once. typedef std::unique_ptr (*TransferManagerCreationFunction)(); static void RegisterTransferManager( - perftools::gputools::Platform::Id platform_id, + se::Platform::Id platform_id, TransferManagerCreationFunction transfer_manager); // Returns the transfer manager singleton pointer if it is available for the // given platform, or an error status if it is not. static StatusOr GetForPlatform( - const perftools::gputools::Platform* platform); + const se::Platform* platform); protected: // Transfer a memory block of the given size from 'source' buffer to the @@ -143,35 +140,32 @@ class TransferManager { // // source is the source data that must be in the target-dependent layout that // the Infeed HLO used in the computation expects. - virtual Status TransferBufferToInfeed( - perftools::gputools::StreamExecutor* executor, int64 size, - const void* source) = 0; + virtual Status TransferBufferToInfeed(se::StreamExecutor* executor, + int64 size, const void* source) = 0; // Transfer a memory block of the given size from the device source into the // 'destination' buffer. // // size is the size to transfer to destination in bytes. - virtual Status TransferBufferFromDevice( - perftools::gputools::StreamExecutor* executor, - const perftools::gputools::DeviceMemoryBase& source, int64 size, - void* destination); + virtual Status TransferBufferFromDevice(se::StreamExecutor* executor, + const se::DeviceMemoryBase& source, + int64 size, void* destination); // Transfer a memory block of the given size from 'source' buffer to the given // destination of the device. // // size is the size to transfer from source in bytes. - virtual Status TransferBufferToDevice( - perftools::gputools::StreamExecutor* executor, int64 size, - const void* source, perftools::gputools::DeviceMemoryBase* destination); + virtual Status TransferBufferToDevice(se::StreamExecutor* executor, + int64 size, const void* source, + se::DeviceMemoryBase* destination); // Writes the given device-memory pointers in 'elements' to the given region // to construct a tuple index table in the platform-specific tuple // representation. virtual Status WriteSingleTupleIndexTable( - perftools::gputools::StreamExecutor* executor, - tensorflow::gtl::ArraySlice - elements, - const Shape& shape, perftools::gputools::DeviceMemoryBase* region) = 0; + se::StreamExecutor* executor, + tensorflow::gtl::ArraySlice elements, + const Shape& shape, se::DeviceMemoryBase* region) = 0; private: // The mutex that guards the platform-to-transfer manager map. @@ -186,8 +180,7 @@ class TransferManager { }; // Map from platform kind to transfer manager singleton. - static std::map* - GetPlatformTransferManagers(); + static std::map* GetPlatformTransferManagers(); }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/transpose_folding.cc b/tensorflow/compiler/xla/service/transpose_folding.cc index 83185ac49e9b7c386d10d1cbc4e20dcdfdfd6cae..3efd38ce0daa3e3f3398b32463019df6cd10a009 100644 --- a/tensorflow/compiler/xla/service/transpose_folding.cc +++ b/tensorflow/compiler/xla/service/transpose_folding.cc @@ -159,6 +159,7 @@ bool FoldTransposeIntoConvolution(InstructionOperandsPair pair) { auto new_conv = HloInstruction::CreateConvolve( convolution.shape(), new_lhs, new_rhs, convolution.window(), new_dnums); + convolution.SetupDerivedInstruction(new_conv.get()); TF_CHECK_OK(convolution.parent()->ReplaceWithNewInstruction( &convolution, std::move(new_conv))); diff --git a/tensorflow/compiler/xla/service/tuple_simplifier.cc b/tensorflow/compiler/xla/service/tuple_simplifier.cc index 113c2e2bd9f73a2b0c783103d7f2da9534bc97c3..d668855084a884518b338cdf396a9330b9f43a2b 100644 --- a/tensorflow/compiler/xla/service/tuple_simplifier.cc +++ b/tensorflow/compiler/xla/service/tuple_simplifier.cc @@ -69,6 +69,7 @@ StatusOr TupleSimplifier::Run(HloModule* module) { // Tuple // HloInstruction* top_tuple = nullptr; + HloInstruction* first_gte = nullptr; bool can_simplify = true; for (int64 operand_number = 0; operand_number < instruction->operand_count(); ++operand_number) { @@ -78,11 +79,17 @@ StatusOr TupleSimplifier::Run(HloModule* module) { can_simplify = false; break; } - + if (first_gte == nullptr) { + first_gte = operand; + } else if (!first_gte->has_compatible_sharding(operand)) { + can_simplify = false; + break; + } if (top_tuple == nullptr) { top_tuple = operand->mutable_operand(0); if (!ShapeUtil::Compatible(top_tuple->shape(), - instruction->shape())) { + instruction->shape()) || + !instruction->has_compatible_sharding(top_tuple)) { can_simplify = false; break; } @@ -108,15 +115,17 @@ StatusOr TupleSimplifier::Run(HloModule* module) { // | // GTE if (instruction->operand(0)->opcode() == HloOpcode::kTuple) { - changed = true; HloInstruction* element_source = instruction->mutable_operand(0)->mutable_operand( instruction->tuple_index()); - TF_RETURN_IF_ERROR(instruction->ReplaceAllUsesWith(element_source)); - for (HloInstruction* user : element_source->users()) { - if (user->opcode() == HloOpcode::kTuple || - user->opcode() == HloOpcode::kGetTupleElement) { - worklist.push(user); + if (instruction->has_compatible_sharding(element_source)) { + changed = true; + TF_RETURN_IF_ERROR(instruction->ReplaceAllUsesWith(element_source)); + for (HloInstruction* user : element_source->users()) { + if (user->opcode() == HloOpcode::kTuple || + user->opcode() == HloOpcode::kGetTupleElement) { + worklist.push(user); + } } } } diff --git a/tensorflow/compiler/xla/service/user_computation.cc b/tensorflow/compiler/xla/service/user_computation.cc index 532f7fd5bfc1dffa86638a6bc51832beebd74e1d..0f16a592b68e20f5dbd1e4655ad5720ecce5a7bd 100644 --- a/tensorflow/compiler/xla/service/user_computation.cc +++ b/tensorflow/compiler/xla/service/user_computation.cc @@ -49,6 +49,8 @@ HloOpcode UnaryOperationToHloOpcode(UnaryOperation unop) { return HloOpcode::kAbs; case UNOP_CEIL: return HloOpcode::kCeil; + case UNOP_CLZ: + return HloOpcode::kClz; case UNOP_COS: return HloOpcode::kCos; case UNOP_EXP: diff --git a/tensorflow/compiler/xla/service_interface.h b/tensorflow/compiler/xla/service_interface.h index 32aae64973dbd7ac2f8d403d8fbd155d432642f9..5b44c26b7c7b082556d9533cf3b3b1b98e5e4b09 100644 --- a/tensorflow/compiler/xla/service_interface.h +++ b/tensorflow/compiler/xla/service_interface.h @@ -112,6 +112,10 @@ class ServiceInterface { virtual tensorflow::Status ComputeConstant( const ComputeConstantRequest* arg, ComputeConstantResponse* result) = 0; + virtual tensorflow::Status ComputeConstantGraph( + const ComputeConstantGraphRequest* arg, + ComputeConstantResponse* result) = 0; + // Methods used by Computation. virtual tensorflow::Status SnapshotComputation( const SnapshotComputationRequest* ag, diff --git a/tensorflow/compiler/xla/shape_util.cc b/tensorflow/compiler/xla/shape_util.cc index 6825d2476587d037aace043230168f78f4e46344..ac7e201bfdceabdd0f11db61bbb3b460017401ca 100644 --- a/tensorflow/compiler/xla/shape_util.cc +++ b/tensorflow/compiler/xla/shape_util.cc @@ -824,6 +824,18 @@ StatusOr ParseShapeStringInternal(tensorflow::StringPiece* s) { return new_shape; } +/* static */ bool ShapeUtil::IndexIsValid(const Shape& shape, + ShapeIndexView index) { + const Shape* subshape = &shape; + for (auto i : index) { + if (!IsTuple(*subshape) || i >= subshape->tuple_shapes_size()) { + return false; + } + subshape = &subshape->tuple_shapes(i); + } + return true; +} + /* static */ const Shape& ShapeUtil::GetSubshape(const Shape& shape, ShapeIndexView index) { const Shape* return_shape = &shape; diff --git a/tensorflow/compiler/xla/shape_util.h b/tensorflow/compiler/xla/shape_util.h index 1375f981a880392e72db9946c5f006d164447baf..63da9154cfc1a5e7e8c0eeaa103d27096540fefe 100644 --- a/tensorflow/compiler/xla/shape_util.h +++ b/tensorflow/compiler/xla/shape_util.h @@ -319,6 +319,11 @@ class ShapeUtil { // Returns an empty tuple shape. Can be used to indicate side-effects. static Shape MakeNil() { return MakeTupleShape({}); } + // Checks whether the shape is initialized. + static bool IsInitialized(const Shape& shape) { + return shape.element_type() != PRIMITIVE_TYPE_INVALID; + } + // Constructs a new shape with the given element type and sequence of // dimensions. static Shape MakeShape(PrimitiveType element_type, @@ -443,6 +448,9 @@ class ShapeUtil { static bool ShapeIs(const Shape& shape, PrimitiveType element_type, std::initializer_list dimensions); + // Returns true if the given shape has a subshape at the given index. + static bool IndexIsValid(const Shape& shape, ShapeIndexView index); + // GetSubshape and GetMutableSubshape return a particular nested Shape within // the given Shape argument. static const Shape& GetSubshape(const Shape& shape, ShapeIndexView index); diff --git a/tensorflow/compiler/xla/tests/BUILD b/tensorflow/compiler/xla/tests/BUILD index 218345772f3e5f26006cb8dfc648c0e64c4318fd..1f90a44d8ba725c1bc7d23b581161f8915ff74fd 100644 --- a/tensorflow/compiler/xla/tests/BUILD +++ b/tensorflow/compiler/xla/tests/BUILD @@ -191,6 +191,8 @@ cc_library( "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client/xla_client:xla_builder", + "//tensorflow/compiler/xla/service:interpreter_plugin", # reference backend + "//tensorflow/compiler/xla/service:platform_util", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:test_utils", "//tensorflow/core:lib", @@ -413,6 +415,8 @@ xla_test( "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client/lib:arithmetic", + "//tensorflow/compiler/xla/client/xla_client:xla_builder", + "//tensorflow/compiler/xla/client/xla_client:xla_computation", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:test_utils", @@ -488,9 +492,10 @@ xla_test( tags = ["enable_for_xla_interpreter"], deps = [ "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/compiler/xla/client/xla_client:xla_builder", + "//tensorflow/compiler/xla/client/xla_client:xla_computation", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:xla_internal_test_main", @@ -639,9 +644,9 @@ xla_test( "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/compiler/xla/client/xla_client:xla_builder", "//tensorflow/compiler/xla/service:reduce_precision_insertion", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", @@ -777,10 +782,10 @@ xla_test( "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:padding", + "//tensorflow/compiler/xla/client/xla_client:xla_builder", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:xla_internal_test_main", @@ -856,11 +861,11 @@ xla_test( "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/client:computation", - "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client/lib:arithmetic", + "//tensorflow/compiler/xla/client/xla_client:xla_builder", + "//tensorflow/compiler/xla/client/xla_client:xla_computation", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:hlo_test_base", @@ -1063,6 +1068,8 @@ xla_test( "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client/lib:arithmetic", + "//tensorflow/compiler/xla/client/xla_client:xla_builder", + "//tensorflow/compiler/xla/client/xla_client:xla_computation", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:xla_internal_test_main", @@ -1085,10 +1092,11 @@ xla_test_library( "//tensorflow/compiler/xla:reference_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:padding", "//tensorflow/compiler/xla/client/lib:arithmetic", + "//tensorflow/compiler/xla/client/xla_client:xla_builder", + "//tensorflow/compiler/xla/client/xla_client:xla_computation", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", @@ -1125,11 +1133,11 @@ xla_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/client:computation", - "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:padding", "//tensorflow/compiler/xla/client/lib:arithmetic", + "//tensorflow/compiler/xla/client/xla_client:xla_builder", + "//tensorflow/compiler/xla/client/xla_client:xla_computation", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:xla_internal_test_main", @@ -1393,8 +1401,8 @@ xla_test( deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array4d", - "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/compiler/xla/client/xla_client:xla_builder", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:xla_internal_test_main", @@ -1482,8 +1490,8 @@ xla_test( deps = [ "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/compiler/xla/client/xla_client:xla_builder", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:xla_internal_test_main", @@ -1547,6 +1555,8 @@ xla_test( "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", + "//tensorflow/compiler/xla/client/xla_client:xla_builder", + "//tensorflow/compiler/xla/client/xla_client:xla_computation", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:test_utils", "//tensorflow/compiler/xla/tests:xla_internal_test_main", @@ -1963,3 +1973,16 @@ tf_cc_test( "//tensorflow/core:test", ], ) + +xla_test( + name = "test_utils_test", + srcs = ["test_utils_test.cc"], + deps = [ + ":local_client_test_base", + ":test_utils", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:test", + ], +) diff --git a/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc b/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc index 03c91745b978f80801e0da5ac44d31959659b20c..4b4dc6dd9d3dddd105c63c7f3ce039bafca65dd7 100644 --- a/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc +++ b/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc @@ -2217,6 +2217,15 @@ XLA_TEST_F(ArrayElementwiseOpTest, LogF32sVector) { error_spec_); } +XLA_TEST_F(ArrayElementwiseOpTest, ClzU32s) { + XlaBuilder builder(TestName()); + auto a = builder.ConstantR1( + {0, 1, 0x10, 0x10000, 0x700000, 0x12345678, 0xF2345678}); + builder.Clz(a); + + ComputeAndCompareR1(&builder, {32, 31, 27, 15, 9, 3, 0}, {}); +} + XLA_TEST_F(ArrayElementwiseOpTest, AddChainFoldLeft) { // a ------ (add) --------- (add) // / / diff --git a/tensorflow/compiler/xla/tests/batch_normalization_test.cc b/tensorflow/compiler/xla/tests/batch_normalization_test.cc index af8af99c791e2a40cfcfa2291b786b33e5652267..f3dac75a44b948c4b45b80b93e7462073010979e 100644 --- a/tensorflow/compiler/xla/tests/batch_normalization_test.cc +++ b/tensorflow/compiler/xla/tests/batch_normalization_test.cc @@ -19,10 +19,10 @@ limitations under the License. #include "tensorflow/compiler/xla/array2d.h" #include "tensorflow/compiler/xla/array4d.h" -#include "tensorflow/compiler/xla/client/computation.h" -#include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/reference_util.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" @@ -69,14 +69,12 @@ class BatchNormalizationTest CHECK_EQ(kY, input_array_.width()); } - ComputationDataHandle CheckShape(ComputationBuilder* b, - const ComputationDataHandle& operand, - const Shape& expected_shape) const { - std::unique_ptr actual_shape = - b->GetShape(operand).ConsumeValueOrDie(); - CHECK(ShapeUtil::Equal(expected_shape, *actual_shape)) + XlaOp CheckShape(XlaBuilder* b, const XlaOp& operand, + const Shape& expected_shape) const { + Shape actual_shape = b->GetShape(operand).ConsumeValueOrDie(); + CHECK(ShapeUtil::Equal(expected_shape, actual_shape)) << "want " << ShapeUtil::HumanString(expected_shape) << " got " - << ShapeUtil::HumanString(*actual_shape); + << ShapeUtil::HumanString(actual_shape); return operand; } @@ -102,7 +100,7 @@ INSTANTIATE_TEST_CASE_P(BatchNormalizationTestInstance, BatchNormalizationTest, #endif XLA_TEST_P(BatchNormalizationTest, SubtractInZ) { - ComputationBuilder builder(client_, "subtract_in_z_one_sample"); + XlaBuilder builder("subtract_in_z_one_sample"); auto x = builder.ConstantLiteral(input_literal_); auto y = builder.ConstantR1({3.14, 4.25}); builder.Sub(x, y, /*broadcast_dimensions=*/{1}); @@ -118,7 +116,7 @@ XLA_TEST_P(BatchNormalizationTest, SubtractInZ) { } XLA_TEST_P(BatchNormalizationTest, SquareTesseractElementwise) { - ComputationBuilder builder(client_, "square_tesseract_elementwise"); + XlaBuilder builder("square_tesseract_elementwise"); auto x = builder.ConstantLiteral(input_literal_); builder.SquareF32(x); @@ -135,9 +133,9 @@ XLA_TEST_P(BatchNormalizationTest, SquareTesseractElementwise) { } XLA_TEST_P(BatchNormalizationTest, SumToZ) { - ComputationBuilder builder(client_, "sum_to_z"); + XlaBuilder builder("sum_to_z"); auto input_activations = builder.ConstantLiteral(input_literal_); - Computation add = CreateScalarAddComputation(F32, &builder); + XlaComputation add = CreateScalarAddComputation(F32, &builder); // Reduce all but the Z dimension. builder.Reduce(input_activations, builder.ConstantR0(0.0f), add, {0, 2, 3}); @@ -147,24 +145,23 @@ XLA_TEST_P(BatchNormalizationTest, SumToZ) { } XLA_TEST_P(BatchNormalizationTest, SquareAndReduce) { - ComputationBuilder builder(client_, "square_and_reduce"); + XlaBuilder builder("square_and_reduce"); auto input_activations = builder.ConstantLiteral(input_literal_); auto set_means = builder.ConstantR1({2.f, 4.2f}); auto activation_deviations = builder.Sub(input_activations, set_means, /*broadcast_dimensions=*/{1}); - Computation add = CreateScalarAddComputation(F32, &builder); + XlaComputation add = CreateScalarAddComputation(F32, &builder); auto dev_squares = builder.SquareF32(activation_deviations); - auto sum_of_squares = builder.Reduce( - dev_squares, builder.ConstantR0(0.0f), add, {0, 2, 3}); + builder.Reduce(dev_squares, builder.ConstantR0(0.0f), add, {0, 2, 3}); std::vector expected = {18, 0.06}; ComputeAndCompareR1(&builder, expected, {}, error_spec_); } XLA_TEST_P(BatchNormalizationTest, VarianceToStddev) { - ComputationBuilder builder(client_, "variance_to_stddev"); + XlaBuilder builder("variance_to_stddev"); auto variance = builder.ConstantR1({6.f, .02f}); - auto sqrt = builder.SqrtF32(variance); + builder.SqrtF32(variance); std::vector expected = {2.44948974f, 0.14142136f}; ComputeAndCompareR1(&builder, expected, {}, error_spec_); @@ -173,13 +170,13 @@ XLA_TEST_P(BatchNormalizationTest, VarianceToStddev) { // Compare against a forward batch normalization example in the NN spec // reference. XLA_TEST_P(BatchNormalizationTest, SpecComparisonForward) { - ComputationBuilder builder(client_, "batch_normalize_per_spec"); + XlaBuilder builder("batch_normalize_per_spec"); auto input_activations = CheckShape(&builder, builder.ConstantLiteral(input_literal_), ShapeUtil::MakeShape(F32, {3, 2, 1, 1})); auto gamma = builder.ConstantR1({1.0, 1.0}); auto beta = builder.ConstantR1({0.0, 0.0}); - Computation add = CreateScalarAddComputation(F32, &builder); + XlaComputation add = CreateScalarAddComputation(F32, &builder); // Reduce all dimensions except dimension 1. Shape TwoElementVectorF32 = ShapeUtil::MakeShape(F32, {2}); auto sum = CheckShape( @@ -189,8 +186,8 @@ XLA_TEST_P(BatchNormalizationTest, SpecComparisonForward) { TwoElementVectorF32); auto input_shape = builder.GetShape(input_activations).ConsumeValueOrDie(); auto sum_shape = builder.GetShape(sum).ConsumeValueOrDie(); - auto count = builder.ConstantR0(ShapeUtil::ElementsIn(*input_shape) / - ShapeUtil::ElementsIn(*sum_shape)); + auto count = builder.ConstantR0(ShapeUtil::ElementsIn(input_shape) / + ShapeUtil::ElementsIn(sum_shape)); auto set_means = builder.Div(sum, count); const float kEpsilon = 1e-9f; @@ -233,7 +230,7 @@ XLA_TEST_P(BatchNormalizationTest, SpecComparisonForward) { XLA_TEST_P(BatchNormalizationTest, BasicTraining) { const int kFeatureIndex = 3; - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto operand = builder.ConstantR4FromArray4D( {{{{1.f, 2.f}}, {{3.f, 4.f}}}, {{{5.f, 6.f}}, {{7.f, 8.f}}}}); @@ -242,8 +239,8 @@ XLA_TEST_P(BatchNormalizationTest, BasicTraining) { auto offset = builder.ConstantR1({1.0f, 2.0f}); - auto tuple = builder.BatchNormTraining(operand, scale, offset, - /*epsilon=*/0.001, kFeatureIndex); + builder.BatchNormTraining(operand, scale, offset, + /*epsilon=*/0.001, kFeatureIndex); auto expected = Literal::MakeTuple( {Literal::CreateR4({{{{-1.6f, -2.0f}}, {{0.1f, 0.6f}}}, @@ -257,7 +254,7 @@ XLA_TEST_P(BatchNormalizationTest, BasicTraining) { XLA_TEST_P(BatchNormalizationTest, BasicTrainingOnSublane) { const int kFeatureIndex = 2; - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto operand = builder.ConstantR4FromArray4D( {{{{1.f}, {2.f}}, {{3.f}, {4.f}}}, {{{5.f}, {6.f}}, {{7.f}, {8.f}}}}); @@ -266,8 +263,8 @@ XLA_TEST_P(BatchNormalizationTest, BasicTrainingOnSublane) { auto offset = builder.ConstantR1({1.0f, 2.0f}); - auto tuple = builder.BatchNormTraining(operand, scale, offset, - /*epsilon=*/0.001, kFeatureIndex); + builder.BatchNormTraining(operand, scale, offset, + /*epsilon=*/0.001, kFeatureIndex); auto expected = Literal::MakeTuple( {Literal::CreateR4({{{{-1.6f}, {-2.0f}}, {{0.1f}, {0.6f}}}, @@ -282,23 +279,23 @@ XLA_TEST_P(BatchNormalizationTest, BasicTrainingOnSublane) { XLA_TEST_P(BatchNormalizationTest, TrainingWithFeatureOnLowDimension) { // Use 0 dimension as feature, tests layout analyzer. const int kFeatureIndex = 0; - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); - ComputationDataHandle h0; + XlaOp h0; auto operand = CreateR3Parameter(Array3D(260, 2, 2, 1.0f), /*parameter_number=*/0, "operand", &builder, &h0); - ComputationDataHandle h1; + XlaOp h1; auto scale = CreateR1Parameter(std::vector(260, 1.0f), /*parameter_number=*/1, "scale", &builder, &h1); - ComputationDataHandle h2; + XlaOp h2; auto offset = CreateR1Parameter(std::vector(260, 1.0f), /*parameter_number=*/2, "offset", &builder, &h2); - auto tuple = builder.BatchNormTraining(h0, h1, h2, - /*epsilon=*/1, kFeatureIndex); + builder.BatchNormTraining(h0, h1, h2, + /*epsilon=*/1, kFeatureIndex); auto expected = Literal::MakeTuple( {Literal::CreateR3FromArray3D(Array3D(260, 2, 2, 1.0f)) @@ -314,24 +311,24 @@ XLA_TEST_P(BatchNormalizationTest, TrainingWithFeatureOnLowDimension) { XLA_TEST_P(BatchNormalizationTest, LargeEpsilonTest) { // Test the correctness of choosing a large epsilon value. const int kFeatureIndex = 2; - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); - ComputationDataHandle h0; + XlaOp h0; auto operand = CreateR3Parameter({{{0.0f}, {10.0f}, {20.0f}, {30.0f}}}, /*parameter_number=*/0, "operand", &builder, &h0); - ComputationDataHandle h1; + XlaOp h1; auto scale = CreateR1Parameter(std::vector(1, 1.0f), /*parameter_number=*/1, "scale", &builder, &h1); - ComputationDataHandle h2; + XlaOp h2; auto offset = CreateR1Parameter(std::vector(1, 0.0f), /*parameter_number=*/2, "offset", &builder, &h2); // var = 125, mean = 15, epsilon = -100 - auto tuple = builder.BatchNormTraining(h0, h1, h2, - /*epsilon=*/-100, kFeatureIndex); + builder.BatchNormTraining(h0, h1, h2, + /*epsilon=*/-100, kFeatureIndex); auto expected = Literal::MakeTuple( {Literal::CreateR3FromArray3D({{{-3.0f}, {-1.0f}, {1.0f}, {3.0f}}}) @@ -346,7 +343,7 @@ XLA_TEST_P(BatchNormalizationTest, LargeEpsilonTest) { XLA_TEST_P(BatchNormalizationTest, BatchNormGradBasic) { const int kFeatureIndex = 2; - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto operand = builder.ConstantR4FromArray4D(Array4D(2, 2, 2, 1, 0.0f)); @@ -453,7 +450,7 @@ INSTANTIATE_TEST_CASE_P(BatchNormTest_Instantiation, BatchNormTestManySizes, XLA_TEST_P(BatchNormTestManySizes, RandomizedTrainingTests) { float epsilon = 0.001; - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); const std::vector& bounds = GetParam().bounds; Array4D input_array(bounds[0], bounds[1], bounds[2], bounds[3]); input_array.FillRandom(GetParam().random_value_var, @@ -553,7 +550,7 @@ XLA_TEST_P(BatchNormTestManySizes, RandomizedTrainingTests) { XLA_TEST_P(BatchNormTestManySizes, RandomizedInferencingTests) { float epsilon = 0.001; - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); const std::vector& bounds = GetParam().bounds; Array4D input_array(bounds[0], bounds[1], bounds[2], bounds[3]); input_array.FillRandom(GetParam().random_value_var, @@ -661,7 +658,7 @@ XLA_TEST_P(BatchNormTestManySizes, RandomizedInferencingTests) { XLA_TEST_P(BatchNormTestManySizes, RandomizedGradTests) { float epsilon = 0.001; - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); const std::vector& bounds = GetParam().bounds; Array4D input_array(bounds[0], bounds[1], bounds[2], bounds[3]); input_array.FillRandom(GetParam().random_value_var, @@ -828,9 +825,9 @@ XLA_TEST_P(BatchNormTestManySizes, RandomizedGradTests) { std::unique_ptr grad_output_data = client_->TransferToServer(*grad_output_literal).ConsumeValueOrDie(); - auto t = builder.BatchNormGrad(input_parameter, scale_parameter, - mean_parameter, var_parameter, - grad_output_parameter, epsilon, feature_index); + builder.BatchNormGrad(input_parameter, scale_parameter, mean_parameter, + var_parameter, grad_output_parameter, epsilon, + feature_index); auto expected = Literal::MakeTuple({expected_grad_activation.get(), diff --git a/tensorflow/compiler/xla/tests/bitcast_convert_test.cc b/tensorflow/compiler/xla/tests/bitcast_convert_test.cc index 0d94d65c1015fb54ada3fdfc95d0c31d0a0f158b..bff60f25ec8f15d372d251ac313200301a04f20f 100644 --- a/tensorflow/compiler/xla/tests/bitcast_convert_test.cc +++ b/tensorflow/compiler/xla/tests/bitcast_convert_test.cc @@ -18,8 +18,8 @@ limitations under the License. #include #include -#include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" @@ -34,7 +34,7 @@ namespace { class BitcastConvertTest : public ClientLibraryTestBase { public: - explicit BitcastConvertTest(perftools::gputools::Platform* platform = nullptr) + explicit BitcastConvertTest(se::Platform* platform = nullptr) : ClientLibraryTestBase(platform) { mutable_debug_options()->add_xla_disable_hlo_passes("algsimp"); mutable_debug_options()->add_xla_disable_hlo_passes("inline"); @@ -42,7 +42,7 @@ class BitcastConvertTest : public ClientLibraryTestBase { }; TEST_F(BitcastConvertTest, ConvertR1S32ToR1S32) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto a = builder.ConstantR1({42, 64}); builder.BitcastConvertType(a, S32); @@ -51,7 +51,7 @@ TEST_F(BitcastConvertTest, ConvertR1S32ToR1S32) { } TEST_F(BitcastConvertTest, ConvertR1F32ToR1F32) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto a = builder.ConstantR1({42.0f, 64.0f}); builder.BitcastConvertType(a, F32); @@ -60,7 +60,7 @@ TEST_F(BitcastConvertTest, ConvertR1F32ToR1F32) { } TEST_F(BitcastConvertTest, BitcastR1S32ToR1F32) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto a = builder.ConstantR1({0, static_cast(0x80000000), 0x3F800000, static_cast(0xBF800000), 0x3F000000, @@ -72,7 +72,7 @@ TEST_F(BitcastConvertTest, BitcastR1S32ToR1F32) { } XLA_TEST_F(BitcastConvertTest, ConvertR1S0S32ToR1S0F32) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto a = builder.ConstantR1({}); builder.BitcastConvertType(a, F32); @@ -81,7 +81,7 @@ XLA_TEST_F(BitcastConvertTest, ConvertR1S0S32ToR1S0F32) { } TEST_F(BitcastConvertTest, ConvertR1F32ToR1S32) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto a = builder.ConstantR1({42.6, 64.4}); builder.BitcastConvertType(a, S32); @@ -90,7 +90,7 @@ TEST_F(BitcastConvertTest, ConvertR1F32ToR1S32) { } TEST_F(BitcastConvertTest, ConvertS32Extremes) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto a = builder.ConstantR1( {std::numeric_limits::min(), std::numeric_limits::max()}); builder.BitcastConvertType(a, F32); @@ -100,7 +100,7 @@ TEST_F(BitcastConvertTest, ConvertS32Extremes) { } TEST_F(BitcastConvertTest, ConvertMapToS32) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto b = builder.CreateSubBuilder("convert"); auto param = b->Parameter(0, ShapeUtil::MakeShape(F32, {}), "in"); b->BitcastConvertType(param, S32); @@ -112,7 +112,7 @@ TEST_F(BitcastConvertTest, ConvertMapToS32) { } TEST_F(BitcastConvertTest, ConvertMapToF32) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto b = builder.CreateSubBuilder("convert"); auto param = b->Parameter(0, ShapeUtil::MakeShape(S32, {}), "in"); b->BitcastConvertType(param, F32); @@ -129,7 +129,7 @@ TEST_F(BitcastConvertTest, ConvertMapToF32) { // input -> convert -> reshape // the new convert should have the same element type as the old convert. TEST_F(BitcastConvertTest, ConvertReshape) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto input = builder.ConstantR1({0x42280000}); auto reshape = builder.Reshape(input, /*dimensions=*/{0}, /*new_sizes=*/{}); builder.BitcastConvertType(reshape, F32); diff --git a/tensorflow/compiler/xla/tests/client_library_test_base.cc b/tensorflow/compiler/xla/tests/client_library_test_base.cc index c2e3cd2350965165dab6e203013eae46c8127753..69389dae3f28a2a20eb8a26e72ab3541eb565305 100644 --- a/tensorflow/compiler/xla/tests/client_library_test_base.cc +++ b/tensorflow/compiler/xla/tests/client_library_test_base.cc @@ -23,6 +23,7 @@ limitations under the License. #include "tensorflow/compiler/xla/execution_options_util.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/ptr_util.h" +#include "tensorflow/compiler/xla/service/platform_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" @@ -31,10 +32,12 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" -namespace se = ::perftools::gputools; - namespace xla { namespace { + +// Name of the interpreter backend. +constexpr char kInterpreter[] = "interpreter"; + // Wrapper function that creates a nicer error message (than a bare // ValueOrDie()) if the platform we intend to test is not available. Client* GetOrCreateLocalClientOrDie(const LocalClientOptions& client_options) { @@ -43,11 +46,18 @@ Client* GetOrCreateLocalClientOrDie(const LocalClientOptions& client_options) { TF_CHECK_OK(result.status()) << " could not create local client for testing"; return result.ValueOrDie(); } + +// Helper functions to get the reference platform. +se::Platform* GetReferencePlatform() { + auto result = PlatformUtil::GetPlatform(kInterpreter); + TF_CHECK_OK(result.status()) << "could not get interpreter platform"; + return result.ValueOrDie(); +} + } // namespace ClientLibraryTestBase::ClientLibraryTestBase( - perftools::gputools::Platform* platform, - const LocalClientOptions& client_options) + se::Platform* platform, const LocalClientOptions& client_options) : client_(GetOrCreateLocalClientOrDie(client_options)), execution_options_(CreateDefaultExecutionOptions()) { CHECK_EQ(platform, client_options.platform()); @@ -66,6 +76,11 @@ ClientLibraryTestBase::ClientLibraryTestBase(se::Platform* platform) LocalClientOptions default_options; default_options.set_platform(platform); client_ = GetOrCreateLocalClientOrDie(default_options); + + LocalClientOptions ref_options; + ref_options.set_platform(GetReferencePlatform()); + ref_client_ = GetOrCreateLocalClientOrDie(ref_options); + execution_options_.mutable_debug_options()->add_xla_disable_hlo_passes( "constant_folding"); } @@ -127,6 +142,20 @@ StatusOr> ClientLibraryTestBase::ExecuteAndTransfer( return ExecuteAndTransfer(computation, arguments, shape_with_output_layout); } +StatusOr> +ClientLibraryTestBase::ExecuteAndTransferReference( + const XlaComputation& computation, + tensorflow::gtl::ArraySlice arguments, + const Shape* shape_with_output_layout) { + ExecutionOptions execution_options = execution_options_; + if (shape_with_output_layout != nullptr) { + *execution_options.mutable_shape_with_output_layout() = + *shape_with_output_layout; + } + return ref_client_->ExecuteAndTransfer(computation, arguments, + &execution_options); +} + std::unique_ptr ClientLibraryTestBase::ExecuteOrDie( ComputationBuilder* builder, tensorflow::gtl::ArraySlice arguments) { @@ -521,6 +550,69 @@ ClientLibraryTestBase::ComputeValueAndReference( return std::make_pair(std::move(reference), std::move(result)); } +void ClientLibraryTestBase::ComputeAndCompare( + XlaBuilder* builder, tensorflow::gtl::ArraySlice arguments) { + auto status_or_data = ComputeValueAndReference(builder, arguments); + EXPECT_IS_OK(status_or_data); + if (!status_or_data.ok()) { + return; + } + std::unique_ptr reference, result; + std::tie(reference, result) = status_or_data.ConsumeValueOrDie(); + LiteralTestUtil::ExpectEqual(*reference, *result); +} + +void ClientLibraryTestBase::ComputeAndCompare( + XlaBuilder* builder, tensorflow::gtl::ArraySlice arguments, + ErrorSpec error) { + auto status_or_data = ComputeValueAndReference(builder, arguments); + EXPECT_IS_OK(status_or_data); + if (!status_or_data.ok()) { + return; + } + std::unique_ptr reference, result; + std::tie(reference, result) = status_or_data.ConsumeValueOrDie(); + LiteralTestUtil::ExpectNear(*reference, *result, error); +} + +StatusOr, std::unique_ptr>> +ClientLibraryTestBase::ComputeValueAndReference( + XlaBuilder* builder, tensorflow::gtl::ArraySlice arguments) { + // Transfer the arguments to the executor service. We put the unique_ptr's + // into a vector to keep the data alive on the service until the end of this + // function. + std::vector> argument_data; + std::vector> ref_argument_data; + for (const auto& arg : arguments) { + TF_ASSIGN_OR_RETURN(auto data, client_->TransferToServer(arg.Clone())); + TF_ASSIGN_OR_RETURN(auto ref_data, ref_client_->TransferToServer(arg)); + argument_data.push_back(std::move(data)); + ref_argument_data.push_back(std::move(ref_data)); + } + + // Create raw pointers to the GlobalData for the rest of the call stack. + std::vector argument_data_ptr; + std::transform( + argument_data.begin(), argument_data.end(), + std::back_inserter(argument_data_ptr), + [](const std::unique_ptr& data) { return data.get(); }); + std::vector ref_argument_data_ptr; + std::transform( + ref_argument_data.begin(), ref_argument_data.end(), + std::back_inserter(ref_argument_data_ptr), + [](const std::unique_ptr& data) { return data.get(); }); + + TF_ASSIGN_OR_RETURN(auto computation, builder->Build()); + + TF_ASSIGN_OR_RETURN(auto result, + ExecuteAndTransfer(computation, argument_data_ptr)); + + TF_ASSIGN_OR_RETURN(auto reference, ExecuteAndTransferReference( + computation, ref_argument_data_ptr)); + + return std::make_pair(std::move(reference), std::move(result)); +} + Computation ClientLibraryTestBase::CreateScalarRelu() { ComputationBuilder builder(client_, "relu"); auto shape = ShapeUtil::MakeShape(use_bfloat16_ ? BF16 : F32, {}); diff --git a/tensorflow/compiler/xla/tests/client_library_test_base.h b/tensorflow/compiler/xla/tests/client_library_test_base.h index 0572acff88313a277fd81c7a13defed4e81a939c..481d7c5c25a7e8ad9d40d11fa3f84e13d6922aa3 100644 --- a/tensorflow/compiler/xla/tests/client_library_test_base.h +++ b/tensorflow/compiler/xla/tests/client_library_test_base.h @@ -64,11 +64,10 @@ std::vector ExpandUseBfloat16( // A client library test establishes an in-process XLA client connection. class ClientLibraryTestBase : public ::testing::Test { protected: - explicit ClientLibraryTestBase( - perftools::gputools::Platform* platform = nullptr); + explicit ClientLibraryTestBase(se::Platform* platform = nullptr); // Creates a new ClientLibraryTestBase with custom client options. - ClientLibraryTestBase(perftools::gputools::Platform* platform, + ClientLibraryTestBase(se::Platform* platform, const LocalClientOptions& client_options); // Returns the name of the test currently being run. @@ -114,6 +113,14 @@ class ClientLibraryTestBase : public ::testing::Test { tensorflow::gtl::ArraySlice arguments, const Shape* shape_with_output_layout = nullptr); + // This executes the computation via the reference client (which connects a + // interpreter backend). The result is used as the expected values of the + // computation. + StatusOr> ExecuteAndTransferReference( + const XlaComputation& computation, + tensorflow::gtl::ArraySlice arguments, + const Shape* shape_with_output_layout = nullptr); + // Convenience OrDie variants of above methods. std::unique_ptr ExecuteOrDie( ComputationBuilder* builder, @@ -236,6 +243,14 @@ class ClientLibraryTestBase : public ::testing::Test { tensorflow::gtl::ArraySlice arguments, ErrorSpec error); + // Convenience method for running a built computation and comparing the result + // with the reference result. + void ComputeAndCompare(XlaBuilder* builder, + tensorflow::gtl::ArraySlice arguments); + void ComputeAndCompare(XlaBuilder* builder, + tensorflow::gtl::ArraySlice arguments, + ErrorSpec error); + // Create scalar operations for use in reductions. Computation CreateScalarRelu(); Computation CreateScalarMax(); @@ -413,6 +428,7 @@ class ClientLibraryTestBase : public ::testing::Test { PrimitiveType FloatType() const { return use_bfloat16_ ? BF16 : F32; } Client* client_; + Client* ref_client_; // To compute reference result. ExecutionOptions execution_options_; private: @@ -444,12 +460,19 @@ class ClientLibraryTestBase : public ::testing::Test { const Shape* output_with_layout = nullptr); // Executes the computation and calculates the expected reference value using - // the HloEvaluator. Returns two literal in the order of (expected, actual). + // the HloEvaluator. Returns two literals in the order of (expected, actual). StatusOr, std::unique_ptr>> ComputeValueAndReference(ComputationBuilder* builder, const ComputationDataHandle& operand, tensorflow::gtl::ArraySlice arguments); + // Executes the computation and calculates the expected reference value using + // the reference client. Returns two literals in the order of (expected, + // actual). + StatusOr, std::unique_ptr>> + ComputeValueAndReference(XlaBuilder* builder, + tensorflow::gtl::ArraySlice arguments); + // Whether to run tests with all float-type input/output converted to // bfloat16. bool use_bfloat16_ = false; diff --git a/tensorflow/compiler/xla/tests/compute_constant_test.cc b/tensorflow/compiler/xla/tests/compute_constant_test.cc index e5a03b49ad259a64b9cbbc88c31d8c6558289d1b..7ea82a791f72ea23be152ce9f548cc6719922236 100644 --- a/tensorflow/compiler/xla/tests/compute_constant_test.cc +++ b/tensorflow/compiler/xla/tests/compute_constant_test.cc @@ -21,6 +21,8 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/global_data.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h" #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -31,6 +33,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/strings/str_util.h" #include "tensorflow/core/platform/types.h" @@ -44,16 +47,14 @@ ClientType client_types[] = {ClientType::kLocal, ClientType::kCompileOnly}; class ComputeConstantTest : public ::testing::Test { public: - explicit ComputeConstantTest( - perftools::gputools::Platform* platform = nullptr) + explicit ComputeConstantTest(se::Platform* platform = nullptr) : platform_(platform) {} string TestName() const { return ::testing::UnitTest::GetInstance()->current_test_info()->name(); } - Client* ClientOrDie(::perftools::gputools::Platform* platform, - ClientType client_type) { + Client* ClientOrDie(se::Platform* platform, ClientType client_type) { if (client_type == ClientType::kLocal) { StatusOr result = ClientLibrary::GetOrCreateLocalClient(platform); @@ -71,39 +72,46 @@ class ComputeConstantTest : public ::testing::Test { } StatusOr> ComputeConstantLiteral( - Client* client, const ComputationDataHandle& operand, - ComputationBuilder* builder, Layout* output_layout = nullptr, - tensorflow::gtl::ArraySlice parameters = {}) { - TF_ASSIGN_OR_RETURN(auto computed, builder->ComputeConstant( - operand, output_layout, parameters)); + Client* client, const XlaOp& operand, XlaBuilder* builder, + Layout* output_layout = nullptr) { + TF_ASSIGN_OR_RETURN(auto subgraph, builder->BuildConstantSubGraph(operand)); + TF_ASSIGN_OR_RETURN(auto computed, + client->ComputeConstant(subgraph, output_layout)); return std::move(computed); } + template + StatusOr ComputeConstantScalar(Client* client, const XlaOp& operand, + XlaBuilder* builder) { + TF_ASSIGN_OR_RETURN(auto literal, ComputeConstantLiteral(client, operand, + builder, nullptr)); + return literal->Get({}); + } + template StatusOr ComputeConstantScalar( Client* client, const ComputationDataHandle& operand, ComputationBuilder* builder, tensorflow::gtl::ArraySlice parameters = {}) { - TF_ASSIGN_OR_RETURN( - auto literal, - ComputeConstantLiteral(client, operand, builder, nullptr, parameters)); + TF_ASSIGN_OR_RETURN(auto literal, + builder->ComputeConstant( + operand, /*output_layout=*/nullptr, parameters)); return literal->Get({}); } - bool IsConstant(const ComputationDataHandle& operand, - ComputationBuilder* builder, int64 num_parameters = 0) { - StatusOr result = builder->IsConstant(operand, num_parameters); + bool IsConstant(const XlaOp& operand, XlaBuilder* builder) { + StatusOr result = builder->IsConstant(operand); EXPECT_TRUE(result.ok()) << result.status(); return result.ok() ? result.ValueOrDie() : false; } - perftools::gputools::Platform* platform_; + se::Platform* platform_; }; TEST_F(ComputeConstantTest, ScalarInt32Literal) { for (ClientType client_type : client_types) { Client* client = ClientOrDie(platform_, client_type); - ComputationBuilder b(client, TestName()); + XlaBuilder b(TestName()); auto computation = b.ConstantR0(42); EXPECT_TRUE(IsConstant(computation, &b)); @@ -116,7 +124,7 @@ TEST_F(ComputeConstantTest, ScalarInt32Literal) { TEST_F(ComputeConstantTest, ScalarFloatAdd) { for (ClientType client_type : client_types) { Client* client = ClientOrDie(platform_, client_type); - ComputationBuilder b(client, TestName()); + XlaBuilder b(TestName()); auto computation = b.Add(b.ConstantR0(42.5f), b.ConstantR0(1.5f)); EXPECT_TRUE(IsConstant(computation, &b)); @@ -130,7 +138,7 @@ TEST_F(ComputeConstantTest, ScalarFloatAdd) { TEST_F(ComputeConstantTest, ScalarRng) { for (ClientType client_type : client_types) { Client* client = ClientOrDie(platform_, client_type); - ComputationBuilder b(client, TestName()); + XlaBuilder b(TestName()); auto computation = b.RngUniform(b.ConstantR0(1.1f), b.ConstantR0(2.1f), ShapeUtil::MakeShape(F32, {})); @@ -151,19 +159,21 @@ TEST_F(ComputeConstantTest, Param) { std::vector arguments; arguments.push_back(std::move(*Literal::CreateR0(42.5f))); - EXPECT_TRUE(IsConstant(computation, &b, arguments.size())); - - auto value = - ComputeConstantScalar(client, computation, &b, arguments); - ASSERT_TRUE(value.ok()) << value.status(); - EXPECT_EQ(value.ValueOrDie(), 44.0f); + TF_ASSERT_OK_AND_ASSIGN(bool is_constant, + b.IsConstant(computation, arguments.size())); + EXPECT_TRUE(is_constant); + + TF_ASSERT_OK_AND_ASSIGN( + auto value, + ComputeConstantScalar(client, computation, &b, arguments)); + EXPECT_EQ(value, 44.0f); } } TEST_F(ComputeConstantTest, DirectParamMissing) { for (ClientType client_type : client_types) { Client* client = ClientOrDie(platform_, client_type); - ComputationBuilder b(client, TestName()); + XlaBuilder b(TestName()); auto computation = b.Parameter(0, ShapeUtil::MakeShape(F32, {}), "param"); EXPECT_FALSE(IsConstant(computation, &b)); @@ -177,7 +187,7 @@ TEST_F(ComputeConstantTest, DirectParamMissing) { TEST_F(ComputeConstantTest, IndirectParamMissing) { for (ClientType client_type : client_types) { Client* client = ClientOrDie(platform_, client_type); - ComputationBuilder b(client, TestName()); + XlaBuilder b(TestName()); auto computation = b.Add(b.ConstantR0(1.0f), b.Parameter(0, ShapeUtil::MakeShape(F32, {}), "param")); @@ -195,7 +205,7 @@ TEST_F(ComputeConstantTest, IndirectParamMissing) { TEST_F(ComputeConstantTest, UnrelatedParam) { for (ClientType client_type : client_types) { Client* client = ClientOrDie(platform_, client_type); - ComputationBuilder b(client, TestName()); + XlaBuilder b(TestName()); auto param_a = b.Parameter(10, ShapeUtil::MakeShape(F32, {}), "param0"); auto constant_4 = @@ -212,64 +222,64 @@ TEST_F(ComputeConstantTest, UnrelatedParam) { EXPECT_TRUE(IsConstant(constant_13, &b)); - auto value = ComputeConstantScalar(client, constant_13, &b); - ASSERT_TRUE(value.ok()) << value.status(); - EXPECT_EQ(value.ValueOrDie(), 13.0f); + TF_ASSERT_OK_AND_ASSIGN( + auto value, ComputeConstantScalar(client, constant_13, &b)); + EXPECT_EQ(value, 13.0f); } } TEST_F(ComputeConstantTest, NonScalarAdd) { for (ClientType client_type : client_types) { Client* client = ClientOrDie(platform_, client_type); - ComputationBuilder b(client, TestName()); + XlaBuilder b(TestName()); auto computation = b.Add(b.ConstantR1({1, 2}), b.ConstantR1({3, 4})); EXPECT_TRUE(IsConstant(computation, &b)); - auto computed = ComputeConstantLiteral(client, computation, &b); - ASSERT_TRUE(computed.ok()) << computed.status(); + TF_ASSERT_OK_AND_ASSIGN(auto computed, + ComputeConstantLiteral(client, computation, &b)); std::unique_ptr expected_literal = Literal::CreateR1({4, 6}); - LiteralTestUtil::ExpectEqual(*expected_literal, *computed.ValueOrDie()); + LiteralTestUtil::ExpectEqual(*expected_literal, *computed); } } TEST_F(ComputeConstantTest, IntegerDivide) { for (ClientType client_type : client_types) { Client* client = ClientOrDie(platform_, client_type); - ComputationBuilder b(client, TestName()); + XlaBuilder b(TestName()); auto computation = b.Div(b.ConstantR0(15), b.ConstantR0(3)); EXPECT_TRUE(IsConstant(computation, &b)); - auto computed = ComputeConstantLiteral(client, computation, &b); - ASSERT_TRUE(computed.ok()) << computed.status(); + TF_ASSERT_OK_AND_ASSIGN(auto computed, + ComputeConstantLiteral(client, computation, &b)); std::unique_ptr expected_literal = Literal::CreateR0(5); - LiteralTestUtil::ExpectEqual(*expected_literal, *computed.ValueOrDie()); + LiteralTestUtil::ExpectEqual(*expected_literal, *computed); } } XLA_TEST_F(ComputeConstantTest, Layout) { for (ClientType client_type : client_types) { Client* client = ClientOrDie(platform_, client_type); - ComputationBuilder b(client, TestName()); + XlaBuilder b(TestName()); std::vector> layouts = {{0, 1}, {1, 0}}; for (const std::vector& layout : layouts) { auto layout_proto = LayoutUtil::MakeLayout(layout); - auto computed = ComputeConstantLiteral( - client, - b.Add(b.ConstantR2({{1, 2}, {3, 4}}), - b.ConstantR2({{10, 20}, {30, 40}})), - &b, &layout_proto); - ASSERT_TRUE(computed.ok()) << computed.status(); + TF_ASSERT_OK_AND_ASSIGN( + auto computed, ComputeConstantLiteral( + client, + b.Add(b.ConstantR2({{1, 2}, {3, 4}}), + b.ConstantR2({{10, 20}, {30, 40}})), + &b, &layout_proto)); std::unique_ptr expected_literal = Literal::CreateR2WithLayout({{11, 22}, {33, 44}}, LayoutUtil::MakeLayout(layout)); - LiteralTestUtil::AssertEqualShapesAndLayouts( - expected_literal->shape(), computed.ValueOrDie()->shape()); - LiteralTestUtil::ExpectEqual(*expected_literal, *computed.ValueOrDie()); + LiteralTestUtil::AssertEqualShapesAndLayouts(expected_literal->shape(), + computed->shape()); + LiteralTestUtil::ExpectEqual(*expected_literal, *computed); } } } diff --git a/tensorflow/compiler/xla/tests/conditional_test.cc b/tensorflow/compiler/xla/tests/conditional_test.cc index b917dee77b5400db8f2c0a6a86258fee64723d71..7ff6706935740c7d76ee5cd03eae292386760397 100644 --- a/tensorflow/compiler/xla/tests/conditional_test.cc +++ b/tensorflow/compiler/xla/tests/conditional_test.cc @@ -13,7 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_computation.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" @@ -23,8 +24,8 @@ namespace { class ConditionalOpTest : public ClientLibraryTestBase { protected: - Computation CreateR0ConstantComputation(float value) { - ComputationBuilder builder(client_, "Constant"); + XlaComputation CreateR0ConstantComputation(float value) { + XlaBuilder builder("Constant"); builder.Parameter(0, empty_tuple_, "tuple"); builder.ConstantR0(value); auto build_status = builder.Build(); @@ -32,16 +33,16 @@ class ConditionalOpTest : public ClientLibraryTestBase { return build_status.ConsumeValueOrDie(); } - Computation CreateR0IdentityComputation() { - ComputationBuilder builder(client_, "Identity"); + XlaComputation CreateR0IdentityComputation() { + XlaBuilder builder("Identity"); builder.Parameter(0, r0f32_, "x"); auto build_status = builder.Build(); EXPECT_IS_OK(build_status.status()); return build_status.ConsumeValueOrDie(); } - Computation CreateCeilComputation(const Shape& shape) { - ComputationBuilder builder(client_, "Ceil"); + XlaComputation CreateCeilComputation(const Shape& shape) { + XlaBuilder builder("Ceil"); auto param = builder.Parameter(0, shape, "param"); builder.Ceil(param); auto build_status = builder.Build(); @@ -49,16 +50,16 @@ class ConditionalOpTest : public ClientLibraryTestBase { return build_status.ConsumeValueOrDie(); } - Computation CreateR0CeilComputation() { + XlaComputation CreateR0CeilComputation() { return CreateCeilComputation(r0f32_); } - Computation CreateR1CeilComputation() { + XlaComputation CreateR1CeilComputation() { return CreateCeilComputation(r1s2f32_); } - Computation CreateFloorComputation(const Shape& shape) { - ComputationBuilder builder(client_, "Floor"); + XlaComputation CreateFloorComputation(const Shape& shape) { + XlaBuilder builder("Floor"); auto param = builder.Parameter(0, shape, "param"); builder.Floor(param); auto build_status = builder.Build(); @@ -66,17 +67,17 @@ class ConditionalOpTest : public ClientLibraryTestBase { return build_status.ConsumeValueOrDie(); } - Computation CreateR0FloorComputation() { + XlaComputation CreateR0FloorComputation() { return CreateFloorComputation(r0f32_); } - Computation CreateR1FloorComputation() { + XlaComputation CreateR1FloorComputation() { return CreateFloorComputation(r1s2f32_); } - Computation CreateTupleCeilComputation(const string& computation_name, - const Shape& tuple_shape) { - ComputationBuilder builder(client_, computation_name); + XlaComputation CreateTupleCeilComputation(const string& computation_name, + const Shape& tuple_shape) { + XlaBuilder builder(computation_name); auto tuple = builder.Parameter(0, tuple_shape, "tuple"); auto x = builder.GetTupleElement(tuple, 0); auto y = builder.GetTupleElement(tuple, 1); @@ -88,17 +89,17 @@ class ConditionalOpTest : public ClientLibraryTestBase { return build_status.ConsumeValueOrDie(); } - Computation CreateR0TupleCeilComputation() { + XlaComputation CreateR0TupleCeilComputation() { return CreateTupleCeilComputation("CeilR0", tuple_2_r0f32_); } - Computation CreateR1TupleCeilComputation() { + XlaComputation CreateR1TupleCeilComputation() { return CreateTupleCeilComputation("CeilR1", tuple_2_r1s2f32_); } - Computation CreateTupleFloorComputation(const string& computation_name, - const Shape& tuple_shape) { - ComputationBuilder builder(client_, computation_name); + XlaComputation CreateTupleFloorComputation(const string& computation_name, + const Shape& tuple_shape) { + XlaBuilder builder(computation_name); auto tuple = builder.Parameter(0, tuple_shape, "tuple"); auto x = builder.GetTupleElement(tuple, 0); auto y = builder.GetTupleElement(tuple, 1); @@ -110,17 +111,17 @@ class ConditionalOpTest : public ClientLibraryTestBase { return build_status.ConsumeValueOrDie(); } - Computation CreateR0TupleFloorComputation() { + XlaComputation CreateR0TupleFloorComputation() { return CreateTupleFloorComputation("FloorR0", tuple_2_r0f32_); } - Computation CreateR1TupleFloorComputation() { + XlaComputation CreateR1TupleFloorComputation() { return CreateTupleFloorComputation("FloorR1", tuple_2_r1s2f32_); } - Computation CreateTupleAddComputation(const string& computation_name, - const Shape& tuple_shape) { - ComputationBuilder builder(client_, computation_name); + XlaComputation CreateTupleAddComputation(const string& computation_name, + const Shape& tuple_shape) { + XlaBuilder builder(computation_name); auto tuple = builder.Parameter(0, tuple_shape, "tuple"); auto x = builder.GetTupleElement(tuple, 0); auto y = builder.GetTupleElement(tuple, 1); @@ -130,17 +131,17 @@ class ConditionalOpTest : public ClientLibraryTestBase { return build_status.ConsumeValueOrDie(); } - Computation CreateR0TupleAddComputation() { + XlaComputation CreateR0TupleAddComputation() { return CreateTupleAddComputation("AddR0", tuple_2_r0f32_); } - Computation CreateR1TupleAddComputation() { + XlaComputation CreateR1TupleAddComputation() { return CreateTupleAddComputation("AddR1", tuple_2_r1s2f32_); } - Computation CreateTupleSubComputation(const string& computation_name, - const Shape& tuple_shape) { - ComputationBuilder builder(client_, computation_name); + XlaComputation CreateTupleSubComputation(const string& computation_name, + const Shape& tuple_shape) { + XlaBuilder builder(computation_name); auto tuple = builder.Parameter(0, tuple_shape, "tuple"); auto x = builder.GetTupleElement(tuple, 0); auto y = builder.GetTupleElement(tuple, 1); @@ -150,11 +151,11 @@ class ConditionalOpTest : public ClientLibraryTestBase { return build_status.ConsumeValueOrDie(); } - Computation CreateR0TupleSubComputation() { + XlaComputation CreateR0TupleSubComputation() { return CreateTupleSubComputation("SubR0", tuple_2_r0f32_); } - Computation CreateR1TupleSubComputation() { + XlaComputation CreateR1TupleSubComputation() { return CreateTupleSubComputation("SubR1", tuple_2_r1s2f32_); } @@ -170,26 +171,25 @@ class ConditionalOpTest : public ClientLibraryTestBase { // Test true and false computations that do not take any parameters. XLA_TEST_F(ConditionalOpTest, Parameters0) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(true); auto operands = builder.Tuple({}); auto true_computation = CreateR0ConstantComputation(56.0f); auto false_computation = CreateR0ConstantComputation(12.0f); - auto result = builder.Conditional(pred, operands, true_computation, operands, - false_computation); + builder.Conditional(pred, operands, true_computation, operands, + false_computation); ComputeAndCompareR0(&builder, 56.0f, {}, error_spec_); } // Test true and false computations that take in 1 parameter. XLA_TEST_F(ConditionalOpTest, Parameters1) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(false); auto operand1 = builder.ConstantR0(56.0f); auto operand2 = builder.ConstantR0(12.0f); auto identity = CreateR0IdentityComputation(); - auto result = - builder.Conditional(pred, operand1, identity, operand2, identity); + builder.Conditional(pred, operand1, identity, operand2, identity); ComputeAndCompareR0(&builder, 12.0f, {}, error_spec_); } @@ -197,12 +197,12 @@ XLA_TEST_F(ConditionalOpTest, Parameters1) { // Test conditional with two different computations in the true and false cases // that take in different arguments. XLA_TEST_F(ConditionalOpTest, DiffComputationsDiffArgs) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(false); auto operand1 = builder.ConstantR0(56.4f); auto operand2 = builder.ConstantR0(12.6f); - auto result = builder.Conditional(pred, operand1, CreateR0CeilComputation(), - operand2, CreateR0FloorComputation()); + builder.Conditional(pred, operand1, CreateR0CeilComputation(), operand2, + CreateR0FloorComputation()); ComputeAndCompareR0(&builder, 12.0f, {}, error_spec_); } @@ -210,11 +210,11 @@ XLA_TEST_F(ConditionalOpTest, DiffComputationsDiffArgs) { // Test conditional with two different computations in the true and false cases // that take in the same arguments. XLA_TEST_F(ConditionalOpTest, DiffComputationsSameArg) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(false); auto operand = builder.ConstantR0(12.6f); - auto result = builder.Conditional(pred, operand, CreateR0CeilComputation(), - operand, CreateR0FloorComputation()); + builder.Conditional(pred, operand, CreateR0CeilComputation(), operand, + CreateR0FloorComputation()); ComputeAndCompareR0(&builder, 12.0f, {}, error_spec_); } @@ -222,12 +222,12 @@ XLA_TEST_F(ConditionalOpTest, DiffComputationsSameArg) { // Test conditional with the same computation in the true and false cases but // take in different arguments. XLA_TEST_F(ConditionalOpTest, SameComputationDiffArgs) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(false); auto operand1 = builder.ConstantR0(56.4f); auto operand2 = builder.ConstantR0(12.6f); auto floor = CreateR0FloorComputation(); - auto result = builder.Conditional(pred, operand1, floor, operand2, floor); + builder.Conditional(pred, operand1, floor, operand2, floor); ComputeAndCompareR0(&builder, 12.0f, {}, error_spec_); } @@ -235,11 +235,11 @@ XLA_TEST_F(ConditionalOpTest, SameComputationDiffArgs) { // Test conditional with the same computation in the true and false cases that // take in the same arguments. XLA_TEST_F(ConditionalOpTest, SameComputationSameArg) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(false); auto operand = builder.ConstantR0(12.6f); auto floor = CreateR0FloorComputation(); - auto result = builder.Conditional(pred, operand, floor, operand, floor); + builder.Conditional(pred, operand, floor, operand, floor); ComputeAndCompareR0(&builder, 12.0f, {}, error_spec_); } @@ -247,12 +247,12 @@ XLA_TEST_F(ConditionalOpTest, SameComputationSameArg) { // Test conditional with different instances of the same computation in the true // and false cases. XLA_TEST_F(ConditionalOpTest, SameComputationDiffInstances) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(false); auto operand1 = builder.ConstantR0(56.4f); auto operand2 = builder.ConstantR0(12.6f); - auto result = builder.Conditional(pred, operand1, CreateR0FloorComputation(), - operand2, CreateR0FloorComputation()); + builder.Conditional(pred, operand1, CreateR0FloorComputation(), operand2, + CreateR0FloorComputation()); ComputeAndCompareR0(&builder, 12.0f, {}, error_spec_); } @@ -260,7 +260,7 @@ XLA_TEST_F(ConditionalOpTest, SameComputationDiffInstances) { // Test the case when a call invokes a computation that contains a conditional. XLA_TEST_F(ConditionalOpTest, ConditionalWithCall) { Shape r0bool = ShapeUtil::MakeShape(PRED, {}); - ComputationBuilder inner_builder(client_, TestName() + ".inner_conditional"); + XlaBuilder inner_builder(TestName() + ".inner_conditional"); auto pred_cond = inner_builder.Parameter(0, r0bool, "param0"); auto true_operand = inner_builder.Parameter(1, r0f32_, "param1"); auto false_operand = inner_builder.Parameter(2, r0f32_, "param2"); @@ -268,7 +268,7 @@ XLA_TEST_F(ConditionalOpTest, ConditionalWithCall) { false_operand, CreateR0FloorComputation()); auto inner_builder_result = inner_builder.Build(); - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(false); auto operand1 = builder.ConstantR0(56.4f); auto operand2 = builder.ConstantR0(12.6f); @@ -281,14 +281,13 @@ XLA_TEST_F(ConditionalOpTest, ConditionalWithCall) { // Test true and false computations that take in 2 parameters and predicate is // true. XLA_TEST_F(ConditionalOpTest, Parameters2TrueBranch) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(true); auto operand1 = builder.ConstantR0(56.0f); auto operand2 = builder.ConstantR0(12.0f); auto operands = builder.Tuple({operand1, operand2}); - auto result = - builder.Conditional(pred, operands, CreateR0TupleAddComputation(), - operands, CreateR0TupleSubComputation()); + builder.Conditional(pred, operands, CreateR0TupleAddComputation(), operands, + CreateR0TupleSubComputation()); ComputeAndCompareR0(&builder, 68.0f, {}, error_spec_); } @@ -296,14 +295,13 @@ XLA_TEST_F(ConditionalOpTest, Parameters2TrueBranch) { // Test true and false computations that take in 2 parameters and predicate is // false. XLA_TEST_F(ConditionalOpTest, Parameters2FalseBranch) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(false); auto operand1 = builder.ConstantR0(56.0f); auto operand2 = builder.ConstantR0(12.0f); auto operands = builder.Tuple({operand1, operand2}); - auto result = - builder.Conditional(pred, operands, CreateR0TupleAddComputation(), - operands, CreateR0TupleSubComputation()); + builder.Conditional(pred, operands, CreateR0TupleAddComputation(), operands, + CreateR0TupleSubComputation()); ComputeAndCompareR0(&builder, 44.0f, {}, error_spec_); } @@ -311,14 +309,13 @@ XLA_TEST_F(ConditionalOpTest, Parameters2FalseBranch) { // Test true and false computations that take in 2 array parameters and // predicate is true. XLA_TEST_F(ConditionalOpTest, Parameters2ArrayTrueBranch) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(true); auto operand1 = builder.ConstantR1({24.0f, 56.0f}); auto operand2 = builder.ConstantR1({10.0f, 11.0f}); auto operands = builder.Tuple({operand1, operand2}); - auto result = - builder.Conditional(pred, operands, CreateR1TupleAddComputation(), - operands, CreateR1TupleSubComputation()); + builder.Conditional(pred, operands, CreateR1TupleAddComputation(), operands, + CreateR1TupleSubComputation()); ComputeAndCompareR1(&builder, {34.0f, 67.0f}, {}, error_spec_); } @@ -326,21 +323,20 @@ XLA_TEST_F(ConditionalOpTest, Parameters2ArrayTrueBranch) { // Test true and false computations that take in 2 array parameters and // predicate is false. XLA_TEST_F(ConditionalOpTest, Parameters2ArrayFalseBranch) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(false); auto operand1 = builder.ConstantR1({24.0f, 56.0f}); auto operand2 = builder.ConstantR1({10.0f, 11.0f}); auto operands = builder.Tuple({operand1, operand2}); - auto result = - builder.Conditional(pred, operands, CreateR1TupleAddComputation(), - operands, CreateR1TupleSubComputation()); + builder.Conditional(pred, operands, CreateR1TupleAddComputation(), operands, + CreateR1TupleSubComputation()); ComputeAndCompareR1(&builder, {14.0f, 45.0f}, {}, error_spec_); } // Test true and false computations that return a tuple of scalars. XLA_TEST_F(ConditionalOpTest, ReturnTupleOfScalars) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(false); auto operands = builder.Tuple( {builder.ConstantR0(12.2f), builder.ConstantR0(25.6f)}); @@ -356,7 +352,7 @@ XLA_TEST_F(ConditionalOpTest, ReturnTupleOfScalars) { // Test true and false computations that return a tuple of arrays. XLA_TEST_F(ConditionalOpTest, ReturnTupleOfArrays) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(true); auto operands = builder.Tuple({builder.ConstantR1({12.2f, 15.8f}), builder.ConstantR1({25.6f, 29.2f})}); @@ -373,7 +369,7 @@ XLA_TEST_F(ConditionalOpTest, ReturnTupleOfArrays) { // Test true and false computations that return a tuple of a predicate, a // scalar, and an array. XLA_TEST_F(ConditionalOpTest, ReturnTupleofPredicateScalarArray) { - ComputationBuilder true_builder(client_, TestName() + ".true"); + XlaBuilder true_builder(TestName() + ".true"); { true_builder.Parameter(0, empty_tuple_, "tuple"); auto true_pred = true_builder.ConstantR0(true); @@ -384,7 +380,7 @@ XLA_TEST_F(ConditionalOpTest, ReturnTupleofPredicateScalarArray) { auto true_builder_result = true_builder.Build(); EXPECT_IS_OK(true_builder_result.status()); - ComputationBuilder false_builder(client_, TestName() + ".false"); + XlaBuilder false_builder(TestName() + ".false"); { false_builder.Parameter(0, empty_tuple_, "tuple"); auto false_pred = false_builder.ConstantR0(false); @@ -395,7 +391,7 @@ XLA_TEST_F(ConditionalOpTest, ReturnTupleofPredicateScalarArray) { auto false_builder_result = false_builder.Build(); EXPECT_IS_OK(false_builder_result.status()); - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(true); auto operands = builder.Tuple({}); builder.Conditional(pred, operands, true_builder_result.ConsumeValueOrDie(), @@ -411,7 +407,7 @@ XLA_TEST_F(ConditionalOpTest, ReturnTupleofPredicateScalarArray) { // Test true and false computations that return a nested tuple. XLA_TEST_F(ConditionalOpTest, ReturnNestedTuple) { - ComputationBuilder true_builder(client_, TestName() + ".true"); + XlaBuilder true_builder(TestName() + ".true"); { true_builder.Parameter(0, empty_tuple_, "tuple"); auto true_constant1 = true_builder.ConstantR0(12.2f); @@ -424,7 +420,7 @@ XLA_TEST_F(ConditionalOpTest, ReturnNestedTuple) { auto true_builder_result = true_builder.Build(); EXPECT_IS_OK(true_builder_result.status()); - ComputationBuilder false_builder(client_, TestName() + ".false"); + XlaBuilder false_builder(TestName() + ".false"); { false_builder.Parameter(0, empty_tuple_, "tuple"); auto false_constant1 = false_builder.ConstantR0(46.6f); @@ -438,7 +434,7 @@ XLA_TEST_F(ConditionalOpTest, ReturnNestedTuple) { auto false_builder_result = false_builder.Build(); EXPECT_IS_OK(false_builder_result.status()); - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(false); auto operands = builder.Tuple({}); builder.Conditional(pred, operands, true_builder_result.ConsumeValueOrDie(), @@ -460,16 +456,16 @@ XLA_TEST_F(ConditionalOpTest, ReturnNestedTuple) { // params. XLA_TEST_F(ConditionalOpTest, ScalarOperandsFromExternalParams) { Shape r0bool = ShapeUtil::MakeShape(PRED, {}); - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); - ComputationDataHandle pred, operand1, operand2; + XlaOp pred, operand1, operand2; auto pred_arg = CreateR0Parameter(true, 0, "pred", &builder, &pred); auto operand1_param = CreateR0Parameter(56.3f, 1, "operand1", &builder, &operand1); auto operand2_param = CreateR0Parameter(12.7f, 2, "operand2", &builder, &operand2); - auto result = builder.Conditional(pred, operand1, CreateR0CeilComputation(), - operand2, CreateR0FloorComputation()); + builder.Conditional(pred, operand1, CreateR0CeilComputation(), operand2, + CreateR0FloorComputation()); ComputeAndCompareR0( &builder, 57.0f, @@ -480,16 +476,16 @@ XLA_TEST_F(ConditionalOpTest, ScalarOperandsFromExternalParams) { // Test conditional that takes in array operands in the form of external params. XLA_TEST_F(ConditionalOpTest, ArrayOperandsFromExternalParams) { Shape r0bool = ShapeUtil::MakeShape(PRED, {}); - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); - ComputationDataHandle pred, operand1, operand2; + XlaOp pred, operand1, operand2; auto pred_arg = CreateR0Parameter(false, 0, "pred", &builder, &pred); auto operand1_param = CreateR1Parameter({24.3f, 56.7f}, 1, "operand1", &builder, &operand1); auto operand2_param = CreateR1Parameter({10.2f, 11.6f}, 2, "operand2", &builder, &operand2); - auto result = builder.Conditional(pred, operand1, CreateR1CeilComputation(), - operand2, CreateR1FloorComputation()); + builder.Conditional(pred, operand1, CreateR1CeilComputation(), operand2, + CreateR1FloorComputation()); ComputeAndCompareR1( &builder, {10.0f, 11.0f}, @@ -499,7 +495,7 @@ XLA_TEST_F(ConditionalOpTest, ArrayOperandsFromExternalParams) { // Test the case where one conditional is nested within another. XLA_TEST_F(ConditionalOpTest, NestedConditionals) { - ComputationBuilder inner_builder(client_, TestName() + ".inner_conditional"); + XlaBuilder inner_builder(TestName() + ".inner_conditional"); { Shape r0bool = ShapeUtil::MakeShape(PRED, {}); Shape tuple_shape = ShapeUtil::MakeTupleShape({r0bool, r0f32_, r0f32_}); @@ -514,7 +510,7 @@ XLA_TEST_F(ConditionalOpTest, NestedConditionals) { auto inner_builder_result = inner_builder.Build(); EXPECT_IS_OK(inner_builder_result.status()); - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred1 = builder.ConstantR0(true); auto pred2 = builder.ConstantR0(false); auto operand1 = builder.ConstantR0(1.1f); @@ -529,7 +525,7 @@ XLA_TEST_F(ConditionalOpTest, NestedConditionals) { } XLA_TEST_F(ConditionalOpTest, ConditionalInNestedComputation) { - ComputationBuilder inner_builder(client_, TestName() + ".inner_conditional"); + XlaBuilder inner_builder(TestName() + ".inner_conditional"); { Shape r0bool = ShapeUtil::MakeShape(PRED, {}); Shape tuple_shape = ShapeUtil::MakeTupleShape({r0bool, r0f32_, r0f32_}); @@ -544,7 +540,7 @@ XLA_TEST_F(ConditionalOpTest, ConditionalInNestedComputation) { auto inner_builder_result = inner_builder.Build(); EXPECT_IS_OK(inner_builder_result.status()); - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred2 = builder.ConstantR0(false); auto operand1 = builder.ConstantR0(1.1f); auto operand2 = builder.ConstantR0(12.2f); @@ -556,7 +552,7 @@ XLA_TEST_F(ConditionalOpTest, ConditionalInNestedComputation) { // Test a mismatch in the shape of the true operand and true computation. XLA_TEST_F(ConditionalOpTest, ShapeMismatch) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto pred = builder.ConstantR0(true); auto operand1 = builder.ConstantR0(56.0f); auto operand2 = builder.ConstantR0(12.0f); @@ -573,27 +569,27 @@ XLA_TEST_F(ConditionalOpTest, ShapeMismatch) { XLA_TEST_F(ConditionalOpTest, SwappedInputsInSequentialConditionals) { Shape tuple_shape = ShapeUtil::MakeTupleShape({r0f32_, r0f32_}); - Computation swapper; + XlaComputation swapper; { - ComputationBuilder builder(client_, TestName() + ".swapper"); + XlaBuilder builder(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; + XlaComputation forwarder; { - ComputationBuilder builder(client_, TestName() + ".forwarder"); + XlaBuilder builder(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; + XlaComputation main; { - ComputationBuilder builder(client_, TestName() + ".main"); + XlaBuilder builder(TestName() + ".main"); auto param0 = builder.Parameter(0, tuple_shape, "mp0"); auto x = builder.GetTupleElement(param0, 0); auto y = builder.GetTupleElement(param0, 1); @@ -605,7 +601,7 @@ XLA_TEST_F(ConditionalOpTest, SwappedInputsInSequentialConditionals) { } auto test_swap = [&](float a, float b) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto x = builder.ConstantR0(a); auto y = builder.ConstantR0(b); auto tuple_operand = builder.Tuple({x, y}); diff --git a/tensorflow/compiler/xla/tests/convert_test.cc b/tensorflow/compiler/xla/tests/convert_test.cc index 0842a8918bcfec037ab0f9aa24014c7d8296cdf8..e67a30d76c2fac1240c5f14a97dea421db2fed04 100644 --- a/tensorflow/compiler/xla/tests/convert_test.cc +++ b/tensorflow/compiler/xla/tests/convert_test.cc @@ -36,7 +36,7 @@ namespace { class ConvertTest : public ClientLibraryTestBase { public: - explicit ConvertTest(perftools::gputools::Platform* platform = nullptr) + explicit ConvertTest(se::Platform* platform = nullptr) : ClientLibraryTestBase(platform) { mutable_debug_options()->add_xla_disable_hlo_passes("algsimp"); mutable_debug_options()->add_xla_disable_hlo_passes("inline"); diff --git a/tensorflow/compiler/xla/tests/convolution_test.cc b/tensorflow/compiler/xla/tests/convolution_test.cc index 72715398dea468d0000144759454c5f8d8673516..947959beb144e1509a77ad2f94b8493de46ba6f2 100644 --- a/tensorflow/compiler/xla/tests/convolution_test.cc +++ b/tensorflow/compiler/xla/tests/convolution_test.cc @@ -20,10 +20,10 @@ limitations under the License. #include "tensorflow/compiler/xla/array2d.h" #include "tensorflow/compiler/xla/array4d.h" -#include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/client/padding.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h" #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/ptr_util.h" @@ -88,12 +88,12 @@ class ForwardPassConvolution_3x3x256_256_OutputZ_Iota : public ConvolutionTest { ASSERT_EQ(2, arhs->width()); ASSERT_EQ(2, arhs->height()); - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto lhs = builder.ConstantR4FromArray4D(*alhs); auto rhs = builder.ConstantR4FromArray4D(*arhs); - auto conv = builder.Conv(lhs, rhs, {1, 1}, Padding::kValid); + builder.Conv(lhs, rhs, {1, 1}, Padding::kValid); - ComputeAndCompare(&builder, conv, {}, error_spec_); + ComputeAndCompare(&builder, {}, error_spec_); } }; @@ -106,12 +106,12 @@ template class Convolve_1x1x1x2_1x1x1x2_Valid : public ConvolutionTest { public: void RunTest() { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(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); + builder.Conv(input, filter, {1, 1}, Padding::kValid); Array4D input_data(1, 1, 1, 2); input_data.FillWithYX(Array2D({ @@ -122,7 +122,7 @@ class Convolve_1x1x1x2_1x1x1x2_Valid : public ConvolutionTest { {5.0f, 6.0f}, })); - ComputeAndCompare(&builder, conv, + ComputeAndCompare(&builder, {std::move(*Literal::CreateFromArray(input_data)), std::move(*Literal::CreateFromArray(filter_data))}, error_spec_); @@ -137,12 +137,12 @@ template class Convolve_1x1x4x4_1x1x2x2_Valid : public ConvolutionTest { public: void RunTest() { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(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); + builder.Conv(input, filter, {1, 1}, Padding::kValid); Array4D input_data(1, 1, 4, 4); input_data.FillWithYX(Array2D({ @@ -156,7 +156,7 @@ class Convolve_1x1x4x4_1x1x2x2_Valid : public ConvolutionTest { {5.0f, 6.0f}, {7.0f, 8.0f}, })); - ComputeAndCompare(&builder, conv, + ComputeAndCompare(&builder, {std::move(*Literal::CreateFromArray(input_data)), std::move(*Literal::CreateFromArray(filter_data))}, error_spec_); @@ -171,12 +171,12 @@ template class Convolve_1x1x4x4_1x1x2x2_Same : public ConvolutionTest { public: void RunTest() { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(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); + builder.Conv(input, filter, {1, 1}, Padding::kSame); Array4D input_data(1, 1, 4, 4); input_data.FillWithYX(Array2D({ @@ -191,7 +191,7 @@ class Convolve_1x1x4x4_1x1x2x2_Same : public ConvolutionTest { {7.0f, 8.0f}, })); - ComputeAndCompare(&builder, conv, + ComputeAndCompare(&builder, {std::move(*Literal::CreateFromArray(input_data)), std::move(*Literal::CreateFromArray(filter_data))}, error_spec_); @@ -207,12 +207,12 @@ template class Convolve_1x1x4x4_1x1x3x3_Same : public ConvolutionTest { public: void RunTest() { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(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); + 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}, @@ -223,7 +223,7 @@ class Convolve_1x1x4x4_1x1x3x3_Same : public ConvolutionTest { 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, + ComputeAndCompare(&builder, {std::move(*Literal::CreateFromArray(input_data)), std::move(*Literal::CreateFromArray(filter_data))}, error_spec_); @@ -234,7 +234,7 @@ 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()); + XlaBuilder builder(TestName()); { Shape input_shape = ShapeUtil::MakeShape(F32, {1, 2, 5}); Shape filter_shape = ShapeUtil::MakeShape(F32, {1, 2, 2}); @@ -264,7 +264,7 @@ template class Convolve1D_1x2x5_1x2x2_WithRHSDilation : public ConvolutionTest { public: void RunTest() { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); { Shape input_shape = ShapeUtil::MakeShapeWithType({1, 2, 5}); Shape filter_shape = ShapeUtil::MakeShapeWithType({1, 2, 2}); @@ -300,7 +300,7 @@ 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()); + XlaBuilder builder(TestName()); { Shape input_shape = ShapeUtil::MakeShape(F32, {1, 2, 5}); Shape filter_shape = ShapeUtil::MakeShape(F32, {1, 2, 2}); @@ -331,7 +331,7 @@ XLA_TEST_F(ConvolutionTest, Convolve1D_1x2x5_1x2x2_WithLHSDilation) { } XLA_TEST_F(ConvolutionTest, Convolve1D_1x2x5_1x2x2_WithLHSAndRHSDilation) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); { Shape input_shape = ShapeUtil::MakeShape(F32, {1, 2, 5}); Shape filter_shape = ShapeUtil::MakeShape(F32, {1, 2, 2}); @@ -365,7 +365,7 @@ template class Convolve1D_1x2x5_1x2x2_WithPadding : public ConvolutionTest { public: void RunTest() { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); { Shape input_shape = ShapeUtil::MakeShapeWithType({1, 2, 5}); Shape filter_shape = ShapeUtil::MakeShapeWithType({1, 2, 2}); @@ -402,7 +402,7 @@ 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()); + XlaBuilder builder(TestName()); std::vector input_dims = {1, 4, 2, 3, 3}; std::vector filter_dims = {2, 2, 2, 3, 3}; Shape input_shape = ShapeUtil::MakeShape(F32, input_dims); @@ -469,7 +469,7 @@ template class Convolve2D_1x3x3x5_3x3x5x5_Valid : public ConvolutionTest { public: void RunTest() { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::vector input_dims = {1, 3, 3, 5}; std::vector filter_dims = {3, 3, 5, 3}; Shape input_shape = ShapeUtil::MakeShapeWithType(input_dims); @@ -537,7 +537,7 @@ XLA_TEST_P(ConvolveWithAndWithoutCanonicalization, execution_options_.mutable_debug_options()->add_xla_disable_hlo_passes( "convolution-canonicalization"); } - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); Shape input_shape = ShapeUtil::MakeShape(F32, {4, 29}); Shape filter_shape = ShapeUtil::MakeShape(F32, {4, 10}); @@ -551,8 +551,7 @@ XLA_TEST_P(ConvolveWithAndWithoutCanonicalization, dnums.set_kernel_output_feature_dimension(1); dnums.set_output_batch_dimension(0); dnums.set_output_feature_dimension(1); - auto conv = builder.ConvWithGeneralDimensions(input, filter, {}, - Padding::kValid, dnums); + builder.ConvWithGeneralDimensions(input, filter, {}, Padding::kValid, dnums); Array2D param0(4, 29); param0.FillUnique(); @@ -563,7 +562,7 @@ XLA_TEST_P(ConvolveWithAndWithoutCanonicalization, Array2D expected_result(29, 10); expected_result.Fill(0); - ComputeAndCompare(&builder, conv, + ComputeAndCompare(&builder, {std::move(*Literal::CreateFromArray(param0)), std::move(*Literal::CreateFromArray(param1))}, error_spec_); @@ -587,7 +586,7 @@ class Convolve1D1WindowTestBase protected: template void TestImpl() { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); int64 input_feature = GetParam().input_feature; int64 output_feature = GetParam().output_feature; int64 batch = GetParam().batch; @@ -724,12 +723,12 @@ INSTANTIATE_TEST_CASE_P( #endif XLA_TEST_F(ConvolutionTest, Convolve_bf16_1x1x1x2_1x1x1x2_Valid) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); Shape input_shape = ShapeUtil::MakeShape(BF16, {1, 1, 1, 2}); Shape filter_shape = ShapeUtil::MakeShape(BF16, {1, 1, 1, 2}); auto input = builder.Parameter(0, input_shape, "input"); auto filter = builder.Parameter(1, filter_shape, "filter"); - auto conv = builder.Conv(input, filter, {1, 1}, Padding::kValid); + builder.Conv(input, filter, {1, 1}, Padding::kValid); Array4D input_data(1, 1, 1, 2); input_data.FillWithYX(Array2D({ @@ -740,11 +739,34 @@ XLA_TEST_F(ConvolutionTest, Convolve_bf16_1x1x1x2_1x1x1x2_Valid) { {bfloat16(5), bfloat16(6)}, })); - ComputeAndCompare(&builder, conv, + ComputeAndCompare(&builder, {std::move(*Literal::CreateFromArray(input_data)), std::move(*Literal::CreateFromArray(filter_data))}, error_spec_); } +// Check that GPU convs still work if the CudnnAlgorithmPicker pass is disabled. +// (We run this test on all platforms, because, what the heck.) +XLA_TEST_F(ConvolutionTest, NoCudnnAlgorithmPicker) { + execution_options_.mutable_debug_options()->add_xla_disable_hlo_passes( + "cudnn-convolution-algorithm-picker"); + + XlaBuilder builder(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"); + builder.Conv(input, filter, {1, 1}, Padding::kValid); + + Array4D input_data(1, 1, 1, 2); + input_data.FillIota(0); + Array4D filter_data(1, 1, 1, 2); + filter_data.FillIota(10); + + ComputeAndCompare(&builder, + {std::move(*Literal::CreateFromArray(input_data)), + std::move(*Literal::CreateFromArray(filter_data))}); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/custom_call_test.cc b/tensorflow/compiler/xla/tests/custom_call_test.cc index 2d847a66b0ae7c8f09fa0cb181a4c84ea99be5b1..b43d5c9ff5d75ee0e1b3c9ceb2bc295e631ac107 100644 --- a/tensorflow/compiler/xla/tests/custom_call_test.cc +++ b/tensorflow/compiler/xla/tests/custom_call_test.cc @@ -134,9 +134,9 @@ class CustomCallClientAPITest : public ClientLibraryTestBase {}; // When using the client API, CustomCall targets can't begin with '$' -- these // are reserved for internal use. XLA_TEST_F(CustomCallClientAPITest, IllegalCustomCallTarget) { - ComputationBuilder builder(client_, TestName()); - auto call = builder.CustomCall("$illegal", /*operands=*/{}, - ShapeUtil::MakeShape(F32, {1})); + XlaBuilder builder(TestName()); + builder.CustomCall("$illegal", /*operands=*/{}, + ShapeUtil::MakeShape(F32, {1})); StatusOr> result = Execute(&builder, /*arguments=*/{}); diff --git a/tensorflow/compiler/xla/tests/dot_operation_test.cc b/tensorflow/compiler/xla/tests/dot_operation_test.cc index 7b994a4c172cafee53ede9bfd4f30b0e0c9888d5..c4031dfee593a13af6a5db15e43ed7bc418603c5 100644 --- a/tensorflow/compiler/xla/tests/dot_operation_test.cc +++ b/tensorflow/compiler/xla/tests/dot_operation_test.cc @@ -50,6 +50,13 @@ 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) && \ + defined(XLA_BACKEND_DOES_NOT_SUPPORT_COMPLEX) +using TypesF16F32 = ::testing::Types; +using TypesF16F32F64 = ::testing::Types; +using TypesF16F32F64CF64 = + ::testing::Types; #else #error "Situation not handled yet" #endif diff --git a/tensorflow/compiler/xla/tests/dynamic_ops_test.cc b/tensorflow/compiler/xla/tests/dynamic_ops_test.cc index 5f00c34002803553b9c17b4fce0abafda7369796..021fbcedb994eb4ea86a9a2da691900a56cd45f2 100644 --- a/tensorflow/compiler/xla/tests/dynamic_ops_test.cc +++ b/tensorflow/compiler/xla/tests/dynamic_ops_test.cc @@ -35,8 +35,6 @@ limitations under the License. #include "tensorflow/core/platform/test_benchmark.h" #include "tensorflow/core/platform/types.h" -namespace se = ::perftools::gputools; - namespace xla { namespace { @@ -737,11 +735,11 @@ void BM_DynamicSlice(int num_iters) { auto start_indices_literal = Literal::CreateR1({0, 1, 2, 3}); ASSERT_IS_OK(transfer_manager->TransferLiteralToDevice( - executors[device_ordinal], *start_indices_literal, *buffer)); + executors[device_ordinal], *start_indices_literal, buffer)); std::unique_ptr executable = client - ->Compile(computation, {&buffer->on_host_shape()}, + ->Compile(computation, {&buffer.on_host_shape()}, ExecutableBuildOptions()) .ConsumeValueOrDie(); @@ -750,14 +748,14 @@ void BM_DynamicSlice(int num_iters) { options.set_allocator(&allocator); const int kWarmups = 2; for (int i = 0; i < kWarmups; ++i) { - auto result = executable->Run({buffer.get()}, options); + auto result = executable->Run({&buffer}, options); ASSERT_TRUE(result.ok()); } // Run benchmark. tensorflow::testing::StartTiming(); for (int i = 0; i < num_iters; ++i) { - auto result = executable->Run({buffer.get()}, options); + auto result = executable->Run({&buffer}, options); ASSERT_TRUE(result.ok()); } } diff --git a/tensorflow/compiler/xla/tests/fusion_test.cc b/tensorflow/compiler/xla/tests/fusion_test.cc index a292eab1d198fbf69c6dc81c780487ea46756f72..c7f64d856097350a948f6b15cefb147aae7cee18 100644 --- a/tensorflow/compiler/xla/tests/fusion_test.cc +++ b/tensorflow/compiler/xla/tests/fusion_test.cc @@ -50,8 +50,6 @@ limitations under the License. using tensorflow::gtl::ArraySlice; -namespace se = ::perftools::gputools; - namespace xla { namespace { @@ -796,19 +794,19 @@ void BM_ParallelFusion(int num_iters) { // Transfer literals to device. auto param0_literal = Literal::CreateR2F32Linspace(1.0, 2.0, param0_dim0, param0_dim1); - std::unique_ptr buffer0 = + ShapedBuffer buffer0 = client->LiteralToShapedBuffer(*param0_literal, device_ordinal) .ConsumeValueOrDie(); auto param1_literal = Literal::CreateR2F32Linspace(1.0, 2.0, param1_dim0, param1_dim1); - std::unique_ptr buffer1 = + ShapedBuffer buffer1 = client->LiteralToShapedBuffer(*param1_literal, device_ordinal) .ConsumeValueOrDie(); auto param2_literal = Literal::CreateR2F32Linspace(1.0, 2.0, param2_dim0, param2_dim1); - std::unique_ptr buffer2 = + ShapedBuffer buffer2 = client->LiteralToShapedBuffer(*param2_literal, device_ordinal) .ConsumeValueOrDie(); @@ -816,8 +814,8 @@ void BM_ParallelFusion(int num_iters) { std::unique_ptr executable = client ->Compile(computation, - {&buffer0->on_host_shape(), &buffer1->on_host_shape(), - &buffer2->on_host_shape()}, + {&buffer0.on_host_shape(), &buffer1.on_host_shape(), + &buffer2.on_host_shape()}, ExecutableBuildOptions()) .ConsumeValueOrDie(); @@ -838,8 +836,7 @@ void BM_ParallelFusion(int num_iters) { // Run some warm-up executions. const int kWarmups = 2; for (int i = 0; i < kWarmups; ++i) { - auto result = - executable->Run({buffer0.get(), buffer1.get(), buffer2.get()}, options); + auto result = executable->Run({&buffer0, &buffer1, &buffer2}, options); ASSERT_TRUE(result.ok()); } @@ -852,8 +849,7 @@ void BM_ParallelFusion(int num_iters) { tensorflow::testing::UseRealTime(); tensorflow::testing::StartTiming(); for (int i = 0; i < num_iters; ++i) { - auto result = - executable->Run({buffer0.get(), buffer1.get(), buffer2.get()}, options); + auto result = executable->Run({&buffer0, &buffer1, &buffer2}, options); ASSERT_TRUE(result.ok()); } } diff --git a/tensorflow/compiler/xla/tests/gather_operation_test.cc b/tensorflow/compiler/xla/tests/gather_operation_test.cc index 9db68ff7a6dcbd9204fb2b3a37734a9aaed35dfd..90496d55e60b4f45fc2d46b2746f94d775cf9f94 100644 --- a/tensorflow/compiler/xla/tests/gather_operation_test.cc +++ b/tensorflow/compiler/xla/tests/gather_operation_test.cc @@ -405,7 +405,7 @@ class GatherClientLibraryTest : public ClientLibraryTestBase {}; // GPU and CPU_PARALLEL. XLA_TEST_F(GatherClientLibraryTest, DISABLED_ON_CPU_PARALLEL(DISABLED_ON_GPU(Basic))) { - // We create this HLO, but using the ComputationBuilder API. + // We create this HLO, but using the XlaBuilder API. // // ENTRY main { // operand = s32[3,3] parameter(0) @@ -418,7 +418,7 @@ XLA_TEST_F(GatherClientLibraryTest, // window_bounds={1, 3} // } - ComputationBuilder builder(client_, "gather_basic"); + XlaBuilder builder("gather_basic"); Shape operand_shape = ShapeUtil::MakeShape(S32, {3, 3}); Shape indices_shape = ShapeUtil::MakeShape(S32, {2}); @@ -443,8 +443,8 @@ XLA_TEST_F(GatherClientLibraryTest, client_->GetDeviceHandles(1)); xla::ExecutionOptions execution_options = CreateDefaultExecutionOptions(); *execution_options.add_device_handles() = devices[0]; - TF_ASSERT_OK_AND_ASSIGN(Computation computation, builder.Build()); - std::vector computation_instances = { + TF_ASSERT_OK_AND_ASSIGN(XlaComputation computation, builder.Build()); + std::vector computation_instances = { {computation, {operand_arg.get(), indices_arg.get()}, execution_options, diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.cc b/tensorflow/compiler/xla/tests/hlo_test_base.cc index 21f71fc91bb84540e5347811cb4643a8aeda445c..9984aba089be89072c5a49f93df68ec805658b68 100644 --- a/tensorflow/compiler/xla/tests/hlo_test_base.cc +++ b/tensorflow/compiler/xla/tests/hlo_test_base.cc @@ -35,8 +35,6 @@ limitations under the License. #include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/types.h" -namespace se = ::perftools::gputools; - namespace xla { namespace { @@ -115,11 +113,13 @@ StatusOr> HloTestBase::Execute( return test_runner_.Execute(std::move(module), arguments); } -StatusOr> HloTestBase::ExecuteNoHloPasses( +std::unique_ptr HloTestBase::ExecuteNoHloPasses( std::unique_ptr module, tensorflow::gtl::ArraySlice arguments) { - return test_runner_.Execute(std::move(module), arguments, - /*run_hlo_passes=*/false); + return test_runner_ + .Execute(std::move(module), arguments, + /*run_hlo_passes=*/false) + .ValueOrDie(); } std::unique_ptr HloTestBase::ExecuteAndTransfer( diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.h b/tensorflow/compiler/xla/tests/hlo_test_base.h index 3e8e2360bb3a87e127920cd222803c0f7b9161f4..79fcea9403e6e2dfc989c86ce8c6609f44acd12b 100644 --- a/tensorflow/compiler/xla/tests/hlo_test_base.h +++ b/tensorflow/compiler/xla/tests/hlo_test_base.h @@ -76,8 +76,7 @@ class HloTestBase : public ::testing::Test { // If your test doesn't use interpreter as the reference backend, you can use // this constructor. Note that your test target is responsible for linking in // both needed backends. - HloTestBase(::perftools::gputools::Platform* test_platform, - ::perftools::gputools::Platform* reference_platform); + HloTestBase(se::Platform* test_platform, se::Platform* reference_platform); ~HloTestBase() override {} @@ -100,7 +99,7 @@ class HloTestBase : public ::testing::Test { // Same as above, except the module will be executed without running any HLO // passes on it. - StatusOr> ExecuteNoHloPasses( + std::unique_ptr ExecuteNoHloPasses( std::unique_ptr module, tensorflow::gtl::ArraySlice arguments); diff --git a/tensorflow/compiler/xla/tests/literal_test_util.cc b/tensorflow/compiler/xla/tests/literal_test_util.cc index 81630df34c58526b6d41492b2b4b3892a02a21c2..c28f79ae386670ca80d603a42f6629dfd30e0bc9 100644 --- a/tensorflow/compiler/xla/tests/literal_test_util.cc +++ b/tensorflow/compiler/xla/tests/literal_test_util.cc @@ -39,6 +39,11 @@ limitations under the License. namespace xla { +using ::tensorflow::strings::Appendf; +using ::tensorflow::strings::Printf; +using ::tensorflow::strings::StrAppend; +using ::tensorflow::strings::StrCat; + /* static */ ::testing::AssertionResult LiteralTestUtil::EqualShapes( const Shape& expected, const Shape& actual) { if (ShapeUtil::IsTuple(expected) != ShapeUtil::IsTuple(actual)) { @@ -173,14 +178,11 @@ template auto lhs_double = static_cast(lhs); auto rhs_double = static_cast(rhs); if (ulhs != urhs) { - return ::testing::AssertionFailure() << tensorflow::strings::Printf( + return ::testing::AssertionFailure() << Printf( "floating values are not bitwise-equal; and equality testing " "was requested: %s=%g=%a vs %s=%g=%a", - tensorflow::strings::StrCat(tensorflow::strings::Hex(ulhs)) - .c_str(), - lhs_double, lhs_double, - tensorflow::strings::StrCat(tensorflow::strings::Hex(urhs)) - .c_str(), + StrCat(tensorflow::strings::Hex(ulhs)).c_str(), lhs_double, + lhs_double, StrCat(tensorflow::strings::Hex(urhs)).c_str(), rhs_double, rhs_double); } return ::testing::AssertionSuccess(); @@ -264,9 +266,7 @@ bool ExpectLiteralsEqual(const Literal& expected, const Literal& actual, << "expected:\n" << expected.ToString() << "\n\tvs actual:\n" << actual.ToString() - << (message.empty() - ? "" - : tensorflow::strings::StrCat("\nmessage: ", message)); + << (message.empty() ? "" : StrCat("\nmessage: ", message)); } /* static */ void LiteralTestUtil::ExpectNotEqual(const Literal& expected, @@ -321,9 +321,8 @@ bool ExpectLiteralsEqual(const Literal& expected, const Literal& actual, case TUPLE: { bool tuple_match = true; for (int i = 0; i < ShapeUtil::TupleElementCount(expected.shape()); ++i) { - SCOPED_TRACE(tensorflow::strings::StrCat( - "Tuple index ", i, " in ", - ShapeUtil::HumanString(expected.shape()))); + SCOPED_TRACE(StrCat("Tuple index ", i, " in ", + ShapeUtil::HumanString(expected.shape()))); // Create LiteralViews of the expected and actual elements. auto result = Equal(LiteralView::Create(expected, {i}), @@ -350,227 +349,301 @@ bool ExpectLiteralsEqual(const Literal& expected, const Literal& actual, namespace { +// Gets the total element count. For tuples, this is not the count of tuple +// elements, but the sum of elements of each tuple element. +int64 RecursiveElementCount(const Shape& shape) { + if (ShapeUtil::IsTuple(shape)) { + const int64 tuple_elements = ShapeUtil::TupleElementCount(shape); + int64 total = 0; + for (int64 i = 0; i < tuple_elements; ++i) { + total += RecursiveElementCount(ShapeUtil::GetTupleElementShape(shape, i)); + } + return total; + } else { + return ShapeUtil::ElementsIn(shape); + } +} + +// Calling ToString on a literal with over 100 million elements takes around +// 3 minutes. The utility of printing a literal with >1000 elements is +// questionable, especially when writing the Literal proto to disk is orders +// of magnitude faster. +string TruncateHugeLiteral(const Literal& literal) { + return RecursiveElementCount(literal.shape()) < 1000 + ? literal.ToString() + : "[TRUNCATED, Literal with more than 1000 values]"; +} + +// Returns whether the actual and expected values are mismatched with respect to +// nans. 'relaxed_nans' is interpreted as in xla::ErrorSpec. +template +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 <> +bool NanMismatch(complex64 expected, complex64 actual, + bool relaxed_nans) { + return NanMismatch(expected.real(), actual.real(), relaxed_nans) || + NanMismatch(expected.imag(), actual.imag(), relaxed_nans); +} + +template <> +bool NanMismatch(half expected, half actual, bool relaxed_nans) { + return NanMismatch(static_cast(expected), + static_cast(actual), relaxed_nans); +} + +// Converts the given floating-point value to a string. +template +string FpValueToString(NativeT value) { + return Printf("%8.4g", static_cast(value)); +} + +template <> +string FpValueToString(complex64 value) { + return Printf("%8.4g + %8.4fi", value.real(), value.imag()); +} + +// Returns the absolute value of the given floating point value. This function +// is used instead of std::abs directly in order to allow type-dependent +// implementations for NearComparator. +template +float FpAbsoluteValue(NativeT value) { + return std::abs(value); +} + +template <> +float FpAbsoluteValue(bfloat16 value) { + return FpAbsoluteValue(static_cast(value)); +} + +template <> +float FpAbsoluteValue(half value) { + return FpAbsoluteValue(static_cast(value)); +} + // Helper class for comparing floating-point literals within an error bound. +template class NearComparator { public: - explicit NearComparator(ErrorSpec error) : error_(error) {} + // Compares the two array literals elementwise and returns an assertion + // result. The assertion result is successful if all actual and expected + // elements are within the given error bound. In case of error, the assertion + // result contains a detailed error message in case of failure. + static ::testing::AssertionResult Compare(const Literal& expected, + const Literal& actual, + ErrorSpec error, + bool detailed_message) { + NearComparator comparator(expected, actual, error, + detailed_message); + return comparator.Run(); + } + + private: + // Data structure encapsulating metadata about a single element mismatch. + struct Mismatch { + NativeT actual; + NativeT expected; + float rel_error; + float abs_error; + + // The linear index of the failure within the shape. This linear index is + // from the 'actual' literal. + int64 linear_index; + + bool operator<(const Mismatch& other) const { + return rel_error < other.rel_error; + } - // Compares the two literals elementwise. EXPECTs each pair of elements to be - // within the error bound. Emits useful log messages and dumps literals to - // temporary files on failure. Returns true if literals match. - bool ExpectNear(const Literal& expected, const Literal& actual) { + string ToString(const Shape& shape) const { + return Printf( + "actual %s, expected %s, index %s, rel error %8.3g, abs error %8.3g", + FpValueToString(actual).c_str(), FpValueToString(expected).c_str(), + LiteralTestUtil::MultiIndexAsString( + IndexUtil::LinearIndexToMultidimensionalIndex(shape, + linear_index)) + .c_str(), + rel_error, abs_error); + } + }; + + explicit NearComparator(const Literal& expected, const Literal& actual, + ErrorSpec error, bool detailed_message) + : expected_(expected), + actual_(actual), + error_(error), + detailed_message_(detailed_message), + abs_value_buckets_(kAbsValueBucketBounds.size() - 1, {0, 0}), + abs_error_buckets_(kErrorBucketBounds.size(), 0), + rel_error_buckets_(kErrorBucketBounds.size(), 0) {} + + // Runs the comparison between expected and actual literals. + ::testing::AssertionResult Run() { VLOG(1) << "expected:"; - XLA_VLOG_LINES(1, TruncateHugeLiteral(expected)); + XLA_VLOG_LINES(1, TruncateHugeLiteral(expected_)); VLOG(1) << "actual:"; - XLA_VLOG_LINES(1, TruncateHugeLiteral(actual)); + XLA_VLOG_LINES(1, TruncateHugeLiteral(actual_)); // If the shapes mismatch, we simply fail the expectation instead of // printing out data, as it's a type error rather than a value error. ::testing::AssertionResult equal_shapes = - LiteralTestUtil::EqualShapes(expected.shape(), actual.shape()); + LiteralTestUtil::EqualShapes(expected_.shape(), actual_.shape()); if (!equal_shapes) { - EXPECT_TRUE(equal_shapes); - return false; + return equal_shapes; } - - // Set up members used during the comparison. - num_miscompares_ = 0; - abs_diff_sum_ = 0.0; - abs_expected_sum_ = 0.0; - abs_diff_miscompare_sum_ = 0.0; - 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); - - switch (expected.shape().element_type()) { - case BF16: - ExpectLiteralsNear(expected, actual, 0); - break; - case F16: - ExpectLiteralsNear(expected, actual, 0); - break; - case F32: - ExpectLiteralsNear(expected, actual, 0); - break; - case F64: - ExpectLiteralsNear(expected, actual, 0); - break; - case C64: - ExpectLiteralsNear(expected, actual, 0); - break; - default: - LOG(FATAL) << "Unsupported primitive type in near comparator: " - << PrimitiveType_Name(expected.shape().element_type()) - << ". Must be floating-point type."; + if (!ShapeUtil::IsArray(expected_.shape())) { + return ::testing::AssertionFailure() << "Expected array shape"; } - if (num_miscompares_ > 0) { - if (!VLOG_IS_ON(1)) { - LOG(INFO) << "expected: " << ShapeUtil::HumanString(expected.shape()) - << " " << TruncateHugeLiteral(expected); - LOG(INFO) << "actual: " << ShapeUtil::HumanString(actual.shape()) - << " " << TruncateHugeLiteral(actual); - LOG(INFO) << "Dumping literals to temp files..."; - WriteLiteralToTempFile(expected, "expected"); - WriteLiteralToTempFile(actual, "actual"); - WriteLiteralToTempFile(miscompares_, "miscompares"); - } - EXPECT_TRUE(num_miscompares_ == 0) - << "\nmax relative mismatch at index " - << LiteralTestUtil::MultiIndexAsString( - IndexUtil::LinearIndexToMultidimensionalIndex( - actual.shape(), max_rel_linear_index_)) - << "\nmaximum relative error " << max_rel_err_ - << "\nmax absolute mismatch at index " - << LiteralTestUtil::MultiIndexAsString( - IndexUtil::LinearIndexToMultidimensionalIndex( - actual.shape(), max_abs_linear_index_)) - << "\nmaximum absolute error " << max_abs_err_ - << "\nfirst mismatch at index " - << LiteralTestUtil::MultiIndexAsString( - IndexUtil::LinearIndexToMultidimensionalIndex( - actual.shape(), first_linear_index_)) - << "\nlast mismatch at index " - << LiteralTestUtil::MultiIndexAsString( - IndexUtil::LinearIndexToMultidimensionalIndex( - actual.shape(), last_linear_index_)) - << "\ntotal absolute error " << abs_diff_sum_ - << "\ntotal absolute error of miscompares " - << abs_diff_miscompare_sum_ << "\ntotal relative error " - << (abs_diff_sum_ / abs_expected_sum_) - << "\ntotal relative error of miscompares " - << (abs_diff_miscompare_sum_ / abs_expected_miscompare_sum_) - << "\nfailure count " << num_miscompares_; + mismatches_ = Literal(ShapeUtil::ChangeElementType(actual_.shape(), PRED)); + mismatches_.PopulateWithValue(false); + + CompareLiterals(); + + if (num_mismatches_ == 0) { + return ::testing::AssertionSuccess(); + } else if (!VLOG_IS_ON(1)) { + LOG(INFO) << "expected: " << ShapeUtil::HumanString(expected_.shape()) + << " " << TruncateHugeLiteral(expected_); + LOG(INFO) << "actual: " << ShapeUtil::HumanString(actual_.shape()) + << " " << TruncateHugeLiteral(actual_); + LOG(INFO) << "Dumping literals to temp files..."; + WriteLiteralToTempFile(expected_, "expected"); + WriteLiteralToTempFile(actual_, "actual"); + WriteLiteralToTempFile(mismatches_, "mismatches"); } - return num_miscompares_ == 0; + return ::testing::AssertionFailure() << ErrorMessage(); } - private: - template - 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); + // Insert the given absolute value into the absolute value bucket vector. The + // bounds of the buckets are given by kAbsValueBucketBounds. + void UpdateAbsValueBucket(NativeT value, bool is_mismatch) { + // Adjust the bucket containing the absolute values of the 'actual' + // elements. + const float abs_value = FpAbsoluteValue(value); + for (int i = 0; i < abs_value_buckets_.size(); ++i) { + if (i == abs_value_buckets_.size() - 1 || + (abs_value >= kAbsValueBucketBounds[i] && + abs_value < kAbsValueBucketBounds[i + 1])) { + // The first value of the pair is the count of elements in the bucket, + // the second is the count of mismatches in the bucket. + abs_value_buckets_[i].first++; + if (is_mismatch) { + abs_value_buckets_[i].second++; + } + return; + } } } - template - void ExpectNear(NativeT expected, NativeT actual, - const ::testing::Message& message) { - EXPECT_NEAR(expected, actual, error_.abs) - << "expected:\n " << expected << "\n\tvs actual:\n " << actual << "\n" - << message; - } - - // EXPECTs that the two given scalar values are within the error bound. Keeps - // track of how many mismatches have occurred to keep the size of the output - // manageable. - template - bool ExpectValuesNear(NativeT expected, NativeT actual) { - if (expected == actual) { - return true; + // Insert the given error into the given error bucket vector. + void UpdateErrorBucket( + float error, tensorflow::gtl::MutableArraySlice error_buckets) { + CHECK_EQ(error_buckets.size(), kErrorBucketBounds.size()); + for (int i = 0; i < error_buckets.size(); ++i) { + if (error >= kErrorBucketBounds[i]) { + error_buckets[i]++; + } } - - const float abs_diff = std::abs(actual - expected); - const float rel_err = abs_diff / std::abs(expected); - const bool nan_mismatch = - NanMismatch(expected, actual, error_.relaxed_nans); - const bool mismatch = - (nan_mismatch || (abs_diff >= error_.abs && rel_err >= error_.rel)); - return !mismatch; } - // Assumes that expected vs actual fail ExpectValuesNear. - template - void UpdateAndLogMiscompares(const NativeT expected, const NativeT actual, - const Shape& shape, const int64 linear_index) { - const float abs_diff = std::abs(actual - expected); - const float rel_err = abs_diff / std::abs(expected); - abs_diff_sum_ += abs_diff; - abs_expected_sum_ += std::abs(expected); - if (rel_err > max_rel_err_ || std::isnan(rel_err)) { - max_rel_err_ = rel_err; - max_rel_linear_index_ = linear_index; + // Compares the two given elements from the expected and actual literals at + // the given literal_index and keeps track of various mismatch statistics. + void CompareValues(NativeT expected, NativeT actual, int64 linear_index) { + const bool is_nan_mismatch = + NanMismatch(expected, actual, error_.relaxed_nans); + float abs_error; + float rel_error; + if (actual == expected) { + abs_error = 0; + rel_error = 0; + } else if (is_nan_mismatch) { + num_nan_mismatches_++; + // A nan mismatch is considered to have infinite error. rel_error is used + // for sorting a std::set of the top mismatchs, and a nan value here will + // result in undefined behavior because nan's do not satisfy the strict + // weak ordering requirement of std containers. + abs_error = std::numeric_limits::infinity(); + rel_error = std::numeric_limits::infinity(); + } else { + abs_error = FpAbsoluteValue(actual - expected); + rel_error = abs_error / FpAbsoluteValue(expected); } - if (abs_diff > max_abs_err_ || std::isnan(abs_diff)) { - max_abs_err_ = abs_diff; - max_abs_linear_index_ = linear_index; + const bool is_abs_mismatch = abs_error > error_.abs; + const bool is_rel_mismatch = rel_error > error_.rel; + const bool is_mismatch = + is_nan_mismatch || (is_abs_mismatch && is_rel_mismatch); + + // Update the error of the relative bucket only if the *absolute* error + // bound is exceeded and vice versa. + if (is_abs_mismatch) { + num_abs_mismatches_++; + UpdateErrorBucket(rel_error, &rel_error_buckets_); } - if (VLOG_IS_ON(10)) { - VLOG(10) << tensorflow::strings::Printf( - "index %s abs_diff %f rel_err %f", - LiteralTestUtil::MultiIndexAsString( - IndexUtil::LinearIndexToMultidimensionalIndex(shape, - linear_index)) - .c_str(), - abs_diff, rel_err); + if (is_rel_mismatch) { + num_rel_mismatches_++; + UpdateErrorBucket(abs_error, &abs_error_buckets_); } - abs_diff_miscompare_sum_ += abs_diff; - abs_expected_miscompare_sum_ += std::abs(expected); - const int64 kMaxFailures = 2; - if (num_miscompares_ < kMaxFailures) { - const auto multi_index = - IndexUtil::LinearIndexToMultidimensionalIndex(shape, linear_index); - ::testing::Message msg; - msg << "mismatch at index " - << LiteralTestUtil::MultiIndexAsString(multi_index) << " abs diff " - << abs_diff << " rel err " << rel_err << " failure #" - << num_miscompares_; - ExpectNear(expected, actual, msg); - } else if (num_miscompares_ == kMaxFailures) { - LOG(ERROR) << "reached max 'loud' failure count; silently proceeding..."; + + UpdateAbsValueBucket(actual, is_mismatch); + + if (!is_mismatch) { + return; } - if (num_miscompares_ == 0) { - first_linear_index_ = linear_index; + + num_mismatches_++; + + // Keep track of the kTopRelativeErrorCount relative error mismatches. + if (top_rel_mismatches_.size() < kTopRelativeErrorCount || + rel_error > top_rel_mismatches_.begin()->rel_error) { + Mismatch mismatch = {actual, expected, rel_error, abs_error, + linear_index}; + top_rel_mismatches_.insert(mismatch); + if (top_rel_mismatches_.size() > kTopRelativeErrorCount) { + top_rel_mismatches_.erase(top_rel_mismatches_.begin()); + } } - num_miscompares_++; - last_linear_index_ = linear_index; - miscompares_.data()[linear_index] = true; + + mismatches_.data()[linear_index] = true; } - // Recursive function which compares the two given literals elementwise. - template - void ExpectLiteralsNear(const Literal& expected, const Literal& actual, - int64 dimension) { + // Compares the two literals elementwise. + void CompareLiterals() { // Fast path optimization for the case were layouts match. - if (LayoutUtil::Equal(actual.shape().layout(), expected.shape().layout())) { + if (LayoutUtil::Equal(actual_.shape().layout(), + expected_.shape().layout())) { tensorflow::gtl::ArraySlice expected_data = - expected.data(); + expected_.data(); tensorflow::gtl::ArraySlice actual_data = - actual.data(); + actual_.data(); const int64 len = expected_data.size(); for (int64 i = 0; i < len; ++i) { - const bool near = ExpectValuesNear(expected_data[i], actual_data[i]); - if (!near) { - UpdateAndLogMiscompares(expected_data[i], actual_data[i], - actual.shape(), i); - } + CompareValues(expected_data[i], actual_data[i], i); } return; } + std::vector multi_index(ShapeUtil::Rank(actual_.shape()), 0); + CompareLiteralsSlow(0, &multi_index); + } - if (dimension == expected.shape().dimensions_size()) { - bool near = ExpectValuesNear(expected.Get(multi_index_), - actual.Get(multi_index_)); - if (!near) { - UpdateAndLogMiscompares( - expected.Get(multi_index_), - actual.Get(multi_index_), actual.shape(), - IndexUtil::MultidimensionalIndexToLinearIndex(actual.shape(), - multi_index_)); - } + // Slow path for CompareLiterals when 'actual' and 'expected' literals have + // different layouts. In this case, multidimensional indices are constructed + // and indexed for each element. + void CompareLiteralsSlow(int64 dimension, std::vector* multi_index) { + if (dimension == multi_index->size()) { + CompareValues(expected_.Get(*multi_index), + actual_.Get(*multi_index), + IndexUtil::MultidimensionalIndexToLinearIndex( + actual_.shape(), *multi_index)); } else { - for (int64 i = 0; i < expected.shape().dimensions(dimension); ++i) { - multi_index_[dimension] = i; - ExpectLiteralsNear(expected, actual, dimension + 1); + for (int64 i = 0; i < expected_.shape().dimensions(dimension); ++i) { + (*multi_index)[dimension] = i; + CompareLiteralsSlow(dimension + 1, multi_index); } } } @@ -580,159 +653,247 @@ class NearComparator { int64 now_usec = tensorflow::Env::Default()->NowMicros(); string filename = tensorflow::io::JoinPath( tensorflow::testing::TmpDir(), - tensorflow::strings::Printf("tempfile-%s-%llx-%s", Hostname().c_str(), - now_usec, name.c_str())); + Printf("tempfile-%s-%llx-%s", Hostname().c_str(), now_usec, + name.c_str())); TF_CHECK_OK(tensorflow::WriteBinaryProto(tensorflow::Env::Default(), filename, literal.ToProto())); LOG(ERROR) << "wrote to " << name << " file: " << filename; } - // Gets the total element count. For tuples, this is not the count of tuple - // elements, but the sum of elements of each tuple element. - int64 RecursiveElementCount(const Shape& shape) { - if (ShapeUtil::IsTuple(shape)) { - const int64 tuple_elements = ShapeUtil::TupleElementCount(shape); - int64 total = 0; - for (int64 i = 0; i < tuple_elements; ++i) { - total += - RecursiveElementCount(ShapeUtil::GetTupleElementShape(shape, i)); - } - return total; - } else { - return ShapeUtil::ElementsIn(shape); + // Returns an error message string with a detailed breakdown of the + // mismatches. Called after calling Run(). + string ErrorMessage() { + string out; + int64 element_count = ShapeUtil::ElementsIn(actual_.shape()); + + auto percent_string = [](float a, float b) { + float pct = b == 0.0 ? 0.0 : 100.0 * a / b; + return Printf("%0.4f%%", pct); + }; + + Appendf(&out, + "\nMismatch count %lld (%s) in shape %s (%lld elements), abs bound " + "%g, rel bound %g\n", + num_mismatches_, + percent_string(num_mismatches_, element_count).c_str(), + ShapeUtil::HumanString(actual_.shape()).c_str(), + ShapeUtil::ElementsIn(actual_.shape()), error_.abs, error_.rel); + if (num_nan_mismatches_ > 0) { + StrAppend(&out, "nan mismatches ", num_nan_mismatches_, "\n"); + } + Appendf(&out, "Top relative error mismatches:\n"); + for (auto it = top_rel_mismatches_.rbegin(); + it != top_rel_mismatches_.rend(); ++it) { + StrAppend(&out, " ", it->ToString(actual_.shape()).c_str(), "\n"); } - } - // Calling ToString on a literal with over 100 million elements takes around - // 3 minutes. The utility of printing a literal with >1000 elements is - // questionable, especially when writing the Literal proto to disk is orders - // of magnitude faster. - string TruncateHugeLiteral(const Literal& literal) { - return RecursiveElementCount(literal.shape()) < 1000 - ? literal.ToString() - : "[TRUNCATED, Literal with more than 1000 values]"; - } + if (!detailed_message_) { + return out; + } - ErrorSpec error_; + StrAppend(&out, "Absolute magnitude breakdown of actual values:\n"); + CHECK_EQ(abs_value_buckets_.size() + 1, kAbsValueBucketBounds.size()); + for (int i = 0; i < abs_value_buckets_.size(); ++i) { + const int64 bucket_size = abs_value_buckets_[i].first; + const int64 bucket_mismatches = abs_value_buckets_[i].second; + string mismatch_str = bucket_mismatches > 0 + ? Printf(", mismatches %lld", bucket_mismatches) + : ""; + Appendf(&out, " %-6g <= x < %-6g : %7lld (%9s)%s\n", + kAbsValueBucketBounds[i], kAbsValueBucketBounds[i + 1], + bucket_size, percent_string(bucket_size, element_count).c_str(), + mismatch_str.c_str()); + } - // Number of element miscomparisons encountered so far. - int64 num_miscompares_; + auto print_accum_buckets = [&](const string& header, int64 total, + tensorflow::gtl::ArraySlice buckets) { + StrAppend(&out, header, ":\n"); + Appendf(&out, " < %-6g : %7lld (%s)\n", kErrorBucketBounds[0], + total - buckets[0], + percent_string(total - buckets[0], total).c_str()); + CHECK_EQ(buckets.size(), kErrorBucketBounds.size()); + for (int i = 0; i < kErrorBucketBounds.size(); ++i) { + Appendf(&out, " >= %-6g : %7lld (%s)\n", kErrorBucketBounds[i], + buckets[i], percent_string(buckets[i], total).c_str()); + } + }; + Appendf(&out, "Elements exceeding abs error bound %g: %lld (%s)\n", + error_.abs, num_abs_mismatches_, + percent_string(num_abs_mismatches_, element_count).c_str()); + print_accum_buckets( + "Relative error breakdown of elements exceeding abs error bound", + num_abs_mismatches_, rel_error_buckets_); + Appendf(&out, "Elements exceeding rel error bound %g: %lld (%s)\n", + error_.rel, num_rel_mismatches_, + percent_string(num_rel_mismatches_, element_count).c_str()); + print_accum_buckets( + "Absolute error breakdown of elements exceeding rel error bound", + num_rel_mismatches_, abs_error_buckets_); + return out; + } - // A Literal containing which elements did not match in the expected and - // actual literals. miscompares_ contains PREDs and is of the same sizes as - // the comparison literals. - Literal miscompares_; - - // A multidimensional index used when performing the recursive comparison. - std::vector multi_index_; - - // Aggregated Statistics on input. - double abs_diff_sum_; - double abs_expected_sum_; - double abs_diff_miscompare_sum_; - double abs_expected_miscompare_sum_; - float max_rel_err_; - float max_abs_err_; - int64 first_linear_index_; - int64 last_linear_index_; - int64 max_rel_linear_index_; - int64 max_abs_linear_index_; -}; + // 'actual' and 'expected' literals being compared. + const Literal& expected_; + const Literal& actual_; -template <> -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); -} + // The error bounds of the comparison. + ErrorSpec error_; -template <> -void NearComparator::ExpectNear(complex64 expected, complex64 actual, - const ::testing::Message& message) { - EXPECT_NEAR(expected.real(), actual.real(), error_.abs) - << "expected:\n " << expected << "\n\tvs actual:\n " << actual << "\n" - << message; - EXPECT_NEAR(expected.imag(), actual.imag(), error_.abs) - << "expected:\n " << expected << "\n\tvs actual:\n " << actual << "\n" - << message; -} + // Whether to include detailed breakdown of mismatches in the error message. + bool detailed_message_; -template <> -bool NearComparator::ExpectValuesNear(bfloat16 expected, - bfloat16 actual) { - return ExpectValuesNear(static_cast(expected), - static_cast(actual)); -} + // Number of element element mismatches encountered so far. + int64 num_mismatches_ = 0; -template <> -bool NearComparator::ExpectValuesNear(half expected, half actual) { - return ExpectValuesNear(static_cast(std::move(expected)), - static_cast(std::move(actual))); -} + // Number of elements with a nan mismatch. + int64 num_nan_mismatches_ = 0; -template <> -void NearComparator::UpdateAndLogMiscompares( - const bfloat16 expected, const bfloat16 actual, const Shape& shape, - const int64 linear_index) { - UpdateAndLogMiscompares(static_cast(expected), - static_cast(actual), shape, linear_index); -} + // Number of elements which exceed the absolute/relative error bound. + int64 num_abs_mismatches_ = 0; + int64 num_rel_mismatches_ = 0; -template <> -void NearComparator::UpdateAndLogMiscompares(half expected, half actual, - const Shape& shape, - const int64 linear_index) { - UpdateAndLogMiscompares(static_cast(std::move(expected)), - static_cast(std::move(actual)), shape, - linear_index); -} - -} // namespace + // A Literal containing which elements did not match in the expected and + // actual literals. mismatches_ contains PREDs and is of the same sizes as + // the comparison literals. + Literal mismatches_; + + // The number of mismatches to report in the output, sorted by relative error + // magnitude. + static constexpr int64 kTopRelativeErrorCount = 5; + + // The set of mismatches with the largest relative error. The size of this set + // is bounded by kTopRelativeErrorCount. + std::multiset top_rel_mismatches_; + + // Actual values are bucketed by absolute value. kAbsValueBucketBounds is the + // bounds of these buckets. abs_value_buckets_ contains a pair for each + // bucket: the element count and failure count. + static constexpr std::array kAbsValueBucketBounds = { + 0.0, 0.0001, 0.001, 0.01, 0.1, 1, std::numeric_limits::infinity()}; + std::vector> abs_value_buckets_; + + // Buckets for relative and absolute errors. The relative error buckets only + // contains those elements which exceed the *absolute* error bound, and vice + // versa. This makes it easy to see the effect of adjusting the relative (or + // absolute) error bound on the success of the comparison. kErrorBucketBounds + // are the lower bounds of the buckets in both vectors. The error buckets are + // a cumulative distribution so an error value may appear in more than one + // bucket. For example an error value of 0.003 may appear in the buckets + // bounded by 0.01, 0.1, and 1.0. + static constexpr std::array kErrorBucketBounds = {0.0001, 0.001, + 0.01, 0.1, 1}; + std::vector abs_error_buckets_; + std::vector rel_error_buckets_; +}; -/* static */ ::testing::AssertionResult LiteralTestUtil::Near( - const Literal& expected, const Literal& actual, const ErrorSpec& error) { +template +constexpr std::array NearComparator::kAbsValueBucketBounds; +template +constexpr std::array NearComparator::kErrorBucketBounds; + +// Helper function for comparing two literals for nearness. Handles tuple-shapes +// via recursion. shape_index is the ShapeIndex of expected (or actual) +// currently being compared. +::testing::AssertionResult NearHelper(const Literal& expected, + const Literal& actual, + const ErrorSpec& error, + bool detailed_message, + const ShapeIndex& shape_index) { ::testing::AssertionResult err = - EqualShapes(expected.shape(), actual.shape()); + LiteralTestUtil::EqualShapes(expected.shape(), actual.shape()); if (!err) { return err; } if (ShapeUtil::IsTuple(expected.shape())) { for (int64 i = 0; i < ShapeUtil::TupleElementCount(expected.shape()); ++i) { - SCOPED_TRACE(tensorflow::strings::StrCat( - "Tuple index ", i, " in ", ShapeUtil::HumanString(expected.shape()))); const auto expected_element = LiteralView::Create(expected, {i}); const auto actual_element = LiteralView::Create(actual, {i}); - + ShapeIndex element_index = shape_index; + element_index.push_back(i); ::testing::AssertionResult res = - Near(expected_element, actual_element, error); - if (err && !res) { - err = res; + NearHelper(expected_element, actual_element, error, detailed_message, + element_index); + if (!res) { + string err_message = + Printf("\nArray at shape index %s%s", + element_index.ToString().c_str(), res.message()); + if (err) { + err = ::testing::AssertionFailure() << err_message; + } else { + err << err_message; + } } } + if (!err && shape_index.empty()) { + // Emit a top-level error message containing the top-level shape in case + // of mismatch. + int64 total_elements = RecursiveElementCount(actual.shape()); + err = ::testing::AssertionFailure() + << Printf("\nMismatches in shape %s (%lld elements):\n%s", + ShapeUtil::HumanString(actual.shape()).c_str(), + total_elements, err.message()); + } return err; } if (ShapeUtil::ElementIsFloating(expected.shape()) || ShapeUtil::ElementIsComplex(expected.shape())) { - NearComparator comparator(error); - return comparator.ExpectNear(expected, actual) - ? ::testing::AssertionSuccess() - : ::testing::AssertionFailure() << "values were not near"; + switch (expected.shape().element_type()) { + case BF16: + return NearComparator::Compare(expected, actual, error, + detailed_message); + break; + case F16: + return NearComparator::Compare(expected, actual, error, + detailed_message); + break; + case F32: + return NearComparator::Compare(expected, actual, error, + detailed_message); + break; + case F64: + return NearComparator::Compare(expected, actual, error, + detailed_message); + break; + case C64: + return NearComparator::Compare(expected, actual, error, + detailed_message); + break; + default: + LOG(FATAL) << "Unsupported primitive type in near comparator: " + << PrimitiveType_Name(expected.shape().element_type()) + << ". Must be floating-point type."; + } } - return Equal(expected, actual); + // Non-floating point literal. + return LiteralTestUtil::Equal(expected, actual); +} + +} // namespace + +/* static */ ::testing::AssertionResult LiteralTestUtil::Near( + const Literal& expected, const Literal& actual, const ErrorSpec& error, + bool detailed_message) { + return NearHelper(expected, actual, error, detailed_message, + /*shape_index=*/{}); } /* static */ void LiteralTestUtil::ExpectNear(const Literal& expected, const Literal& actual, const ErrorSpec& error, const string& message) { - EXPECT_TRUE(Near(expected, actual, error)) - << (message.empty() - ? "" - : tensorflow::strings::StrCat("\nmessage: ", message)); + ::testing::AssertionResult res = + Near(expected, actual, error, /*detailed_message=*/false); + if (!res) { + res << "Expected: " << TruncateHugeLiteral(expected) << "\n"; + res << "Actual: " << TruncateHugeLiteral(actual) << "\n"; + if (!message.empty()) { + res << StrCat("\nmessage: ", message); + } + } + EXPECT_TRUE(res); } /*static*/ ::testing::AssertionResult LiteralTestUtil::NearOrEqual( @@ -754,8 +915,7 @@ void NearComparator::UpdateAndLogMiscompares(half expected, half actual, /* static */ string LiteralTestUtil::MultiIndexAsString( tensorflow::gtl::ArraySlice multi_index) { - return tensorflow::strings::StrCat( - "{", tensorflow::str_util::Join(multi_index, ","), "}"); + return StrCat("{", tensorflow::str_util::Join(multi_index, ","), "}"); } /* static */ std::unique_ptr LiteralTestUtil::Reshape( diff --git a/tensorflow/compiler/xla/tests/literal_test_util.h b/tensorflow/compiler/xla/tests/literal_test_util.h index 7b757a4bd7e7592583b7596b4305ddb7e6c52d75..a755568c0f098e15512bd1d3720269c867bc9c49 100644 --- a/tensorflow/compiler/xla/tests/literal_test_util.h +++ b/tensorflow/compiler/xla/tests/literal_test_util.h @@ -122,16 +122,19 @@ class LiteralTestUtil { // bounds are equivalent. // // Tuples are matched recursively. When comparing tensors of - // non-floating-point type, checks for exact equality, ignoring the ErroSpec. + // non-floating-point type, checks for exact equality, ignoring the ErrorSpec. // // If the shape of the literals is neither a complex/floating-point tensor nor // a tuple which contains a complex/floating-point tensor, Near() is // equivalent to Equal(). We don't raise an error in this case, because we // want to allow callers to call Near() even if they have no preconceptions // about the shapes being compared. + // + // If detailed_message is true, then the error message in the assertion result + // will contain a more detailed breakdown of mismatches. static ::testing::AssertionResult Near( - const Literal& expected, const Literal& actual, - const ErrorSpec& error) TF_MUST_USE_RESULT; + const Literal& expected, const Literal& actual, const ErrorSpec& error, + bool detailed_message = false) TF_MUST_USE_RESULT; // Expects expected and actual to be Near with the given error. static void ExpectNear(const Literal& expected, const Literal& actual, diff --git a/tensorflow/compiler/xla/tests/literal_test_util_test.cc b/tensorflow/compiler/xla/tests/literal_test_util_test.cc index 3a421f8458268a14dcdd84889bcae4990c095ea4..9d619a77c7e8d6398b559e8f562cd7f8194e0811 100644 --- a/tensorflow/compiler/xla/tests/literal_test_util_test.cc +++ b/tensorflow/compiler/xla/tests/literal_test_util_test.cc @@ -89,7 +89,7 @@ TEST(LiteralTestUtilTest, ExpectNearFailurePlacesResultsInTemporaryDirectory) { EXPECT_EQ("2", literal->ToString()); } else if (result.find("actual") != string::npos) { EXPECT_EQ("4", literal->ToString()); - } else if (result.find("miscompares") != string::npos) { + } else if (result.find("mismatches") != string::npos) { EXPECT_EQ("true", literal->ToString()); } else { FAIL() << "unknown file in temporary directory: " << result; diff --git a/tensorflow/compiler/xla/tests/llvm_compiler_test.cc b/tensorflow/compiler/xla/tests/llvm_compiler_test.cc index 7e92439c494b677f718a63c71c20828d65bebef4..2f46ee0be216d7dabf1c476d3cfb7d528f8ab6a4 100644 --- a/tensorflow/compiler/xla/tests/llvm_compiler_test.cc +++ b/tensorflow/compiler/xla/tests/llvm_compiler_test.cc @@ -43,7 +43,7 @@ class LLVMCompilerTest : public ::testing::Test { ~LLVMCompilerTest() override {} protected: - using Platform = ::perftools::gputools::Platform; + using Platform = se::Platform; explicit LLVMCompilerTest(string platform_name) : platform_name_(std::move(platform_name)) {} @@ -95,7 +95,7 @@ class LLVMCompilerTest : public ::testing::Test { modules.push_back(hlo_module->Clone()); modules.push_back(std::move(hlo_module)); - std::vector> executors; + std::vector> executors; executors.push_back({backend_->default_stream_executor()}); executors.push_back({backend_->default_stream_executor()}); diff --git a/tensorflow/compiler/xla/tests/local_client_allocation_test.cc b/tensorflow/compiler/xla/tests/local_client_allocation_test.cc index 3d30ceeaf1b0369b6fdc0cd9620c04aae287941c..7209f91639b5f24e7da569135c3b9444ed522d6e 100644 --- a/tensorflow/compiler/xla/tests/local_client_allocation_test.cc +++ b/tensorflow/compiler/xla/tests/local_client_allocation_test.cc @@ -25,6 +25,7 @@ limitations under the License. #include "tensorflow/compiler/xla/tests/literal_test_util.h" #include "tensorflow/compiler/xla/tests/local_client_test_base.h" #include "tensorflow/compiler/xla/tests/test_macros.h" +#include "tensorflow/core/lib/gtl/optional.h" #include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/types.h" @@ -53,7 +54,7 @@ XLA_TEST_F(LocalClientAllocationTest, AddVectors) { // deallocation happen on the right allocator. ExecutableRunOptions options; options.set_allocator(allocator); - std::unique_ptr result = + tensorflow::gtl::optional result = ExecuteLocallyOrDie(builder.Build().ValueOrDie(), {}, DefaultExecutableBuildOptions(), options); @@ -66,7 +67,7 @@ XLA_TEST_F(LocalClientAllocationTest, AddVectors) { // Deallocate result and verify that deallocate was called once. int64 deallocation_count_before = allocator_->deallocation_count(); - result = nullptr; + result.reset(); EXPECT_EQ(deallocation_count_before + 1, allocator_->deallocation_count()); } @@ -92,7 +93,7 @@ XLA_TEST_F(LocalClientAllocationTest, RunOnDevices) { computation, {}, ExecutableBuildOptions().set_device_ordinal(d), ExecutableRunOptions().set_device_ordinal(d).set_allocator(allocator)); LiteralTestUtil::ExpectR1Near( - {2.0f, 4.0f, 6.0f}, *ShapedBufferToLiteral(*result), error_spec_); + {2.0f, 4.0f, 6.0f}, *ShapedBufferToLiteral(result), error_spec_); // At least one allocation should have been performed when executing the // computation. diff --git a/tensorflow/compiler/xla/tests/local_client_execute_test.cc b/tensorflow/compiler/xla/tests/local_client_execute_test.cc index 2462ea39f914b1dbb525ea777a48d9ce66035638..7e14e77366d1d4d7f21a585883bb2fc750f4247e 100644 --- a/tensorflow/compiler/xla/tests/local_client_execute_test.cc +++ b/tensorflow/compiler/xla/tests/local_client_execute_test.cc @@ -43,8 +43,6 @@ limitations under the License. #include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/test_benchmark.h" -namespace se = ::perftools::gputools; - namespace xla { namespace { @@ -59,10 +57,9 @@ XLA_TEST_F(LocalClientExecuteTest, Constant) { ComputationBuilder builder(local_client_, TestName()); auto y = builder.ConstantR0(123.0f); - std::unique_ptr result = + ScopedShapedBuffer result = ExecuteLocallyOrDie(builder.Build().ValueOrDie(), {}); - - LiteralTestUtil::ExpectR0Near(123.f, *ShapedBufferToLiteral(*result), + LiteralTestUtil::ExpectR0Near(123.f, *ShapedBufferToLiteral(result), error_spec_); } @@ -73,10 +70,9 @@ XLA_TEST_F(LocalClientExecuteTest, AddScalars) { builder.Add(x, y); auto x_value = LiteralToShapedBuffer(*Literal::CreateR0(42.0f)); - std::unique_ptr result = - ExecuteLocallyOrDie(builder.Build().ValueOrDie(), {x_value.get()}); - - LiteralTestUtil::ExpectR0Near(165.f, *ShapedBufferToLiteral(*result), + ScopedShapedBuffer result = + ExecuteLocallyOrDie(builder.Build().ValueOrDie(), {&x_value}); + LiteralTestUtil::ExpectR0Near(165.f, *ShapedBufferToLiteral(result), error_spec_); } @@ -87,10 +83,9 @@ XLA_TEST_F(LocalClientExecuteTest, AddZeroElementVectors) { builder.Add(x, y); auto x_array = LiteralToShapedBuffer(*Literal::CreateR1({})); - std::unique_ptr result = - ExecuteLocallyOrDie(builder.Build().ValueOrDie(), {x_array.get()}); - - LiteralTestUtil::ExpectR1Near({}, *ShapedBufferToLiteral(*result), + ScopedShapedBuffer result = + ExecuteLocallyOrDie(builder.Build().ValueOrDie(), {&x_array}); + LiteralTestUtil::ExpectR1Near({}, *ShapedBufferToLiteral(result), error_spec_); } @@ -102,11 +97,10 @@ XLA_TEST_F(LocalClientExecuteTest, AddVectors) { auto x_array = LiteralToShapedBuffer(*Literal::CreateR1({0.0f, 1.0f, 2.0f})); - std::unique_ptr result = - ExecuteLocallyOrDie(builder.Build().ValueOrDie(), {x_array.get()}); - + ScopedShapedBuffer result = + ExecuteLocallyOrDie(builder.Build().ValueOrDie(), {&x_array}); LiteralTestUtil::ExpectR1Near( - {2.0f, 4.0f, 6.0f}, *ShapedBufferToLiteral(*result), error_spec_); + {2.0f, 4.0f, 6.0f}, *ShapedBufferToLiteral(result), error_spec_); } XLA_TEST_F(LocalClientExecuteTest, AddVectorsWithProfile) { @@ -118,13 +112,12 @@ XLA_TEST_F(LocalClientExecuteTest, AddVectorsWithProfile) { auto x_array = LiteralToShapedBuffer(*Literal::CreateR1({0.0f, 1.0f, 2.0f})); ExecutionProfile profile; - std::unique_ptr result = ExecuteLocallyOrDie( - builder.Build().ValueOrDie(), {x_array.get()}, - DefaultExecutableBuildOptions(), + ScopedShapedBuffer result = ExecuteLocallyOrDie( + builder.Build().ValueOrDie(), {&x_array}, DefaultExecutableBuildOptions(), DefaultExecutableRunOptions().set_execution_profile(&profile)); LiteralTestUtil::ExpectR1Near( - {2.0f, 4.0f, 6.0f}, *ShapedBufferToLiteral(*result), error_spec_); + {2.0f, 4.0f, 6.0f}, *ShapedBufferToLiteral(result), error_spec_); EXPECT_GT(profile.compute_and_transfer_time_ns(), 0); } @@ -138,27 +131,27 @@ XLA_TEST_F(LocalClientExecuteTest, AddArraysWithDifferentInputLayouts) { // Create x as a col-major array. auto x_array = LiteralToShapedBuffer(*Literal::CreateR2WithLayout( {{1.0f, 2.0f}, {3.0f, 4.0f}}, LayoutUtil::MakeLayout({0, 1}))); - EXPECT_TRUE(LayoutUtil::Equal(x_array->on_device_shape().layout(), + EXPECT_TRUE(LayoutUtil::Equal(x_array.on_device_shape().layout(), LayoutUtil::MakeLayout({0, 1}))); // Create y as a row-major array. auto y_array = LiteralToShapedBuffer(*Literal::CreateR2WithLayout( {{10.0f, 20.0f}, {30.0f, 40.0f}}, LayoutUtil::MakeLayout({1, 0}))); - EXPECT_TRUE(LayoutUtil::Equal(y_array->on_device_shape().layout(), + EXPECT_TRUE(LayoutUtil::Equal(y_array.on_device_shape().layout(), LayoutUtil::MakeLayout({1, 0}))); - std::unique_ptr result_colmaj = - ExecuteLocallyOrDie(computation, {x_array.get(), y_array.get()}); + ScopedShapedBuffer result_colmaj = + ExecuteLocallyOrDie(computation, {&x_array, &y_array}); LiteralTestUtil::ExpectR2Near({{11.0f, 22.0f}, {33.0f, 44.0f}}, - *ShapedBufferToLiteral(*result_colmaj), + *ShapedBufferToLiteral(result_colmaj), error_spec_); // Run with the parameter values in a different order. - std::unique_ptr result_param_swap = - ExecuteLocallyOrDie(computation, {y_array.get(), x_array.get()}); + ScopedShapedBuffer result_param_swap = + ExecuteLocallyOrDie(computation, {&y_array, &x_array}); LiteralTestUtil::ExpectR2Near( {{11.0f, 22.0f}, {33.0f, 44.0f}}, - *ShapedBufferToLiteral(*result_param_swap), error_spec_); + *ShapedBufferToLiteral(result_param_swap), error_spec_); } XLA_TEST_F(LocalClientExecuteTest, AddArraysWithDifferentOutputLayouts) { @@ -174,27 +167,27 @@ XLA_TEST_F(LocalClientExecuteTest, AddArraysWithDifferentOutputLayouts) { *Literal::CreateR2({{10.0f, 20.0f}, {30.0f, 40.0f}})); // Run with col-major result layout. - std::unique_ptr result_colmaj = ExecuteLocallyOrDie( - computation, {x_array.get(), y_array.get()}, + ScopedShapedBuffer result_colmaj = ExecuteLocallyOrDie( + computation, {&x_array, &y_array}, DefaultExecutableBuildOptions().set_result_layout( ShapeUtil::MakeShapeWithLayout(F32, /*dimensions=*/{2, 2}, {0, 1})), DefaultExecutableRunOptions()); - EXPECT_TRUE(LayoutUtil::Equal(result_colmaj->on_device_shape().layout(), + EXPECT_TRUE(LayoutUtil::Equal(result_colmaj.on_device_shape().layout(), LayoutUtil::MakeLayout({0, 1}))); LiteralTestUtil::ExpectR2Near({{11.0f, 22.0f}, {33.0f, 44.0f}}, - *ShapedBufferToLiteral(*result_colmaj), + *ShapedBufferToLiteral(result_colmaj), error_spec_); // Run with row-major result layout. - std::unique_ptr result_rowmaj = ExecuteLocallyOrDie( - computation, {x_array.get(), y_array.get()}, + ScopedShapedBuffer result_rowmaj = ExecuteLocallyOrDie( + computation, {&x_array, &y_array}, DefaultExecutableBuildOptions().set_result_layout( ShapeUtil::MakeShapeWithLayout(F32, /*dimensions=*/{2, 2}, {1, 0})), DefaultExecutableRunOptions()); - EXPECT_TRUE(LayoutUtil::Equal(result_rowmaj->on_device_shape().layout(), + EXPECT_TRUE(LayoutUtil::Equal(result_rowmaj.on_device_shape().layout(), LayoutUtil::MakeLayout({1, 0}))); LiteralTestUtil::ExpectR2Near({{11.0f, 22.0f}, {33.0f, 44.0f}}, - *ShapedBufferToLiteral(*result_rowmaj), + *ShapedBufferToLiteral(result_rowmaj), error_spec_); } @@ -210,13 +203,13 @@ XLA_TEST_F(LocalClientExecuteTest, TupleResult) { auto y_array = LiteralToShapedBuffer( *Literal::CreateR2({{10.0f, 20.0f}, {30.0f, 40.0f}})); - std::unique_ptr result = - ExecuteLocallyOrDie(computation, {x_array.get(), y_array.get()}); + ScopedShapedBuffer result = + ExecuteLocallyOrDie(computation, {&x_array, &y_array}); - EXPECT_TRUE(ShapeUtil::IsTuple(result->on_host_shape())); - EXPECT_EQ(3, ShapeUtil::TupleElementCount(result->on_host_shape())); + EXPECT_TRUE(ShapeUtil::IsTuple(result.on_host_shape())); + EXPECT_EQ(3, ShapeUtil::TupleElementCount(result.on_host_shape())); - std::unique_ptr result_literal = ShapedBufferToLiteral(*result); + std::unique_ptr result_literal = ShapedBufferToLiteral(result); LiteralTestUtil::ExpectR2Equal( {{1.0f, 2.0f}, {3.0f, 4.0f}}, LiteralView::Create(*result_literal, {0})); LiteralTestUtil::ExpectR2Equal( @@ -239,13 +232,13 @@ XLA_TEST_F(LocalClientExecuteTest, NestedTupleResult) { auto y_array = LiteralToShapedBuffer( *Literal::CreateR2({{10.0f, 20.0f}, {30.0f, 40.0f}})); - std::unique_ptr result = - ExecuteLocallyOrDie(computation, {x_array.get(), y_array.get()}); + ScopedShapedBuffer result = + ExecuteLocallyOrDie(computation, {&x_array, &y_array}); - EXPECT_TRUE(ShapeUtil::IsTuple(result->on_host_shape())); - EXPECT_EQ(2, ShapeUtil::TupleElementCount(result->on_host_shape())); + EXPECT_TRUE(ShapeUtil::IsTuple(result.on_host_shape())); + EXPECT_EQ(2, ShapeUtil::TupleElementCount(result.on_host_shape())); - std::unique_ptr result_literal = ShapedBufferToLiteral(*result); + std::unique_ptr result_literal = ShapedBufferToLiteral(result); LiteralTestUtil::ExpectR2Equal( {{1.0f, 2.0f}, {3.0f, 4.0f}}, LiteralView::Create(*result_literal, {1})); LiteralTestUtil::ExpectR2Equal( @@ -276,11 +269,11 @@ XLA_TEST_F(LocalClientExecuteTest, TupleResultWithLayout) { ShapeUtil::MakeShapeWithLayout(F32, /*dimensions=*/{2, 2}, /*minor_to_major=*/{1, 0})}); options.set_result_layout(shape_with_layout); - std::unique_ptr result = ExecuteLocallyOrDie( - builder.Build().ValueOrDie(), {array.get(), array.get()}, options, - DefaultExecutableRunOptions()); + ScopedShapedBuffer result = + ExecuteLocallyOrDie(builder.Build().ValueOrDie(), {&array, &array}, + options, DefaultExecutableRunOptions()); - std::unique_ptr result_literal = ShapedBufferToLiteral(*result); + std::unique_ptr result_literal = ShapedBufferToLiteral(result); LiteralTestUtil::ExpectR2Equal( {{1.0f, 2.0f}, {3.0f, 4.0f}}, LiteralView::Create(*result_literal, {0})); LiteralTestUtil::ExpectR2Equal( @@ -320,13 +313,13 @@ XLA_TEST_F(LocalClientExecuteTest, TupleArguments) { auto x_buffer = LiteralToShapedBuffer(*x_literal); auto y_buffer = LiteralToShapedBuffer(*y_literal); - std::unique_ptr result = - ExecuteLocallyOrDie(computation, {x_buffer.get(), y_buffer.get()}); + ScopedShapedBuffer result = + ExecuteLocallyOrDie(computation, {&x_buffer, &y_buffer}); - EXPECT_TRUE(ShapeUtil::IsTuple(result->on_host_shape())); - EXPECT_EQ(2, ShapeUtil::TupleElementCount(result->on_host_shape())); + EXPECT_TRUE(ShapeUtil::IsTuple(result.on_host_shape())); + EXPECT_EQ(2, ShapeUtil::TupleElementCount(result.on_host_shape())); - std::unique_ptr result_literal = ShapedBufferToLiteral(*result); + std::unique_ptr result_literal = ShapedBufferToLiteral(result); LiteralTestUtil::ExpectR2Equal( {{56.0f, 46.0f}, {36.0f, 26.0f}}, LiteralView::Create(*result_literal, {0})); @@ -365,10 +358,9 @@ XLA_TEST_F(LocalClientExecuteTest, NestedTupleArgument) { Literal::CreateR1({222.0, -2.0, 10.0}).get()}); auto arg_buffer = LiteralToShapedBuffer(*arg_literal); - std::unique_ptr result = - ExecuteLocallyOrDie(computation, {arg_buffer.get()}); + ScopedShapedBuffer result = ExecuteLocallyOrDie(computation, {&arg_buffer}); - std::unique_ptr result_literal = ShapedBufferToLiteral(*result); + std::unique_ptr result_literal = ShapedBufferToLiteral(result); LiteralTestUtil::ExpectR2Equal( {{-1.0, -2.0}, {-3.0, -4}}, LiteralView::Create(*result_literal, {0})); LiteralTestUtil::ExpectR1Equal( @@ -396,18 +388,16 @@ XLA_TEST_F(LocalClientExecuteTest, PassingTupleResultBackIntoComputation) { Literal::CreateR2({{11.0, 3.0}, {4.0, 5.0}}).get()}); auto arg_buffer = LiteralToShapedBuffer(*arg_literal); - std::unique_ptr result_0 = - ExecuteLocallyOrDie(computation, {arg_buffer.get()}); - std::unique_ptr result_0_literal = ShapedBufferToLiteral(*result_0); + ScopedShapedBuffer result_0 = ExecuteLocallyOrDie(computation, {&arg_buffer}); + std::unique_ptr result_0_literal = ShapedBufferToLiteral(result_0); LiteralTestUtil::ExpectR2Equal( {{-1.0, -2.0}, {-3.0, -4.0}}, LiteralView::Create(*result_0_literal, {0})); LiteralTestUtil::ExpectR2Equal( {{22.0, 6.0}, {8.0, 10}}, LiteralView::Create(*result_0_literal, {1})); - std::unique_ptr result_1 = - ExecuteLocallyOrDie(computation, {result_0.get()}); - std::unique_ptr result_1_literal = ShapedBufferToLiteral(*result_1); + ScopedShapedBuffer result_1 = ExecuteLocallyOrDie(computation, {&result_0}); + std::unique_ptr result_1_literal = ShapedBufferToLiteral(result_1); LiteralTestUtil::ExpectR2Equal( {{1.0, 2.0}, {3.0, 4.0}}, LiteralView::Create(*result_1_literal, {0})); LiteralTestUtil::ExpectR2Equal( @@ -453,10 +443,8 @@ XLA_TEST_F(LocalClientExecuteTest, LargeTuple) { Literal::MakeTupleOwned(std::move(arg_elements)); auto arg_buffer = LiteralToShapedBuffer(*arg_literal); - std::unique_ptr result = - ExecuteLocallyOrDie(computation, {arg_buffer.get()}); - - std::unique_ptr result_literal = ShapedBufferToLiteral(*result); + ScopedShapedBuffer result = ExecuteLocallyOrDie(computation, {&arg_buffer}); + std::unique_ptr result_literal = ShapedBufferToLiteral(result); for (int i = 0; i < kElementCount; ++i) { LiteralTestUtil::ExpectR1Near( @@ -511,9 +499,8 @@ XLA_TEST_F(LocalClientExecuteTest, DISABLED_ON_CPU_PARALLEL(LargeNestedTuple)) { auto arg_literal = Literal::MakeTupleOwned(std::move(outer_tuple_elements)); auto arg_buffer = LiteralToShapedBuffer(*arg_literal); - std::unique_ptr result = - ExecuteLocallyOrDie(computation, {arg_buffer.get()}); - std::unique_ptr result_literal = ShapedBufferToLiteral(*result); + ScopedShapedBuffer result = ExecuteLocallyOrDie(computation, {&arg_buffer}); + std::unique_ptr result_literal = ShapedBufferToLiteral(result); for (int i = 0; i < kFanout; ++i) { for (int j = 0; j < kFanout; ++j) { @@ -556,9 +543,8 @@ XLA_TEST_F(LocalClientExecuteTest, DeepTuple) { } auto arg_buffer = LiteralToShapedBuffer(*arg_literal); - std::unique_ptr result = - ExecuteLocallyOrDie(computation, {arg_buffer.get()}); - std::unique_ptr result_literal = ShapedBufferToLiteral(*result); + ScopedShapedBuffer result = ExecuteLocallyOrDie(computation, {&arg_buffer}); + std::unique_ptr result_literal = ShapedBufferToLiteral(result); ShapeIndex index; for (int i = 0; i < kTupleDepth; ++i) { @@ -578,7 +564,7 @@ XLA_TEST_F(LocalClientExecuteTest, InvalidNumberOfArguments) { auto x_array = LiteralToShapedBuffer(*Literal::CreateR1({1.0f, 2.0f, 3.0f})); auto execute_status = - ExecuteLocally(builder.Build().ValueOrDie(), {x_array.get()}); + ExecuteLocally(builder.Build().ValueOrDie(), {&x_array}); EXPECT_FALSE(execute_status.ok()); EXPECT_THAT(execute_status.status().error_message(), @@ -594,7 +580,7 @@ XLA_TEST_F(LocalClientExecuteTest, IncorrectArgumentShape) { auto x_array = LiteralToShapedBuffer( *Literal::CreateR2({{0.0f, 1.0f}, {2.0f, 3.0f}})); auto execute_status = - ExecuteLocally(builder.Build().ValueOrDie(), {x_array.get()}); + ExecuteLocally(builder.Build().ValueOrDie(), {&x_array}); EXPECT_FALSE(execute_status.ok()); EXPECT_THAT(execute_status.status().error_message(), @@ -611,7 +597,7 @@ XLA_TEST_F(LocalClientExecuteTest, InvalidResultLayout) { auto x_array = LiteralToShapedBuffer( *Literal::CreateR2({{0.0f, 1.0f}, {2.0f, 3.0f}})); auto execute_status = ExecuteLocally( - builder.Build().ValueOrDie(), {x_array.get()}, + builder.Build().ValueOrDie(), {&x_array}, DefaultExecutableBuildOptions().set_result_layout( ShapeUtil::MakeShapeWithLayout(F32, /*dimensions=*/{1, 2, 3, 4}, @@ -644,9 +630,9 @@ XLA_TEST_F(LocalClientExecuteTest, RunOnAllDeviceOrdinals) { computation, {}, DefaultExecutableBuildOptions().set_device_ordinal(d), DefaultExecutableRunOptions().set_device_ordinal(d)); - EXPECT_EQ(d, result->device_ordinal()); + EXPECT_EQ(d, result.device_ordinal()); LiteralTestUtil::ExpectR0Equal(42.0f, - *ShapedBufferToLiteral(*result)); + *ShapedBufferToLiteral(result)); } } } @@ -689,9 +675,9 @@ XLA_TEST_F(LocalClientExecuteTest, RunOnStream) { DefaultExecutableRunOptions().set_stream(&stream)); // As a check to verify that the computation ran of the device associated // with the stream. This is a weak check, but stronger verification is hard. - EXPECT_EQ(d, result->device_ordinal()); + EXPECT_EQ(d, result.device_ordinal()); LiteralTestUtil::ExpectR0Equal(42.0f, - *ShapedBufferToLiteral(*result)); + *ShapedBufferToLiteral(result)); } } @@ -767,9 +753,9 @@ XLA_TEST_F(LocalClientExecuteTest, SelectBetweenTuples) { {builder.ConstantR1(vec2), builder.ConstantR1(vec1)}); builder.Select(builder.ConstantR0(false), tuple12, tuple21); - std::unique_ptr result = + ScopedShapedBuffer result = ExecuteLocallyOrDie(builder.Build().ValueOrDie(), {}); - std::unique_ptr tuple_literal = ShapedBufferToLiteral(*result); + std::unique_ptr tuple_literal = ShapedBufferToLiteral(result); LiteralTestUtil::ExpectR1Equal( {2.0f, 4.0f, 6.0f}, LiteralView::Create(*tuple_literal, {0})); LiteralTestUtil::ExpectR1Equal( @@ -793,12 +779,12 @@ XLA_TEST_F(LocalClientExecuteTest, CompileExecutable) { auto x_array = LiteralToShapedBuffer(*Literal::CreateR1({0.0f, 1.0f, 2.0f})); - std::unique_ptr result = - executable->Run({x_array.get()}, DefaultExecutableRunOptions()) + ScopedShapedBuffer result = + executable->Run({&x_array}, DefaultExecutableRunOptions()) .ConsumeValueOrDie(); LiteralTestUtil::ExpectR1Near( - {2.0f, 4.0f, 6.0f}, *ShapedBufferToLiteral(*result), error_spec_); + {2.0f, 4.0f, 6.0f}, *ShapedBufferToLiteral(result), error_spec_); } XLA_TEST_F(LocalClientExecuteTest, ShapeBufferToLiteralConversion) { @@ -811,7 +797,7 @@ XLA_TEST_F(LocalClientExecuteTest, ShapeBufferToLiteralConversion) { literal, local_client_->default_device_ordinal(), allocator_)); TF_ASSERT_OK_AND_ASSIGN( auto transferred_literal, - local_client_->ShapedBufferToLiteral(*shaped_buffer)); + local_client_->ShapedBufferToLiteral(shaped_buffer)); EXPECT_EQ(literal, *transferred_literal); }; @@ -851,7 +837,7 @@ XLA_TEST_F(LocalClientExecuteTest, ShapeBufferToLiteralConversion64bit) { literal, local_client_->default_device_ordinal(), allocator_)); TF_ASSERT_OK_AND_ASSIGN( auto transferred_literal, - local_client_->ShapedBufferToLiteral(*shaped_buffer)); + local_client_->ShapedBufferToLiteral(shaped_buffer)); EXPECT_EQ(literal, *transferred_literal); }; @@ -919,12 +905,12 @@ void BM_LocalClientOverhead(int num_iters) { .ConsumeValueOrDie(); auto literal = Literal::CreateR2({{0, 0, 0}, {0, 0, 0}}); ASSERT_IS_OK(transfer_manager->TransferLiteralToDevice( - executors[device_ordinal], *literal, *buffer)); + executors[device_ordinal], *literal, buffer)); const int kWarmups = 2; auto executable_status = client->Compile( - computation, {&buffer->on_host_shape()}, ExecutableBuildOptions()); + computation, {&buffer.on_host_shape()}, ExecutableBuildOptions()); ASSERT_IS_OK(executable_status); std::unique_ptr executable = executable_status.ConsumeValueOrDie(); @@ -936,13 +922,13 @@ void BM_LocalClientOverhead(int num_iters) { run_options.set_allocator(&allocator).set_stream(&stream); for (int i = 0; i < kWarmups; ++i) { - auto result = executable->Run({buffer.get()}, run_options); + auto result = executable->Run({&buffer}, run_options); ASSERT_IS_OK(result); } tensorflow::testing::StartTiming(); for (int i = 0; i < num_iters; ++i) { - auto result = executable->Run({buffer.get()}, run_options); + auto result = executable->Run({&buffer}, run_options); ASSERT_IS_OK(result); } } diff --git a/tensorflow/compiler/xla/tests/local_client_test_base.cc b/tensorflow/compiler/xla/tests/local_client_test_base.cc index 96b976d25d75d35f46adfd104a03aceb363661eb..bb5aabb214da41f8bcd5da7d23d6d732e1277133 100644 --- a/tensorflow/compiler/xla/tests/local_client_test_base.cc +++ b/tensorflow/compiler/xla/tests/local_client_test_base.cc @@ -35,19 +35,21 @@ namespace xla { /* static */ TestAllocator* LocalClientTestBase::allocator_; -StatusOr TestAllocator::Allocate( - int device_ordinal, uint64 size, bool retry_on_failure) { +StatusOr TestAllocator::Allocate(int device_ordinal, + uint64 size, + bool retry_on_failure) { VLOG(2) << "Allocate(" << device_ordinal << ", " << size << ")"; { tensorflow::mutex_lock lock(count_mutex_); allocation_count_++; device_allocation_count_[device_ordinal]++; } - return StreamExecutorMemoryAllocator::Allocate(device_ordinal, size); + return StreamExecutorMemoryAllocator::Allocate(device_ordinal, size, + retry_on_failure); } -tensorflow::Status TestAllocator::Deallocate( - int device_ordinal, perftools::gputools::DeviceMemoryBase* mem) { +tensorflow::Status TestAllocator::Deallocate(int device_ordinal, + se::DeviceMemoryBase* mem) { VLOG(2) << "Deallocate(" << device_ordinal << ")"; { tensorflow::mutex_lock lock(count_mutex_); @@ -88,7 +90,7 @@ int64 TestAllocator::deallocation_count(int device_ordinal) const { } /* static */ TestAllocator* LocalClientTestBase::GetOrCreateAllocator( - perftools::gputools::Platform* platform) { + se::Platform* platform) { static tensorflow::mutex mu(tensorflow::LINKER_INITIALIZED); tensorflow::mutex_lock lock(mu); @@ -115,8 +117,7 @@ struct LocalClientTestBase::EigenThreadPoolWrapper { std::unique_ptr device; }; -LocalClientTestBase::LocalClientTestBase( - perftools::gputools::Platform* platform) +LocalClientTestBase::LocalClientTestBase(se::Platform* platform) : local_client_( ClientLibrary::GetOrCreateLocalClient(platform).ValueOrDie()), thread_pool_wrapper_(new EigenThreadPoolWrapper()) { @@ -128,7 +129,7 @@ LocalClientTestBase::LocalClientTestBase( LocalClientTestBase::~LocalClientTestBase() {} -std::unique_ptr LocalClientTestBase::LiteralToShapedBuffer( +ScopedShapedBuffer LocalClientTestBase::LiteralToShapedBuffer( const Literal& literal) { return local_client_ ->LiteralToShapedBuffer(literal, local_client_->default_device_ordinal()) @@ -155,7 +156,7 @@ ExecutableRunOptions LocalClientTestBase::DefaultExecutableRunOptions() const { return run_options; } -std::unique_ptr LocalClientTestBase::ExecuteLocallyOrDie( +ScopedShapedBuffer LocalClientTestBase::ExecuteLocallyOrDie( const Computation& computation, tensorflow::gtl::ArraySlice arguments) { return ExecuteLocally(computation, arguments, DefaultExecutableBuildOptions(), @@ -163,7 +164,7 @@ std::unique_ptr LocalClientTestBase::ExecuteLocallyOrDie( .ConsumeValueOrDie(); } -std::unique_ptr LocalClientTestBase::ExecuteLocallyOrDie( +ScopedShapedBuffer LocalClientTestBase::ExecuteLocallyOrDie( const Computation& computation, tensorflow::gtl::ArraySlice arguments, const ExecutableBuildOptions& build_options, @@ -172,16 +173,14 @@ std::unique_ptr LocalClientTestBase::ExecuteLocallyOrDie( .ConsumeValueOrDie(); } -StatusOr> -LocalClientTestBase::ExecuteLocally( +StatusOr LocalClientTestBase::ExecuteLocally( const Computation& computation, tensorflow::gtl::ArraySlice arguments) { return ExecuteLocally(computation, arguments, DefaultExecutableBuildOptions(), DefaultExecutableRunOptions()); } -StatusOr> -LocalClientTestBase::ExecuteLocally( +StatusOr LocalClientTestBase::ExecuteLocally( const Computation& computation, tensorflow::gtl::ArraySlice arguments, const ExecutableBuildOptions& build_options, diff --git a/tensorflow/compiler/xla/tests/local_client_test_base.h b/tensorflow/compiler/xla/tests/local_client_test_base.h index f0c73f04f6eb67b2e9cb5e111eccdc3818059b2b..4ee56a05ec6ec1ddf3fba0b357e2c3b586eb456e 100644 --- a/tensorflow/compiler/xla/tests/local_client_test_base.h +++ b/tensorflow/compiler/xla/tests/local_client_test_base.h @@ -41,15 +41,15 @@ namespace xla { class TestAllocator : public StreamExecutorMemoryAllocator { public: - explicit TestAllocator(perftools::gputools::Platform* platform) + explicit TestAllocator(se::Platform* platform) : StreamExecutorMemoryAllocator( platform, PlatformUtil::GetStreamExecutors(platform).ValueOrDie()) { } - StatusOr Allocate( - int device_ordinal, uint64 size, bool retry_on_failure) override; - tensorflow::Status Deallocate( - int device_ordinal, perftools::gputools::DeviceMemoryBase* mem) override; + StatusOr Allocate(int device_ordinal, uint64 size, + bool retry_on_failure) override; + tensorflow::Status Deallocate(int device_ordinal, + se::DeviceMemoryBase* mem) override; // Return the number of allocations that have been performed. int64 allocation_count() const; @@ -75,18 +75,15 @@ class TestAllocator : public StreamExecutorMemoryAllocator { class LocalClientTestBase : public ::testing::Test { protected: struct EigenThreadPoolWrapper; - explicit LocalClientTestBase( - perftools::gputools::Platform* platform = nullptr); + explicit LocalClientTestBase(se::Platform* platform = nullptr); virtual ~LocalClientTestBase(); - static TestAllocator* GetOrCreateAllocator( - perftools::gputools::Platform* platform); + static TestAllocator* GetOrCreateAllocator(se::Platform* platform); // Copy the given literal onto the default device and return a // ScopedShapedBuffer. Convenience wrapper around // LocalClient::LiteralToShapedBuffer. - std::unique_ptr LiteralToShapedBuffer( - const Literal& literal); + ScopedShapedBuffer LiteralToShapedBuffer(const Literal& literal); // Construct and return a literal containing the array represented by // shaped_buffer. @@ -95,19 +92,19 @@ class LocalClientTestBase : public ::testing::Test { // Execute the given computation on the local client. With and without // options. - StatusOr> ExecuteLocally( + StatusOr ExecuteLocally( const Computation& computation, tensorflow::gtl::ArraySlice arguments); - StatusOr> ExecuteLocally( + StatusOr ExecuteLocally( const Computation& computation, tensorflow::gtl::ArraySlice arguments, const ExecutableBuildOptions& build_options, const ExecutableRunOptions& run_options); - std::unique_ptr ExecuteLocallyOrDie( + ScopedShapedBuffer ExecuteLocallyOrDie( const Computation& computation, tensorflow::gtl::ArraySlice arguments); - std::unique_ptr ExecuteLocallyOrDie( + ScopedShapedBuffer ExecuteLocallyOrDie( const Computation& computation, tensorflow::gtl::ArraySlice arguments, const ExecutableBuildOptions& build_options, @@ -128,7 +125,7 @@ class LocalClientTestBase : public ::testing::Test { // of the process. So make the allocator static. static TestAllocator* allocator_; - perftools::gputools::StreamExecutor* stream_executor_; + se::StreamExecutor* stream_executor_; TransferManager* transfer_manager_; LocalClient* local_client_; diff --git a/tensorflow/compiler/xla/tests/map_test.cc b/tensorflow/compiler/xla/tests/map_test.cc index 0cd812fd1b4bc69c34b70d3ca0fd0aa6cf57fa4c..8fabcaca1b968dbd0546e16b97f48c43dbb60e3c 100644 --- a/tensorflow/compiler/xla/tests/map_test.cc +++ b/tensorflow/compiler/xla/tests/map_test.cc @@ -21,6 +21,8 @@ limitations under the License. #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/statusor.h" @@ -39,7 +41,7 @@ namespace { class MapTest : public ClientLibraryTestBase { public: - explicit MapTest(perftools::gputools::Platform* platform = nullptr) + explicit MapTest(se::Platform* platform = nullptr) : ClientLibraryTestBase(platform) { mutable_debug_options()->add_xla_disable_hlo_passes("algsimp"); mutable_debug_options()->add_xla_disable_hlo_passes("inline"); @@ -50,18 +52,18 @@ class MapTest : public ClientLibraryTestBase { // x {R0F32} ----> (add) // / // 1.0f ---------/ - Computation CreateAdderToOne() { - ComputationBuilder mapped_builder(client_, TestName()); + XlaComputation CreateAdderToOne() { + XlaBuilder mapped_builder(TestName()); auto x = mapped_builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x"); auto one = mapped_builder.ConstantR0(1.0); - auto adder_to_one = mapped_builder.Add(x, one); + mapped_builder.Add(x, one); auto computation_status = mapped_builder.Build(); TF_CHECK_OK(computation_status.status()); return computation_status.ConsumeValueOrDie(); } - Computation CreateMax() { - ComputationBuilder b(client_, TestName()); + XlaComputation CreateMax() { + XlaBuilder b(TestName()); auto lhs = b.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x"); auto rhs = b.Parameter(1, ShapeUtil::MakeShape(F32, {}), "y"); b.Max(lhs, rhs); @@ -73,8 +75,8 @@ class MapTest : public ClientLibraryTestBase { // Creates a computation that accepts an F32 and returns T(1) (ignoring the // argument). template - Computation CreateScalarOne() { - ComputationBuilder mapped_builder(client_, "scalar_one"); + XlaComputation CreateScalarOne() { + XlaBuilder mapped_builder("scalar_one"); (void)mapped_builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x"); mapped_builder.ConstantR0(1); auto computation_status = mapped_builder.Build(); @@ -87,11 +89,11 @@ class MapTest : public ClientLibraryTestBase { // x {R0F32} ----> (mul) // / // 2.0f ---------/ - Computation CreateMulByTwo() { - ComputationBuilder mapped_builder(client_, TestName()); + XlaComputation CreateMulByTwo() { + XlaBuilder mapped_builder(TestName()); auto x = mapped_builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x"); auto two = mapped_builder.ConstantR0(2.0); - auto mul_by_two = mapped_builder.Mul(x, two); + mapped_builder.Mul(x, two); auto computation_status = mapped_builder.Build(); TF_CHECK_OK(computation_status.status()); return computation_status.ConsumeValueOrDie(); @@ -105,12 +107,12 @@ class MapTest : public ClientLibraryTestBase { // x {R0F32} ----> (add) ----> (mul) // / // 1.0f ---------/ - Computation CreateAdderToOneTimesItself() { - ComputationBuilder mapped_builder(client_, TestName()); + XlaComputation CreateAdderToOneTimesItself() { + XlaBuilder mapped_builder(TestName()); auto x = mapped_builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x"); auto one = mapped_builder.ConstantR0(1.0); auto adder_to_one = mapped_builder.Add(x, one); - auto result = mapped_builder.Mul(x, adder_to_one); + mapped_builder.Mul(x, adder_to_one); auto computation_status = mapped_builder.Build(); TF_CHECK_OK(computation_status.status()); return computation_status.ConsumeValueOrDie(); @@ -122,12 +124,13 @@ class MapTest : public ClientLibraryTestBase { // x {R0F32} -----------> (map) ----> (add) // / / // embedded_computation --/ n --/ - Computation CreateMapPlusN(const Computation& embedded_computation, float n) { - ComputationBuilder builder(client_, TestName()); + XlaComputation CreateMapPlusN(const XlaComputation& embedded_computation, + float n) { + XlaBuilder builder(TestName()); auto x = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x"); auto map = builder.Map({x}, embedded_computation, {}); auto constant_n = builder.ConstantR0(n); - auto add = builder.Add(map, constant_n); + builder.Add(map, constant_n); auto computation_status = builder.Build(); TF_CHECK_OK(computation_status.status()); return computation_status.ConsumeValueOrDie(); @@ -135,11 +138,11 @@ class MapTest : public ClientLibraryTestBase { // Creates a binary function with signature (F32, F32) -> Pred // defined by (x, y) -> x > y. - Computation CreateGt() { - ComputationBuilder b(client_, "Gt"); + XlaComputation CreateGt() { + XlaBuilder b("Gt"); auto x = b.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x"); auto y = b.Parameter(1, ShapeUtil::MakeShape(F32, {}), "y"); - auto gt = b.Gt(x, y); + b.Gt(x, y); auto computation_status = b.Build(); TF_CHECK_OK(computation_status.status()); return computation_status.ConsumeValueOrDie(); @@ -152,13 +155,13 @@ class MapTest : public ClientLibraryTestBase { // y {R0F32} ----> (add) ---> (add) // / // z {R0F32} ---------------/ - Computation CreateTernaryAdder() { - ComputationBuilder mapped_builder(client_, "TernaryAdder"); + XlaComputation CreateTernaryAdder() { + XlaBuilder mapped_builder("TernaryAdder"); auto x = mapped_builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x"); auto y = mapped_builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "y"); auto z = mapped_builder.Parameter(2, ShapeUtil::MakeShape(F32, {}), "z"); auto xy = mapped_builder.Add(x, y); - auto xyz = mapped_builder.Add(xy, z); + mapped_builder.Add(xy, z); auto computation_status = mapped_builder.Build(); TF_CHECK_OK(computation_status.status()); return computation_status.ConsumeValueOrDie(); @@ -167,13 +170,13 @@ class MapTest : public ClientLibraryTestBase { TEST_F(MapTest, MapEachElemPlusOneR0) { // Applies lambda (x) (+ x 1)) to an input scalar. - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr param0_literal = Literal::CreateR0(42.0); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); auto param = builder.Parameter(0, param0_literal->shape(), "param0"); - auto map = builder.Map({param}, CreateAdderToOne(), {}); + builder.Map({param}, CreateAdderToOne(), {}); ComputeAndCompareR0(&builder, 43.0, {param0_data.get()}, ErrorSpec(0.01f)); @@ -181,13 +184,13 @@ TEST_F(MapTest, MapEachElemPlusOneR0) { XLA_TEST_F(MapTest, MapEachElemPlusOneR1S0) { // Maps (lambda (x) (+ x 1)) onto an input R1F32 vector of length 0. - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr param0_literal = Literal::CreateR1({}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); auto param = builder.Parameter(0, param0_literal->shape(), "param0"); - auto map = builder.Map({param}, CreateAdderToOne(), {0}); + builder.Map({param}, CreateAdderToOne(), {0}); ComputeAndCompareR1(&builder, {}, {param0_data.get()}, ErrorSpec(0.01f)); @@ -195,55 +198,55 @@ XLA_TEST_F(MapTest, MapEachElemPlusOneR1S0) { TEST_F(MapTest, MapEachElemPlusOneR1S4) { // Maps (lambda (x) (+ x 1)) onto an input R1F32 vector of length 4. - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr param0_literal = Literal::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); auto param = builder.Parameter(0, param0_literal->shape(), "param0"); - auto map = builder.Map({param}, CreateAdderToOne(), {0}); + builder.Map({param}, CreateAdderToOne(), {0}); ComputeAndCompareR1(&builder, {3.2f, 4.3f, 5.4f, 6.5f}, {param0_data.get()}, ErrorSpec(0.01f)); } TEST_F(MapTest, MapEachF32ElementToS32Constant) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr param0_literal = Literal::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); auto param = builder.Parameter(0, param0_literal->shape(), "param0"); - auto map = builder.Map({param}, CreateScalarOne(), {0}); + builder.Map({param}, CreateScalarOne(), {0}); ComputeAndCompareR1(&builder, {1, 1, 1, 1}, {param0_data.get()}); } TEST_F(MapTest, MapEachF32ElementToU32Constant) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr param0_literal = Literal::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); auto param = builder.Parameter(0, param0_literal->shape(), "param0"); - auto map = builder.Map({param}, CreateScalarOne(), {0}); + builder.Map({param}, CreateScalarOne(), {0}); ComputeAndCompareR1(&builder, {1, 1, 1, 1}, {param0_data.get()}); } TEST_F(MapTest, MapEachElemLongerChainR1) { // Maps (lambda (x) (* (+ x 1) x)) onto an input R1F32 vector. - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr param0_literal = Literal::CreateR1({2.6f, -5.1f, 0.1f, 0.2f, 999.0f, 255.5f}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); auto param = builder.Parameter(0, param0_literal->shape(), "param0"); - auto map = builder.Map({param}, CreateAdderToOneTimesItself(), {0}); + builder.Map({param}, CreateAdderToOneTimesItself(), {0}); ComputeAndCompareR1( &builder, {9.36f, 20.91f, 0.11f, 0.24f, 999000.0f, 65535.75f}, @@ -253,14 +256,14 @@ TEST_F(MapTest, MapEachElemLongerChainR1) { XLA_TEST_F(MapTest, MapMultipleMapsR1S0) { // Maps (lambda (x) (+ x 1)) onto an input R1F32 vector of length 0, and then // maps (lambda (x) (* x 2)) on the result. - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr param0_literal = Literal::CreateR1({}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); auto param = builder.Parameter(0, param0_literal->shape(), "param0"); auto map1 = builder.Map({param}, CreateAdderToOne(), {0}); - auto map2 = builder.Map({map1}, CreateMulByTwo(), {0}); + builder.Map({map1}, CreateMulByTwo(), {0}); ComputeAndCompareR1(&builder, {}, {param0_data.get()}, ErrorSpec(0.01f)); @@ -269,7 +272,7 @@ XLA_TEST_F(MapTest, MapMultipleMapsR1S0) { TEST_F(MapTest, MapMultipleMapsR1S4) { // Maps (lambda (x) (+ x 1)) onto an input R1F32 vector of length 4, and then // maps (lambda (x) (* x 2)) on the result. - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr param0_literal = Literal::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); std::unique_ptr param0_data = @@ -277,7 +280,7 @@ TEST_F(MapTest, MapMultipleMapsR1S4) { auto param = builder.Parameter(0, param0_literal->shape(), "param0"); auto map1 = builder.Map({param}, CreateAdderToOne(), {0}); - auto map2 = builder.Map({map1}, CreateMulByTwo(), {0}); + builder.Map({map1}, CreateMulByTwo(), {0}); ComputeAndCompareR1(&builder, {6.4f, 8.6f, 10.8f, 13.0f}, {param0_data.get()}, ErrorSpec(0.01f)); @@ -285,14 +288,14 @@ TEST_F(MapTest, MapMultipleMapsR1S4) { TEST_F(MapTest, MapEachElemPlusOneR2) { // Maps (lambda (x) (+ x 1)) onto an input R2F32 vector. - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr param0_literal = Literal::CreateR2( {{13.25f, 14.0f}, {-7.1f, -7.2f}, {-8.8f, 8.8f}}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); auto param = builder.Parameter(0, param0_literal->shape(), "param0"); - auto map = builder.Map({param}, CreateAdderToOne(), {0, 1}); + builder.Map({param}, CreateAdderToOne(), {0, 1}); Array2D expected_array( {{14.25f, 15.0f}, {-6.1f, -6.2f}, {-7.8f, 9.8f}}); @@ -317,18 +320,18 @@ XLA_TEST_F(MapTest, ComplexNestedMaps) { auto embed2 = CreateMapPlusN(embed1, 2.0); auto embed3 = CreateMapPlusN(embed1, 4.0); - ComputationBuilder embed4_builder(client_, "embed4"); + XlaBuilder embed4_builder("embed4"); auto embed4_param = embed4_builder.Parameter(0, scalar_shape, "x"); auto embed4_map_lhs = embed4_builder.Map({embed4_param}, embed2, {}); auto embed4_map_rhs = embed4_builder.Map({embed4_param}, embed3, {}); - auto embed4_add = embed4_builder.Add(embed4_map_lhs, embed4_map_rhs); + embed4_builder.Add(embed4_map_lhs, embed4_map_rhs); auto embed4_status = embed4_builder.Build(); ASSERT_IS_OK(embed4_status.status()); auto embed4 = embed4_status.ConsumeValueOrDie(); auto embed5 = CreateMapPlusN(embed2, 6.0); - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto constant_42 = builder.ConstantR0(42.0); auto constant_7 = builder.ConstantR0(7.0); auto map_42 = builder.Map({constant_42}, embed5, {}); @@ -359,7 +362,8 @@ TEST_F(MapTest, VersionedEmbeddedComputation) { // Add another Add(1) operation to the existing embedded computation. This // requires using the stub interface because the ComputationBuilder does not - // allow modification to the Computation objects after they have been built. + // allow modification to the XlaComputation objects after they have been + // built. BinaryOpRequest request; request.set_binop(BINOP_ADD); *request.mutable_lhs() = adder_to_one; @@ -381,7 +385,7 @@ TEST_F(MapTest, VersionedEmbeddedComputation) { TEST_F(MapTest, MapBinaryAdder) { // Maps (lambda (x y) (+ x y)) onto two R1F32 vectors. - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr param0_literal = Literal::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); std::unique_ptr param0_data = @@ -393,8 +397,7 @@ TEST_F(MapTest, MapBinaryAdder) { auto param0 = builder.Parameter(0, param0_literal->shape(), "param0"); auto param1 = builder.Parameter(1, param1_literal->shape(), "param1"); - auto map = builder.Map({param0, param1}, - CreateScalarAddComputation(F32, &builder), {0}); + builder.Map({param0, param1}, CreateScalarAddComputation(F32, &builder), {0}); ComputeAndCompareR1(&builder, {7.3f, 7.7, 4.3f, 0}, {param0_data.get(), param1_data.get()}, @@ -404,7 +407,7 @@ TEST_F(MapTest, MapBinaryAdder) { // Adds two rank-2 arrays with different layouts. This test exercises a path // for Map that used to fail in shape inference (b/28989438). XLA_TEST_F(MapTest, AddWithMixedLayouts) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr param0_literal = Literal::CreateR2WithLayout( {{1, 2}, {3, 4}}, LayoutUtil::MakeLayout({1, 0})); std::unique_ptr param0_data = @@ -417,8 +420,8 @@ XLA_TEST_F(MapTest, AddWithMixedLayouts) { auto param0 = builder.Parameter(0, param0_literal->shape(), "param0"); auto param1 = builder.Parameter(1, param1_literal->shape(), "param1"); - auto map = builder.Map({param0, param1}, - CreateScalarAddComputation(S32, &builder), {0, 1}); + builder.Map({param0, param1}, CreateScalarAddComputation(S32, &builder), + {0, 1}); Array2D expected(2, 2); expected(0, 0) = 11; @@ -430,7 +433,7 @@ XLA_TEST_F(MapTest, AddWithMixedLayouts) { } XLA_TEST_F(MapTest, AddR3_3x0x2) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr param0_literal = Literal::CreateR3FromArray3D(Array3D(3, 0, 2)); std::unique_ptr param0_data = @@ -443,8 +446,8 @@ XLA_TEST_F(MapTest, AddR3_3x0x2) { auto param0 = builder.Parameter(0, param0_literal->shape(), "param0"); auto param1 = builder.Parameter(1, param1_literal->shape(), "param1"); - auto map = builder.Map({param0, param1}, - CreateScalarAddComputation(S32, &builder), {0, 1, 2}); + builder.Map({param0, param1}, CreateScalarAddComputation(S32, &builder), + {0, 1, 2}); ComputeAndCompareR3(&builder, Array3D(3, 0, 2), {param0_data.get(), param1_data.get()}); @@ -452,7 +455,7 @@ XLA_TEST_F(MapTest, AddR3_3x0x2) { TEST_F(MapTest, MapTernaryAdder) { // Maps (lambda (x y z) (+ x y z)) onto three R1F32 vectors. - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr param0_literal = Literal::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); std::unique_ptr param0_data = @@ -469,7 +472,7 @@ TEST_F(MapTest, MapTernaryAdder) { auto param0 = builder.Parameter(0, param0_literal->shape(), "param0"); auto param1 = builder.Parameter(1, param1_literal->shape(), "param1"); auto param2 = builder.Parameter(2, param2_literal->shape(), "param2"); - auto map = builder.Map({param0, param1, param2}, CreateTernaryAdder(), {0}); + builder.Map({param0, param1, param2}, CreateTernaryAdder(), {0}); ComputeAndCompareR1( &builder, {-2.7f, -92.3f, -895.7f, -400.0f}, @@ -479,24 +482,24 @@ TEST_F(MapTest, MapTernaryAdder) { TEST_F(MapTest, MapGt) { // Maps (x,y) -> x > y onto two R1F32 vectors. - ComputationBuilder b(client_, TestName()); + XlaBuilder b(TestName()); auto gt = CreateGt(); b.Map({b.ConstantR1({1, 20}), b.ConstantR1({10, 2})}, gt, {0}); ComputeAndCompareR1(&b, {false, true}, {}); } TEST_F(MapTest, NestedBinaryMap) { - Computation max_with_square; + XlaComputation max_with_square; { // max_with_square(x) = do max(x, x^2) via a map. - ComputationBuilder b(client_, "max_with_square"); + XlaBuilder b("max_with_square"); auto x = b.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x"); b.Map({x, b.Mul(x, x)}, CreateMax(), {}); auto computation_status = b.Build(); ASSERT_IS_OK(computation_status.status()); max_with_square = computation_status.ConsumeValueOrDie(); } - ComputationBuilder b(client_, TestName()); + XlaBuilder b(TestName()); auto input = b.ConstantR1({0.1f, 0.5f, -0.5f, 1.0f, 2.0f}); b.Map({input}, max_with_square, {0}); ComputeAndCompareR1(&b, {0.1f, 0.5f, 0.25f, 1.0f, 4.0f}, {}); @@ -505,13 +508,13 @@ TEST_F(MapTest, NestedBinaryMap) { TEST_F(MapTest, MapOperantionWithBuildError) { // Maps (lambda (x y) (+ x y)) onto two R1F32 vectors but uses an unsupported // type combination (F32 + U16) to test that the error is reported to the - // outermost ComputationBuilder. - ComputationBuilder builder(client_, TestName()); + // outermost XlaBuilder. + XlaBuilder builder(TestName()); auto sub_builder = builder.CreateSubBuilder("ErrorAdd"); auto x = sub_builder->Parameter(0, ShapeUtil::MakeShape(F32, {}), "x"); auto y = sub_builder->Parameter(1, ShapeUtil::MakeShape(U16, {}), "y"); - auto adder = sub_builder->Add(x, y); + sub_builder->Add(x, y); auto error_add = sub_builder->BuildAndNoteError(); std::unique_ptr param0_literal = @@ -525,9 +528,9 @@ TEST_F(MapTest, MapOperantionWithBuildError) { auto param0 = builder.Parameter(0, param0_literal->shape(), "param0"); auto param1 = builder.Parameter(1, param1_literal->shape(), "param1"); - auto map = builder.Map({param0, param1}, error_add, {0}); + builder.Map({param0, param1}, error_add, {0}); - StatusOr computation_status = builder.Build(); + StatusOr computation_status = builder.Build(); ASSERT_TRUE(!computation_status.ok()); EXPECT_THAT( computation_status.status().ToString(), @@ -545,7 +548,7 @@ using MapTestWithFullOpt = ClientLibraryTestBase; // to have issues with such patterns and maybe invalidate the pointer to entry // computation. TEST_F(MapTestWithFullOpt, MapScalarPower) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto sub_builder = builder.CreateSubBuilder("power"); auto x = sub_builder->Parameter(0, ShapeUtil::MakeShape(F32, {}), "x"); @@ -572,7 +575,7 @@ TEST_F(MapTestWithFullOpt, MapScalarPower) { // Regression test for b/35786417, where the inliner would not notice the change // of parameter order inside the map. TEST_F(MapTestWithFullOpt, MapSubtractOppositeOrder) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto sub_builder = builder.CreateSubBuilder("power"); auto x = sub_builder->Parameter(0, ShapeUtil::MakeShape(F32, {}), "x"); @@ -598,7 +601,7 @@ TEST_F(MapTestWithFullOpt, MapSubtractOppositeOrder) { // Regression test for b/35786417, where the inliner would CHECK-fail due to the // mul inside the map having more parameters than the map does. TEST_F(MapTestWithFullOpt, MapSquare) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto sub_builder = builder.CreateSubBuilder("power"); auto x = sub_builder->Parameter(0, ShapeUtil::MakeShape(F32, {}), "x"); diff --git a/tensorflow/compiler/xla/tests/reduce_precision_test.cc b/tensorflow/compiler/xla/tests/reduce_precision_test.cc index dc7ce3253cee255a7949326fa5b49fc8917432b8..b311785449f1774c3bc1e4d7ad35c2866e3b4061 100644 --- a/tensorflow/compiler/xla/tests/reduce_precision_test.cc +++ b/tensorflow/compiler/xla/tests/reduce_precision_test.cc @@ -20,9 +20,9 @@ limitations under the License. #include #include "tensorflow/compiler/xla/array2d.h" -#include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h" #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/service/reduce_precision_insertion.h" @@ -228,15 +228,14 @@ XLA_TEST_P(ReducePrecisionAccuracyTest, ReducePrecisionF32) { // This is required for proper handling of NaN values. SetFastMathDisabled(true); - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr a_literal = Literal::CreateR1({input_values}); std::unique_ptr a_data = client_->TransferToServer(*a_literal).ConsumeValueOrDie(); auto a = builder.Parameter(0, a_literal->shape(), "a"); - auto reduce_precision = - builder.ReducePrecision(a, exponent_bits, mantissa_bits); + builder.ReducePrecision(a, exponent_bits, mantissa_bits); ComputeAndCompareR1(&builder, expected_values, {a_data.get()}); } @@ -252,7 +251,7 @@ class ReducePrecisionInsertionTest : public ClientLibraryTestBase {}; // The interpreter has no fusion pass, so skip this test. XLA_TEST_F(ReducePrecisionInsertionTest, DISABLED_ON_INTERPRETER(ReducePrecisionBeforeFusion)) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr a_literal = Literal::CreateR1({1.00001}); std::unique_ptr a_data = @@ -265,7 +264,7 @@ XLA_TEST_F(ReducePrecisionInsertionTest, // Near 1.0, Log(x) approximates x - 1; this lets us confirm that the // reduce-precision operation showed up in the correct place in the // graph. - auto log = builder.Log(abs); + builder.Log(abs); // Insert precision-reduction after the Abs(x) operation, rounding that // result to exactly 1.0f. @@ -281,7 +280,7 @@ XLA_TEST_F(ReducePrecisionInsertionTest, // The interpreter has no fusion pass, so skip this test. XLA_TEST_F(ReducePrecisionInsertionTest, DISABLED_ON_INTERPRETER(ReducePrecisionSkippedAfterFusion)) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr a_literal = Literal::CreateR1({1.00001}); std::unique_ptr a_data = @@ -290,7 +289,7 @@ XLA_TEST_F(ReducePrecisionInsertionTest, // These two operations should be fused by any reasonable backend. auto abs = builder.Abs(a); - auto neg = builder.Neg(abs); + builder.Neg(abs); // Add a pass after operation fusion, suffixing kAbs operations. This // should not see into the fusion nodes and thus should not affect the @@ -307,7 +306,7 @@ XLA_TEST_F(ReducePrecisionInsertionTest, // The interpreter has no fusion pass, so skip this test. XLA_TEST_F(ReducePrecisionInsertionTest, DISABLED_ON_INTERPRETER(ReducePrecisionAddedAfterFusion)) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr a_literal = Literal::CreateR1({1.00001}); std::unique_ptr a_data = @@ -316,7 +315,7 @@ XLA_TEST_F(ReducePrecisionInsertionTest, // These two operations should be fused by any reasonable backend. auto abs = builder.Abs(a); - auto neg = builder.Neg(abs); + builder.Neg(abs); // Add a pass after operation fusion, suffixing kFusion operations. auto reduce_precision_pass = execution_options_.mutable_debug_options() @@ -331,7 +330,7 @@ XLA_TEST_F(ReducePrecisionInsertionTest, // The interpreter has no fusion pass, so skip this test. XLA_TEST_F(ReducePrecisionInsertionTest, DISABLED_ON_INTERPRETER(ReducePrecisionSkippedFusionContains)) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr a_literal = Literal::CreateR1({1.00001}); std::unique_ptr a_data = @@ -340,7 +339,7 @@ XLA_TEST_F(ReducePrecisionInsertionTest, // These two operations should be fused by any reasonable backend. auto abs = builder.Abs(a); - auto neg = builder.Neg(abs); + builder.Neg(abs); // Add a pass suffixing fusion nodes containing kCos operations. This // should have no effect. @@ -356,7 +355,7 @@ XLA_TEST_F(ReducePrecisionInsertionTest, // The interpreter has no fusion pass, so skip this test. XLA_TEST_F(ReducePrecisionInsertionTest, DISABLED_ON_INTERPRETER(ReducePrecisionAddedFusionContains)) { - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); std::unique_ptr a_literal = Literal::CreateR1({1.00001}); std::unique_ptr a_data = @@ -365,7 +364,7 @@ XLA_TEST_F(ReducePrecisionInsertionTest, // These two operations should be fused by any reasonable backend. auto abs = builder.Abs(a); - auto neg = builder.Neg(abs); + builder.Neg(abs); // Add a pass suffixing fusion nodes containing kAbs operations. This // should see the kAbs operation within the above fusion node. diff --git a/tensorflow/compiler/xla/tests/reduce_test.cc b/tensorflow/compiler/xla/tests/reduce_test.cc index d24927d22b6534b46e711cd442f19a3e5cfcebdf..423ccadb5b3b7df950824349737a833c08870d77 100644 --- a/tensorflow/compiler/xla/tests/reduce_test.cc +++ b/tensorflow/compiler/xla/tests/reduce_test.cc @@ -39,6 +39,8 @@ limitations under the License. #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h" #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/reference_util.h" @@ -50,6 +52,7 @@ limitations under the License. #include "tensorflow/compiler/xla/tests/test_macros.h" #include "tensorflow/compiler/xla/util.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/test.h" #include "tensorflow/core/platform/types.h" @@ -502,21 +505,18 @@ XLA_TEST_F(ReduceTest, TransposeAndReduceElementwiseR2_111x50_To_R1) { // 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); + XlaBuilder builder(TestName()); + XlaComputation 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}); + XlaOp input = builder.Parameter(0, input_shape, "input"); + XlaOp zero = builder.ConstantR0(0.0); + XlaOp transpose = builder.Transpose(input, /*permutation=*/{1, 0, 2}); + 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)); + ComputeAndCompare(&builder, {std::move(*input_data)}, ErrorSpec(0.01, 1e-4)); } XLA_TEST_F(ReduceTest, Reshape_111x2x25Reduce_111x50_To_R1) { @@ -935,5 +935,36 @@ XLA_TEST_F(ReduceInitializerTest, U64InitializerBigValue) { DoTest(1234556789123, 1024); } +// Test the operational semantic that the init value is passed on the lhs for +// reduces. Can be tested by performing an "identity" reduce (that simply +// returns one of the parameters). In this case, we return the rhs, which for +// a 1D array with one element, should not be the init value. +XLA_TEST_F(ReduceTest, ReduceIdentity) { + ComputationBuilder builder(client_, TestName()); + Shape single_float = ShapeUtil::MakeShape(F32, {}); + builder.Parameter(0, single_float, "lhs-unused"); + builder.Parameter(1, single_float, "rhs-used"); + auto computation_status = builder.Build(); + TF_ASSERT_OK(computation_status.status()); + + Shape operand_shape = ShapeUtil::MakeShape(F32, {1}); + builder.Reduce(builder.Parameter(0, operand_shape, "operand"), + builder.Parameter(1, single_float, "init"), + computation_status.ValueOrDie(), {0}); + + float operand[] = {42.0f}; + float init = 58.5f; + float expected = 42.0f; + std::unique_ptr input_literal = Literal::CreateR1(operand); + std::unique_ptr input_global_data = + client_->TransferToServer(*input_literal).ConsumeValueOrDie(); + std::unique_ptr input_literal2 = Literal::CreateR0(init); + std::unique_ptr input_global_data2 = + client_->TransferToServer(*input_literal2).ConsumeValueOrDie(); + ComputeAndCompareR0( + &builder, expected, {input_global_data.get(), input_global_data2.get()}, + ErrorSpec(0.0001)); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/reduce_window_test.cc b/tensorflow/compiler/xla/tests/reduce_window_test.cc index 8dd24f1237136e2807cea8a261ead25f5c7adbb2..0a097667222d8c315cb9943688d6dbcb4426a8b3 100644 --- a/tensorflow/compiler/xla/tests/reduce_window_test.cc +++ b/tensorflow/compiler/xla/tests/reduce_window_test.cc @@ -21,10 +21,11 @@ limitations under the License. #include "tensorflow/compiler/xla/array2d.h" #include "tensorflow/compiler/xla/array3d.h" #include "tensorflow/compiler/xla/array4d.h" -#include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/client/padding.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h" #include "tensorflow/compiler/xla/reference_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" @@ -63,11 +64,9 @@ class ReduceWindowTestBase : public ClientLibraryTestBase { class ReduceWindowTest : public ::testing::WithParamInterface, public ReduceWindowTestBase { public: - ReduceWindowTest() : builder_(client_, TestName()) { - set_use_bfloat16(GetParam()); - } + ReduceWindowTest() : builder_(TestName()) { set_use_bfloat16(GetParam()); } - void ReduceWindowAdd(const ComputationDataHandle& input, + void ReduceWindowAdd(const XlaOp& input, tensorflow::gtl::ArraySlice window_dimensions, tensorflow::gtl::ArraySlice window_strides, Padding padding) { @@ -78,16 +77,17 @@ class ReduceWindowTest : public ::testing::WithParamInterface, window_dimensions, window_strides, padding); } - void ReduceWindowMax(const ComputationDataHandle& input, + void ReduceWindowMax(const XlaOp& input, tensorflow::gtl::ArraySlice window_dimensions, tensorflow::gtl::ArraySlice window_strides, Padding padding) { auto init = CreateConstantFromLiteral(Literal::MinValue(F32), &builder_); - builder_.ReduceWindow(input, init, CreateScalarMax(), window_dimensions, - window_strides, padding); + builder_.ReduceWindow(input, init, + CreateScalarMaxComputation(FloatType(), &builder_), + window_dimensions, window_strides, padding); } - void ReduceWindowMin(const ComputationDataHandle& input, + void ReduceWindowMin(const XlaOp& input, tensorflow::gtl::ArraySlice window_dimensions, tensorflow::gtl::ArraySlice window_strides, Padding padding) { @@ -97,7 +97,7 @@ class ReduceWindowTest : public ::testing::WithParamInterface, window_dimensions, window_strides, padding); } - ComputationBuilder builder_; + XlaBuilder builder_; }; TEST_P(ReduceWindowTest, MismatchedRanksGivesErrorStatus) { @@ -310,7 +310,7 @@ XLA_TEST_P(ReduceWindowTest, NonstandardReduceFunction) { auto rhs = b->Parameter(1, scalar, "rhs"); b->Min(b->Add(lhs, rhs), CreateConstantFromLiteral(*Literal::CreateR0(8.0f), b.get())); - Computation reduce_fn = b->BuildAndNoteError(); + XlaComputation reduce_fn = b->BuildAndNoteError(); builder_.ReduceWindow( input, @@ -338,7 +338,7 @@ TEST_P(ReduceWindowTest, R4UnitWindow) { std::unique_ptr input_literal = Literal::CreateR4FromArray4DWithLayout( input_array, LayoutUtil::MakeLayout({0, 3, 2, 1})); - ComputationDataHandle input; + XlaOp input; auto input_data = CreateParameterAndTransferLiteral( 0, *input_literal, "parameter", &builder_, &input); @@ -406,7 +406,7 @@ XLA_TEST_P(ReduceWindowTest, R4SecondMinorStride) { std::unique_ptr input_literal = Literal::CreateR4FromArray4DWithLayout( input_array, LayoutUtil::MakeLayout({3, 2, 1, 0})); - ComputationDataHandle input; + XlaOp input; auto input_data = CreateParameterAndTransferLiteral( 0, *input_literal, "parameter", &builder_, &input); @@ -428,7 +428,7 @@ XLA_TEST_P(ReduceWindowTest, R4SecondMinorUnitStride) { std::unique_ptr input_literal = Literal::CreateR4FromArray4DWithLayout( input_array, LayoutUtil::MakeLayout({3, 2, 1, 0})); - ComputationDataHandle input; + XlaOp input; auto input_data = CreateParameterAndTransferLiteral( 0, *input_literal, "parameter", &builder_, &input); @@ -450,7 +450,7 @@ XLA_TEST_P(ReduceWindowTest, R4SecondMinorWin) { std::unique_ptr input_literal = Literal::CreateR4FromArray4DWithLayout( input_array, LayoutUtil::MakeLayout({3, 2, 1, 0})); - ComputationDataHandle input; + XlaOp input; auto input_data = CreateParameterAndTransferLiteral( 0, *input_literal, "parameter", &builder_, &input); @@ -551,7 +551,7 @@ TEST_P(ReduceWindowTest, R2ReduceWindowInceptionFromBroadcast) { TEST_P(ReduceWindowTest, R2ReduceWindowNonOverlappingFromBroadcast) { Array2D input_array(6, 4, 1.0f); - ComputationDataHandle input = builder_.Broadcast( + XlaOp input = builder_.Broadcast( CreateConstantFromLiteral(Literal::One(F32), &builder_), {6, 4}); Padding padding = Padding::kSame; @@ -610,7 +610,7 @@ class R4ReduceWindowTest : public ReduceWindowTestBase, R4ReduceWindowTest() { set_use_bfloat16(::testing::get<1>(GetParam())); } void DoIt() { - ComputationBuilder b(client_, TestName()); + XlaBuilder b(TestName()); const auto& param = ::testing::get<0>(GetParam()); const float kInitValue = 0.0f; @@ -621,7 +621,7 @@ class R4ReduceWindowTest : public ReduceWindowTestBase, std::unique_ptr input_literal = Literal::CreateR4FromArray4DWithLayout( input, LayoutUtil::MakeLayout(param.layout)); - ComputationDataHandle parameter; + XlaOp parameter; auto input_arg = CreateParameterAndTransferLiteral(0, *input_literal, "p0", &b, ¶meter); @@ -962,7 +962,7 @@ class R3ReduceWindowTest : public ReduceWindowTestBase, }; TEST_P(R3ReduceWindowTest, Add) { - ComputationBuilder b(client_, TestName()); + XlaBuilder b(TestName()); const auto& param = ::testing::get<0>(GetParam()); CHECK(param.reducer == kAdd); @@ -973,7 +973,7 @@ TEST_P(R3ReduceWindowTest, Add) { Literal::CreateR3FromArray3DWithLayout( input, LayoutUtil::MakeLayout(param.layout)); - ComputationDataHandle parameter; + XlaOp parameter; auto input_arg = CreateParameterAndTransferLiteral(0, *input_literal, "p0", &b, ¶meter); auto init_value = @@ -1063,15 +1063,12 @@ struct R2ReduceWindowTestData { /*strides=*/{1, 1}, /*pad_low=*/{0, 130}, /*pad_high=*/{0, 0}, /*layout=*/{1, 0}, /*reducer=*/Reducer::kAdd}, -// TODO(b/76025683): These tests fail on TPU. -#if defined(XLA_TEST_BACKEND_CPU) || defined(XLA_TEST_BACKEND_GPU) - {/*base_bounds=*/{4096, 4096}, /*window_bounds=*/{1, 4}, - /*strides=*/{1, 1024}, /*pad_low=*/{0, 0}, /*pad-high=*/{0, 0}, - /*layout=*/{1, 0}, /*reducer=*/Reducer::kAdd}, {/*base_bounds=*/{8, 256}, /*window_bounds=*/{1, 4}, /*strides=*/{1, 64}, /*pad_low=*/{0, 0}, /*pad_high=*/{0, 0}, /*layout=*/{1, 0}, /*reducer=*/Reducer::kAdd}, -#endif + {/*base_bounds=*/{4096, 4096}, /*window_bounds=*/{1, 4}, + /*strides=*/{1, 1024}, /*pad_low=*/{0, 0}, /*pad-high=*/{0, 0}, + /*layout=*/{1, 0}, /*reducer=*/Reducer::kAdd}, }; string R2ReduceWindowTestDataToString( @@ -1100,7 +1097,7 @@ class R2ReduceWindowTest : public ReduceWindowTestBase, R2ReduceWindowTest() { set_use_bfloat16(::testing::get<1>(GetParam())); } void DoIt() { - ComputationBuilder b(client_, TestName()); + XlaBuilder b(TestName()); const auto& param = ::testing::get<0>(GetParam()); CHECK(param.reducer == kAdd); @@ -1110,7 +1107,7 @@ class R2ReduceWindowTest : public ReduceWindowTestBase, Literal::CreateR2FromArray2DWithLayout( input, LayoutUtil::MakeLayout(param.layout)); - ComputationDataHandle parameter; + XlaOp parameter; auto input_arg = CreateParameterAndTransferLiteral(0, *input_literal, "p0", &b, ¶meter); std::vector> padding(2); @@ -1298,7 +1295,7 @@ class R1ReduceWindowTest : public ReduceWindowTestBase, }; TEST_P(R1ReduceWindowTest, DoIt) { - ComputationBuilder b(client_, TestName()); + XlaBuilder b(TestName()); const auto& param = ::testing::get<0>(GetParam()); CHECK(param.reducer == kAdd || param.reducer == kMax); @@ -1307,7 +1304,7 @@ TEST_P(R1ReduceWindowTest, DoIt) { std::iota(std::begin(input_vector), std::end(input_vector), 0); std::unique_ptr input_literal = Literal::CreateR1(tensorflow::gtl::ArraySlice(input_vector)); - ComputationDataHandle parameter; + XlaOp parameter; auto input_arg = CreateParameterAndTransferLiteral(0, *input_literal, "p0", &b, ¶meter); @@ -1438,5 +1435,22 @@ ENTRY R3Window { EXPECT_TRUE(RunAndCompare(hlo_string, ErrorSpec{0.001})); } +TEST_F(HloTestBase, ReduceWindowIdentity) { + const string& hlo_string = R"( +HloModule ReduceWindowIdentity +identity.pad_to_reduce_window { + param0 = f32[] parameter(0) + ROOT param1 = f32[] parameter(1) +} +ENTRY reduce-window-identity { + operand = f32[1,32,64]{2,1,0} parameter(0) + constant.4466 = f32[] constant(0) + ROOT reduce-window = f32[1,33,64]{2,1,0} reduce-window(operand, constant.4466), window={size=1x1x1 pad=0_0x1_0x0_0}, to_apply=identity.pad_to_reduce_window +} + +)"; + EXPECT_TRUE(RunAndCompare(hlo_string, tensorflow::gtl::nullopt)); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/reverse_test.cc b/tensorflow/compiler/xla/tests/reverse_test.cc index 8fc841f14087cdea02fe44cdaea521ff92122aec..6959c95502cb7af6b720592e7836c6789719a528 100644 --- a/tensorflow/compiler/xla/tests/reverse_test.cc +++ b/tensorflow/compiler/xla/tests/reverse_test.cc @@ -17,8 +17,8 @@ limitations under the License. #include "tensorflow/compiler/xla/array2d.h" #include "tensorflow/compiler/xla/array4d.h" -#include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_builder.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" @@ -85,7 +85,7 @@ TEST_P(FloatReverseTest, Reverses) { auto r1_literal = Literal::CreateR1(input_vector); auto input_literal = r1_literal->Reshape(spec.input_dims).ConsumeValueOrDie(); - ComputationBuilder builder(client_, TestName()); + XlaBuilder builder(TestName()); auto a = AddParam(*input_literal, &builder); builder.Rev(a, spec.reversal); diff --git a/tensorflow/compiler/xla/tests/select_and_scatter_test.cc b/tensorflow/compiler/xla/tests/select_and_scatter_test.cc index d268fdcacebcb162bf61bc7dd4b208f4db6c4a5f..7015e5a6a31f506d30c2629d7735482cf354455a 100644 --- a/tensorflow/compiler/xla/tests/select_and_scatter_test.cc +++ b/tensorflow/compiler/xla/tests/select_and_scatter_test.cc @@ -19,11 +19,11 @@ limitations under the License. #include #include "tensorflow/compiler/xla/array2d.h" -#include "tensorflow/compiler/xla/client/computation.h" -#include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/client/padding.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h" +#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h" #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/reference_util.h" @@ -50,7 +50,7 @@ class SelectAndScatterTest : public ClientLibraryTestBase, public ::testing::WithParamInterface { public: - SelectAndScatterTest() : builder_(client_, TestName()) { + SelectAndScatterTest() : builder_(TestName()) { // Create S32 GE and ADD computations for select and scatter respectively. ge_s32_ = CreateScalarGeComputation(S32, &builder_); add_s32_ = CreateScalarAddComputation(S32, &builder_); @@ -60,13 +60,13 @@ class SelectAndScatterTest min_f32_ = CreateScalarMinComputation(F32, &builder_); } - ComputationBuilder builder_; - Computation ge_s32_; - Computation add_s32_; - Computation ge_f32_; - Computation add_f32_; - Computation max_f32_; - Computation min_f32_; + XlaBuilder builder_; + XlaComputation ge_s32_; + XlaComputation add_s32_; + XlaComputation ge_f32_; + XlaComputation add_f32_; + XlaComputation max_f32_; + XlaComputation min_f32_; }; XLA_TEST_P(SelectAndScatterTest, ParamTest) { @@ -80,12 +80,11 @@ XLA_TEST_P(SelectAndScatterTest, ParamTest) { s.FillRandom(12.0f); auto source = builder_.ConstantFromArray(s); - auto select_and_scatter = builder_.SelectAndScatter( - operand, ge_f32_, GetParam().window_dimensions, GetParam().window_strides, - GetParam().padding_type, source, builder_.ConstantR0(0.0f), - add_f32_); + builder_.SelectAndScatter(operand, ge_f32_, GetParam().window_dimensions, + GetParam().window_strides, GetParam().padding_type, + source, builder_.ConstantR0(0.0f), add_f32_); - ComputeAndCompare(&builder_, select_and_scatter, {}, ErrorSpec(1e-5)); + ComputeAndCompare(&builder_, {}, ErrorSpec(1e-5)); } INSTANTIATE_TEST_CASE_P( diff --git a/tensorflow/compiler/xla/tests/slice_test.cc b/tensorflow/compiler/xla/tests/slice_test.cc index 69fbe98bd63661322d37936c90a5fe3580efc2de..52195db2aa74710b901dd7744a670764a034e96b 100644 --- a/tensorflow/compiler/xla/tests/slice_test.cc +++ b/tensorflow/compiler/xla/tests/slice_test.cc @@ -299,9 +299,12 @@ INSTANTIATE_TEST_CASE_P( SliceR1TestBigSlicesInstantiation, SliceR1LargeTest, ::testing::Values( - 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} + 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} ), SliceR1TestDataToString ); diff --git a/tensorflow/compiler/xla/tests/test_utils.cc b/tensorflow/compiler/xla/tests/test_utils.cc index 68f75d50cbaa9b6eab56714987196cc9aced60a9..997a1d8273736af31994ebbd07ff3857d1e8e0b5 100644 --- a/tensorflow/compiler/xla/tests/test_utils.cc +++ b/tensorflow/compiler/xla/tests/test_utils.cc @@ -165,7 +165,7 @@ enum class ConstantType { kUnknown, kZero, kOne }; // Return the constant type required by this computation, if known. ConstantType GetInitValue(const HloComputation& computation) { const HloInstruction* const root = computation.root_instruction(); - if (computation.num_parameters() != 2 || + if (computation.num_parameters() != 2 || root->operand_count() != 2 || root->operand(0)->opcode() != HloOpcode::kParameter || root->operand(1)->opcode() != HloOpcode::kParameter || root->operand(0) == root->operand(1)) { @@ -339,9 +339,9 @@ StatusOr>> MakeFakeArguments( return std::move(arguments); } -Status VerifyHloModule(const perftools::gputools::Platform& platform, - HloModule* const module) { - return HloVerifier().Run(module).status(); +Status VerifyHloModule(const se::Platform& platform, HloModule* const module, + bool allow_mixed_precision) { + return HloVerifier(allow_mixed_precision).Run(module).status(); } } // namespace xla diff --git a/tensorflow/compiler/xla/tests/test_utils.h b/tensorflow/compiler/xla/tests/test_utils.h index 0fb024ffb074f1c90b75022bc7f5a8b58b03c0c2..30c147910cae85e1ebdddc22e637a6c1fd577c20 100644 --- a/tensorflow/compiler/xla/tests/test_utils.h +++ b/tensorflow/compiler/xla/tests/test_utils.h @@ -68,8 +68,8 @@ StatusOr>> MakeFakeArguments( // Check that a given module satisfies various constraints before trying to // execute it. -Status VerifyHloModule(const perftools::gputools::Platform& platform, - HloModule* const module); +Status VerifyHloModule(const se::Platform& platform, HloModule* const module, + bool allow_mixed_precision = false); } // namespace xla diff --git a/tensorflow/compiler/xla/tests/test_utils_test.cc b/tensorflow/compiler/xla/tests/test_utils_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..e8efc6e2a83f42bf81fc1261ba508632cf3f85b3 --- /dev/null +++ b/tensorflow/compiler/xla/tests/test_utils_test.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/compiler/xla/tests/test_utils.h" + +#include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/tests/local_client_test_base.h" +#include "tensorflow/compiler/xla/tests/test_macros.h" +#include "tensorflow/core/lib/core/status_test_util.h" + +namespace xla { +namespace { + +// A test fixture is used because we need a client for our computation builder. +class TestUtilsTest : public LocalClientTestBase {}; + +XLA_TEST_F(TestUtilsTest, UnusedParam) { + ComputationBuilder builder(local_client_, TestName()); + // Make the reduction lambda. + Shape single_float = ShapeUtil::MakeShape(F32, {}); + builder.Parameter(0, single_float, "unused"); + builder.Parameter(1, single_float, "used"); + auto computation_status = builder.Build(); + TF_ASSERT_OK(computation_status.status()); + + // Make the reduction. + Shape pair_float = ShapeUtil::MakeShape(F32, {2}); + builder.Reduce(builder.Parameter(0, pair_float, "operand"), + builder.Parameter(1, single_float, "init"), + computation_status.ValueOrDie(), {0}); + computation_status = builder.Build(); + TF_ASSERT_OK(computation_status.status()); + + auto executable_status = local_client_->Compile( + computation_status.ValueOrDie(), {&pair_float, &single_float}, + ExecutableBuildOptions()); + TF_ASSERT_OK(executable_status.status()); + HloModule& module = const_cast( + executable_status.ValueOrDie()->executable()->module()); + TF_ASSERT_OK(MakeFakeArguments(&module).status()); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/transfer_manager_test.cc b/tensorflow/compiler/xla/tests/transfer_manager_test.cc index 268ba338f2e6740a1d1a046d5a85494f3cf2e9f8..e2067bc1b835a946fc56801cbf227e05ef0686b4 100644 --- a/tensorflow/compiler/xla/tests/transfer_manager_test.cc +++ b/tensorflow/compiler/xla/tests/transfer_manager_test.cc @@ -45,7 +45,7 @@ class TransferManagerTest : public LocalClientTestBase { ~TransferManagerTest() override = default; - std::unique_ptr AllocateDeviceBuffer(const Shape& shape) { + ScopedShapedBuffer AllocateDeviceBuffer(const Shape& shape) { return transfer_manager_ ->AllocateScopedShapedBuffer( shape, GetOrCreateAllocator(local_client_->platform()), @@ -64,10 +64,10 @@ XLA_TEST_F(TransferManagerTest, TransferR0U32) { // Round trip literal through device. ASSERT_IS_OK(transfer_manager_->TransferLiteralToDevice( - stream_executor_, *literal, *device_buffer)); + stream_executor_, *literal, device_buffer)); TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result, transfer_manager_->TransferLiteralFromDevice( - stream_executor_, *device_buffer)); + stream_executor_, device_buffer)); LiteralTestUtil::ExpectR0Equal(42, *result); } @@ -80,10 +80,10 @@ XLA_TEST_F(TransferManagerTest, TransferR1F32) { // Round trip literal through device. ASSERT_IS_OK(transfer_manager_->TransferLiteralToDevice( - stream_executor_, *literal, *device_buffer)); + stream_executor_, *literal, device_buffer)); TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result, transfer_manager_->TransferLiteralFromDevice( - stream_executor_, *device_buffer)); + stream_executor_, device_buffer)); LiteralTestUtil::ExpectR1Equal({1.25f, 2.5f, -17.0f, -20.125f}, *result); @@ -98,10 +98,10 @@ XLA_TEST_F(TransferManagerTest, TransferR1LargeF32) { // Round trip literal through device. ASSERT_IS_OK(transfer_manager_->TransferLiteralToDevice( - stream_executor_, *literal, *device_buffer)); + stream_executor_, *literal, device_buffer)); TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result, transfer_manager_->TransferLiteralFromDevice( - stream_executor_, *device_buffer)); + stream_executor_, device_buffer)); LiteralTestUtil::ExpectR1Equal(test_vector, *result); } @@ -114,10 +114,10 @@ XLA_TEST_F(TransferManagerTest, TransferR1U8) { // Round trip literal through device. ASSERT_IS_OK(transfer_manager_->TransferLiteralToDevice( - stream_executor_, *literal, *device_buffer)); + stream_executor_, *literal, device_buffer)); TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result, transfer_manager_->TransferLiteralFromDevice( - stream_executor_, *device_buffer)); + stream_executor_, device_buffer)); EXPECT_EQ(result->GetR1U8AsString(), test_string); } @@ -130,10 +130,10 @@ XLA_TEST_F(TransferManagerTest, TransferR2F32) { // Round trip literal through device. ASSERT_IS_OK(transfer_manager_->TransferLiteralToDevice( - stream_executor_, *literal, *device_buffer)); + stream_executor_, *literal, device_buffer)); TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result, transfer_manager_->TransferLiteralFromDevice( - stream_executor_, *device_buffer)); + stream_executor_, device_buffer)); LiteralTestUtil::ExpectR2Equal( {{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}}, *result); @@ -150,10 +150,10 @@ XLA_TEST_F(TransferManagerTest, // Round trip literal through device. Set the on-device layout to something // different than the literal layout. ASSERT_IS_OK(transfer_manager_->TransferLiteralToDevice( - stream_executor_, *literal, *device_buffer)); + stream_executor_, *literal, device_buffer)); TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result, transfer_manager_->TransferLiteralFromDevice( - stream_executor_, *device_buffer)); + stream_executor_, device_buffer)); EXPECT_FALSE( LayoutUtil::Equal(result->shape().layout(), literal->shape().layout())); @@ -170,10 +170,10 @@ XLA_TEST_F(TransferManagerTest, TransferTuple) { // Round trip literal through device. ASSERT_IS_OK(transfer_manager_->TransferLiteralToDevice( - stream_executor_, *literal, *device_buffer)); + stream_executor_, *literal, device_buffer)); TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result, transfer_manager_->TransferLiteralFromDevice( - stream_executor_, *device_buffer)); + stream_executor_, device_buffer)); LiteralTestUtil::ExpectEqual(*literal, *result); } @@ -184,10 +184,10 @@ XLA_TEST_F(TransferManagerTest, TransferEmptyTuple) { // Round trip literal through device. ASSERT_IS_OK(transfer_manager_->TransferLiteralToDevice( - stream_executor_, *literal, *device_buffer)); + stream_executor_, *literal, device_buffer)); TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result, transfer_manager_->TransferLiteralFromDevice( - stream_executor_, *device_buffer)); + stream_executor_, device_buffer)); LiteralTestUtil::ExpectEqual(*literal, *result); } @@ -204,10 +204,10 @@ XLA_TEST_F(TransferManagerTest, TransferNestedTuple) { // Round trip literal through device. ASSERT_IS_OK(transfer_manager_->TransferLiteralToDevice( - stream_executor_, *literal, *device_buffer)); + stream_executor_, *literal, device_buffer)); TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result, transfer_manager_->TransferLiteralFromDevice( - stream_executor_, *device_buffer)); + stream_executor_, device_buffer)); LiteralTestUtil::ExpectEqual(*literal, *result); } @@ -219,10 +219,10 @@ XLA_TEST_F(TransferManagerTest, TransferComplexValue) { // Round trip literal through device. ASSERT_IS_OK(transfer_manager_->TransferLiteralToDevice( - stream_executor_, *literal, *device_buffer)); + stream_executor_, *literal, device_buffer)); TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result, transfer_manager_->TransferLiteralFromDevice( - stream_executor_, *device_buffer)); + stream_executor_, device_buffer)); LiteralTestUtil::ExpectEqual(*literal, *result); } @@ -238,10 +238,10 @@ XLA_TEST_F(TransferManagerTest, TransferComplexValueInTuple) { // Round trip literal through device. ASSERT_IS_OK(transfer_manager_->TransferLiteralToDevice( - stream_executor_, *literal, *device_buffer)); + stream_executor_, *literal, device_buffer)); TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr result, transfer_manager_->TransferLiteralFromDevice( - stream_executor_, *device_buffer)); + stream_executor_, device_buffer)); LiteralTestUtil::ExpectEqual(*literal, *result); } diff --git a/tensorflow/compiler/xla/tests/tuple_test.cc b/tensorflow/compiler/xla/tests/tuple_test.cc index 098be6d7aabe88d0deef600716229ddbd0bcae2f..61d0fa02aba80ddf60ab83f8355a5e3c1ef44518 100644 --- a/tensorflow/compiler/xla/tests/tuple_test.cc +++ b/tensorflow/compiler/xla/tests/tuple_test.cc @@ -535,8 +535,7 @@ TEST_F(TupleHloTest, HloRunner::CreateModuleFromString(testcase, GetDebugOptionsForTest()) .ValueOrDie(); auto param = Literal::MakeTupleOwned(Literal::CreateR1({1, 2, 3})); - TF_ASSERT_OK_AND_ASSIGN(auto result, - ExecuteNoHloPasses(std::move(module), {param.get()})); + auto result = ExecuteNoHloPasses(std::move(module), {param.get()}); EXPECT_TRUE(LiteralTestUtil::Equal( *result, *Literal::MakeTupleOwned(Literal::CreateR2({{1, 2, 3}})))); diff --git a/tensorflow/compiler/xla/tests/vector_ops_simple_test.cc b/tensorflow/compiler/xla/tests/vector_ops_simple_test.cc index b52c718814d4ffeff68c60588a6637a2159d57e5..697d78fe6e9abf6136f9ec4a708ef983732d0013 100644 --- a/tensorflow/compiler/xla/tests/vector_ops_simple_test.cc +++ b/tensorflow/compiler/xla/tests/vector_ops_simple_test.cc @@ -39,7 +39,7 @@ namespace { class VecOpsSimpleTest : public ClientLibraryTestBase { public: - explicit VecOpsSimpleTest(perftools::gputools::Platform* platform = nullptr) + explicit VecOpsSimpleTest(se::Platform* platform = nullptr) : ClientLibraryTestBase(platform) { mutable_debug_options()->add_xla_disable_hlo_passes("algsimp"); mutable_debug_options()->add_xla_disable_hlo_passes("inline"); diff --git a/tensorflow/compiler/xla/tests/while_test.cc b/tensorflow/compiler/xla/tests/while_test.cc index 89ce2ce797f979b8668fbdb172a4a3abc5922b9f..1e18b5679956710d7e43c6209be3fdfe4877b4d3 100644 --- a/tensorflow/compiler/xla/tests/while_test.cc +++ b/tensorflow/compiler/xla/tests/while_test.cc @@ -37,8 +37,6 @@ limitations under the License. #include "tensorflow/core/platform/test_benchmark.h" #include "tensorflow/core/platform/types.h" -namespace se = ::perftools::gputools; - namespace xla { namespace { diff --git a/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc b/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc index ff3418a128eed82b730a6602d6e3faba4ad7be32..837a01e873e1cfbcf7c52c13833bbe4366cc9163 100644 --- a/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc +++ b/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc @@ -34,7 +34,7 @@ limitations under the License. namespace xla { namespace { -namespace se = ::perftools::gputools; + namespace gtl = ::tensorflow::gtl; class HloProfileTest : public ClientLibraryTestBase {}; @@ -129,18 +129,18 @@ void ExecuteAndFetchProfile(string* profile_output, LocalClient* client, auto* transfer_manager = backend->transfer_manager(); TF_ASSERT_OK_AND_ASSIGN( - std::unique_ptr lhs_arg, + ScopedShapedBuffer lhs_arg, transfer_manager->AllocateScopedShapedBuffer( lhs_arg_shape, allocator, backend->default_device_ordinal())); TF_ASSERT_OK(transfer_manager->TransferLiteralToDevice( - executor, *Literal::CreateFromShape(lhs_arg_shape), *lhs_arg)); + executor, *Literal::CreateFromShape(lhs_arg_shape), lhs_arg)); TF_ASSERT_OK_AND_ASSIGN( - std::unique_ptr rhs_arg, + ScopedShapedBuffer rhs_arg, transfer_manager->AllocateScopedShapedBuffer( rhs_arg_shape, allocator, backend->default_device_ordinal())); TF_ASSERT_OK(transfer_manager->TransferLiteralToDevice( - executor, *Literal::CreateFromShape(rhs_arg_shape), *rhs_arg)); + executor, *Literal::CreateFromShape(rhs_arg_shape), rhs_arg)); TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr local_executable, @@ -165,7 +165,7 @@ void ExecuteAndFetchProfile(string* profile_output, LocalClient* client, backend->eigen_intra_op_thread_pool()); TF_ASSERT_OK_AND_ASSIGN( auto execution_result, - executable->ExecuteOnStream(&run_options, {lhs_arg.get(), rhs_arg.get()}, + executable->ExecuteOnStream(&run_options, {&lhs_arg, &rhs_arg}, &hlo_execution_profile)); (void)execution_result; diff --git a/tensorflow/compiler/xla/tools/parser/hlo_parser.cc b/tensorflow/compiler/xla/tools/parser/hlo_parser.cc index b2f122982adf750106f034e7e786367720ebafcf..95d3fd28b38a5945f5ed685f07db046701250273 100644 --- a/tensorflow/compiler/xla/tools/parser/hlo_parser.cc +++ b/tensorflow/compiler/xla/tools/parser/hlo_parser.cc @@ -470,6 +470,7 @@ bool HloParser::ParseInstruction(HloComputation::Builder* builder, case HloOpcode::kRoundNearestAfz: case HloOpcode::kBitcast: case HloOpcode::kCeil: + case HloOpcode::kClz: case HloOpcode::kCopy: case HloOpcode::kCos: case HloOpcode::kExp: @@ -724,15 +725,6 @@ bool HloParser::ParseInstruction(HloComputation::Builder* builder, shape, operands[0], *broadcast_dimensions)); break; } - case HloOpcode::kBroadcastDimOne: { - if (!ParseOperands(&operands, /*expected_size=*/1) || - !ParseAttributes(attrs)) { - return false; - } - instruction = builder->AddInstruction( - HloInstruction::CreateBroadcastDimOne(shape, operands[0])); - break; - } case HloOpcode::kConcatenate: { optional> dimensions; attrs["dimensions"] = {/*required=*/true, AttrTy::kBracedInt64List, diff --git a/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc b/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc index 57684b58346166f7e3ef9576f6cd8f70ab9dc389..adc8b1d620eb65fdca19072831360b71847abf9e 100644 --- a/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc +++ b/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc @@ -57,18 +57,6 @@ ENTRY %axpy.v5 (alpha: f32[], x: f32[2,4], y: f32[2,4]) -> f32[2,4] { ROOT %add = f32[2,4]{1,0} add(f32[2,4]{1,0} %multiply, f32[2,4]{1,0} %y) } -)" -}, -// broadcast size-one dimensions -{ -"BroadcastDimOne", -R"(HloModule broadcast_dim_one_module - -ENTRY %broadcast-dim-one () -> f32[2,2] { - %constant = f32[1,2]{1,0} constant(f32[1,2] { { 1.1, 2.2 } }) - ROOT %broadcast-dim-one = f32[2,2]{1,0} broadcast-dim-one(f32[1,2]{1,0} %constant) -} - )" }, // pred constant diff --git a/tensorflow/compiler/xla/types.h b/tensorflow/compiler/xla/types.h index 9fa4297523bab0748863479be52dff1b7b523a8b..b645acb700b0f168112a40c9c72b4669435f717d 100644 --- a/tensorflow/compiler/xla/types.h +++ b/tensorflow/compiler/xla/types.h @@ -46,4 +46,10 @@ using ::Eigen::half; } // namespace xla +// Alias namespace ::stream_executor as ::xla::se. +namespace stream_executor {} +namespace xla { +namespace se = ::stream_executor; +} // namespace xla + #endif // TENSORFLOW_COMPILER_XLA_TYPES_H_ diff --git a/tensorflow/compiler/xla/util.h b/tensorflow/compiler/xla/util.h index 2da9f9ed6f40fcf5b2512f974519df0b355da10f..be33bd6dd1304fa8fc6e5aed1d4c4d65bf97e692 100644 --- a/tensorflow/compiler/xla/util.h +++ b/tensorflow/compiler/xla/util.h @@ -528,6 +528,16 @@ bool IsInt32(T x) { // value is implementation-defined." return static_cast(x) == x; } + +template +Status EraseElementFromVector(std::vector* container, const T& value) { + // c_find returns a const_iterator which does not seem to work on gcc 4.8.4, + // and this breaks the ubuntu/xla_gpu build bot. + auto it = std::find(container->begin(), container->end(), value); + TF_RET_CHECK(it != container->end()); + container->erase(it); + return Status::OK(); +} } // namespace xla #define XLA_LOG_LINES(SEV, STRING) \ diff --git a/tensorflow/compiler/xla/xla.bzl b/tensorflow/compiler/xla/xla.bzl index 6b136d333bbf079efd314833f46fe3b98743fbac..1439f1bcc5cec39203a7cb4b1f8604e7349382c6 100644 --- a/tensorflow/compiler/xla/xla.bzl +++ b/tensorflow/compiler/xla/xla.bzl @@ -6,7 +6,9 @@ load("//tensorflow/core:platform/default/build_config_root.bzl", "if_static") # xla_proto_library() is a convenience wrapper around cc_proto_library. -def xla_proto_library(name, srcs=[], deps=[], visibility=None, testonly=0): +def xla_proto_library(name, srcs=[], deps=[], visibility=None, testonly=0, **kwargs): + if kwargs.get('use_grpc_plugin'): + kwargs['use_grpc_namespace'] = True cc_proto_library(name=name, srcs=srcs, deps=deps, @@ -16,6 +18,13 @@ def xla_proto_library(name, srcs=[], deps=[], visibility=None, testonly=0): ), protoc="@protobuf_archive//:protoc", testonly=testonly, - visibility=visibility,) + visibility=visibility, + **kwargs) + +def xla_py_grpc_library(**kwargs): + # Note: we don't currently define any special targets for Python GRPC in OSS. + _ignore = kwargs + pass + ORC_JIT_MEMORY_MAPPER_TARGETS = [] diff --git a/tensorflow/compiler/xla/xla.proto b/tensorflow/compiler/xla/xla.proto index f9943f71d31627a95ccb2dbbf77a1b203eff79b4..f619b8dc24038af64a27fc0565c74447ca9d09cf 100644 --- a/tensorflow/compiler/xla/xla.proto +++ b/tensorflow/compiler/xla/xla.proto @@ -192,6 +192,9 @@ message DebugOptions { // Generate calls to MKL-DNN in the CPU backend. bool xla_cpu_use_mkl_dnn = 97; + // Maximum kernel unroll factor for the GPU backend. + int32 xla_gpu_max_kernel_unroll_factor = 98; + // Extra options to pass to the compilation backend; specific interpretation // of these values is left to the backend. map xla_backend_extra_options = 500; @@ -417,6 +420,11 @@ message ComputeConstantRequest { repeated LiteralProto parameters = 4; } +message ComputeConstantGraphRequest { + HloModuleProto computation = 1; + Layout output_layout = 2; +} + message ComputeConstantResponse { // A LiteralProto is returned directly for this request, instead of a // ComputationDataHandle. diff --git a/tensorflow/compiler/xla/xla_data.proto b/tensorflow/compiler/xla/xla_data.proto index 1f16e6d25178fd9c10a30b0c500e090ee2e08117..d23f9e5918f54c4f385f3b16fd84bbee51ed5a95 100644 --- a/tensorflow/compiler/xla/xla_data.proto +++ b/tensorflow/compiler/xla/xla_data.proto @@ -355,17 +355,19 @@ message WindowDimension { // positions of the window in this dimension. int64 stride = 2; - // If positive, means the amount of padding with zeroes to add to the base - // area at the low end of this dimension; if negative, its negative means the - // number of elements removed from the low end of this dimension. For example, - // in the horizontal dimension of a rectangle, this would be the number of - // zeroes to pad on the left, given that indices increase when going right. + // If positive, means the amount of padding to add to the base area at the low + // end of this dimension; if negative, its negative means the number of + // elements removed from the low end of this dimension. For example, in the + // horizontal dimension of a rectangle, this would be the number of padding + // values to pad on the left, given that indices increase when going right. + // The actual padding value depends upon the context. Convolution pads with + // zeros. ReduceWindow and SelectAndScatter pads with the reduce function's + // init value. int64 padding_low = 3; - // As padding_low, but on the high end of this dimension. For - // example, in the horizontal dimension of a rectangle, this would - // be the number of zeroes to pad on the right, given that indices - // increase when going right. + // As padding_low, but on the high end of this dimension. For example, in the + // horizontal dimension of a rectangle, this would be the number of values to + // pad on the right, given that indices increase when going right. int64 padding_high = 4; // Dilation factor of the sliding window in this dimension. A dilation factor @@ -799,6 +801,9 @@ enum UnaryOperation { // Elementwise, extract real component of complex x. UNOP_IMAG = 16; + + // Elementwise, computes clz(x). + UNOP_CLZ = 17; } message UnaryOpRequest { diff --git a/tensorflow/contrib/BUILD b/tensorflow/contrib/BUILD index 7e47516550068e8f38d2155e48e229c2ab77b488..d28392a62c2fa9e34ec7912c77d29ac28adb5dee 100644 --- a/tensorflow/contrib/BUILD +++ b/tensorflow/contrib/BUILD @@ -25,6 +25,7 @@ py_library( "//tensorflow/contrib/batching:batch_py", "//tensorflow/contrib/bayesflow:bayesflow_py", "//tensorflow/contrib/boosted_trees:init_py", + "//tensorflow/contrib/checkpoint/python:checkpoint", "//tensorflow/contrib/cloud:cloud_py", "//tensorflow/contrib/cluster_resolver:cluster_resolver_pip", "//tensorflow/contrib/cluster_resolver:cluster_resolver_py", diff --git a/tensorflow/contrib/__init__.py b/tensorflow/contrib/__init__.py index 36cc5144d072893a950c1701fc46fb55510d1fd2..0d163daa6e23435fe42aaa046860ea434039a619 100644 --- a/tensorflow/contrib/__init__.py +++ b/tensorflow/contrib/__init__.py @@ -24,6 +24,7 @@ import os # Add projects here, they will show up under tf.contrib. from tensorflow.contrib import batching from tensorflow.contrib import bayesflow +from tensorflow.contrib import checkpoint from tensorflow.contrib import cloud from tensorflow.contrib import cluster_resolver from tensorflow.contrib import coder diff --git a/tensorflow/contrib/autograph/__init__.py b/tensorflow/contrib/autograph/__init__.py index a39f44b21aa0ddf683b30c18bbe15a43262f7db2..3386c4eca4b93e850f6fe3c6239d29c61d787ece 100644 --- a/tensorflow/contrib/autograph/__init__.py +++ b/tensorflow/contrib/autograph/__init__.py @@ -21,6 +21,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +# TODO(mdan): Bring only the relevant symbols to the top level. from tensorflow.contrib.autograph import utils from tensorflow.contrib.autograph.impl.api import convert from tensorflow.contrib.autograph.impl.api import converted_call diff --git a/tensorflow/contrib/autograph/converters/asserts.py b/tensorflow/contrib/autograph/converters/asserts.py index f011a97ade94f2979486ef6329673a0160dd9bac..2d9e2c58e3afcef5c18f477a7a29e518e98e672e 100644 --- a/tensorflow/contrib/autograph/converters/asserts.py +++ b/tensorflow/contrib/autograph/converters/asserts.py @@ -27,8 +27,6 @@ from tensorflow.contrib.autograph.pyct import transformer class AssertsTransformer(transformer.Base): """Transforms Print nodes to Call so they can be handled as functions.""" - # pylint:disable=invalid-name - def visit_Assert(self, node): self.generic_visit(node) @@ -44,9 +42,7 @@ class AssertsTransformer(transformer.Base): elif isinstance(node.msg, gast.Str): return templates.replace(template, test=node.test, msg=node.msg) else: - raise NotImplementedError('Can only convert string messages for now.') - - # pylint:enable=invalid-name + raise NotImplementedError('can only convert string messages for now.') def transform(node, context): diff --git a/tensorflow/contrib/autograph/converters/break_statements.py b/tensorflow/contrib/autograph/converters/break_statements.py index 62115d4005cb80af1bc2e916c2f3b78f0cc91044..91de82f0a78ccae711298d78364810dd099a5c38 100644 --- a/tensorflow/contrib/autograph/converters/break_statements.py +++ b/tensorflow/contrib/autograph/converters/break_statements.py @@ -18,104 +18,108 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import gast - from tensorflow.contrib.autograph.pyct import anno from tensorflow.contrib.autograph.pyct import templates from tensorflow.contrib.autograph.pyct import transformer from tensorflow.contrib.autograph.pyct.static_analysis.annos import NodeAnno -class BreakCanonicalizationTransformer(transformer.Base): - """Canonicalizes break statements into additional conditionals.""" +# Tags for local state. +BREAK_USED = 'break_used' +CONTROL_VAR_NAME = 'control_var_name' - def __init__(self, context): - super(BreakCanonicalizationTransformer, self).__init__(context) - # This is a stack structure, to correctly process nested loops. - # Each item is a list [break_used, break_variable_name] - self.break_uses = [] - def _create_break_check(self): - template = """ - (not var_name) - """ - expr, = templates.replace(template, var_name=self.break_uses[-1][1]) - return expr.value +class BreakStatementTransformer(transformer.Base): + """Canonicalizes break statements into additional conditionals.""" + + def _track_body(self, nodes, break_var): + self.enter_local_scope() + self.set_local(CONTROL_VAR_NAME, break_var) + nodes = self.visit_block(nodes) + break_used = self.get_local(BREAK_USED, False) + self.exit_local_scope() + return nodes, break_used - def _create_break_trigger(self): + def visit_Break(self, node): + self.set_local(BREAK_USED, True) + var_name = self.get_local(CONTROL_VAR_NAME) + # TODO(mdan): This will fail when expanded inside a top-level else block. template = """ var_name = True + continue """ - block = templates.replace(template, var_name=self.break_uses[-1][1]) - block.append(gast.Continue()) - return block + return templates.replace(template, var_name=var_name) - def _create_break_init(self): + def _guard_if_present(self, block, var_name): + """Prevents the block from executing if var_name is set.""" + if not block: + return block template = """ - var_name = False - """ - assign, = templates.replace(template, var_name=self.break_uses[-1][1]) - return assign - - # TODO(mdan): Surely the transformer supports this better? - def _manual_visit_list(self, block): - new_block = [] - for n in block: - new_n = self.visit(n) - if isinstance(new_n, list): - new_block.extend(new_n) - else: - new_block.append(new_n) - return new_block + if not var_name: + block + """ + node = templates.replace( + template, + var_name=var_name, + block=block) + return node def visit_While(self, node): - self.generic_visit(node.test) scope = anno.getanno(node, NodeAnno.BODY_SCOPE) - - break_var = self.context.namer.new_symbol('break_requested', - scope.referenced) - self.break_uses.append([False, break_var]) - node.body = self._manual_visit_list(node.body) - if self.break_uses[-1][0]: - node.test = gast.BoolOp(gast.And(), [ - node.test, - gast.UnaryOp(gast.Not(), gast.Name(break_var, gast.Load(), None)) - ]) - final_nodes = [self._create_break_init(), node] - else: - final_nodes = node - self.break_uses.pop() - - for n in node.orelse: - self.generic_visit(n) - return final_nodes + break_var = self.context.namer.new_symbol('break__', scope.referenced) + + node.test = self.visit(node.test) + node.body, break_used = self._track_body(node.body, break_var) + # A break in the else clause applies to the containing scope. + node.orelse = self.visit_block(node.orelse) + + if break_used: + template = """ + var_name = False + while test and not var_name: + body + else: + orelse + """ + # Python's else clause only triggers if the loop exited cleanly (e.g. + # break did not trigger). + node = templates.replace( + template, + var_name=break_var, + test=node.test, + body=node.body, + orelse=self._guard_if_present(node.orelse, break_var)) + + return node def visit_For(self, node): - self.generic_visit(node.target) - self.generic_visit(node.iter) scope = anno.getanno(node, NodeAnno.BODY_SCOPE) - - break_var = self.context.namer.new_symbol('break_requested', - scope.referenced) - self.break_uses.append([False, break_var]) - node.body = self._manual_visit_list(node.body) - if self.break_uses[-1][0]: + break_var = self.context.namer.new_symbol('break__', scope.referenced) + + node.target = self.visit(node.target) + node.iter = self.visit(node.iter) + node.body, break_used = self._track_body(node.body, break_var) + # A break in the else clause applies to the containing scope. + node.orelse = self.visit_block(node.orelse) + + if break_used: + node.orelse = self._guard_if_present(node.orelse, break_var) + template = """ + var_name = False + for_stmt + """ + # Python's else clause only triggers if the loop exited cleanly (e.g. + # break did not trigger). + node = templates.replace( + template, + var_name=break_var, + for_stmt=node) extra_cond = templates.replace_as_expression( 'not var_name', var_name=break_var) - anno.setanno(node, 'extra_cond', extra_cond) - final_nodes = [self._create_break_init(), node] - else: - final_nodes = node - self.break_uses.pop() - - for n in node.orelse: - self.generic_visit(n) - return final_nodes + anno.setanno(node[1], 'extra_cond', extra_cond) - def visit_Break(self, node): - self.break_uses[-1][0] = True - return self._create_break_trigger() + return node def transform(node, context): - return BreakCanonicalizationTransformer(context).visit(node) + return BreakStatementTransformer(context).visit(node) diff --git a/tensorflow/contrib/autograph/converters/break_statements_test.py b/tensorflow/contrib/autograph/converters/break_statements_test.py index dd4914a022f57b3bb4a19ec132f311f12269fa9e..1af59e9b5260fe0d3a3ef72c7a003dc451e230f3 100644 --- a/tensorflow/contrib/autograph/converters/break_statements_test.py +++ b/tensorflow/contrib/autograph/converters/break_statements_test.py @@ -25,7 +25,7 @@ from tensorflow.python.platform import test class BreakCanonicalizationTest(converter_test_base.TestCase): - def test_basic_break(self): + def test_basic_while(self): def test_fn(x): v = [] @@ -40,13 +40,11 @@ class BreakCanonicalizationTest(converter_test_base.TestCase): node = break_statements.transform(node, self.ctx) with self.compiled(node) as result: - self.assertEqual(test_fn(0), result.test_fn(0)) - self.assertEqual(test_fn(1), result.test_fn(1)) - self.assertEqual(test_fn(2), result.test_fn(2)) - self.assertEqual(test_fn(3), result.test_fn(3)) - self.assertEqual(test_fn(4), result.test_fn(4)) + self.assertEqual([], result.test_fn(0)) + self.assertEqual([], result.test_fn(1)) + self.assertEqual([3], result.test_fn(4)) - def test_basic_break_for_loop(self): + def test_basic_for(self): def test_fn(a): v = [] @@ -57,30 +55,18 @@ class BreakCanonicalizationTest(converter_test_base.TestCase): v.append(x) return v - # The break is incompletely canonicalized for for loops. Everything is - # in place except for the condition verification. - def test_equiv_fn(a): - v = [] - for x in a: - x -= 1 - if x % 2 == 0: - continue - v.append(x) - return v - node = self.parse_and_analyze(test_fn, {}) node = break_statements.transform(node, self.ctx) with self.compiled(node) as result: - # The break is incompletely canonicalized. Everything is in place, but - # the loop does not break. - self.assertEqual(test_equiv_fn([]), result.test_fn([])) - self.assertEqual(test_equiv_fn([1]), result.test_fn([1])) - self.assertEqual(test_equiv_fn([2]), result.test_fn([2])) - self.assertEqual( - test_equiv_fn([1, 2, 3, 4]), result.test_fn([1, 2, 3, 4])) + # The break is incompletely canonicalized. The loop will not interrupt, + # but the section following the break will be skipped. + self.assertEqual([], result.test_fn([])) + self.assertEqual([3, 3], result.test_fn([4, 4])) + self.assertEqual([3], result.test_fn([4, 5])) + self.assertEqual([3], result.test_fn([5, 4])) - def test_continue_deeply_nested(self): + def test_deeply_nested(self): def test_fn(x): v = [] @@ -93,7 +79,7 @@ class BreakCanonicalizationTest(converter_test_base.TestCase): u.append(x) else: w.append(x) - continue + break v.append(x) return v, u, w @@ -101,11 +87,60 @@ class BreakCanonicalizationTest(converter_test_base.TestCase): node = break_statements.transform(node, self.ctx) with self.compiled(node) as result: - self.assertEqual(test_fn(0), result.test_fn(0)) - self.assertEqual(test_fn(1), result.test_fn(1)) - self.assertEqual(test_fn(2), result.test_fn(2)) - self.assertEqual(test_fn(3), result.test_fn(3)) - self.assertEqual(test_fn(4), result.test_fn(4)) + self.assertEqual(([], [], []), result.test_fn(0)) + self.assertEqual(([2, 1], [2], [0]), result.test_fn(3)) + self.assertEqual(([10, 9, 8, 7], [10, 8], [6]), result.test_fn(11)) + + def test_nested_loops(self): + + def test_fn(x): + v = [] + u = [] + while x > 0: + x -= 1 + y = x + while y > 0: + y -= 1 + if y % 2 == 0: + break + u.append(y) + if x == 0: + break + v.append(x) + return v, u + + node = self.parse_and_analyze(test_fn, {}) + node = break_statements.transform(node, self.ctx) + + with self.compiled(node) as result: + self.assertEqual(([], []), result.test_fn(0)) + self.assertEqual(([1], []), result.test_fn(2)) + self.assertEqual(([2, 1], [1]), result.test_fn(3)) + self.assertEqual(([4, 3, 2, 1], [3, 1]), result.test_fn(5)) + + def test_loop_else(self): + + def test_fn(x): + v = [] + u = [] + while x > 0: + x -= 1 + y = x + while y > 1: + break + else: + u.append(y) + break + v.append(x) + return v, u + + node = self.parse_and_analyze(test_fn, {}) + node = break_statements.transform(node, self.ctx) + + with self.compiled(node) as result: + self.assertEqual(([], []), result.test_fn(0)) + self.assertEqual(([], [1]), result.test_fn(2)) + self.assertEqual(([2], [1]), result.test_fn(3)) if __name__ == '__main__': diff --git a/tensorflow/contrib/autograph/converters/builtin_functions.py b/tensorflow/contrib/autograph/converters/builtin_functions.py index 0349ce29ceb097fbebc36a0378b9072750772416..317711a866f731de1b497295a2752dee0eb544f5 100644 --- a/tensorflow/contrib/autograph/converters/builtin_functions.py +++ b/tensorflow/contrib/autograph/converters/builtin_functions.py @@ -34,24 +34,24 @@ class BuiltinFunctionTransformer(transformer.Base): def __init__(self, context): super(BuiltinFunctionTransformer, self).__init__(context) - # pylint:disable=invalid-name - def _convert_builtin(self, node): template = """ - autograph_utils.dynamic_builtin(func, args) + ag__.utils.dynamic_builtin(func, args) """ return templates.replace(template, func=node.func, args=node.args)[0].value def _convert_print(self, node): template = """ - autograph_utils.dynamic_print(args) + ag__.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 in ('len', 'range'): + # TODO(mdan): Rely on the live_val and use inspect_utils.is_builtin instead. + if (isinstance(node.func, gast.Name) and + node.func.id in ('len', 'range', 'xrange')): 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': @@ -70,8 +70,6 @@ class BuiltinFunctionTransformer(transformer.Base): function_call = templates.replace(template, fname='print', args=args)[0] return self.visit(function_call) - # pylint:enable=invalid-name - def transform(node, context): return BuiltinFunctionTransformer(context).visit(node) diff --git a/tensorflow/contrib/autograph/converters/builtin_functions_test.py b/tensorflow/contrib/autograph/converters/builtin_functions_test.py index ac7e756c47c31816ad34a7ea6926917712afa6c3..30272409df322560b04ba75b3e1cb6f9ad5ff0af 100644 --- a/tensorflow/contrib/autograph/converters/builtin_functions_test.py +++ b/tensorflow/contrib/autograph/converters/builtin_functions_test.py @@ -26,8 +26,6 @@ from tensorflow.contrib.autograph.converters import builtin_functions from tensorflow.contrib.autograph.converters import converter_test_base from tensorflow.python.framework import constant_op from tensorflow.python.ops import array_ops -from tensorflow.python.ops import logging_ops -from tensorflow.python.ops import script_ops from tensorflow.python.platform import test @@ -49,7 +47,7 @@ class BuiltinFunctionsTest(converter_test_base.TestCase): self.assertEqual(3, result.test_fn([0, 0, 0])) - def test_print_with_op(self): + def test_print(self): def test_fn(a): print(a) @@ -57,14 +55,12 @@ class BuiltinFunctionsTest(converter_test_base.TestCase): 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.compiled(node) as result: with self.test_session() as sess: try: out_capturer = six.StringIO() sys.stdout = out_capturer - result.test_fn('a') + result.test_fn(constant_op.constant('a')) sess.run(sess.graph.get_operations()) self.assertEqual(out_capturer.getvalue(), 'a\n') finally: @@ -72,41 +68,19 @@ class BuiltinFunctionsTest(converter_test_base.TestCase): 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.compiled(node) as result: with self.test_session() as sess: try: out_capturer = six.StringIO() sys.stdout = out_capturer - result.test_fn('a', 1, [2, 3]) + result.test_fn( + constant_op.constant('a'), constant_op.constant(1), [2, 3]) sess.run(sess.graph.get_operations()) self.assertEqual(out_capturer.getvalue(), 'a 1 [2, 3]\n') finally: diff --git a/tensorflow/contrib/autograph/converters/call_trees.py b/tensorflow/contrib/autograph/converters/call_trees.py index 61f6bfd7e733fc3e2e0bea35a955509c39d57bc9..554f0471d44d54194c45c3855b1483796ae65a6a 100644 --- a/tensorflow/contrib/autograph/converters/call_trees.py +++ b/tensorflow/contrib/autograph/converters/call_trees.py @@ -23,7 +23,6 @@ from __future__ import division from __future__ import print_function from collections import namedtuple -import types import gast @@ -114,7 +113,7 @@ class CallTreeTransformer(transformer.Base): 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) + return not inspect_utils.isbuiltin(target_entity) def _should_compile(self, node, fqn): """Determines whether an entity should be compiled in the context.""" @@ -147,7 +146,7 @@ class CallTreeTransformer(transformer.Base): # Inspect the target function decorators. If any include a @convert # or @graph_ready annotation, then they must be called as they are. # TODO(mdan): This may be quite heavy. - # To parse and re-analize each function for every call site could be quite + # To parse and re-analyze each function for every call site could be quite # wasteful. Maybe we could cache the parsed AST? try: target_node, _ = parser.parse_entity(target_entity) @@ -199,7 +198,7 @@ class CallTreeTransformer(transformer.Base): def _wrap_to_py_func_no_return(self, node): # TODO(mdan): Properly handle varargs, etc. template = """ - autograph_utils.wrap_py_func(func, None, (args,), kwargs, True) + ag__.utils.wrap_py_func(func, None, (args,), kwargs, True) """ return templates.replace( template, @@ -210,7 +209,7 @@ class CallTreeTransformer(transformer.Base): def _wrap_to_py_func_single_return(self, node, dtype): # TODO(mdan): Properly handle varargs, etc. template = """ - autograph_utils.wrap_py_func(func, dtype, (args,), kwargs, False) + ag__.utils.wrap_py_func(func, dtype, (args,), kwargs, False) """ return templates.replace_as_expression( template, @@ -238,7 +237,7 @@ class CallTreeTransformer(transformer.Base): # Before we could convert all the time though, we'd need a reasonable # caching mechanism. template = """ - autograph_api.converted_call(func, True, False, {}, args) + ag__.converted_call(func, True, False, {}, args) """ call_expr = templates.replace(template, func=node.func, args=node.args) new_call = call_expr[0].value @@ -246,8 +245,6 @@ class CallTreeTransformer(transformer.Base): new_call.keywords = node.keywords return new_call - # pylint:disable=invalid-name - def visit_Expr(self, node): if isinstance(node.value, gast.Call): if anno.hasanno(node.value.func, 'live_val'): @@ -295,15 +292,17 @@ class CallTreeTransformer(transformer.Base): raise NotImplementedError( 'py_func with return values (unknown function)') else: - if self.context.recursive: + if ast_util.matches(node, 'super(_)'): + # super() calls are preserved. The class conversion mechanism will + # ensure that they return the correct value. + pass + elif self.context.recursive: node = self._insert_dynamic_conversion(node) else: # Unresolved functions are allowed in non-recursive mode. pass return node - # pylint:enable=invalid-name - def transform(node, context, uncompiled_modules, nocompile_decorators): """Transform function call to the compiled counterparts. diff --git a/tensorflow/contrib/autograph/converters/call_trees_test.py b/tensorflow/contrib/autograph/converters/call_trees_test.py index c666dcb73b232ce443898cfe3359f74605af98f2..303dd54a4ee49de27fad0c5cdc2d6274abfe0fa8 100644 --- a/tensorflow/contrib/autograph/converters/call_trees_test.py +++ b/tensorflow/contrib/autograph/converters/call_trees_test.py @@ -34,7 +34,7 @@ class CallTreesTest(converter_test_base.TestCase): def test_basic(self): def test_fn_1(_): - raise ValueError('This should not be called in the compiled verison.') + raise ValueError('This should not be called in the compiled version.') def renamed_test_fn_1(a): return a + 1 diff --git a/tensorflow/contrib/autograph/converters/control_flow.py b/tensorflow/contrib/autograph/converters/control_flow.py index 55a28e8ac30397d317f1e6548909501d473eb4e8..2e26cdb3d9387d358e0225555506f199e9945d0b 100644 --- a/tensorflow/contrib/autograph/converters/control_flow.py +++ b/tensorflow/contrib/autograph/converters/control_flow.py @@ -78,7 +78,7 @@ class ControlFlowTransformer(transformer.Base): def _create_cond_expr(self, results, test, body_name, orelse_name): if results is not None: template = """ - results = autograph_utils.run_cond(test, body_name, orelse_name) + results = ag__.utils.run_cond(test, body_name, orelse_name) """ return templates.replace( template, @@ -88,7 +88,7 @@ class ControlFlowTransformer(transformer.Base): orelse_name=orelse_name) else: template = """ - autograph_utils.run_cond(test, body_name, orelse_name) + ag__.utils.run_cond(test, body_name, orelse_name) """ return templates.replace( template, test=test, body_name=body_name, orelse_name=orelse_name) @@ -207,7 +207,7 @@ class ControlFlowTransformer(transformer.Base): def body_name(state_ssf): body return state_ssf, - state_ast_tuple = __ops.while_loop( + state_ast_tuple = ag__.while_loop( test_name, body_name, (state,), (extra_deps,)) """ node = templates.replace( @@ -264,7 +264,7 @@ class ControlFlowTransformer(transformer.Base): def body_name(iterate, state_ssf): body return state_ssf, - state_ast_tuple = __ops.for_loop( + state_ast_tuple = ag__.for_loop( iterated, extra_cond_name, body_name, (state,)) """ node = templates.replace( diff --git a/tensorflow/contrib/autograph/converters/converter_test_base.py b/tensorflow/contrib/autograph/converters/converter_test_base.py index 6f75e9a529b60d08873b4b90046e1cabc274140e..41c2e71702e7e3ee3811a2cbee27c8c988eb3a5c 100644 --- a/tensorflow/contrib/autograph/converters/converter_test_base.py +++ b/tensorflow/contrib/autograph/converters/converter_test_base.py @@ -35,14 +35,17 @@ from tensorflow.python.platform import test class FakeNamer(object): + """A fake namer that uses a global counter to generate unique names.""" + + def __init__(self): + self.i = 0 def new_symbol(self, name_root, used): - i = 0 while True: - name = '%s%d' % (name_root, i) + self.i += 1 + name = '%s%d' % (name_root, self.i) if name not in used: return name - i += 1 def compiled_function_name(self, original_fqn, @@ -76,9 +79,10 @@ class TestCase(test.TestCase): try: result, source = compiler.ast_to_object(node) result.tf = self.make_fake_mod('fake_tf', *symbols) - result.autograph_utils = utils - result.autograph_api = self.make_fake_mod('fake_api', converted_call) - result.__dict__['__ops'] = operators + fake_ag = self.make_fake_mod('fake_ag', converted_call) + fake_ag.__dict__.update(operators.__dict__) + fake_ag.__dict__['utils'] = utils + result.__dict__['ag__'] = fake_ag yield result except Exception: # pylint:disable=broad-except if source is None: diff --git a/tensorflow/contrib/autograph/converters/decorators_test.py b/tensorflow/contrib/autograph/converters/decorators_test.py index e67ab1cd6a15ceb66fe75140419c7abca9653ae4..9c01f689127dbedad7669c65b03e7da071b2d64d 100644 --- a/tensorflow/contrib/autograph/converters/decorators_test.py +++ b/tensorflow/contrib/autograph/converters/decorators_test.py @@ -28,7 +28,7 @@ from tensorflow.python.platform import test # 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 +# trace of the decorator (which is normally 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 diff --git a/tensorflow/contrib/autograph/converters/ifexp.py b/tensorflow/contrib/autograph/converters/ifexp.py index bb0c0a36a7827e5c73e0fa67f09aa4f54d497a2c..616d222762e09feeba1809f119d915dfbe522283 100644 --- a/tensorflow/contrib/autograph/converters/ifexp.py +++ b/tensorflow/contrib/autograph/converters/ifexp.py @@ -27,7 +27,7 @@ class IfExp(transformer.Base): def visit_IfExp(self, node): template = """ - autograph_utils.run_cond(test, lambda: (body,), lambda: (orelse,)) + ag__.utils.run_cond(test, lambda: (body,), lambda: (orelse,)) """ desugared_ifexp = templates.replace_as_expression( template, test=node.test, body=node.body, orelse=node.orelse) diff --git a/tensorflow/contrib/autograph/converters/lists.py b/tensorflow/contrib/autograph/converters/lists.py index 234a0a7487d5fc9e068acf4a19af3bac84f4737e..b49521b2c328f418828a5e92890aa1b169384b70 100644 --- a/tensorflow/contrib/autograph/converters/lists.py +++ b/tensorflow/contrib/autograph/converters/lists.py @@ -45,7 +45,7 @@ class ListTransformer(transformer.Base): 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') + 'use set_element_type(, ) to continue') dtype = anno.getanno(node, 'element_type') if not isinstance(dtype, dtypes.DType): # TODO(mdan): Allow non-TF dtypes? @@ -74,7 +74,7 @@ class ListTransformer(transformer.Base): if qn.qn[-1] == 'append' and (len(call_node.args) == 1): template = """ - target = autograph_utils.dynamic_list_append(target, element) + target = ag__.utils.dynamic_list_append(target, element) """ node = templates.replace( template, @@ -82,23 +82,33 @@ class ListTransformer(transformer.Base): element=call_node.args[0]) return node + def _replace_list_constructors(self, targets, values): + for target in targets: + if (isinstance(target, (gast.Tuple, gast.List)) and + isinstance(values, (gast.Tuple, gast.List))): + n_targets = len(target.elts) + for i in range(n_targets): + target_el, value_el = target.elts[i], values.elts[i] + values.elts[i] = self._replace_list_constructors( + (target_el,), value_el) + return values + if isinstance(values, gast.List): + if values.elts: + return self._pre_populated_list(values) + else: + return self._empty_list(values) + return values + 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) + # TODO(mdan): A similar pattern exists in type_info.py + # We should add a generic "unpack_assignment" function to the base + # transformer, that has the same effect as applying some logic to the SSA + # form. + node.value = self._replace_list_constructors(node.targets, node.value) return node diff --git a/tensorflow/contrib/autograph/converters/lists_test.py b/tensorflow/contrib/autograph/converters/lists_test.py index 749ba14347314f975c5a6e1111133336e2f5c5e6..74c6dc64f197f75eb3e66c01fb078467e8e8ea89 100644 --- a/tensorflow/contrib/autograph/converters/lists_test.py +++ b/tensorflow/contrib/autograph/converters/lists_test.py @@ -45,7 +45,51 @@ class ListTest(converter_test_base.TestCase): result.utils = utils result.dtypes = dtypes with self.test_session() as sess: - self.assertEqual(test_fn(), sess.run(result.test_fn().stack())) + self.assertAllEqual([1], sess.run(result.test_fn().stack())) + + def test_empty_annotated_lists_unpacked(self): + + def test_fn(): + l, m = [], [] + utils.set_element_type(l, dtypes.int32) + utils.set_element_type(m, dtypes.int32) + l.append(1) + m.append(2) + return l, m + + 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: + result.utils = utils + result.dtypes = dtypes + with self.test_session() as sess: + res_l, res_m = result.test_fn() + self.assertEqual([1], sess.run(res_l.stack())) + self.assertEqual([2], sess.run(res_m.stack())) + + def test_empty_annotated_lists_list_unpacked(self): + + def test_fn(): + [l, m] = [], [] + utils.set_element_type(l, dtypes.int32) + utils.set_element_type(m, dtypes.int32) + l.append(1) + m.append(2) + return l, m + + 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: + result.utils = utils + result.dtypes = dtypes + with self.test_session() as sess: + res_l, res_m = result.test_fn() + self.assertEqual([1], sess.run(res_l.stack())) + self.assertEqual([2], sess.run(res_m.stack())) if __name__ == '__main__': diff --git a/tensorflow/contrib/autograph/converters/name_scopes.py b/tensorflow/contrib/autograph/converters/name_scopes.py index 2a3f474360e94635470bf9581222e4c79f46b7a1..dfee529abaa8c14d9b408819b32c5199500a2c2f 100644 --- a/tensorflow/contrib/autograph/converters/name_scopes.py +++ b/tensorflow/contrib/autograph/converters/name_scopes.py @@ -12,8 +12,7 @@ # 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. -""" +"""Wraps a function body with a `name_scope` of the function name.""" from __future__ import absolute_import from __future__ import division @@ -28,23 +27,46 @@ from tensorflow.contrib.autograph.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 _name_for_current_scope(self): + innermost = self.enclosing_entities[-1] + if len(self.enclosing_entities) > 1: + parent = self.enclosing_entities[-2] + if isinstance(parent, gast.ClassDef): + # Methods also take the name of their class. + name = '%s/%s' % (parent.name, innermost.name) + else: + name = innermost.name + else: + name = innermost.name + + # Sanitize the name. + # See https://www.tensorflow.org/api_docs/python/tf/Graph#name_scope + # TensorFlow doesn't like leading underscores at the top level. + while name[0] == '_': + name = name[1:] + return name 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) + node = self.generic_visit(node) + + unscoped_body = [] + scoped_body = node.body + if scoped_body: + first = scoped_body[0] + if isinstance(first, gast.Expr) and isinstance(first.value, gast.Str): + # Skip any docstring. + unscoped_body = scoped_body[:1] + scoped_body = scoped_body[1:] + + template = """ + with tf.name_scope(scope_name): + body + """ + scoped_body = templates.replace( + template, + scope_name=gast.Str(self._name_for_current_scope()), + body=scoped_body) + node.body = unscoped_body + scoped_body return node diff --git a/tensorflow/contrib/autograph/converters/name_scopes_test.py b/tensorflow/contrib/autograph/converters/name_scopes_test.py index 61e5db2af826d0c2238f1af0f3240411596f7429..17692cbd880dbc1db4bb40ad7345e27907499f9d 100644 --- a/tensorflow/contrib/autograph/converters/name_scopes_test.py +++ b/tensorflow/contrib/autograph/converters/name_scopes_test.py @@ -27,9 +27,10 @@ from tensorflow.python.platform import test class FunctionNameScopeTransformer(converter_test_base.TestCase): - def test_basic_name(self): + def test_basic(self): def test_fn(l): + """This should stay here.""" a = 5 l += a return l @@ -38,41 +39,62 @@ class FunctionNameScopeTransformer(converter_test_base.TestCase): 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])) + result_op = result.test_fn(constant_op.constant(1)) self.assertIn('test_fn/', result_op.op.name) - def test_nested_name(self): + self.assertEqual('This should stay here.', result.test_fn.__doc__) + + def test_long_docstring(self): def test_fn(l): + """Multi-line docstring. + + Args: + l: A thing. + Returns: + l + """ + return l - def body(i): - return i**2 + node = self.parse_and_analyze(test_fn, {}) + node = name_scopes.transform(node, self.ctx) + + with self.compiled(node, ops.name_scope) as result: + self.assertIn('Multi-line', result.test_fn.__doc__) + self.assertIn('Returns:', result.test_fn.__doc__) - l += [4] - return body(l) + def test_nested_functions(self): + + def test_fn(l): + + def inner_fn(i): + return i ** 2 + + l += 4 + return inner_fn(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])) + result_op = result.test_fn(constant_op.constant(1)) 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) + self.assertNotIn('inner_fn', first_result_input_name) + self.assertIn('test_fn/inner_fn/', second_result_input_name) - def test_class_name(self): + def test_method(self): class TestClass(object): def test_fn(self, l): - def body(i): - return i**2 + def inner_fn(i): + return i ** 2 - l += [4] - return body(l) + l += 4 + return inner_fn(l) # Note that 'TestClass' was needed in the namespace here. node = self.parse_and_analyze( @@ -80,12 +102,37 @@ class FunctionNameScopeTransformer(converter_test_base.TestCase): 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])) + result_op = result.TestClass().test_fn(constant_op.constant(1)) 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) + self.assertNotIn('inner_fn', first_result_input_name) + self.assertIn('TestClass/test_fn/inner_fn/', second_result_input_name) + + def test_operator(self): + + class TestClass(object): + + def __call__(self, l): + + def inner_fn(i): + return i ** 2 + + l += 4 + return inner_fn(l) + + # Note that 'TestClass' was needed in the namespace here. + node = self.parse_and_analyze( + TestClass.__call__, {'TestClass': TestClass}, owner_type=TestClass) + node = name_scopes.transform(node, self.ctx) + + with self.compiled(node, ops.name_scope) as result: + result_op = result.__call__(TestClass(), constant_op.constant(1)) + first_result_input_name = result_op.op.inputs[0].name + second_result_input_name = result_op.op.inputs[1].name + self.assertIn('call__/', first_result_input_name) + self.assertNotIn('inner_fn', first_result_input_name) + self.assertIn('call__/inner_fn/', second_result_input_name) if __name__ == '__main__': diff --git a/tensorflow/contrib/autograph/converters/side_effect_guards.py b/tensorflow/contrib/autograph/converters/side_effect_guards.py index 1c1293d2c411b51b563ac3965284a48725ed3278..3bcb2d3c42c6e0663c8f78523199a364b6ac231f 100644 --- a/tensorflow/contrib/autograph/converters/side_effect_guards.py +++ b/tensorflow/contrib/autograph/converters/side_effect_guards.py @@ -160,8 +160,8 @@ class SideEffectGuardTransformer(transformer.Base): [alias_map.get(s, s).ast() for s in guarded_args], None) template = """ - with autograph_utils.control_dependency_on_returns(call): - aliased_guarded_args = autograph_utils.alias_tensors(guarded_args) + with ag__.utils.control_dependency_on_returns(call): + aliased_guarded_args = ag__.utils.alias_tensors(guarded_args) """ control_deps_guard = templates.replace( template, @@ -172,7 +172,7 @@ class SideEffectGuardTransformer(transformer.Base): alias_map = {} template = """ - with autograph_utils.control_dependency_on_returns(call): + with ag__.utils.control_dependency_on_returns(call): pass """ control_deps_guard = templates.replace(template, call=node.value)[-1] diff --git a/tensorflow/contrib/autograph/examples/notebooks/rnn_keras_estimator.ipynb b/tensorflow/contrib/autograph/examples/notebooks/rnn_keras_estimator.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..324b23c24b5a7970d7f20ed955839ba1cf1774fc --- /dev/null +++ b/tensorflow/contrib/autograph/examples/notebooks/rnn_keras_estimator.ipynb @@ -0,0 +1,1078 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "colab": { + "autoexec": { + "startup": false, + "wait_interval": 0 + } + }, + "colab_type": "code", + "id": "LqNpENf-ec0X", + "slideshow": { + "slide_type": "slide" + } + }, + "outputs": [], + "source": [ + "!pip install -U tf-nightly" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "colab": { + "autoexec": { + "startup": false, + "wait_interval": 0 + } + }, + "colab_type": "code", + "id": "Pa2qpEmoVOGe", + "slideshow": { + "slide_type": "-" + } + }, + "outputs": [], + "source": [ + "import os\n", + "import time\n", + "\n", + "import tensorflow as tf\n", + "from tensorflow.contrib import autograph\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import six\n", + "\n", + "from google.colab import widgets" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "colab_type": "text", + "id": "HNqUFL4deCsL", + "slideshow": { + "slide_type": "slide" + } + }, + "source": [ + "# Case study: training a custom RNN, using Keras and Estimators\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "colab_type": "text", + "id": "YkC1k4HEQ7rw", + "slideshow": { + "slide_type": "-" + } + }, + "source": [ + "In this section, we show how you can use AutoGraph to build RNNColorbot, an RNN that takes as input names of colors and predicts their corresponding RGB tuples. The model will be trained by a [custom Estimator](https://www.tensorflow.org/get_started/custom_estimators)." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "colab_type": "text", + "id": "7nkPDl5CTCNb", + "slideshow": { + "slide_type": "-" + } + }, + "source": [ + "To get started, set up the dataset. The following cells defines methods that download and format the data needed for RNNColorbot; the details aren't important (read them in the privacy of your own home if you so wish), but make sure to run the cells before proceeding." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "colab": { + "autoexec": { + "startup": false, + "wait_interval": 0 + } + }, + "colab_type": "code", + "id": "A0uREmVXCQEw", + "slideshow": { + "slide_type": "-" + } + }, + "outputs": [], + "source": [ + "def parse(line):\n", + " \"\"\"Parses a line from the colors dataset.\"\"\"\n", + " items = tf.string_split([line], \",\").values\n", + " rgb = tf.string_to_number(items[1:], out_type=tf.float32) / 255.0\n", + " color_name = items[0]\n", + " chars = tf.one_hot(tf.decode_raw(color_name, tf.uint8), depth=256)\n", + " length = tf.cast(tf.shape(chars)[0], dtype=tf.int64)\n", + " return rgb, chars, length\n", + "\n", + "\n", + "def set_static_batch_shape(batch_size):\n", + " def apply(rgb, chars, length):\n", + " rgb.set_shape((batch_size, None))\n", + " chars.set_shape((batch_size, None, 256))\n", + " length.set_shape((batch_size,))\n", + " return rgb, chars, length\n", + " return apply\n", + "\n", + "\n", + "def load_dataset(data_dir, url, batch_size, training=True):\n", + " \"\"\"Loads the colors data at path into a tf.PaddedDataset.\"\"\"\n", + " path = tf.keras.utils.get_file(os.path.basename(url), url, cache_dir=data_dir)\n", + " dataset = tf.data.TextLineDataset(path)\n", + " dataset = dataset.skip(1)\n", + " dataset = dataset.map(parse)\n", + " dataset = dataset.cache()\n", + " dataset = dataset.repeat()\n", + " if training:\n", + " dataset = dataset.shuffle(buffer_size=3000)\n", + " dataset = dataset.padded_batch(\n", + " batch_size, padded_shapes=((None,), (None, 256), ()))\n", + " # To simplify the model code, we statically set as many of the shapes that we\n", + " # know.\n", + " dataset = dataset.map(set_static_batch_shape(batch_size))\n", + " return dataset" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "colab_type": "text", + "id": "waZ89t3DTUla", + "slideshow": { + "slide_type": "-" + } + }, + "source": [ + "To show the use of control flow, we write the RNN loop by hand, rather than using a pre-built RNN model.\n", + "\n", + "Note how we write the model code in Eager style, with regular `if` and `while` statements. Then, we annotate the functions with `@autograph.convert` to have them automatically compiled to run in graph mode.\n", + "We use Keras to define the model, and we will train it using Estimators." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "colab": { + "autoexec": { + "startup": false, + "wait_interval": 0 + } + }, + "colab_type": "code", + "id": "9v8AJouiC44V", + "slideshow": { + "slide_type": "slide" + } + }, + "outputs": [], + "source": [ + "@autograph.convert()\n", + "class RnnColorbot(tf.keras.Model):\n", + " \"\"\"RNN Colorbot model.\"\"\"\n", + "\n", + " def __init__(self):\n", + " super(RnnColorbot, self).__init__()\n", + " self.lower_cell = tf.contrib.rnn.LSTMBlockCell(256)\n", + " self.upper_cell = tf.contrib.rnn.LSTMBlockCell(128)\n", + " self.relu_layer = tf.layers.Dense(3, activation=tf.nn.relu)\n", + "\n", + "\n", + " def _rnn_layer(self, chars, cell, batch_size, training):\n", + " \"\"\"A single RNN layer.\n", + "\n", + " Args:\n", + " chars: A Tensor of shape (max_sequence_length, batch_size, input_size)\n", + " cell: An object of type tf.contrib.rnn.LSTMBlockCell\n", + " batch_size: Int, the batch size to use\n", + " training: Boolean, whether the layer is used for training\n", + "\n", + " Returns:\n", + " A Tensor of shape (max_sequence_length, batch_size, output_size).\n", + " \"\"\"\n", + " hidden_outputs = []\n", + " autograph.utils.set_element_type(hidden_outputs, tf.float32)\n", + " state, output = cell.zero_state(batch_size, tf.float32)\n", + " for ch in chars:\n", + " cell_output, (state, output) = cell.call(ch, (state, output))\n", + " hidden_outputs.append(cell_output)\n", + " hidden_outputs = hidden_outputs.stack()\n", + " if training:\n", + " hidden_outputs = tf.nn.dropout(hidden_outputs, 0.5)\n", + " return hidden_outputs\n", + "\n", + " def build(self, _):\n", + " \"\"\"Creates the model variables. See keras.Model.build().\"\"\"\n", + " self.lower_cell.build(tf.TensorShape((None, 256)))\n", + " self.upper_cell.build(tf.TensorShape((None, 256)))\n", + " self.relu_layer.build(tf.TensorShape((None, 128))) \n", + " self.built = True\n", + "\n", + "\n", + " def call(self, inputs, training=False):\n", + " \"\"\"The RNN model code. Uses Eager and \n", + "\n", + " The model consists of two RNN layers (made by lower_cell and upper_cell),\n", + " followed by a fully connected layer with ReLU activation.\n", + "\n", + " Args:\n", + " inputs: A tuple (chars, length)\n", + " training: Boolean, whether the layer is used for training\n", + "\n", + " Returns:\n", + " A Tensor of shape (batch_size, 3) - the model predictions.\n", + " \"\"\"\n", + " chars, length = inputs\n", + " batch_size = chars.shape[0]\n", + " seq = tf.transpose(chars, (1, 0, 2))\n", + "\n", + " seq = self._rnn_layer(seq, self.lower_cell, batch_size, training)\n", + " seq = self._rnn_layer(seq, self.upper_cell, batch_size, training)\n", + "\n", + " # Grab just the end-of-sequence from each output.\n", + " indices = tf.stack([length - 1, range(batch_size)], axis=1)\n", + " sequence_ends = tf.gather_nd(seq, indices)\n", + " return self.relu_layer(sequence_ends)\n", + "\n", + "@autograph.convert()\n", + "def loss_fn(labels, predictions):\n", + " return tf.reduce_mean((predictions - labels) ** 2)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "colab_type": "text", + "id": "JjK4gXFvFsf4", + "slideshow": { + "slide_type": "slide" + } + }, + "source": [ + "We will now create the model function for the custom Estimator.\n", + "\n", + "In the model function, we simply use the model class we defined above - that's it!" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "colab": { + "autoexec": { + "startup": false, + "wait_interval": 0 + } + }, + "colab_type": "code", + "id": "-yso_Nx23Gy1", + "slideshow": { + "slide_type": "-" + } + }, + "outputs": [], + "source": [ + "def model_fn(features, labels, mode, params):\n", + " \"\"\"Estimator model function.\"\"\"\n", + " chars = features['chars']\n", + " sequence_length = features['sequence_length']\n", + " inputs = (chars, sequence_length)\n", + "\n", + " # Create the model. Simply using the AutoGraph-ed class just works!\n", + " colorbot = RnnColorbot()\n", + " colorbot.build(None)\n", + "\n", + " if mode == tf.estimator.ModeKeys.TRAIN:\n", + " predictions = colorbot(inputs, training=True)\n", + " loss = loss_fn(labels, predictions)\n", + "\n", + " learning_rate = params['learning_rate']\n", + " optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)\n", + " global_step = tf.train.get_global_step()\n", + " train_op = optimizer.minimize(loss, global_step=global_step)\n", + " return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)\n", + "\n", + " elif mode == tf.estimator.ModeKeys.EVAL:\n", + " predictions = colorbot(inputs)\n", + " loss = loss_fn(labels, predictions)\n", + "\n", + " return tf.estimator.EstimatorSpec(mode, loss=loss)\n", + "\n", + " elif mode == tf.estimator.ModeKeys.PREDICT:\n", + " predictions = colorbot(inputs)\n", + "\n", + " predictions = tf.minimum(predictions, 1.0)\n", + " return tf.estimator.EstimatorSpec(mode, predictions=predictions)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "colab_type": "text", + "id": "HOQfoBnHC9CP", + "slideshow": { + "slide_type": "-" + } + }, + "source": [ + "We'll create an input function that will feed our training and eval data." + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": { + "colab": { + "autoexec": { + "startup": false, + "wait_interval": 0 + } + }, + "colab_type": "code", + "id": "FJZlx7yG2MP0", + "slideshow": { + "slide_type": "slide" + } + }, + "outputs": [], + "source": [ + "def input_fn(data_dir, data_url, params, training=True):\n", + " \"\"\"An input function for training\"\"\"\n", + " batch_size = params['batch_size']\n", + " \n", + " # load_dataset defined above\n", + " dataset = load_dataset(data_dir, data_url, batch_size, training=training)\n", + "\n", + " # Package the pipeline end in a format suitable for the estimator.\n", + " labels, chars, sequence_length = dataset.make_one_shot_iterator().get_next()\n", + " features = {\n", + " 'chars': chars,\n", + " 'sequence_length': sequence_length\n", + " }\n", + "\n", + " return features, labels" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "colab_type": "text", + "id": "qsvv-lzbDqXd", + "slideshow": { + "slide_type": "-" + } + }, + "source": [ + "We now have everything in place to build our custom estimator and use it for training and eval!" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": { + "colab": { + "autoexec": { + "startup": false, + "wait_interval": 0 + }, + "height": 35 + }, + "colab_type": "code", + "executionInfo": { + "elapsed": 10604, + "status": "ok", + "timestamp": 1524095272039, + "user": { + "displayName": "", + "photoUrl": "", + "userId": "" + }, + "user_tz": 240 + }, + "id": "2pg1AfbxBJQq", + "outputId": "9c924b4f-06e1-4538-976c-a3e1ddac5660", + "slideshow": { + "slide_type": "-" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Eval loss at step 100: 0.0674834\n" + ] + } + ], + "source": [ + "params = {\n", + " 'batch_size': 64,\n", + " 'learning_rate': 0.01,\n", + "}\n", + "\n", + "train_url = \"https://raw.githubusercontent.com/random-forests/tensorflow-workshop/master/extras/colorbot/data/train.csv\"\n", + "test_url = \"https://raw.githubusercontent.com/random-forests/tensorflow-workshop/master/extras/colorbot/data/test.csv\"\n", + "data_dir = \"tmp/rnn/data\"\n", + "\n", + "regressor = tf.estimator.Estimator(\n", + " model_fn=model_fn,\n", + " params=params)\n", + "\n", + "regressor.train(\n", + " input_fn=lambda: input_fn(data_dir, train_url, params),\n", + " steps=100)\n", + "eval_results = regressor.evaluate(\n", + " input_fn=lambda: input_fn(data_dir, test_url, params, training=False),\n", + " steps=2\n", + ")\n", + "\n", + "print('Eval loss at step %d: %s' % (eval_results['global_step'], eval_results['loss']))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "colab_type": "text", + "id": "zG1YAjB_cUnQ", + "slideshow": { + "slide_type": "slide" + } + }, + "source": [ + "And here's the same estimator used for inference." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": { + "colab": { + "autoexec": { + "startup": false, + "wait_interval": 0 + }, + "height": 343 + }, + "colab_type": "code", + "executionInfo": { + "elapsed": 7990, + "status": "ok", + "timestamp": 1524095280105, + "user": { + "displayName": "", + "photoUrl": "", + "userId": "" + }, + "user_tz": 240 + }, + "id": "dxHex2tUN_10", + "outputId": "2b889e5a-b9ed-4645-bf03-d98f26c72101", + "slideshow": { + "slide_type": "slide" + } + }, + "outputs": [ + { + "data": { + "text/html": [ + "\u003clink rel=stylesheet type=text/css href='/nbextensions/google.colab/tabbar.css'\u003e\u003c/link\u003e" + ], + "text/plain": [ + "\u003cIPython.core.display.HTML at 0x7f3f36aa6cd0\u003e" + ] + }, + "metadata": { + "tags": [ + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "\u003cscript src='/nbextensions/google.colab/tabbar_main.min.js'\u003e\u003c/script\u003e" + ], + "text/plain": [ + "\u003cIPython.core.display.HTML at 0x7f3eca67f7d0\u003e" + ] + }, + "metadata": { + "tags": [ + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "\u003cdiv id=\"id1\"\u003e\u003c/div\u003e" + ], + "text/plain": [ + "\u003cIPython.core.display.HTML at 0x7f3eca67f8d0\u003e" + ] + }, + "metadata": { + "tags": [ + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"e8ddfa22-4362-11e8-91ec-c8d3ffb5fbe0\"] = colab_lib.createTabBar({\"contentBorder\": [\"0px\"], \"elementId\": \"id1\", \"borderColor\": [\"#a7a7a7\"], \"contentHeight\": [\"initial\"], \"tabNames\": [\"RNN Colorbot\"], \"location\": \"top\", \"initialSelection\": 0});\n", + "//# sourceURL=js_71b9087b6d" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3eca67f950\u003e" + ] + }, + "metadata": { + "tags": [ + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"e8ddfa23-4362-11e8-91ec-c8d3ffb5fbe0\"] = window[\"id1\"].setSelectedTabIndex(0);\n", + "//# sourceURL=js_e390445f33" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3eca67f990\u003e" + ] + }, + "metadata": { + "tags": [ + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"e8ddfa24-4362-11e8-91ec-c8d3ffb5fbe0\"] = google.colab.output.getActiveOutputArea();\n", + "//# sourceURL=js_241dd76d85" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3eca67fc50\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"e8ddfa25-4362-11e8-91ec-c8d3ffb5fbe0\"] = document.querySelector(\"#id1_content_0\");\n", + "//# sourceURL=js_60c64e3d50" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3eca67fd90\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"e8ddfa26-4362-11e8-91ec-c8d3ffb5fbe0\"] = google.colab.output.setActiveOutputArea(window[\"e8ddfa25-4362-11e8-91ec-c8d3ffb5fbe0\"]);\n", + "//# sourceURL=js_14ea437cbd" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3eca67fe10\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"e8ddfa27-4362-11e8-91ec-c8d3ffb5fbe0\"] = window[\"id1\"].setSelectedTabIndex(0);\n", + "//# sourceURL=js_09294c2226" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3eca67fcd0\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ec965514-4362-11e8-91ec-c8d3ffb5fbe0\"] = google.colab.output.setActiveOutputArea(window[\"e8ddfa24-4362-11e8-91ec-c8d3ffb5fbe0\"]);\n", + "//# sourceURL=js_e5e8266997" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3eca67fe10\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ec965515-4362-11e8-91ec-c8d3ffb5fbe0\"] = google.colab.output.getActiveOutputArea();\n", + "//# sourceURL=js_07a097f0ee" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3eca67fc90\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ec965516-4362-11e8-91ec-c8d3ffb5fbe0\"] = document.querySelector(\"#id1_content_0\");\n", + "//# sourceURL=js_790d669ca8" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3eca67f8d0\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ec965517-4362-11e8-91ec-c8d3ffb5fbe0\"] = google.colab.output.setActiveOutputArea(window[\"ec965516-4362-11e8-91ec-c8d3ffb5fbe0\"]);\n", + "//# sourceURL=js_d30df771f0" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3eca67fd90\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ec965518-4362-11e8-91ec-c8d3ffb5fbe0\"] = window[\"id1\"].setSelectedTabIndex(0);\n", + "//# sourceURL=js_8a43a2da4b" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3eca67fc50\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "image/png": "iVBORw0KGgoAAAANSUhEUgAAAQwAAAENCAYAAAD60Fs2AAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAACMBJREFUeJzt3F+I1XX+x/G32zjiFERUpgaFd2JBzOg5joX4h0SiMgmM\n/uhVGIlgFBlERGB3hUEkhkRdtDfRP1ACL6KpLBqcguxCjEAkmGamQcSohFHzsxe7O6zssvsydtff\n+ns8rs758j3f8z7fiyef7/k3o7XWCiDwh4s9APC/QzCAmGAAMcEAYoIBxAQDiAkGF8XTTz9d3W63\n7rvvvhoZGakVK1Zc7JEICMYlbvXq1TU8PHyxxzjPV199VcPDw/XZZ5/V22+/XVVVM2bMuMhTkRAM\n/qt+++23+uGHH+r666+vWbNmXexxuECCcQl76qmnanx8vLZs2VIDAwP1+uuv1zfffFP3339/dTqd\nWr9+fY2MjEzvv2nTpnr55ZfrgQceqIGBgXr44Yfr5MmTVVV1+vTp2r59ey1durQ6nU5t2LChTpw4\nUVVVk5OTtWXLllq6dGmtXbu23nnnnelj7tq1q7Zt21bbt2+vJUuW1HvvvVfPPvtsHTp0qAYGBmrX\nrl1/N/fRo0dr06ZN1el06u67766hoaGqqhodHa1OpzO93zPPPFO33nrr9P3t27fXm2+++e89iZyv\ncUlbtWpVGx4ebq21NjEx0brdbjtw4EBrrbUvvviidbvdduLEidZaaxs3bmxr1qxp33//fZuammob\nN25sO3fubK219tZbb7VHH320TU1NtXPnzrXDhw+3X375pbXW2kMPPdR27NjRTp8+3Y4cOdIGBwen\nn/OVV15pN910U/voo49aa61NTU21999/vz344IPTMx48eLCtWLGitdbamTNn2po1a9qePXvamTNn\n2vDwcOvv72/Hjh2bfj2HDx9urbW2du3advvtt7ejR4+21lpbuXJlO3LkyH/qVNJas8L4f6D95edC\n+/btq5UrV9by5curqmrZsmV1880316effjq977333ls33HBD9fb21h133FFHjhypqqqenp46efJk\nHTt2rGbMmFGLFi2qyy+/vCYmJurrr7+uJ598smbOnFkLFy6sDRs21N69e6eP2d/fX6tXr66qqt7e\n3n8666FDh+rUqVP1yCOPVE9PTw0ODtaqVavqgw8+qKqqJUuW1MjISB0/fryqqtauXVtffvlljY6O\n1q+//loLFy78N501/pGeiz0A/z1jY2O1f//++vjjj6vqzyE5e/ZsLVu2bHqfa665Zvr27Nmz69Sp\nU1VVdc8999TExEQ98cQT9fPPP9e6devq8ccfr8nJybryyitr9uzZ04+bP39+HT58ePr+3Llz4xkn\nJydr3rx5522bP39+TU5OVlVVp9OpoaGhuu6666rb7Va32629e/dWb29vLV68+ALOBr+HYFzi/vbT\nh3nz5tX69etrx44dF3ycnp6e2rp1a23durXGxsZq8+bNtWDBgrrtttvqp59+qlOnTlVfX19VVY2P\nj9ecOXP+4Qz/ypw5c2p8fPy8bWNjY7VgwYKqqup2u/Xiiy/WvHnzqtPp1MDAQD333HPV29tb3W73\ngl8XF8YlySXu2muvrdHR0aqqWrduXQ0NDdXnn39e586dq6mpqRoZGakff/zxXx7n4MGD9d1339W5\nc+eqr6+venp66rLLLqu5c+dWf39/vfTSS3X69On69ttv6913361169b9rnlvueWW6uvrq9dee63O\nnj1bBw8erE8++aTuvPPOqqq68cYba9asWbVv377qdDp1xRVX1NVXX10ffvjheW+I8p8hGJe4zZs3\n1+7du6vb7db+/ftr9+7dtWfPnlq2bFmtWrWq3njjjen3OP7ZSuD48eO1bdu2Wrx4cd111121dOnS\n6Sjs3LmzRkdHa/ny5bVt27Z67LHHzrvMuRAzZ86sV199tQ4cOFCDg4P1/PPP1wsvvDC9wqj68yrj\nqquumr7U+WsoFi1a9Luek9yM1vyBDpCxwgBiggHEBAOICQYQ+z/7PYzjf/QRGVxM12z68u+2WWEA\nMcEAYoIBxAQDiAkGEBMMICYYQEwwgJhgADHBAGKCAcQEA4gJBhATDCAmGEBMMICYYAAxwQBiggHE\nBAOICQYQEwwgJhhATDCAmGAAMcEAYoIBxAQDiAkGEBMMICYYQEwwgJhgADHBAGKCAcQEA4gJBhAT\nDCAmGEBMMICYYAAxwQBiggHEBAOICQYQEwwgJhhATDCAmGAAMcEAYoIBxAQDiAkGEBMMICYYQEww\ngJhgADHBAGKCAcQEA4gJBhATDCAmGEBMMICYYAAxwQBiggHEBAOICQYQEwwgJhhATDCAmGAAMcEA\nYoIBxAQDiAkGEBMMICYYQEwwgJhgADHBAGKCAcQEA4gJBhATDCAmGEBMMICYYAAxwQBiggHEBAOI\nCQYQEwwgJhhATDCAmGAAMcEAYoIBxAQDiAkGEBMMICYYQEwwgJhgADHBAGKCAcQEA4gJBhATDCAm\nGEBMMICYYAAxwQBiggHEBAOICQYQEwwgJhhATDCAmGAAMcEAYoIBxAQDiAkGEBMMICYYQEwwgJhg\nADHBAGKCAcQEA4gJBhATDCAmGEBMMICYYAAxwQBiggHEBAOICQYQEwwgJhhATDCAmGAAMcEAYoIB\nxAQDiAkGEBMMICYYQEwwgJhgADHBAGKCAcQEA4gJBhATDCAmGEBMMICYYAAxwQBiggHEBAOICQYQ\nEwwgJhhATDCAmGAAMcEAYoIBxAQDiAkGEBMMICYYQEwwgJhgADHBAGKCAcQEA4gJBhATDCAmGEBM\nMICYYAAxwQBiggHEBAOICQYQEwwgJhhATDCAmGAAMcEAYoIBxAQDiAkGEBMMICYYQEwwgJhgADHB\nAGKCAcQEA4gJBhATDCAmGEBMMICYYAAxwQBiggHEBAOICQYQEwwgJhhATDCAmGAAMcEAYoIBxAQD\niAkGEBMMIDajtdYu9hDA/wYrDCAmGEBMMICYYAAxwQBiggHEBAOICQYQEwwgJhhATDCAmGAAMcEA\nYoIBxAQDiAkGEBMMICYYQEwwgJhgADHBAGKCAcQEA4j9CY2LTAbbRbWuAAAAAElFTkSuQmCC\n", + "text/plain": [ + "\u003cmatplotlib.figure.Figure at 0x7f3ecc00bf10\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1", + "user_output" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ec965519-4362-11e8-91ec-c8d3ffb5fbe0\"] = google.colab.output.setActiveOutputArea(window[\"ec965515-4362-11e8-91ec-c8d3ffb5fbe0\"]);\n", + "//# sourceURL=js_893ad561f4" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3f31b55c90\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ec96551a-4362-11e8-91ec-c8d3ffb5fbe0\"] = google.colab.output.getActiveOutputArea();\n", + "//# sourceURL=js_2d99e0ac17" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3eca67fe50\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ec96551b-4362-11e8-91ec-c8d3ffb5fbe0\"] = document.querySelector(\"#id1_content_0\");\n", + "//# sourceURL=js_5c19462e32" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3f31b55dd0\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ec96551c-4362-11e8-91ec-c8d3ffb5fbe0\"] = google.colab.output.setActiveOutputArea(window[\"ec96551b-4362-11e8-91ec-c8d3ffb5fbe0\"]);\n", + "//# sourceURL=js_b9c8b7567b" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3f31b55a50\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ec96551d-4362-11e8-91ec-c8d3ffb5fbe0\"] = window[\"id1\"].setSelectedTabIndex(0);\n", + "//# sourceURL=js_fd05186348" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3f31b55810\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "\u003cdiv class=id_888646481 style=\"margin-right:10px; display:flex;align-items:center;\"\u003e\u003cspan style=\"margin-right: 3px;\"\u003e\u003c/span\u003e\u003c/div\u003e" + ], + "text/plain": [ + "\u003cIPython.core.display.HTML at 0x7f3f32414810\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1", + "user_output" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ec96551e-4362-11e8-91ec-c8d3ffb5fbe0\"] = jQuery(\".id_888646481 span\");\n", + "//# sourceURL=js_efef96e882" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3f31b55710\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1", + "user_output" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ec96551f-4362-11e8-91ec-c8d3ffb5fbe0\"] = window[\"ec96551e-4362-11e8-91ec-c8d3ffb5fbe0\"].text(\"Give me a color name (or press 'enter' to exit): \");\n", + "//# sourceURL=js_6eca889864" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3eca67f990\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1", + "user_output" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ed8ea972-4362-11e8-91ec-c8d3ffb5fbe0\"] = jQuery(\".id_888646481 input\");\n", + "//# sourceURL=js_f02070cc60" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3f31b553d0\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1", + "user_output" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ed8ea973-4362-11e8-91ec-c8d3ffb5fbe0\"] = window[\"ed8ea972-4362-11e8-91ec-c8d3ffb5fbe0\"].remove();\n", + "//# sourceURL=js_ed9faba660" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3f31a95450\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1", + "user_output" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ed8ea974-4362-11e8-91ec-c8d3ffb5fbe0\"] = jQuery(\".id_888646481 span\");\n", + "//# sourceURL=js_f3458d7074" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3f31a95250\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1", + "user_output" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ed8ea975-4362-11e8-91ec-c8d3ffb5fbe0\"] = window[\"ed8ea974-4362-11e8-91ec-c8d3ffb5fbe0\"].text(\"Give me a color name (or press 'enter' to exit): \");\n", + "//# sourceURL=js_3ffd97bd6f" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3f31a953d0\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1", + "user_output" + ] + }, + "output_type": "display_data" + }, + { + "data": { + "application/javascript": [ + "window[\"ed8ea976-4362-11e8-91ec-c8d3ffb5fbe0\"] = google.colab.output.setActiveOutputArea(window[\"ec96551a-4362-11e8-91ec-c8d3ffb5fbe0\"]);\n", + "//# sourceURL=js_7f73e8bcca" + ], + "text/plain": [ + "\u003cIPython.core.display.Javascript at 0x7f3f31b55710\u003e" + ] + }, + "metadata": { + "tags": [ + "id1_content_0", + "outputarea_id1" + ] + }, + "output_type": "display_data" + } + ], + "source": [ + "def predict_input_fn(color_name):\n", + " \"\"\"An input function for prediction.\"\"\"\n", + " _, chars, sequence_length = parse(color_name)\n", + "\n", + " # We create a batch of a single element.\n", + " features = {\n", + " 'chars': tf.expand_dims(chars, 0),\n", + " 'sequence_length': tf.expand_dims(sequence_length, 0)\n", + " }\n", + " return features, None\n", + "\n", + "\n", + "def draw_prediction(color_name, pred):\n", + " pred = pred * 255\n", + " pred = pred.astype(np.uint8)\n", + " plt.axis('off')\n", + " plt.imshow(pred)\n", + " plt.title(color_name)\n", + " plt.show()\n", + "\n", + "\n", + "def predict_with_estimator(color_name, regressor):\n", + " predictions = regressor.predict(\n", + " input_fn=lambda:predict_input_fn(color_name))\n", + " pred = next(predictions)\n", + " predictions.close()\n", + " pred = np.minimum(pred, 1.0)\n", + " pred = np.expand_dims(np.expand_dims(pred, 0), 0)\n", + "\n", + " draw_prediction(color_name, pred)\n", + "\n", + "tb = widgets.TabBar([\"RNN Colorbot\"])\n", + "while True:\n", + " with tb.output_to(0):\n", + " try:\n", + " color_name = six.moves.input(\"Give me a color name (or press 'enter' to exit): \")\n", + " except (EOFError, KeyboardInterrupt):\n", + " break\n", + " if not color_name:\n", + " break\n", + " with tb.output_to(0):\n", + " tb.clear_tab()\n", + " predict_with_estimator(color_name, regressor)\n", + " " + ] + } + ], + "metadata": { + "colab": { + "collapsed_sections": [], + "default_view": {}, + "last_runtime": { + "build_target": "", + "kind": "local" + }, + "name": "RNN Colorbot using Keras and Estimators", + "provenance": [ + { + "file_id": "1CtzefX39ffFibX_BqE6cRbT0UW_DdVKl", + "timestamp": 1523579810961 + }, + { + "file_id": "1DcfimonWU11tmyivKBGVrbpAl3BIOaRG", + "timestamp": 1523016192637 + }, + { + "file_id": "1wCZUh73zTNs1jzzYjqoxMIdaBWCdKJ2K", + "timestamp": 1522238054357 + }, + { + "file_id": "1_HpC-RrmIv4lNaqeoslUeWaX8zH5IXaJ", + "timestamp": 1521743157199 + }, + { + "file_id": "1mjO2fQ2F9hxpAzw2mnrrUkcgfb7xSGW-", + "timestamp": 1520522344607 + } + ], + "version": "0.3.2", + "views": {} + }, + "kernelspec": { + "display_name": "Python 2", + "name": "python2" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/tensorflow/contrib/autograph/impl/api.py b/tensorflow/contrib/autograph/impl/api.py index dce994e50df60d8bd419f62207d77035beac9f5a..24f87b2c14da4a3523f1e580d4362cbd3679a2cd 100644 --- a/tensorflow/contrib/autograph/impl/api.py +++ b/tensorflow/contrib/autograph/impl/api.py @@ -49,7 +49,7 @@ def convert(recursive=False, verbose=False, arg_types=None): function is called. This means the parameter values are known at compilation. Args: - recursive: Whether to recusrively convert any functions that the decorator + recursive: Whether to recursively convert any functions that the decorator function may call. verbose: Whether to output the compiled code in the logs. arg_types: See to_graph. @@ -137,7 +137,7 @@ def converted_call(f, recursive, verbose, arg_types, *args, **kwargs): unknown_arg_value = object() # Sentinel for arguments of unknown value - if tf_inspect.isbuiltin(f): + if inspect_utils.isbuiltin(f): return builtins.dynamic_builtin(f, *args, **kwargs) if tf_inspect.isfunction(f) or tf_inspect.ismethod(f): @@ -156,7 +156,7 @@ def converted_call(f, recursive, verbose, arg_types, *args, **kwargs): # Constructors target_entity = f arg_map_target = f.__init__ - effective_args = (unknown_arg_value,) + args + effective_args = args partial_types = () elif hasattr(f, '__call__') and hasattr(f, '__class__'): @@ -215,7 +215,7 @@ def to_graph(e, Args: e: A Python entity. - recursive: Whether to recusrively convert any functions that the decorator + recursive: Whether to recursively 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 @@ -235,22 +235,22 @@ def to_graph(e, nocompile_decorators=(convert, do_not_convert, converted_call), partial_types=partial_types, api_module=tf_inspect.getmodule(to_graph)) - _, name = conversion.entity_to_graph(e, conversion_map, arg_values, arg_types) + _, name, namespace = conversion.entity_to_graph(e, conversion_map, arg_values, + arg_types) module = gast.Module([]) for import_line in config.COMPILED_IMPORT_STATEMENTS: module.body.extend(parser.parse_str(import_line).body) - for dep in conversion_map.dependency_cache.values(): + for dep in reversed(conversion_map.dependency_cache.values()): module.body.append(dep) compiled_node, compiled_src = compiler.ast_to_object(module) - # The compiled code should see everything the entry function saw. + # The compiled code should see everything the entry entity saw. # TODO(mdan): This might not work well if the call tree spans modules? - if tf_inspect.isfunction(e): - for key, val in inspect_utils.getnamespace(e).items(): - # Avoid overwriting entities that have been transformed. - if key not in compiled_node.__dict__: - compiled_node.__dict__[key] = val + for key, val in namespace.items(): + # Avoid overwriting entities that have been transformed. + if key not in compiled_node.__dict__: + compiled_node.__dict__[key] = val compiled_fn = getattr(compiled_node, name) if verbose: diff --git a/tensorflow/contrib/autograph/impl/api_test.py b/tensorflow/contrib/autograph/impl/api_test.py index f9db07778a33498f699923a9e0a193c843bfefd8..a7737b7f448131b1c54951efa719b481e1f4d0c9 100644 --- a/tensorflow/contrib/autograph/impl/api_test.py +++ b/tensorflow/contrib/autograph/impl/api_test.py @@ -39,8 +39,6 @@ class ApiTest(test.TestCase): 'from __future__ import print_function', 'from tensorflow.contrib.autograph import utils' ' as autograph_utils', - 'from tensorflow.contrib.autograph import operators' - ' as __ops', 'tf = autograph_utils.fake_tf()', ) @@ -179,6 +177,92 @@ class ApiTest(test.TestCase): constant_op.constant(-2)) self.assertListEqual([0, 1], sess.run(x).tolist()) + def test_converted_call_builtin(self): + x = api.converted_call(range, False, False, {}, 3) + self.assertEqual((0, 1, 2), tuple(x)) + + def test_converted_call_function(self): + + def test_fn(x): + if x < 0: + return -x + return x + + with self.test_session() as sess: + x = api.converted_call( + test_fn, False, False, {}, constant_op.constant(-1)) + self.assertEqual(1, sess.run(x)) + + def test_converted_call_method(self): + + class TestClass(object): + + def __init__(self, x): + self.x = x + + def test_method(self): + if self.x < 0: + return -self.x + return self.x + + with self.test_session() as sess: + tc = TestClass(constant_op.constant(-1)) + x = api.converted_call(tc.test_method, False, False, {}, tc) + self.assertEqual(1, sess.run(x)) + + def test_converted_call_method_by_class(self): + + class TestClass(object): + + def __init__(self, x): + self.x = x + + def test_method(self): + if self.x < 0: + return -self.x + return self.x + + with self.test_session() as sess: + tc = TestClass(constant_op.constant(-1)) + x = api.converted_call(TestClass.test_method, False, False, {}, tc) + self.assertEqual(1, sess.run(x)) + + def test_converted_call_callable_object(self): + + class TestClass(object): + + def __init__(self, x): + self.x = x + + def __call__(self): + if self.x < 0: + return -self.x + return self.x + + with self.test_session() as sess: + tc = TestClass(constant_op.constant(-1)) + x = api.converted_call(tc, False, False, {}) + self.assertEqual(1, sess.run(x)) + + def test_converted_call_constructor(self): + + class TestClass(object): + + def __init__(self, x): + self.x = x + + def test_method(self): + if self.x < 0: + return -self.x + return self.x + + with self.test_session() as sess: + tc = api.converted_call( + TestClass, False, False, {}, constant_op.constant(-1)) + # tc is now a converted object. + x = tc.test_method() + self.assertEqual(1, sess.run(x)) + def test_to_graph_basic(self): def test_fn(x, s): diff --git a/tensorflow/contrib/autograph/impl/config.py b/tensorflow/contrib/autograph/impl/config.py index 26326465e265f5b40c3badedc0ea2813248ef60f..2600088595a12761b1138c4649c06882bd8fd000 100644 --- a/tensorflow/contrib/autograph/impl/config.py +++ b/tensorflow/contrib/autograph/impl/config.py @@ -46,10 +46,4 @@ NO_SIDE_EFFECT_CONSTRUCTORS = set(('tensorflow',)) COMPILED_IMPORT_STATEMENTS = ( 'from __future__ import print_function', 'import tensorflow as tf', - 'from tensorflow.contrib.autograph.impl import api' - ' as autograph_api', - 'from tensorflow.contrib.autograph import utils' - ' as autograph_utils', - 'from tensorflow.contrib.autograph import operators' - ' as __ops', ) diff --git a/tensorflow/contrib/autograph/impl/conversion.py b/tensorflow/contrib/autograph/impl/conversion.py index 3bacc9430098d9cebf1726074524731899cdd965..55a30dc127957b2a9caa053db843380c94bacfbf 100644 --- a/tensorflow/contrib/autograph/impl/conversion.py +++ b/tensorflow/contrib/autograph/impl/conversion.py @@ -18,8 +18,12 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import collections +import imp + import gast +from tensorflow.contrib.autograph import operators from tensorflow.contrib.autograph import utils from tensorflow.contrib.autograph.converters import asserts from tensorflow.contrib.autograph.converters import break_statements @@ -36,6 +40,7 @@ from tensorflow.contrib.autograph.converters import side_effect_guards from tensorflow.contrib.autograph.converters import single_return from tensorflow.contrib.autograph.impl import config from tensorflow.contrib.autograph.impl import naming +from tensorflow.contrib.autograph.pyct import ast_util from tensorflow.contrib.autograph.pyct import context from tensorflow.contrib.autograph.pyct import inspect_utils from tensorflow.contrib.autograph.pyct import parser @@ -56,7 +61,7 @@ class ConversionMap(object): This object is mutable, and is updated as functions are converted. Attributes: - recursive: Whether to recusrively convert any functions that the decorator + recursive: Whether to recursively convert any functions that the decorator function may call. nocompile_decorators: tuple of decorator functions that toggle compilation off. @@ -78,7 +83,9 @@ class ConversionMap(object): self.recursive = recursive self.nocompile_decorators = nocompile_decorators self.partial_types = partial_types if partial_types else () - self.dependency_cache = {} + # Required to output dependencies in discovery order, which should match + # the reverse dependency order. + self.dependency_cache = collections.OrderedDict() self.additional_imports = set() self.name_map = {} self.api_module = api_module @@ -138,20 +145,31 @@ def entity_to_graph(o, conversion_map, arg_values, arg_types): parameters. Returns: - A tuple (ast, new_name): + A tuple (ast, new_name, namespace): * ast: An AST representing an entity with interface equivalent to `o`, but which when executed it creates TF a graph. * new_name: The symbol name under which the new entity can be found. + * namespace: A dict mapping all symbols visible to the converted entity, + keyed by their symbol name. Raises: ValueError: if the entity type is not supported. """ if tf_inspect.isclass(o): - node, new_name = class_to_graph(o, conversion_map) + node, name, ns = class_to_graph(o, conversion_map) elif tf_inspect.isfunction(o): - node, new_name = function_to_graph(o, conversion_map, arg_values, arg_types) + # TODO(mdan): This is not a reliable mechanism. + # The most reliable way is to check the source code, the AST will contain + # a Lambda node instead of a FunctionDef + if o.__name__ == '': + raise NotImplementedError( + 'lambda functions are not yet supported; declare the function' + ' using def instead: %s' % o) + else: + node, name, ns = function_to_graph(o, conversion_map, arg_values, + arg_types) elif tf_inspect.ismethod(o): - node, new_name = function_to_graph(o, conversion_map, arg_values, arg_types) + node, name, ns = function_to_graph(o, conversion_map, arg_values, arg_types) else: raise ValueError( 'Entity "%s" has unsupported type "%s". Only functions and classes are ' @@ -174,7 +192,7 @@ def entity_to_graph(o, conversion_map, arg_values, arg_types): continue entity_to_graph(candidate, conversion_map, {}, {}) - return node, new_name + return node, name, ns def class_to_graph(c, conversion_map): @@ -185,46 +203,96 @@ def class_to_graph(c, conversion_map): if not members: raise ValueError('Cannot convert %s: it has no member methods.' % c) - class_namespace = None + class_namespace = {} for _, m in members: - node, _ = function_to_graph( + # Only convert the members that are directly defined by the class. + if inspect_utils.getdefiningclass(m, c) is not c: + continue + node, _, namespace = function_to_graph( m, conversion_map=conversion_map, arg_values={}, arg_types={'self': (c.__name__, c)}, owner_type=c) - # TODO(mdan): Do not assume all members have the same view of globals. if class_namespace is None: - class_namespace = inspect_utils.getnamespace(m) + class_namespace = namespace + else: + class_namespace.update(namespace) converted_members[m] = node namer = conversion_map.new_namer(class_namespace) class_name = namer.compiled_class_name(c.__name__, c) - node = gast.ClassDef( - class_name, - bases=[], - keywords=[], - body=list(converted_members.values()), - decorator_list=[]) - return node, class_name + # TODO(mdan): This needs to be explained more thoroughly. + # Process any base classes: if the sueprclass if of a whitelisted type, an + # absolute import line is generated. Otherwise, it is marked for conversion + # (as a side effect of the call to namer.compiled_class_name() followed by + # conversion_map.update_name_map(namer)). + output_nodes = [] + renames = {} + bases = [] + for base in c.__bases__: + if isinstance(object, base): + bases.append('object') + continue + if is_whitelisted_for_graph(base): + alias = namer.new_symbol(base.__name__, ()) + output_nodes.append( + gast.ImportFrom( + module=base.__module__, + names=[gast.alias(name=base.__name__, asname=alias)], + level=0)) + else: + # This will trigger a conversion into a class with this name. + alias = namer.compiled_class_name(base.__name__, base) + bases.append(alias) + renames[qual_names.QN(base.__name__)] = qual_names.QN(alias) + conversion_map.update_name_map(namer) + # Generate the definition of the converted class. + output_nodes.append( + gast.ClassDef( + class_name, + bases=bases, + keywords=[], + body=list(converted_members.values()), + decorator_list=[])) + node = gast.Module(output_nodes) + + # Make a final pass to replace references to the class or its base classes. + # Most commonly, this occurs when making super().__init__() calls. + # TODO(mdan): Making direct references to superclass' superclass will fail. + node = qual_names.resolve(node) + renames[qual_names.QN(c.__name__)] = qual_names.QN(class_name) + node = ast_util.rename_symbols(node, renames) -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 'autograph_util' not in namespace: - namespace['autograph_utils'] = utils - elif namespace['autograph_utils'] != utils: - raise ValueError( - 'The module name "autograph_utils" is reserved and may not be used.') + return node, class_name, class_namespace - # We also make reference to the api module for dynamic conversion, but - # to avoid circular references we don't import it here. - if 'autograph_api' not in namespace: - namespace['autograph_api'] = api_module - elif namespace['autograph_api'] != api_module: - raise ValueError( - 'The module name "autograph_api" is reserved and may not be used.') + +def _add_reserved_symbol(namespace, name, entity): + if name not in namespace: + namespace[name] = entity + elif namespace[name] != entity: + raise ValueError('The name "%s" is reserved and may not be used.' % name) + + +ag_internal = None + + +def _add_self_references(namespace, api_module): + """Adds namespace references to the module that exposes the api itself.""" + global ag_internal + if ag_internal is None: + # Craft a module that exposes parts of the external API as well as certain + # internal modules. + ag_internal = imp.new_module('autograph') + ag_internal.converted_call = api_module.converted_call + ag_internal.utils = utils + # TODO(mdan): Add safeguards against name clashes. + # We don't want to create a submodule because we want the operators to be + # accessible as ag__. + ag_internal.__dict__.update(operators.__dict__) + + _add_reserved_symbol(namespace, 'ag__', ag_internal) def function_to_graph(f, conversion_map, arg_values, arg_types, @@ -261,7 +329,7 @@ def function_to_graph(f, conversion_map, arg_values, arg_types, # TODO(mdan): Use this at compilation. conversion_map.additional_imports.update(deps) - return node, new_name + return node, new_name, namespace def _static_analysis_pass(node, ctx): @@ -310,6 +378,8 @@ def node_to_graph(node, ctx, nocompile_decorators): node = ifexp.transform(node, ctx) node, deps = decorators.transform(node, nocompile_decorators) node = break_statements.transform(node, ctx) + node = _static_analysis_pass(node, ctx) + node = asserts.transform(node, ctx) # Note: sequencing continue canonicalization before for loop one avoids diff --git a/tensorflow/contrib/autograph/impl/conversion_test.py b/tensorflow/contrib/autograph/impl/conversion_test.py index 7066739eb87f89ab98e906b10dab62baeaa2de8e..5edd8e74a8899a25fb51e2a4e133f3cb7933fa26 100644 --- a/tensorflow/contrib/autograph/impl/conversion_test.py +++ b/tensorflow/contrib/autograph/impl/conversion_test.py @@ -21,13 +21,18 @@ from __future__ import print_function import gast from tensorflow.contrib.autograph import utils +from tensorflow.contrib.autograph.impl import api from tensorflow.contrib.autograph.impl import conversion from tensorflow.python.framework import constant_op +from tensorflow.python.keras._impl.keras.engine import training from tensorflow.python.platform import test class ConversionTest(test.TestCase): + def _simple_conversion_map(self): + return conversion.ConversionMap(True, (), (), api) + def test_is_whitelisted_for_graph(self): def test_fn(): @@ -39,18 +44,19 @@ class ConversionTest(test.TestCase): def test_entity_to_graph_unsupported_types(self): with self.assertRaises(ValueError): - conversion_map = conversion.ConversionMap(True, (), (), None) + conversion_map = self._simple_conversion_map() conversion.entity_to_graph('dummy', conversion_map, None, None) def test_entity_to_graph_callable(self): - + b = 2 def f(a): - return a + return a + b - conversion_map = conversion.ConversionMap(True, (), (), None) - ast, new_name = conversion.entity_to_graph(f, conversion_map, None, None) + conversion_map = self._simple_conversion_map() + ast, name, ns = conversion.entity_to_graph(f, conversion_map, None, None) self.assertTrue(isinstance(ast, gast.FunctionDef), ast) - self.assertEqual('tf__f', new_name) + self.assertEqual('tf__f', name) + self.assertTrue(ns['b'] is b) def test_entity_to_graph_call_tree(self): @@ -60,7 +66,7 @@ class ConversionTest(test.TestCase): def f(a): return g(a) - conversion_map = conversion.ConversionMap(True, (), (), None) + conversion_map = self._simple_conversion_map() conversion.entity_to_graph(f, conversion_map, None, None) self.assertTrue(f in conversion_map.dependency_cache) @@ -73,6 +79,87 @@ class ConversionTest(test.TestCase): conversion_map.dependency_cache[f].body[0].body[0].value.func.id) self.assertEqual('tf__g', conversion_map.dependency_cache[g].name) + def test_entity_to_graph_class_hierarchy(self): + + class TestBase(object): + + def __init__(self, x='base'): + self.x = x + + def foo(self): + return self.x + + def bar(self): + return self.x + + class TestSubclass(TestBase): + + def __init__(self, y): + super(TestSubclass, self).__init__('sub') + self.y = y + + def foo(self): + return self.y + + def baz(self): + return self.y + + conversion_map = self._simple_conversion_map() + conversion.entity_to_graph(TestSubclass, conversion_map, None, None) + + self.assertTrue(TestBase in conversion_map.dependency_cache) + self.assertTrue(TestSubclass in conversion_map.dependency_cache) + self.assertEqual('TfTestBase', + conversion_map.dependency_cache[TestBase].body[-1].name) + self.assertEqual( + 'TfTestSubclass', + conversion_map.dependency_cache[TestSubclass].body[-1].name) + + def test_entity_to_graph_class_hierarchy_whitelisted(self): + + class TestSubclass(training.Model): + + def __init__(self, y): + super(TestSubclass, self).__init__() + self.built = False + + def call(self, x): + return 3 * x + + conversion_map = self._simple_conversion_map() + conversion.entity_to_graph(TestSubclass, conversion_map, None, None) + + self.assertTrue(TestSubclass in conversion_map.dependency_cache) + self.assertFalse(training.Model in conversion_map.dependency_cache) + self.assertEqual( + 'Model', + conversion_map.dependency_cache[TestSubclass].body[0].names[0].name) + self.assertEqual( + 'TfTestSubclass', + conversion_map.dependency_cache[TestSubclass].body[-1].name) + + def test_entity_to_graph_lambda(self): + f = lambda a: a + + with self.assertRaises(NotImplementedError): + conversion_map = self._simple_conversion_map() + conversion.entity_to_graph(f, conversion_map, None, None) + + def test_ag_module_cached(self): + def callee(): + return range(3) + + def caller(a): + return a() + + conversion_map = self._simple_conversion_map() + _, _, callee_ns = conversion.entity_to_graph( + callee, conversion_map, None, None) + _, _, caller_ns = conversion.entity_to_graph( + caller, conversion_map, None, None) + + self.assertTrue(callee_ns['ag__'] is caller_ns['ag__']) + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/autograph/impl/naming.py b/tensorflow/contrib/autograph/impl/naming.py index 1facaa0ca0ebcc6d4281e7c92a462ceeb00b453a..b1d3f76be7763fada88fd0a1da9d3aa43b67ddfa 100644 --- a/tensorflow/contrib/autograph/impl/naming.py +++ b/tensorflow/contrib/autograph/impl/naming.py @@ -62,8 +62,6 @@ class Namer(object): n += 1 new_name = '%s_%d' % (new_name_root, n) - if live_entity is not None: - self.renamed_calls[live_entity] = new_name self.generated_names.add(new_name) if live_entity is not None: self.renamed_calls[live_entity] = new_name diff --git a/tensorflow/contrib/autograph/operators/BUILD b/tensorflow/contrib/autograph/operators/BUILD index 4c624685751f7de551f6a9e891a96a2d141e9f3e..efb8d441dd839bd34dcbc18f701c7993a2f03906 100644 --- a/tensorflow/contrib/autograph/operators/BUILD +++ b/tensorflow/contrib/autograph/operators/BUILD @@ -21,11 +21,24 @@ py_library( srcs = [ "__init__.py", "control_flow.py", + "data_structures.py", ], srcs_version = "PY2AND3", visibility = ["//tensorflow:__subpackages__"], deps = [ "//tensorflow/contrib/autograph/utils", + "//tensorflow/python:tensor_array_ops", + "//tensorflow/python/data/ops:dataset_ops", + ], +) + +py_test( + name = "data_structures_test", + srcs = ["data_structures_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":operators", + "//tensorflow/python:client_testlib", ], ) diff --git a/tensorflow/contrib/autograph/operators/control_flow.py b/tensorflow/contrib/autograph/operators/control_flow.py index 81ae64f110924cb9f8f089ced2f44bb8e3aa5135..d9d8b0d593e5372942ca6423d10022f0f56d78ce 100644 --- a/tensorflow/contrib/autograph/operators/control_flow.py +++ b/tensorflow/contrib/autograph/operators/control_flow.py @@ -25,6 +25,9 @@ from tensorflow.python.framework import tensor_util from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import gen_math_ops +# TODO(mdan): Rename _loop to _stmt to follow Python nomenclature. +# TODO(mdan): Rename arguments to match the AST names. + def for_loop(iterated, extra_cond, loop_body, init_state): """Functional form of a for statement. @@ -182,3 +185,32 @@ def _py_while_loop(loop_cond, loop_body, init_state, opts): while loop_cond(*state): state = loop_body(*state) return state + + +def if_stmt(cond, body, orelse): + """Functional form of an if statement. + + Args: + cond: Boolean. + body: Callable with no arguments, and outputs of the positive (if) branch + as return type. + orelse: Callable with no arguments, and outputs of the negative (else) + branch as return type. + + Returns: + Tuple containing the statement outputs. + """ + if tensor_util.is_tensor(cond): + return _tf_if_stmt(cond, body, orelse) + else: + return _py_if_stmt(cond, body, orelse) + + +def _tf_if_stmt(cond, body, orelse): + """Overload of if_stmt that stages a TF cond.""" + return control_flow_ops.cond(cond, body, orelse) + + +def _py_if_stmt(cond, body, orelse): + """Overload of if_stmt that executes a Python if statement.""" + return body() if cond else orelse() diff --git a/tensorflow/contrib/autograph/operators/control_flow_test.py b/tensorflow/contrib/autograph/operators/control_flow_test.py index 9112b1627fccc0e34216c6710c782e510be29d8b..a0cd0bfa82bb052d55dfe30f8700fc33a794a59f 100644 --- a/tensorflow/contrib/autograph/operators/control_flow_test.py +++ b/tensorflow/contrib/autograph/operators/control_flow_test.py @@ -18,7 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.autograph import operators +from tensorflow.contrib.autograph.operators import control_flow from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes @@ -29,7 +29,7 @@ from tensorflow.python.platform import test class ForLoopTest(test.TestCase): def test_tensor(self): - s = operators.for_loop( + s = control_flow.for_loop( constant_op.constant([1, 2, 3, 4]), extra_cond=lambda s: True, loop_body=lambda i, s: (s + i,), @@ -38,7 +38,7 @@ class ForLoopTest(test.TestCase): self.assertEqual((10,), sess.run(s)) def test_python(self): - s = operators.for_loop( + s = control_flow.for_loop( range(5), extra_cond=lambda s: True, loop_body=lambda i, s: (s + i,), @@ -47,7 +47,7 @@ class ForLoopTest(test.TestCase): def test_dataset(self): to_int32 = lambda i: math_ops.cast(i, dtypes.int32) - s = operators.for_loop( + s = control_flow.for_loop( dataset_ops.Dataset.range(5).map(to_int32), extra_cond=lambda s: True, loop_body=lambda i, s: (s + i,), @@ -60,7 +60,7 @@ class WhileLoopTest(test.TestCase): def test_tensor(self): n = constant_op.constant(5) - results = operators.while_loop( + results = control_flow.while_loop( loop_cond=lambda i, s: i < n, loop_body=lambda i, s: (i + 1, s + i,), init_state=(0, 0), @@ -70,7 +70,7 @@ class WhileLoopTest(test.TestCase): def test_python(self): n = 5 - results = operators.while_loop( + results = control_flow.while_loop( loop_cond=lambda i, s: i < n, loop_body=lambda i, s: (i + 1, s + i), init_state=(0, 0), @@ -78,5 +78,22 @@ class WhileLoopTest(test.TestCase): self.assertEqual((5, 10), results) +class IfStmtTest(test.TestCase): + + def test_tensor(self): + def test_if_stmt(cond): + return control_flow.if_stmt( + cond=cond, + body=lambda: 1, + orelse=lambda: -1) + with self.test_session() as sess: + self.assertEqual(1, sess.run(test_if_stmt(constant_op.constant(True)))) + self.assertEqual(-1, sess.run(test_if_stmt(constant_op.constant(False)))) + + def test_python(self): + self.assertEqual(1, control_flow.if_stmt(True, lambda: 1, lambda: -1)) + self.assertEqual(-1, control_flow.if_stmt(False, lambda: 1, lambda: -1)) + + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/autograph/operators/data_structures.py b/tensorflow/contrib/autograph/operators/data_structures.py new file mode 100644 index 0000000000000000000000000000000000000000..c862306baa9e8114a71a26323ddcbd35c8592c55 --- /dev/null +++ b/tensorflow/contrib/autograph/operators/data_structures.py @@ -0,0 +1,56 @@ +# 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. +# ============================================================================== +"""Operators specific to data structures: list append, subscripts, etc.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.ops import tensor_array_ops + +# TODO(mdan): Add support for TensorList once functional. +# TODO(mdan): Add primitives for empty list, list with elements. + + +def append(target, element): + """The list append function. + + Note: it is unspecified where target will be mutated or not. If target is + a TensorFlow entity, it will not be typically mutated. If target is a plain + list, it will be. In general, if the target is mutated then the return value + should point to the original entity. + + Args: + target: An entity that supports append semantics. + element: The element to append. + + Returns: + Same as target, after the append was performed. + """ + if isinstance(target, tensor_array_ops.TensorArray): + return _tf_tensorarray_append(target, element) + else: + return _py_append(target, element) + + +def _tf_tensorarray_append(target, element): + """Overload of append that stages a TensorArray write at the last position.""" + return target.write(target.size(), element) + + +def _py_append(target, element): + """Overload of append that executes a Python list append.""" + target.append(element) + return target diff --git a/tensorflow/contrib/autograph/operators/data_structures_test.py b/tensorflow/contrib/autograph/operators/data_structures_test.py new file mode 100644 index 0000000000000000000000000000000000000000..577d28c34da39f1216669513c29a00ac07bec126 --- /dev/null +++ b/tensorflow/contrib/autograph/operators/data_structures_test.py @@ -0,0 +1,44 @@ +# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for data_structures module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.autograph.operators import data_structures +from tensorflow.python.framework import dtypes +from tensorflow.python.ops import tensor_array_ops +from tensorflow.python.platform import test + + +class AppendTest(test.TestCase): + + def test_tf_tensorarray(self): + l = tensor_array_ops.TensorArray(dtypes.int32, size=0, dynamic_size=True) + l1 = data_structures.append(l, 1) + l2 = data_structures.append(l1, 2) + with self.test_session() as sess: + self.assertAllEqual(sess.run(l1.stack()), [1]) + self.assertAllEqual(sess.run(l2.stack()), [1, 2]) + + def test_python(self): + l = [] + self.assertAllEqual(data_structures.append(l, 1), [1]) + self.assertAllEqual(data_structures.append(l, 2), [1, 2]) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/autograph/pyct/BUILD b/tensorflow/contrib/autograph/pyct/BUILD index c483ff68c4b7c6d9a3315f569b62b8f253079f00..796ab445c74128e1123e24b67c288e0e3c5ca24c 100644 --- a/tensorflow/contrib/autograph/pyct/BUILD +++ b/tensorflow/contrib/autograph/pyct/BUILD @@ -125,3 +125,14 @@ py_test( "@gast_archive//:gast", ], ) + +py_test( + name = "transformer_test", + srcs = ["transformer_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":pyct", + "//tensorflow/python:client_testlib", + "@gast_archive//:gast", + ], +) diff --git a/tensorflow/contrib/autograph/pyct/ast_util.py b/tensorflow/contrib/autograph/pyct/ast_util.py index 4a70bab4402a940dec6a8b183daf7406a7e34131..c4f82d11708393a6029d3f17be428b47eb9342ff 100644 --- a/tensorflow/contrib/autograph/pyct/ast_util.py +++ b/tensorflow/contrib/autograph/pyct/ast_util.py @@ -23,10 +23,11 @@ import ast import gast from tensorflow.contrib.autograph.pyct import anno +from tensorflow.contrib.autograph.pyct import parser class CleanCopier(gast.NodeVisitor): - """Copy AST nodes. + """Copies AST nodes. The copied nodes will ignore almost all fields that are prefixed by '__'. Exceptions make some annotations. @@ -106,3 +107,79 @@ def keywords_to_dict(keywords): keys.append(gast.Str(kw.arg)) values.append(kw.value) return gast.Dict(keys=keys, values=values) + + +class PatternMatcher(gast.NodeVisitor): + """Matches a node against a pattern represented by a node. + + The pattern may contain wildcards represented by the symbol '_'. + """ + + def __init__(self, pattern): + self.pattern = pattern + self.pattern_stack = [] + self.matches = True + + def compare_and_visit(self, node, pattern): + self.pattern_stack.append(self.pattern) + self.pattern = pattern + self.generic_visit(node) + self.pattern = self.pattern_stack.pop() + + def no_match(self): + self.matches = False + return False + + def is_wildcard(self, p): + if isinstance(p, (list, tuple)) and len(p) == 1: + p, = p + if isinstance(p, gast.Name) and p.id == '_': + return True + if p == '_': + return True + return False + + def generic_visit(self, node): + if not self.matches: + return + + pattern = self.pattern + for f in node._fields: + if f.startswith('__'): + continue + + if not hasattr(node, f): + if hasattr(pattern, f) and getattr(pattern, f): + return self.no_match() + else: + continue + if not hasattr(pattern, f): + return self.no_match() + + v = getattr(node, f) + p = getattr(pattern, f) + + if self.is_wildcard(p): + continue + if isinstance(v, (list, tuple)): + if not isinstance(p, (list, tuple)) or len(v) != len(p): + return self.no_match() + for v_item, p_item in zip(v, p): + self.compare_and_visit(v_item, p_item) + elif isinstance(v, (gast.AST, ast.AST)): + if not isinstance(v, type(p)) and not isinstance(p, type(v)): + return self.no_match() + self.compare_and_visit(v, p) + else: + # Assume everything else is a value type. + if v != p: + return self.no_match() + + +def matches(node, pattern): + if isinstance(pattern, str): + pattern = parser.parse_expression(pattern) + matcher = PatternMatcher(pattern) + matcher.visit(node) + return matcher.matches + diff --git a/tensorflow/contrib/autograph/pyct/ast_util_test.py b/tensorflow/contrib/autograph/pyct/ast_util_test.py index 8faf92c705d997db298dbb1115981fd9da26372d..3afa04a50685d19c90944c14ed39f9d3ad35e486 100644 --- a/tensorflow/contrib/autograph/pyct/ast_util_test.py +++ b/tensorflow/contrib/autograph/pyct/ast_util_test.py @@ -85,7 +85,33 @@ class AstUtilTest(test.TestCase): output.body += (ast.Assign([ast.Name(id='d', ctx=ast.Store())], d),) result, _ = compiler.ast_to_object(output) self.assertDictEqual(result.d, {'a': 3, 'c': 1, 'd': 'e'}) - print(d) + + def assertMatch(self, target_str, pattern_str): + node = parser.parse_expression(target_str) + pattern = parser.parse_expression(pattern_str) + self.assertTrue(ast_util.matches(node, pattern)) + + def assertNoMatch(self, target_str, pattern_str): + node = parser.parse_expression(target_str) + pattern = parser.parse_expression(pattern_str) + self.assertFalse(ast_util.matches(node, pattern)) + + def test_matches_symbols(self): + self.assertMatch('foo', '_') + self.assertNoMatch('foo()', '_') + self.assertMatch('foo + bar', 'foo + _') + self.assertNoMatch('bar + bar', 'foo + _') + self.assertNoMatch('foo - bar', 'foo + _') + + def test_matches_function_args(self): + self.assertMatch('super(Foo, self).__init__(arg1, arg2)', + 'super(_).__init__(_)') + self.assertMatch('super().__init__()', 'super(_).__init__(_)') + self.assertNoMatch('super(Foo, self).bar(arg1, arg2)', + 'super(_).__init__(_)') + self.assertMatch('super(Foo, self).__init__()', 'super(Foo, _).__init__(_)') + self.assertNoMatch('super(Foo, self).__init__()', + 'super(Bar, _).__init__(_)') if __name__ == '__main__': diff --git a/tensorflow/contrib/autograph/pyct/inspect_utils.py b/tensorflow/contrib/autograph/pyct/inspect_utils.py index 386a6d21ec9ecfe9c5c70ddcd1027dacf6860cea..eef74599a7d5415b4b05d2f05fb094b1dcd33323 100644 --- a/tensorflow/contrib/autograph/pyct/inspect_utils.py +++ b/tensorflow/contrib/autograph/pyct/inspect_utils.py @@ -22,12 +22,25 @@ from __future__ import division from __future__ import print_function import itertools +import types import six from tensorflow.python.util import tf_inspect +def isbuiltin(f): + # Note these return false for isinstance(f, types.BuiltinFunctionType) so we + # need to specifically check for them. + if f in (range, int, float): + return True + if isinstance(f, types.BuiltinFunctionType): + return True + if tf_inspect.isbuiltin(f): + return True + return False + + def getnamespace(f): """Returns the complete namespace of a function. @@ -50,16 +63,27 @@ def getnamespace(f): return namespace +def _get_unbound_function(m): + # TODO(mdan): Figure out why six.get_unbound_function fails in some cases. + # The failure case is for tf.keras.Model. + if hasattr(m, 'im_func'): + return m.im_func + return m + + def getdefiningclass(m, owner_class): """Resolves the class (e.g. one of the superclasses) that defined a method.""" - m = six.get_unbound_function(m) - last_defining = owner_class - for superclass in tf_inspect.getmro(owner_class): + # Normalize bound functions to their respective unbound versions. + m = _get_unbound_function(m) + for superclass in owner_class.__bases__: if hasattr(superclass, m.__name__): superclass_m = getattr(superclass, m.__name__) - if six.get_unbound_function(superclass_m) == m: - last_defining = superclass - return last_defining + if _get_unbound_function(superclass_m) is m: + return superclass + elif hasattr(m, '__self__') and m.__self__ == owner_class: + # Python 3 class methods only work this way it seems :S + return superclass + return owner_class def getmethodclass(m): diff --git a/tensorflow/contrib/autograph/pyct/inspect_utils_test.py b/tensorflow/contrib/autograph/pyct/inspect_utils_test.py index 58f827b79a943a276c2f330f9cfd26e8bcb36119..1a212f676a616307b41feafafda9d1d794ba3d2d 100644 --- a/tensorflow/contrib/autograph/pyct/inspect_utils_test.py +++ b/tensorflow/contrib/autograph/pyct/inspect_utils_test.py @@ -243,6 +243,10 @@ class InspectUtilsTest(test.TestCase): def bar(self): pass + @classmethod + def class_method(cls): + pass + class Subclass(Superclass): def foo(self): @@ -257,6 +261,16 @@ class InspectUtilsTest(test.TestCase): inspect_utils.getdefiningclass(Subclass.bar, Subclass) is Superclass) self.assertTrue( inspect_utils.getdefiningclass(Subclass.baz, Subclass) is Subclass) + self.assertTrue( + inspect_utils.getdefiningclass(Subclass.class_method, Subclass) is + Superclass) + + def test_isbuiltin(self): + self.assertTrue(inspect_utils.isbuiltin(range)) + self.assertTrue(inspect_utils.isbuiltin(float)) + self.assertTrue(inspect_utils.isbuiltin(int)) + self.assertTrue(inspect_utils.isbuiltin(len)) + self.assertFalse(inspect_utils.isbuiltin(function_decorator)) if __name__ == '__main__': diff --git a/tensorflow/contrib/autograph/pyct/static_analysis/activity.py b/tensorflow/contrib/autograph/pyct/static_analysis/activity.py index 6dd53091fa3a4d874e9133520a2299b8a93ba231..2c14c2c8c23810c64446eb9e7ffc5402ce9a2298 100644 --- a/tensorflow/contrib/autograph/pyct/static_analysis/activity.py +++ b/tensorflow/contrib/autograph/pyct/static_analysis/activity.py @@ -133,18 +133,18 @@ class Scope(object): def mark_param(self, name): self.params.add(name) - def mark_creation(self, name): + def mark_creation(self, name, writes_create_symbol=False): 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 + if not writes_create_symbol: + return else: raise ValueError('Unknown symbol "%s".' % parent) self.created.add(name) def mark_write(self, name): + """Marks the given symbol as modified in the current scope.""" self.modified.add(name) if self.isolated: self.mark_creation(name) @@ -162,23 +162,45 @@ class Scope(object): self.parent.mark_returned(name) -class ActivityAnalizer(transformer.Base): +class ActivityAnalyzer(transformer.Base): """Annotates nodes with local scope information. See Scope.""" def __init__(self, context, parent_scope): - super(ActivityAnalizer, self).__init__(context) + super(ActivityAnalyzer, self).__init__(context) self.scope = Scope(parent_scope) self._in_return_statement = False - def _track_symbol(self, node): - # This can happen when we have an attribute (or subscript) on a function - # call. Example: a().b + @property + def _in_constructor(self): + innermost = self.enclosing_entities[-1] + if len(self.enclosing_entities) > 1: + parent = self.enclosing_entities[-2] + return isinstance(parent, gast.ClassDef) and innermost.name == '__init__' + return False + + def _node_sets_self_attribute(self, node): + if anno.hasanno(node, anno.Basic.QN): + qn = anno.getanno(node, anno.Basic.QN) + # TODO(mdan): The 'self' argument is not guaranteed to be called 'self'. + if qn.has_attr and qn.parent.qn == ('self',): + return True + + def _track_symbol(self, + node, + composite_writes_alter_parent=False, + writes_create_symbol=False): + # A QN may be missing when we have an attribute (or subscript) on a function + # call. Example: a().b if not anno.hasanno(node, anno.Basic.QN): return qn = anno.getanno(node, anno.Basic.QN) if isinstance(node.ctx, gast.Store): self.scope.mark_write(qn) + if qn.is_composite and composite_writes_alter_parent: + self.scope.mark_write(qn.parent) + if writes_create_symbol: + self.scope.mark_creation(qn, writes_create_symbol=True) elif isinstance(node.ctx, gast.Load): self.scope.mark_read(qn) elif isinstance(node.ctx, gast.Param): @@ -207,7 +229,18 @@ class ActivityAnalizer(transformer.Base): def visit_Attribute(self, node): self.generic_visit(node) - self._track_symbol(node) + if self._in_constructor and self._node_sets_self_attribute(node): + self._track_symbol( + node, composite_writes_alter_parent=True, writes_create_symbol=True) + else: + self._track_symbol(node) + return node + + def visit_Subscript(self, node): + self.generic_visit(node) + # Subscript writes (e.g. a[b] = "value") are considered to modify + # both the element itself (a[b]) and its parent (a). + self._track_symbol(node, composite_writes_alter_parent=True) return node def visit_Print(self, node): @@ -323,4 +356,4 @@ class ActivityAnalizer(transformer.Base): def resolve(node, context, parent_scope=None): - return ActivityAnalizer(context, parent_scope).visit(node) + return ActivityAnalyzer(context, parent_scope).visit(node) diff --git a/tensorflow/contrib/autograph/pyct/static_analysis/activity_test.py b/tensorflow/contrib/autograph/pyct/static_analysis/activity_test.py index 1e6c686b01445a86499d4f5254ea7e139e450843..ef79a295bfa3940705d2f341edd4eda74d7d7068 100644 --- a/tensorflow/contrib/autograph/pyct/static_analysis/activity_test.py +++ b/tensorflow/contrib/autograph/pyct/static_analysis/activity_test.py @@ -108,7 +108,7 @@ class ScopeTest(test.TestCase): self.assertFalse(QN('a') in child.referenced) -class ActivityAnalizerTest(test.TestCase): +class ActivityAnalyzerTest(test.TestCase): def _parse_and_analyze(self, test_fn): node, source = parser.parse_entity(test_fn) @@ -144,10 +144,21 @@ class ActivityAnalizerTest(test.TestCase): anno.getanno(node.body[0].body[2].value, NodeAnno.IS_LOCAL)) # b in return b + def assertSymbolSetsAre(self, expected, actual, name): + expected = set(expected) + actual = set(str(s) for s in actual) + self.assertSetEqual( + expected, actual, 'for symbol set: %s\n' + ' Expected: %s\n' + ' Got: %s\n' + ' Missing: %s\n' + ' Extra: %s\n' % (name.upper(), expected, actual, + expected - actual, actual - expected)) + 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)) + self.assertSymbolSetsAre(used, scope.used, 'read') + self.assertSymbolSetsAre(modified, scope.modified, 'modified') + self.assertSymbolSetsAre(created, scope.created, 'created') def test_print_statement(self): @@ -172,7 +183,7 @@ class ActivityAnalizerTest(test.TestCase): # arguments. self.assertScopeIsRmc(print_args_scope, ('a', 'b'), (), ()) - def test_call(self): + def test_call_args(self): def test_fn(a): b = 0 @@ -187,6 +198,57 @@ class ActivityAnalizerTest(test.TestCase): self.assertScopeIsRmc( anno.getanno(call_node, NodeAnno.ARGS_SCOPE), ('a', 'b'), (), ()) + def test_call_args_attributes(self): + + def foo(*_): + pass + + def test_fn(a): + a.c = 0 + foo(a.b, a.c) + return a.d + + node = self._parse_and_analyze(test_fn) + call_node = node.body[0].body[1].value + self.assertScopeIsRmc( + anno.getanno(call_node, NodeAnno.ARGS_SCOPE), + ('a', 'a.b', 'a.c'), + (), + (), + ) + self.assertScopeIsRmc( + anno.getanno(call_node, NodeAnno.ARGS_SCOPE).parent, + ('a', 'a.b', 'a.c', 'a.d', 'foo'), + ('a.c',), + ('a',), + ) + + def test_call_args_subscripts(self): + + def foo(*_): + pass + + def test_fn(a): + b = 1 + c = 2 + foo(a[0], a[b]) + return a[c] + + node = self._parse_and_analyze(test_fn) + call_node = node.body[0].body[2].value + self.assertScopeIsRmc( + anno.getanno(call_node, NodeAnno.ARGS_SCOPE), + ('a', 'a[0]', 'a[b]', 'b'), + (), + (), + ) + self.assertScopeIsRmc( + anno.getanno(call_node, NodeAnno.ARGS_SCOPE).parent, + ('a', 'a[0]', 'a[b]', 'a[c]', 'b', 'c', 'foo'), + ('b', 'c'), + ('a', 'b', 'c'), + ) + def test_while(self): def test_fn(a): @@ -253,7 +315,72 @@ class ActivityAnalizerTest(test.TestCase): 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_if_attributes(self): + + def test_fn(a): + if a > 0: + a.b = -a.c + d = 2 * a + else: + a.b = a.c + d = 1 + return d + + node = self._parse_and_analyze(test_fn) + if_node = node.body[0].body[0] + self.assertScopeIsRmc( + anno.getanno(if_node, NodeAnno.BODY_SCOPE), + ('a', 'a.c'), + ('a.b', 'd'), + ('d',), + ) + self.assertScopeIsRmc( + anno.getanno(if_node, NodeAnno.ORELSE_SCOPE), + ('a', 'a.c'), + ('a.b', 'd'), + ('d',), + ) + self.assertScopeIsRmc( + anno.getanno(if_node, NodeAnno.BODY_SCOPE).parent, + ('a', 'a.c', 'd'), + ('a.b', 'd'), + ('a', 'd'), + ) + + def test_if_subscripts(self): + + def test_fn(a, b, c, e): + if a > 0: + a[b] = -a[c] + d = 2 * a + else: + a[0] = e + d = 1 + return d + + node = self._parse_and_analyze(test_fn) + if_node = node.body[0].body[0] + self.assertScopeIsRmc( + anno.getanno(if_node, NodeAnno.BODY_SCOPE), + ('a', 'b', 'c', 'a[c]'), + ('a', 'a[b]', 'd'), + ('d',), + ) + # TODO(mdan): Should subscript writes (a[0] = 1) be considered to read "a"? + self.assertScopeIsRmc( + anno.getanno(if_node, NodeAnno.ORELSE_SCOPE), + ('a', 'e'), + ('a', 'a[0]', 'd'), + ('d',), + ) + self.assertScopeIsRmc( + anno.getanno(if_node, NodeAnno.ORELSE_SCOPE).parent, + ('a', 'b', 'c', 'd', 'e', 'a[c]'), + ('a', 'd', 'a[b]', 'a[0]'), + ('a', 'b', 'c', 'd', 'e'), + ) + + def test_nested_if(self): def test_fn(b): if b > 0: @@ -272,7 +399,7 @@ class ActivityAnalizerTest(test.TestCase): anno.getanno(inner_if_node, NodeAnno.ORELSE_SCOPE), ('b',), ('a',), ('a',)) - def test_function_def(self): + def test_nested_function(self): def test_fn(a): @@ -287,44 +414,48 @@ class ActivityAnalizerTest(test.TestCase): return b, c node = self._parse_and_analyze(test_fn) - fndef_node = node.body[0].body[0] + fn_def_node = node.body[0].body[0] self.assertScopeIsRmc( - anno.getanno(fndef_node, + anno.getanno(fn_def_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',), ( + anno.getanno(fn_def_node, NodeAnno.BODY_SCOPE), ('x', 'y'), ('y',), ( 'x', 'y', )) - def test_call_with_composite_names(self): + def test_constructor_attributes(self): - def foo(*_): - pass + class TestClass(object): + + def __init__(self, a): + self.b = a + self.b.c = 1 + + node = self._parse_and_analyze(TestClass) + init_node = node.body[0].body[0] + self.assertScopeIsRmc( + anno.getanno(init_node, NodeAnno.BODY_SCOPE), + ('self', 'a', 'self.b'), + ('self', 'self.b', 'self.b.c'), + ('self', 'a', 'self.b'), + ) + + def test_aug_assign_subscripts(self): 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 + a[0] += 1 node = self._parse_and_analyze(test_fn) - call_node = node.body[0].body[0].value + fn_node = node.body[0] 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')) + anno.getanno(fn_node, NodeAnno.BODY_SCOPE), + ('a',), + ('a', 'a[0]'), + ('a',), + ) if __name__ == '__main__': diff --git a/tensorflow/contrib/autograph/pyct/static_analysis/annos.py b/tensorflow/contrib/autograph/pyct/static_analysis/annos.py index d6d9f7e1a6028d1ce8eee6c3e250a260c3bf827f..b929b35b79200b0968c9c4f26b10cda28763773a 100644 --- a/tensorflow/contrib/autograph/pyct/static_analysis/annos.py +++ b/tensorflow/contrib/autograph/pyct/static_analysis/annos.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Annotations used by the static analizer.""" +"""Annotations used by the static analyzer.""" from __future__ import absolute_import from __future__ import division @@ -28,15 +28,15 @@ class NoValue(Enum): class NodeAnno(NoValue): - """Additionnal annotations used by the static analyzer. + """Additional 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_LOCAL = 'Symbol is local to the function scope being analyzed.' + IS_PARAM = 'Symbol is a parameter to the function being analyzed.' IS_MODIFIED_SINCE_ENTRY = ( 'Symbol has been explicitly replaced in the current function scope.') diff --git a/tensorflow/contrib/autograph/pyct/static_analysis/type_info.py b/tensorflow/contrib/autograph/pyct/static_analysis/type_info.py index 203aa3c3d18ab15300bbf424adeece6e74d9c994..c00946f9c41bc68d5c638d71f356b484db1286d1 100644 --- a/tensorflow/contrib/autograph/pyct/static_analysis/type_info.py +++ b/tensorflow/contrib/autograph/pyct/static_analysis/type_info.py @@ -48,6 +48,9 @@ from tensorflow.contrib.autograph.pyct import transformer from tensorflow.python.util import tf_inspect +# TODO(mdan): Remove the duplication between this and activity.py. +# In particular, the symbol definitions we track here could as well be tracked +# there because they follow the same rules for visibility. class Scope(object): """Tracks symbol value references. @@ -99,20 +102,16 @@ class TypeInfoResolver(transformer.Base): def __init__(self, context): super(TypeInfoResolver, self).__init__(context) self.scope = Scope(None) - self.function_level = 0 def visit_FunctionDef(self, node): self.scope = Scope(self.scope) - self.function_level += 1 - self.generic_visit(node) - self.function_level -= 1 + node = self.generic_visit(node) self.scope = self.scope.parent return node def _visit_block(self, block): self.scope = Scope(self.scope) - for i, n in enumerate(block): - block[i] = self.generic_visit(n) + block = self.visit_block(block) self.scope = self.scope.parent return block @@ -137,7 +136,7 @@ class TypeInfoResolver(transformer.Base): def _process_function_arg(self, arg_name): str_name = str(arg_name) - if self.function_level == 1 and str_name in self.context.arg_types: + if len(self.enclosing_entities) == 1 and str_name in self.context.arg_types: # Forge a node to hold the type information, so that method calls on # it can resolve the type. type_holder = arg_name.ast() @@ -168,16 +167,8 @@ class TypeInfoResolver(transformer.Base): anno.getanno(definition, 'element_type')) return node - def _process_tuple_assignment(self, source, t): - for i, e in enumerate(t.elts): - if isinstance(e, gast.Tuple): - self._process_tuple_assignment(source, e) - else: - self.scope.setval( - anno.getanno(e, anno.Basic.QN), - gast.Subscript(source, gast.Index(i), ctx=gast.Store())) - def _process_variable_assignment(self, source, targets): + # Special case: constructors. if isinstance(source, gast.Call): func = source.func if anno.hasanno(func, 'live_val'): @@ -190,15 +181,25 @@ class TypeInfoResolver(transformer.Base): # We can have a whitelist of no-side-effects constructors. # We can also step inside the constructor and further analyze. - for t in targets: - if isinstance(t, gast.Tuple): - # need to recurse on the case of assigning nested tuples, - # ex. a, (b, c) = f() - self._process_tuple_assignment(source, t) - elif isinstance(t, (gast.Name, gast.Attribute)): - self.scope.setval(anno.getanno(t, anno.Basic.QN), source) + # Multiple targets mean multiple assignment. + for target in targets: + # Tuple target means unpacking. + if isinstance(target, (gast.Tuple, gast.List)): + for i, target_item in enumerate(target.elts): + # Two cases here: + # 1. Static unpacking, e.g. a, b = c, d + # 2. Dynamic unpacking, e.g. a, b = c + # The former case is optimized away. + if isinstance(source, (gast.Tuple, gast.List)): + source_item = source.elts[i] + else: + source_item = gast.Subscript(source, gast.Index(i), ctx=None) + self._process_variable_assignment(source_item, (target_item,)) + elif isinstance(target, (gast.Name, gast.Attribute)): + target_symbol = anno.getanno(target, anno.Basic.QN) + self.scope.setval(target_symbol, source) else: - raise ValueError('Dont know how to handle assignment to %s' % t) + raise ValueError('assignment target has unknown type: %s' % target) def visit_With(self, node): for wi in node.items: @@ -218,19 +219,26 @@ class TypeInfoResolver(transformer.Base): # 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): + target_arg, type_arg = node.args + if not anno.hasanno(target_arg, 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') + if isinstance(type_arg, gast.Str): + element_type = type_arg.s + elif isinstance(type_arg, gast.Num): + element_type = type_arg.n + else: + if not anno.hasanno(type_arg, 'live_val'): + raise ValueError( + 'the second argument of "%s" must be statically resolvable' % + self.context.type_annotation_func) + element_type = anno.getanno(type_arg, 'live_val') + + target_symbol = anno.getanno(target_arg, anno.Basic.QN) # 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. diff --git a/tensorflow/contrib/autograph/pyct/static_analysis/type_info_test.py b/tensorflow/contrib/autograph/pyct/static_analysis/type_info_test.py index c0de4a604301b6e9f80ee83e4797b9ac7e558a48..46b7701624a43073fb7cc612d2678ab851513d91 100644 --- a/tensorflow/contrib/autograph/pyct/static_analysis/type_info_test.py +++ b/tensorflow/contrib/autograph/pyct/static_analysis/type_info_test.py @@ -196,23 +196,46 @@ class TypeInfoResolverTest(test.TestCase): f_ref = node.body[0].body[1].value self.assertEqual(anno.getanno(f_ref, 'element_type'), Foo) - def test_nested_assignment(self): + def test_nested_unpacking(self): - def test_fn(foo): - a, (b, c) = foo + class Foo(object): + pass + + class Bar(object): + pass + + def test_fn(): + a, (b, c) = (Foo(), (Bar(), Foo())) return a, b, c - node = self._parse_and_analyze(test_fn, {'foo': (1, 2, 3)}) - lhs = node.body[0].body[1].value.elts - a = lhs[0] - b = lhs[1] - c = lhs[2] - # TODO(mdan): change these once we have the live values propagating - # correctly + node = self._parse_and_analyze(test_fn, {'Foo': Foo, 'Bar': Bar}) + a, b, c = node.body[0].body[1].value.elts + self.assertEquals(Foo, anno.getanno(a, 'type')) + self.assertEquals(Bar, anno.getanno(b, 'type')) + self.assertEquals(Foo, anno.getanno(c, 'type')) self.assertFalse(anno.hasanno(a, 'live_val')) self.assertFalse(anno.hasanno(b, 'live_val')) self.assertFalse(anno.hasanno(c, 'live_val')) + def test_inner_scope(self): + + def test_fn(): + a = [] + utils.set_element_type(a, 1) + for _ in a: + b = [] + utils.set_element_type(b, 2) + return a, b + + node = self._parse_and_analyze(test_fn, {'utils': utils}) + a, b = node.body[0].body[2].body[2].value.elts + self.assertEquals(1, anno.getanno(a, 'element_type')) + self.assertEquals(2, anno.getanno(b, 'element_type')) + self.assertFalse(anno.hasanno(a, 'type')) + self.assertFalse(anno.hasanno(b, 'type')) + self.assertFalse(anno.hasanno(a, 'live_val')) + self.assertFalse(anno.hasanno(b, 'live_val')) + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/autograph/pyct/transformer.py b/tensorflow/contrib/autograph/pyct/transformer.py index 35f114b6e11901a854c1d631061ae42285c0e261..4db6cc0adfad90ffc1a6bbcadfc80215688d271e 100644 --- a/tensorflow/contrib/autograph/pyct/transformer.py +++ b/tensorflow/contrib/autograph/pyct/transformer.py @@ -40,7 +40,13 @@ def try_ast_to_source(node): class Base(gast.NodeTransformer): - """Base class for specialized transformers.""" + """Base class for specialized transformers. + + Scope-local state tracking: to keep state across nodes, at the level of + (possibly nested) scopes, use enter/exit_local_scope and set/get_local. + You must call enter/exit_local_scope manually, but the transformer detects + when they are not properly paired. + """ def __init__(self, context): """Initialize the transformer. Subclasses should call this. @@ -51,6 +57,33 @@ class Base(gast.NodeTransformer): self._lineno = 0 self._col_offset = 0 self.context = context + self._enclosing_entities = [] + + # A stack that allows keeping mutable, scope-local state where scopes may be + # nested. For example, it can be used to track the usage of break + # statements in each loop, where loops may be nested. + self._local_scope_state = [] + self.enter_local_scope() + + @property + def enclosing_entities(self): + return tuple(self._enclosing_entities) + + @property + def locel_scope_level(self): + return len(self._local_scope_state) + + def enter_local_scope(self): + self._local_scope_state.append({}) + + def exit_local_scope(self): + return self._local_scope_state.pop() + + def set_local(self, name, value): + self._local_scope_state[-1][name] = value + + def get_local(self, name, default=None): + return self._local_scope_state[-1].get(name, default) def debug_print(self, node): """Helper method useful for debugging.""" @@ -58,16 +91,36 @@ class Base(gast.NodeTransformer): print(pretty_printer.fmt(node)) return node + def visit_block(self, nodes): + """Helper equivalent to generic_visit, but for node lists.""" + results = [] + for node in nodes: + replacement = self.visit(node) + if replacement: + if isinstance(replacement, (list, tuple)): + results.extend(replacement) + else: + results.append(replacement) + return results + def visit(self, node): source_code = self.context.source_code source_file = self.context.source_file + did_enter_function = False + local_scope_state_size = len(self._local_scope_state) + try: + if isinstance(node, (gast.FunctionDef, gast.ClassDef, gast.Lambda)): + self._enclosing_entities.append(node) + did_enter_function = True + if source_code and hasattr(node, 'lineno'): self._lineno = node.lineno self._col_offset = node.col_offset if anno.hasanno(node, anno.Basic.SKIP_PROCESSING): return node return super(Base, self).visit(node) + except (ValueError, AttributeError, KeyError, NotImplementedError, AssertionError) as e: msg = '%s: %s\nOffending source:\n%s\n\nOccurred at node:\n%s' % ( @@ -82,3 +135,13 @@ class Base(gast.NodeTransformer): msg, (source_file, self._lineno, self._col_offset + 1, line)), sys.exc_info()[2]) + finally: + if did_enter_function: + self._enclosing_entities.pop() + + if local_scope_state_size != len(self._local_scope_state): + raise AssertionError( + 'Inconsistent local scope stack. Before entering node %s, the' + ' stack had length %d, after exit it has length %d. This' + ' indicates enter_local_scope and exit_local_scope are not' + ' well paired.') diff --git a/tensorflow/contrib/autograph/pyct/transformer_test.py b/tensorflow/contrib/autograph/pyct/transformer_test.py new file mode 100644 index 0000000000000000000000000000000000000000..f96b0dc377521a482d347436caa98633a0a32c8a --- /dev/null +++ b/tensorflow/contrib/autograph/pyct/transformer_test.py @@ -0,0 +1,179 @@ +# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for templates module.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.autograph.pyct import anno +from tensorflow.contrib.autograph.pyct import context +from tensorflow.contrib.autograph.pyct import parser +from tensorflow.contrib.autograph.pyct import transformer +from tensorflow.python.platform import test + + +class TransformerTest(test.TestCase): + + def _context_for_nodetesting(self): + return context.EntityContext( + namer=None, + source_code=None, + source_file=None, + namespace=None, + arg_values=None, + arg_types=None, + owner_type=None, + recursive=False) + + def test_entity_scope_tracking(self): + + class TestTransformer(transformer.Base): + + # The choice of note to assign to is arbitrary. Using Assign because it's + # easy to find in the tree. + def visit_Assign(self, node): + anno.setanno(node, 'enclosing_entities', self.enclosing_entities) + return self.generic_visit(node) + + # This will show up in the lambda function. + def visit_BinOp(self, node): + anno.setanno(node, 'enclosing_entities', self.enclosing_entities) + return self.generic_visit(node) + + tr = TestTransformer(self._context_for_nodetesting()) + + def test_function(): + a = 0 + + class TestClass(object): + + def test_method(self): + b = 0 + def inner_function(x): + c = 0 + d = lambda y: (x + y) + return c, d + return b, inner_function + return a, TestClass + + node, _ = parser.parse_entity(test_function) + node = tr.visit(node) + + test_function_node = node.body[0] + test_class = test_function_node.body[1] + test_method = test_class.body[0] + inner_function = test_method.body[1] + lambda_node = inner_function.body[1].value + + a = test_function_node.body[0] + b = test_method.body[0] + c = inner_function.body[0] + lambda_expr = lambda_node.body + + self.assertEqual( + (test_function_node,), anno.getanno(a, 'enclosing_entities')) + self.assertEqual((test_function_node, test_class, test_method), + anno.getanno(b, 'enclosing_entities')) + self.assertEqual( + (test_function_node, test_class, test_method, inner_function), + anno.getanno(c, 'enclosing_entities')) + self.assertEqual((test_function_node, test_class, test_method, + inner_function, lambda_node), + anno.getanno(lambda_expr, 'enclosing_entities')) + + def test_statement_info_stack(self): + + class TestTransformer(transformer.Base): + + # Extract all string constants from the block. + def visit_Str(self, node): + self.set_local('string', self.get_local('string', default='') + node.s) + return self.generic_visit(node) + + def _annotate_result(self, node): + self.enter_local_scope() + node = self.generic_visit(node) + anno.setanno(node, 'test', self.get_local('string')) + self.exit_local_scope() + return node + + def visit_While(self, node): + return self._annotate_result(node) + + def visit_For(self, node): + return self._annotate_result(node) + + tr = TestTransformer(self._context_for_nodetesting()) + + def test_function(a): + """Docstring.""" + assert a == 'This should not be counted' + for i in range(3): + _ = 'a' + if i > 2: + return 'b' + else: + _ = 'c' + while True: + raise '1' + return 'nor this' + + node, _ = parser.parse_entity(test_function) + node = tr.visit(node) + + for_node = node.body[0].body[2] + while_node = for_node.body[1].orelse[1] + + self.assertFalse(anno.hasanno(for_node, 'string')) + self.assertEqual('abc', anno.getanno(for_node, 'test')) + self.assertFalse(anno.hasanno(while_node, 'string')) + self.assertEqual('1', anno.getanno(while_node, 'test')) + + def test_statement_info_stack_checks_integrity(self): + + class TestTransformer(transformer.Base): + + def visit_If(self, node): + self.enter_local_scope() + return self.generic_visit(node) + + def visit_For(self, node): + node = self.generic_visit(node) + self.exit_local_scope() + return node + + tr = TestTransformer(self._context_for_nodetesting()) + + def no_exit(a): + if a > 0: + print(a) + return None + + node, _ = parser.parse_entity(no_exit) + with self.assertRaises(AssertionError): + tr.visit(node) + + def no_entry(a): + for _ in a: + print(a) + + node, _ = parser.parse_entity(no_entry) + with self.assertRaises(AssertionError): + tr.visit(node) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/autograph/utils/builtins.py b/tensorflow/contrib/autograph/utils/builtins.py index 7fbb7c09d81ff3d3916d48d3c48e377b90a25907..211e8eaee9082dd3e4f035e4379871cd2e154a39 100644 --- a/tensorflow/contrib/autograph/utils/builtins.py +++ b/tensorflow/contrib/autograph/utils/builtins.py @@ -28,24 +28,17 @@ 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__) + raise ValueError('%s is not supported' % f) def dynamic_len(list_or_tensor): @@ -84,7 +77,7 @@ def is_tf_print_compatible(value): def dynamic_print(*values): - """Implementartion of print using dynamic dispatch. + """Implementation 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. @@ -98,9 +91,15 @@ def dynamic_print(*values): if all(map(is_tf_print_compatible, values)): return logging_ops.Print(1, values) - def flushed_print(*vals): + def print_wrapper(*vals): + if six.PY3: + # TensorFlow doesn't seem to generate Unicode when passing strings to + # py_func. This causes the print to add a "b'" wrapper to the output, + # which is probably never what you want. + vals = tuple(v.decode() if isinstance(v, bytes) else v for v in vals) print(*vals) + # The flush helps avoid garbled output in IPython. sys.stdout.flush() return py_func.wrap_py_func( - flushed_print, None, values, use_dummy_return=True) + print_wrapper, None, values, use_dummy_return=True) diff --git a/tensorflow/contrib/autograph/utils/builtins_test.py b/tensorflow/contrib/autograph/utils/builtins_test.py index d9f7913d89a5471c76eb7ae484674bd7a1853ac9..163e6984079fea5c3b3d9aeda0ec8048d651686f 100644 --- a/tensorflow/contrib/autograph/utils/builtins_test.py +++ b/tensorflow/contrib/autograph/utils/builtins_test.py @@ -76,8 +76,9 @@ class BuiltinsTest(test.TestCase): 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) + # Functions that just have the names of builtins are rejected. + with self.assertRaises(ValueError): + self.assertEqual(builtins.dynamic_builtin(range, 1), 1) if six.PY2: self.assertListEqual( list(builtins.dynamic_builtin(xrange, 3)), [0, 1, 2]) diff --git a/tensorflow/contrib/bayesflow/python/ops/monte_carlo_impl.py b/tensorflow/contrib/bayesflow/python/ops/monte_carlo_impl.py index d193a8459d00b83580509c8de25d5f7801b195fe..032b859d469ee5039e08e4af4c2f4ebf35c2ff19 100644 --- a/tensorflow/contrib/bayesflow/python/ops/monte_carlo_impl.py +++ b/tensorflow/contrib/bayesflow/python/ops/monte_carlo_impl.py @@ -44,15 +44,13 @@ def expectation_importance_sampler(f, n=None, seed=None, name='expectation_importance_sampler'): - r"""Monte Carlo estimate of `\\(E_p[f(Z)] = E_q[f(Z) p(Z) / q(Z)]\\)`. + r"""Monte Carlo estimate of \\(E_p[f(Z)] = E_q[f(Z) p(Z) / q(Z)]\\). - With `\\(p(z) := exp^{log_p(z)}\\)`, this `Op` returns + With \\(p(z) := exp^{log_p(z)}\\), this `Op` returns - ``` \\(n^{-1} sum_{i=1}^n [ f(z_i) p(z_i) / q(z_i) ], z_i ~ q,\\) \\(\approx E_q[ f(Z) p(Z) / q(Z) ]\\) \\(= E_p[f(Z)]\\) - ``` This integral is done in log-space with max-subtraction to better handle the often extreme values that `f(z) p(z) / q(z)` can take on. @@ -121,14 +119,12 @@ def expectation_importance_sampler_logspace( name='expectation_importance_sampler_logspace'): r"""Importance sampling with a positive function, in log-space. - With `\\(p(z) := exp^{log_p(z)}\\)`, and `\\(f(z) = exp{log_f(z)}\\)`, + With \\(p(z) := exp^{log_p(z)}\\), and \\(f(z) = exp{log_f(z)}\\), this `Op` returns - ``` \\(Log[ n^{-1} sum_{i=1}^n [ f(z_i) p(z_i) / q(z_i) ] ], z_i ~ q,\\) \\(\approx Log[ E_q[ f(Z) p(Z) / q(Z) ] ]\\) \\(= Log[E_p[f(Z)]]\\) - ``` This integral is done in log-space with max-subtraction to better handle the often extreme values that `f(z) p(z) / q(z)` can take on. @@ -196,13 +192,11 @@ def _logspace_mean(log_values): def expectation(f, samples, log_prob=None, use_reparametrization=True, axis=0, keep_dims=False, name=None): - """Computes the Monte-Carlo approximation of `\\(E_p[f(X)]\\)`. + """Computes the Monte-Carlo approximation of \\(E_p[f(X)]\\). This function computes the Monte-Carlo approximation of an expectation, i.e., - ```none \\(E_p[f(X)] \approx= m^{-1} sum_i^m f(x_j), x_j\ ~iid\ p(X)\\) - ``` where: @@ -216,8 +210,8 @@ def expectation(f, samples, log_prob=None, use_reparametrization=True, 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 ~ p\\)`. + 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 ~ p\\). However, if p is not reparameterized, TensorFlow's gradient will be incorrect since the chain-rule stops at samples of non-reparameterized distributions. @@ -296,7 +290,7 @@ def expectation(f, samples, log_prob=None, use_reparametrization=True, Args: f: Python callable which can return `f(samples)`. samples: `Tensor` of samples used to form the Monte-Carlo approximation of - `\\(E_p[f(X)]\\)`. A batch of samples should be indexed by `axis` + \\(E_p[f(X)]\\). A batch of samples should be indexed by `axis` dimensions. log_prob: Python callable which can return `log_prob(samples)`. Must correspond to the natural-logarithm of the pdf/pmf of each sample. Only @@ -317,7 +311,7 @@ def expectation(f, samples, log_prob=None, use_reparametrization=True, Returns: approx_expectation: `Tensor` corresponding to the Monte-Carlo approximation - of `\\(E_p[f(X)]\\)`. + of \\(E_p[f(X)]\\). Raises: ValueError: if `f` is not a Python `callable`. @@ -329,7 +323,7 @@ def expectation(f, samples, log_prob=None, use_reparametrization=True, if not callable(f): raise ValueError('`f` must be a callable function.') if use_reparametrization: - return math_ops.reduce_mean(f(samples), axis=axis, keep_dims=keep_dims) + return math_ops.reduce_mean(f(samples), axis=axis, keepdims=keep_dims) else: if not callable(log_prob): raise ValueError('`log_prob` must be a callable function.') @@ -349,7 +343,7 @@ def expectation(f, samples, log_prob=None, use_reparametrization=True, # "Is there a floating point value of x, for which x-x == 0 is false?" # http://stackoverflow.com/q/2686644 fx += stop(fx) * (logpx - stop(logpx)) # Add zeros_like(logpx). - return math_ops.reduce_mean(fx, axis=axis, keep_dims=keep_dims) + return math_ops.reduce_mean(fx, axis=axis, keepdims=keep_dims) def _sample_mean(values): diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/BUILD b/tensorflow/contrib/boosted_trees/estimator_batch/BUILD index 17e20c4b315bab8852c90788567a2f2f92119f40..8cff1a3bb1d11aff6a264636291a7149b40de516 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/BUILD +++ b/tensorflow/contrib/boosted_trees/estimator_batch/BUILD @@ -28,12 +28,13 @@ py_library( srcs = ["model.py"], srcs_version = "PY2AND3", deps = [ + ":estimator_utils", ":trainer_hooks", "//tensorflow/contrib/boosted_trees:gbdt_batch", "//tensorflow/contrib/boosted_trees:model_ops_py", "//tensorflow/python:framework_ops", "//tensorflow/python:state_ops", - "//tensorflow/python:training", + "//tensorflow/python:training_util", ], ) @@ -51,6 +52,18 @@ py_library( ], ) +py_library( + name = "estimator_utils", + srcs = ["estimator_utils.py"], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/contrib/learn", + "//tensorflow/core:protos_all_py", + "//tensorflow/python:control_flow_ops", + "//tensorflow/python:framework_ops", + ], +) + py_test( name = "trainer_hooks_test", size = "small", @@ -118,6 +131,7 @@ py_library( srcs = ["estimator.py"], srcs_version = "PY2AND3", deps = [ + ":estimator_utils", ":model", "//tensorflow/contrib/boosted_trees:losses", "//tensorflow/contrib/learn", @@ -130,6 +144,7 @@ py_library( srcs = ["dnn_tree_combined_estimator.py"], srcs_version = "PY2AND3", deps = [ + ":estimator_utils", ":trainer_hooks", "//tensorflow/contrib/boosted_trees:gbdt_batch", "//tensorflow/contrib/boosted_trees:model_ops_py", @@ -159,3 +174,22 @@ py_test( "//tensorflow/python:framework_for_generated_wrappers", ], ) + +py_test( + name = "estimator_test", + size = "medium", + srcs = ["estimator_test.py"], + srcs_version = "PY2AND3", + tags = [ + "no_gpu", + "no_pip_gpu", + "notsan", + ], + deps = [ + ":estimator", + "//tensorflow/contrib/boosted_trees:gbdt_batch", + "//tensorflow/contrib/layers:layers_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_for_generated_wrappers", + ], +) 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 d9b0d89a03dce40d34f76bb1262d26bb587a2dc7..62f1f4122b05b56a708823df4246d618bd3fa5d4 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/custom_export_strategy.py +++ b/tensorflow/contrib/boosted_trees/estimator_batch/custom_export_strategy.py @@ -39,7 +39,8 @@ _SPARSE_FLOAT_FEATURE_NAME_TEMPLATE = "%s_%d" def make_custom_export_strategy(name, convert_fn, feature_columns, - export_input_fn): + export_input_fn, + use_core_columns=False): """Makes custom exporter of GTFlow tree format. Args: @@ -58,7 +59,7 @@ def make_custom_export_strategy(name, input_fn = export_input_fn() (sorted_feature_names, dense_floats, sparse_float_indices, _, _, sparse_int_indices, _, _) = gbdt_batch.extract_features( - input_fn.features, feature_columns) + input_fn.features, feature_columns, use_core_columns) def export_fn(estimator, export_dir, checkpoint_path=None, eval_result=None): """A wrapper to export to SavedModel, and convert it to other formats.""" diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/dnn_tree_combined_estimator.py b/tensorflow/contrib/boosted_trees/estimator_batch/dnn_tree_combined_estimator.py index 2e7b8cba05b89feaac3f47e13d26e7ae37a7b0ae..9994c84ebdb930eea0818188225488eb5eca84eb 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/dnn_tree_combined_estimator.py +++ b/tensorflow/contrib/boosted_trees/estimator_batch/dnn_tree_combined_estimator.py @@ -19,25 +19,20 @@ logits of the DNN. The input layer of the DNN (including the embeddings learned over sparse features) can optionally be provided to the boosted trees as an additional input feature. """ - from __future__ import absolute_import from __future__ import division from __future__ import print_function import six from tensorflow.contrib import layers +from tensorflow.contrib.boosted_trees.estimator_batch import estimator_utils from tensorflow.contrib.boosted_trees.estimator_batch import trainer_hooks from tensorflow.contrib.boosted_trees.python.ops import model_ops from tensorflow.contrib.boosted_trees.python.training.functions import gbdt_batch from tensorflow.contrib.layers.python.layers import optimizers -from tensorflow.contrib.learn.python.learn.estimators import constants from tensorflow.contrib.learn.python.learn.estimators import estimator from tensorflow.contrib.learn.python.learn.estimators import head as head_lib from tensorflow.contrib.learn.python.learn.estimators import model_fn -from tensorflow.contrib.learn.python.learn.estimators import model_fn as contrib_model_fn_lib -from tensorflow.contrib.learn.python.learn.estimators import prediction_key -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 feature_column_lib from tensorflow.python.framework import ops from tensorflow.python.ops import control_flow_ops @@ -48,56 +43,8 @@ from tensorflow.python.ops import variable_scope from tensorflow.python.summary import summary from tensorflow.python.training import training_util - _DNN_LEARNING_RATE = 0.001 -_CORE_MODE_TO_CONTRIB_MODE_ = { - model_fn_lib.ModeKeys.TRAIN: contrib_model_fn_lib.ModeKeys.TRAIN, - model_fn_lib.ModeKeys.EVAL: contrib_model_fn_lib.ModeKeys.EVAL, - model_fn_lib.ModeKeys.PREDICT: contrib_model_fn_lib.ModeKeys.INFER -} - - -def _core_mode_to_contrib_mode(mode): - return _CORE_MODE_TO_CONTRIB_MODE_[mode] - - -def _export_outputs_to_output_alternatives(export_outputs): - """Converts EstimatorSpec.export_outputs to output_alternatives. - - Args: - export_outputs: export_outputs created by create_estimator_spec. - Returns: - converted output_alternatives. - """ - output = dict() - if export_outputs is not None: - for key, value in export_outputs.items(): - if isinstance(value, export_output.ClassificationOutput): - exported_predictions = { - prediction_key.PredictionKey.SCORES: value.scores, - prediction_key.PredictionKey.CLASSES: value.classes - } - output[key] = (constants.ProblemType.CLASSIFICATION, - exported_predictions) - return output - return None - - -def _estimator_spec_to_model_fn_ops(estimator_spec, is_regression): - alternatives = [] - if not is_regression: - _export_outputs_to_output_alternatives(estimator_spec.export_outputs) - - return model_fn.ModelFnOps( - mode=_core_mode_to_contrib_mode(estimator_spec.mode), - predictions=estimator_spec.predictions, - loss=estimator_spec.loss, - train_op=estimator_spec.train_op, - eval_metric_ops=estimator_spec.eval_metric_ops, - output_alternatives=alternatives) - - def _get_optimizer(optimizer): if callable(optimizer): return optimizer() @@ -128,8 +75,7 @@ def _dnn_tree_combined_model_fn(features, dnn_steps_to_train=10000, tree_feature_columns=None, tree_center_bias=False, - use_core_versions=False, - is_regression=False): + use_core_versions=False): """DNN and GBDT combined model_fn. Args: @@ -169,7 +115,6 @@ def _dnn_tree_combined_model_fn(features, first fitting the bias. use_core_versions: Whether feature columns and loss are from the core (as opposed to contrib) version of tensorflow. - is_regression: Whether the problem is regression or not. Returns: A `ModelFnOps` object. @@ -305,8 +250,8 @@ def _dnn_tree_combined_model_fn(features, labels=labels, train_op_fn=_dnn_train_op_fn, logits=dnn_logits) - dnn_train_op = _estimator_spec_to_model_fn_ops(dnn_train_op, - is_regression).train_op + dnn_train_op = estimator_utils.estimator_spec_to_model_fn_ops( + dnn_train_op).train_op tree_train_op = head.create_estimator_spec( features=tree_features, @@ -314,10 +259,10 @@ def _dnn_tree_combined_model_fn(features, labels=labels, train_op_fn=_tree_train_op_fn, logits=tree_train_logits) - tree_train_op = _estimator_spec_to_model_fn_ops(tree_train_op, - is_regression).train_op + tree_train_op = estimator_utils.estimator_spec_to_model_fn_ops( + tree_train_op).train_op - model_fn_ops = _estimator_spec_to_model_fn_ops(model_fn_ops, is_regression) + model_fn_ops = estimator_utils.estimator_spec_to_model_fn_ops(model_fn_ops) else: model_fn_ops = head.create_model_fn_ops( features=features, @@ -529,26 +474,12 @@ class DNNBoostedTreeCombinedRegressor(estimator.Estimator): def _model_fn(features, labels, mode, config): return _dnn_tree_combined_model_fn( - features, - labels, - mode, - head, - dnn_hidden_units, - dnn_feature_columns, - tree_learner_config, - num_trees, - tree_examples_per_layer, - config, - dnn_optimizer, - dnn_activation_fn, - dnn_dropout, - dnn_input_layer_partitioner, - dnn_input_layer_to_tree, - dnn_steps_to_train, - tree_feature_columns, - tree_center_bias, - use_core_versions, - is_regression=True) + features, labels, mode, head, dnn_hidden_units, dnn_feature_columns, + tree_learner_config, num_trees, tree_examples_per_layer, config, + dnn_optimizer, dnn_activation_fn, dnn_dropout, + dnn_input_layer_partitioner, dnn_input_layer_to_tree, + dnn_steps_to_train, tree_feature_columns, tree_center_bias, + use_core_versions) super(DNNBoostedTreeCombinedRegressor, self).__init__( model_fn=_model_fn, model_dir=model_dir, diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/estimator.py b/tensorflow/contrib/boosted_trees/estimator_batch/estimator.py index 70454aa6dbdb19297028a3f80822719bef5a0f72..89d0d611d2905492cec09e033b8cbc238ec7fac6 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/estimator.py +++ b/tensorflow/contrib/boosted_trees/estimator_batch/estimator.py @@ -40,7 +40,8 @@ class GradientBoostedDecisionTreeClassifier(estimator.Estimator): label_keys=None, feature_engineering_fn=None, logits_modifier_function=None, - center_bias=True): + center_bias=True, + use_core_libs=False): """Initializes a GradientBoostedDecisionTreeClassifier estimator instance. Args: @@ -63,7 +64,8 @@ class GradientBoostedDecisionTreeClassifier(estimator.Estimator): logits_modifier_function: A modifier function for the logits. center_bias: Whether a separate tree should be created for first fitting the bias. - + use_core_libs: Whether feature columns and loss are from the core (as + opposed to contrib) version of tensorflow. Raises: ValueError: If learner_config is not valid. """ @@ -99,6 +101,7 @@ class GradientBoostedDecisionTreeClassifier(estimator.Estimator): 'examples_per_layer': examples_per_layer, 'center_bias': center_bias, 'logits_modifier_function': logits_modifier_function, + 'use_core_libs': use_core_libs, }, model_dir=model_dir, config=config, @@ -120,7 +123,8 @@ class GradientBoostedDecisionTreeRegressor(estimator.Estimator): config=None, feature_engineering_fn=None, logits_modifier_function=None, - center_bias=True): + center_bias=True, + use_core_libs=False): """Initializes a GradientBoostedDecisionTreeRegressor estimator instance. Args: @@ -145,6 +149,8 @@ class GradientBoostedDecisionTreeRegressor(estimator.Estimator): logits_modifier_function: A modifier function for the logits. center_bias: Whether a separate tree should be created for first fitting the bias. + use_core_libs: Whether feature columns and loss are from the core (as + opposed to contrib) version of tensorflow. """ head = head_lib.regression_head( label_name=label_name, @@ -166,6 +172,7 @@ class GradientBoostedDecisionTreeRegressor(estimator.Estimator): 'examples_per_layer': examples_per_layer, 'logits_modifier_function': logits_modifier_function, 'center_bias': center_bias, + 'use_core_libs': use_core_libs, }, model_dir=model_dir, config=config, @@ -189,7 +196,8 @@ class GradientBoostedDecisionTreeEstimator(estimator.Estimator): config=None, feature_engineering_fn=None, logits_modifier_function=None, - center_bias=True): + center_bias=True, + use_core_libs=False): """Initializes a GradientBoostedDecisionTreeEstimator estimator instance. Args: @@ -210,6 +218,8 @@ class GradientBoostedDecisionTreeEstimator(estimator.Estimator): logits_modifier_function: A modifier function for the logits. center_bias: Whether a separate tree should be created for first fitting the bias. + use_core_libs: Whether feature columns and loss are from the core (as + opposed to contrib) version of tensorflow. """ super(GradientBoostedDecisionTreeEstimator, self).__init__( model_fn=model.model_builder, @@ -222,6 +232,7 @@ class GradientBoostedDecisionTreeEstimator(estimator.Estimator): 'examples_per_layer': examples_per_layer, 'logits_modifier_function': logits_modifier_function, 'center_bias': center_bias, + 'use_core_libs': use_core_libs, }, model_dir=model_dir, config=config, diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/estimator_test.py b/tensorflow/contrib/boosted_trees/estimator_batch/estimator_test.py new file mode 100644 index 0000000000000000000000000000000000000000..0d58317bd59331cfcde0e12aeb3a3a03fc45d89b --- /dev/null +++ b/tensorflow/contrib/boosted_trees/estimator_batch/estimator_test.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. +# ============================================================================== +"""Tests for GBDT estimator.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function +import tempfile +from tensorflow.contrib.boosted_trees.estimator_batch import estimator +from tensorflow.contrib.boosted_trees.proto import learner_pb2 +from tensorflow.contrib.layers.python.layers import feature_column as contrib_feature_column +from tensorflow.contrib.learn.python.learn.estimators import run_config +from tensorflow.python.estimator.canned import head as head_lib +from tensorflow.python.feature_column import feature_column_lib as core_feature_column +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import test_util +from tensorflow.python.ops.losses import losses +from tensorflow.python.platform import gfile +from tensorflow.python.platform import googletest + + +def _train_input_fn(): + features = {"x": constant_op.constant([[2.], [1.], [1.]])} + label = constant_op.constant([[1], [0], [0]], dtype=dtypes.int32) + return features, label + + +def _eval_input_fn(): + features = {"x": constant_op.constant([[1.], [2.], [2.]])} + label = constant_op.constant([[0], [1], [1]], dtype=dtypes.int32) + return features, label + + +class BoostedTreeEstimatorTest(test_util.TensorFlowTestCase): + + def setUp(self): + self._export_dir_base = tempfile.mkdtemp() + "export/" + gfile.MkDir(self._export_dir_base) + + def testFitAndEvaluateDontThrowException(self): + learner_config = learner_pb2.LearnerConfig() + learner_config.num_classes = 2 + learner_config.constraints.max_tree_depth = 1 + model_dir = tempfile.mkdtemp() + config = run_config.RunConfig() + + classifier = estimator.GradientBoostedDecisionTreeClassifier( + learner_config=learner_config, + num_trees=1, + examples_per_layer=3, + model_dir=model_dir, + config=config, + feature_columns=[contrib_feature_column.real_valued_column("x")]) + + classifier.fit(input_fn=_train_input_fn, steps=15) + classifier.evaluate(input_fn=_eval_input_fn, steps=1) + classifier.export(self._export_dir_base) + + def testFitAndEvaluateDontThrowExceptionWithCoreForEstimator(self): + learner_config = learner_pb2.LearnerConfig() + learner_config.num_classes = 2 + learner_config.constraints.max_tree_depth = 1 + model_dir = tempfile.mkdtemp() + config = run_config.RunConfig() + + # Use core head + head_fn = head_lib._binary_logistic_head_with_sigmoid_cross_entropy_loss( + loss_reduction=losses.Reduction.SUM_OVER_BATCH_SIZE) + + model = estimator.GradientBoostedDecisionTreeEstimator( + head=head_fn, + learner_config=learner_config, + num_trees=1, + examples_per_layer=3, + model_dir=model_dir, + config=config, + feature_columns=[core_feature_column.numeric_column("x")], + use_core_libs=True) + + model.fit(input_fn=_train_input_fn, steps=15) + model.evaluate(input_fn=_eval_input_fn, steps=1) + model.export(self._export_dir_base) + + def testFitAndEvaluateDontThrowExceptionWithCoreForClassifier(self): + learner_config = learner_pb2.LearnerConfig() + learner_config.num_classes = 2 + learner_config.constraints.max_tree_depth = 1 + model_dir = tempfile.mkdtemp() + config = run_config.RunConfig() + + classifier = estimator.GradientBoostedDecisionTreeClassifier( + learner_config=learner_config, + num_trees=1, + examples_per_layer=3, + model_dir=model_dir, + config=config, + feature_columns=[core_feature_column.numeric_column("x")], + use_core_libs=True) + + classifier.fit(input_fn=_train_input_fn, steps=15) + classifier.evaluate(input_fn=_eval_input_fn, steps=1) + classifier.export(self._export_dir_base) + + def testFitAndEvaluateDontThrowExceptionWithCoreForRegressor(self): + learner_config = learner_pb2.LearnerConfig() + learner_config.num_classes = 2 + learner_config.constraints.max_tree_depth = 1 + model_dir = tempfile.mkdtemp() + config = run_config.RunConfig() + + regressor = estimator.GradientBoostedDecisionTreeRegressor( + learner_config=learner_config, + num_trees=1, + examples_per_layer=3, + model_dir=model_dir, + config=config, + feature_columns=[core_feature_column.numeric_column("x")], + use_core_libs=True) + + regressor.fit(input_fn=_train_input_fn, steps=15) + regressor.evaluate(input_fn=_eval_input_fn, steps=1) + regressor.export(self._export_dir_base) + + +if __name__ == "__main__": + googletest.main() diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/estimator_utils.py b/tensorflow/contrib/boosted_trees/estimator_batch/estimator_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..48a7f85eada8c72de83b814af2f00e97a62a073e --- /dev/null +++ b/tensorflow/contrib/boosted_trees/estimator_batch/estimator_utils.py @@ -0,0 +1,74 @@ +# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Utilities for converting between core and contrib feature columns.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.learn.python.learn.estimators import constants +from tensorflow.contrib.learn.python.learn.estimators import model_fn +from tensorflow.contrib.learn.python.learn.estimators import model_fn as contrib_model_fn_lib +from tensorflow.contrib.learn.python.learn.estimators import prediction_key +from tensorflow.python.estimator import model_fn as model_fn_lib +from tensorflow.python.estimator.export import export_output + +_CORE_MODE_TO_CONTRIB_MODE_ = { + model_fn_lib.ModeKeys.TRAIN: contrib_model_fn_lib.ModeKeys.TRAIN, + model_fn_lib.ModeKeys.EVAL: contrib_model_fn_lib.ModeKeys.EVAL, + model_fn_lib.ModeKeys.PREDICT: contrib_model_fn_lib.ModeKeys.INFER +} + + +def _core_mode_to_contrib_mode(mode): + return _CORE_MODE_TO_CONTRIB_MODE_[mode] + + +def _export_outputs_to_output_alternatives(export_outputs): + """Converts EstimatorSpec.export_outputs to output_alternatives. + + Args: + export_outputs: export_outputs created by create_estimator_spec. + Returns: + converted output_alternatives. + """ + output = dict() + if export_outputs is not None: + for key, value in export_outputs.items(): + if isinstance(value, export_output.ClassificationOutput): + exported_predictions = { + prediction_key.PredictionKey.SCORES: value.scores, + prediction_key.PredictionKey.CLASSES: value.classes + } + output[key] = (constants.ProblemType.CLASSIFICATION, + exported_predictions) + return output + return None + + +def estimator_spec_to_model_fn_ops(estimator_spec, export_alternatives=False): + if export_alternatives: + alternatives = _export_outputs_to_output_alternatives( + estimator_spec.export_outputs) + else: + alternatives = [] + + return model_fn.ModelFnOps( + mode=_core_mode_to_contrib_mode(estimator_spec.mode), + predictions=estimator_spec.predictions, + loss=estimator_spec.loss, + train_op=estimator_spec.train_op, + eval_metric_ops=estimator_spec.eval_metric_ops, + output_alternatives=alternatives) diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/model.py b/tensorflow/contrib/boosted_trees/estimator_batch/model.py index c6455a7ea3d18eb358edee034cee58b2bed21024..15ab6d814522ab1dee58dcd71246354fc4d8a483 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/model.py +++ b/tensorflow/contrib/boosted_trees/estimator_batch/model.py @@ -20,6 +20,7 @@ from __future__ import print_function import copy +from tensorflow.contrib.boosted_trees.estimator_batch import estimator_utils from tensorflow.contrib.boosted_trees.estimator_batch import trainer_hooks from tensorflow.contrib.boosted_trees.python.ops import model_ops from tensorflow.contrib.boosted_trees.python.training.functions import gbdt_batch @@ -60,6 +61,7 @@ def model_builder(features, labels, mode, params, config): feature_columns = params["feature_columns"] weight_column_name = params["weight_column_name"] num_trees = params["num_trees"] + use_core_libs = params["use_core_libs"] logits_modifier_function = params["logits_modifier_function"] if features is None: raise ValueError("At least one feature must be specified.") @@ -93,7 +95,8 @@ def model_builder(features, labels, mode, params, config): learner_config=learner_config, feature_columns=feature_columns, logits_dimension=head.logits_dimension, - features=training_features) + features=training_features, + use_core_columns=use_core_libs) with ops.name_scope("gbdt", "gbdt_optimizer"): predictions_dict = gbdt_model.predict(mode) logits = predictions_dict["predictions"] @@ -108,12 +111,22 @@ def model_builder(features, labels, mode, params, config): update_op = state_ops.assign_add(global_step, 1).op return update_op - model_fn_ops = head.create_model_fn_ops( - features=features, - mode=mode, - labels=labels, - train_op_fn=_train_op_fn, - logits=logits) + create_estimator_spec_op = getattr(head, "create_estimator_spec", None) + if use_core_libs and callable(create_estimator_spec_op): + model_fn_ops = head.create_estimator_spec( + features=features, + mode=mode, + labels=labels, + train_op_fn=_train_op_fn, + logits=logits) + model_fn_ops = estimator_utils.estimator_spec_to_model_fn_ops(model_fn_ops) + else: + model_fn_ops = head.create_model_fn_ops( + features=features, + mode=mode, + labels=labels, + train_op_fn=_train_op_fn, + logits=logits) if num_trees: if center_bias: num_trees += 1 diff --git a/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_stream_test.cc b/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_stream_test.cc index 4481c0d0e4400acd93c9a277de185db7aaf9bcb0..67ac9bf387ae9b3ca29e610c2c4138c28302ca33 100644 --- a/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_stream_test.cc +++ b/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_stream_test.cc @@ -138,6 +138,12 @@ void GenerateOneValue(int32 worker_id, int64 max_elements, double *total_weight, stream->Finalize(); } +void GenerateOneZeroWeightedValue(int32 worker_id, int64 max_elements, + double *total_weight, Stream *stream) { + stream->PushEntry(10, 0); + stream->Finalize(); +} + TEST(WeightedQuantilesStreamTest, OneValue) { const double eps = 0.01; const int64 max_elements = 1 << 16; @@ -145,6 +151,13 @@ TEST(WeightedQuantilesStreamTest, OneValue) { {10.0, 10.0, 10.0, 10.0, 10.0}, 1e-2); } +TEST(WeightedQuantilesStreamTest, OneZeroWeightValue) { + const double eps = 0.01; + const int64 max_elements = 1 << 16; + TestSingleWorkerStreams(eps, max_elements, GenerateOneZeroWeightedValue, {}, + 1e-2); +} + TEST(WeightedQuantilesStreamTest, FixedUniform) { const double eps = 0.01; const int64 max_elements = 1 << 16; diff --git a/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_summary.h b/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_summary.h index aec232f3cbb096f0aa51e4362a821882391f8027..7576856dc3a6d0b6681ee9745c875cf46d1e2960 100644 --- a/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_summary.h +++ b/tensorflow/contrib/boosted_trees/lib/quantiles/weighted_quantiles_summary.h @@ -235,6 +235,11 @@ class WeightedQuantilesSummary { // The resulting boundaries are guaranteed to both contain at least // num_boundaries unique elements and maintain approximation bounds. std::vector GenerateBoundaries(int64 num_boundaries) const { + std::vector output; + if (entries_.empty()) { + return output; + } + // Generate soft compressed summary. WeightedQuantilesSummary compressed_summary; @@ -246,7 +251,6 @@ class WeightedQuantilesSummary { compressed_summary.Compress(num_boundaries, compression_eps); // Return boundaries. - std::vector output; output.reserve(compressed_summary.entries_.size()); for (const auto& entry : compressed_summary.entries_) { output.push_back(entry.value); @@ -260,6 +264,9 @@ class WeightedQuantilesSummary { // full rank queries O(nlogn). std::vector GenerateQuantiles(int64 num_quantiles) const { std::vector output; + if (entries_.empty()) { + return output; + } num_quantiles = std::max(num_quantiles, 2LL); output.reserve(num_quantiles + 1); 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 85b909e4f2556c520a5bffe46d5954683d9dda5a..4bde7f3e33d6f8b295cd35cb32bbbccecf8a2b87 100644 --- a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py +++ b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py @@ -23,7 +23,6 @@ import copy from tensorflow.contrib import learn from tensorflow.contrib import stateless - from tensorflow.contrib.boosted_trees.lib.learner.batch import categorical_split_handler from tensorflow.contrib.boosted_trees.lib.learner.batch import ordinal_split_handler from tensorflow.contrib.boosted_trees.proto import learner_pb2 @@ -141,7 +140,7 @@ class _OpRoundRobinStrategy(object): return task -def extract_features(features, feature_columns): +def extract_features(features, feature_columns, use_core_columns): """Extracts columns from a dictionary of features. Args: @@ -174,7 +173,11 @@ def extract_features(features, feature_columns): transformed_features = collections.OrderedDict() for fc in feature_columns: # pylint: disable=protected-access - if isinstance(fc, feature_column_lib._EmbeddingColumn): + if use_core_columns: + # pylint: disable=protected-access + tensor = fc_core._transform_features(features, [fc])[fc] + transformed_features[fc.name] = tensor + elif isinstance(fc, feature_column_lib._EmbeddingColumn): # pylint: enable=protected-access transformed_features[fc.name] = fc_core.input_layer( features, [fc], @@ -265,7 +268,8 @@ class GradientBoostedDecisionTreeModel(object): learner_config, features, logits_dimension, - feature_columns=None): + feature_columns=None, + use_core_columns=False): """Construct a new GradientBoostedDecisionTreeModel function. Args: @@ -338,8 +342,9 @@ class GradientBoostedDecisionTreeModel(object): if not features: raise ValueError("Features dictionary must be specified.") (fc_names, dense_floats, sparse_float_indices, sparse_float_values, - sparse_float_shapes, sparse_int_indices, sparse_int_values, - sparse_int_shapes) = extract_features(features, self._feature_columns) + sparse_float_shapes, sparse_int_indices, + sparse_int_values, sparse_int_shapes) = extract_features( + features, self._feature_columns, use_core_columns) logging.info("Active Feature Columns: " + str(fc_names)) self._fc_names = fc_names self._dense_floats = dense_floats 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 6411f57a5419123e799af9231a04fce8ae7724d4..f9c22283b7f5136777bfa60a12c94974adfbd245 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 @@ -27,9 +27,11 @@ from tensorflow.contrib.boosted_trees.python.ops import model_ops from tensorflow.contrib.boosted_trees.python.training.functions import gbdt_batch from tensorflow.contrib.boosted_trees.python.utils import losses +from tensorflow.python.feature_column import feature_column_lib as core_feature_column from tensorflow.contrib.layers.python.layers import feature_column as feature_column_lib from tensorflow.contrib.learn.python.learn.estimators import model_fn + from tensorflow.python.framework import dtypes from tensorflow.python.framework import sparse_tensor from tensorflow.python.framework import test_util @@ -43,7 +45,7 @@ from tensorflow.python.platform import googletest def _squared_loss(label, unused_weights, predictions): """Unweighted loss implementation.""" loss = math_ops.reduce_sum( - math_ops.square(predictions - label), 1, keep_dims=True) + math_ops.square(predictions - label), 1, keepdims=True) return loss @@ -99,7 +101,8 @@ class GbdtTest(test_util.TensorFlowTestCase): array_ops.zeros([2], dtypes.int64)) (fc_names, dense_floats, sparse_float_indices, sparse_float_values, sparse_float_shapes, sparse_int_indices, sparse_int_values, - sparse_int_shapes) = (gbdt_batch.extract_features(features, None)) + sparse_int_shapes) = ( + gbdt_batch.extract_features(features, None, use_core_columns=False)) self.assertEqual(len(fc_names), 3) self.assertAllEqual(fc_names, ["dense_float", "sparse_float", "sparse_int"]) @@ -148,8 +151,9 @@ class GbdtTest(test_util.TensorFlowTestCase): "sparse_categorical", hash_bucket_size=1000000)) (fc_names, dense_floats, sparse_float_indices, sparse_float_values, sparse_float_shapes, sparse_int_indices, sparse_int_values, - sparse_int_shapes) = (gbdt_batch.extract_features( - features, feature_columns)) + sparse_int_shapes) = ( + gbdt_batch.extract_features( + features, feature_columns, use_core_columns=False)) self.assertEqual(len(fc_names), 3) self.assertAllEqual(fc_names, ["dense_float", "sparse_float", "sparse_categorical"]) @@ -174,6 +178,41 @@ class GbdtTest(test_util.TensorFlowTestCase): self.assertAllEqual(sparse_int_shapes[0].eval(), features["sparse_categorical"].dense_shape.eval()) + def testExtractFeaturesFromCoreFeatureColumns(self): + """Tests feature extraction when using core columns.""" + with self.test_session(): + features = {} + # Sparse float column does not exist in core, so only dense numeric and + # categorical. + features["dense_float"] = array_ops.zeros([2, 1], dtypes.float32) + features["sparse_categorical"] = sparse_tensor.SparseTensor( + array_ops.zeros([2, 2], dtypes.int64), + array_ops.zeros([2], dtypes.string), array_ops.zeros([2], + dtypes.int64)) + + feature_columns = set() + feature_columns.add(core_feature_column.numeric_column("dense_float")) + feature_columns.add( + core_feature_column.categorical_column_with_hash_bucket( + "sparse_categorical", hash_bucket_size=1000000)) + (fc_names, dense_floats, _, _, _, sparse_int_indices, sparse_int_values, + sparse_int_shapes) = ( + gbdt_batch.extract_features( + features, feature_columns, use_core_columns=True)) + self.assertEqual(len(fc_names), 2) + self.assertAllEqual(fc_names, ["dense_float", "sparse_categorical"]) + self.assertEqual(len(dense_floats), 1) + self.assertEqual(len(sparse_int_indices), 1) + self.assertEqual(len(sparse_int_values), 1) + self.assertEqual(len(sparse_int_shapes), 1) + self.assertAllEqual(dense_floats[0].eval(), + features["dense_float"].eval()) + self.assertAllEqual(sparse_int_indices[0].eval(), + features["sparse_categorical"].indices.eval()) + self.assertAllEqual(sparse_int_values[0].eval(), [397263, 397263]) + self.assertAllEqual(sparse_int_shapes[0].eval(), + features["sparse_categorical"].dense_shape.eval()) + def testTrainFnChiefNoBiasCentering(self): """Tests the train function running on chief without bias centering.""" with self.test_session() as sess: diff --git a/tensorflow/contrib/checkpoint/README.md b/tensorflow/contrib/checkpoint/README.md new file mode 100644 index 0000000000000000000000000000000000000000..d35c5bae3b702c0fea5194e5e653660e319e38c5 --- /dev/null +++ b/tensorflow/contrib/checkpoint/README.md @@ -0,0 +1,2 @@ +Tools for working with object-based checkpoints produced by +`tf.train.Checkpoint`. diff --git a/tensorflow/python/ops/distributions/bijectors.py b/tensorflow/contrib/checkpoint/__init__.py similarity index 63% rename from tensorflow/python/ops/distributions/bijectors.py rename to tensorflow/contrib/checkpoint/__init__.py index 69c3a5d4c0ba86586ccb6e55e71d898b1bf7c035..1192cc44a17823f69db28947308a8b839a83e57e 100644 --- a/tensorflow/python/ops/distributions/bijectors.py +++ b/tensorflow/contrib/checkpoint/__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,20 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Core module for TensorFlow distribution bijectors.""" +"""Tools for working with object-based checkpoints. + + +For creating and managing dependencies: +@@dot_graph_from_checkpoint +@@split_dependency +""" + from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.contrib.checkpoint.python.split_dependency import split_dependency +from tensorflow.contrib.checkpoint.python.visualize import dot_graph_from_checkpoint -# go/tf-wildcard-import -# pylint: disable=wildcard-import,unused-import -from tensorflow.python.ops.distributions.bijector import Bijector -from tensorflow.python.ops.distributions.identity_bijector import Identity - -# pylint: enable=wildcard-import,unused-import from tensorflow.python.util.all_util import remove_undocumented -_allowed_symbols = ["Bijector", "Identity"] - -remove_undocumented(__name__, _allowed_symbols) +remove_undocumented(module_name=__name__) diff --git a/tensorflow/contrib/checkpoint/python/BUILD b/tensorflow/contrib/checkpoint/python/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..a5681ffa61d07ef29d0a0862db9736a210c8e26e --- /dev/null +++ b/tensorflow/contrib/checkpoint/python/BUILD @@ -0,0 +1,61 @@ +licenses(["notice"]) # Apache 2.0 + +package(default_visibility = ["//tensorflow:internal"]) + +load("//tensorflow:tensorflow.bzl", "py_test") + +py_library( + name = "checkpoint", + srcs_version = "PY2AND3", + deps = [ + ":split_dependency", + ":visualize", + ], +) + +py_library( + name = "split_dependency", + srcs = ["split_dependency.py"], + srcs_version = "PY2AND3", + visibility = ["//tensorflow:internal"], + deps = [ + "//tensorflow/python:control_flow_ops", + "//tensorflow/python:training", + ], +) + +py_test( + name = "split_dependency_test", + srcs = ["split_dependency_test.py"], + deps = [ + ":split_dependency", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:resource_variable_ops", + "//tensorflow/python:training", + "//tensorflow/python/eager:test", + ], +) + +py_library( + name = "visualize", + srcs = ["visualize.py"], + srcs_version = "PY2AND3", + visibility = ["//tensorflow:internal"], + deps = [ + "//tensorflow/python:pywrap_tensorflow", + ], +) + +py_test( + name = "visualize_test", + srcs = ["visualize_test.py"], + deps = [ + ":visualize", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:resource_variable_ops", + "//tensorflow/python:training", + "//tensorflow/python/eager:test", + ], +) diff --git a/tensorflow/contrib/checkpoint/python/split_dependency.py b/tensorflow/contrib/checkpoint/python/split_dependency.py new file mode 100644 index 0000000000000000000000000000000000000000..3aec8c96e90440d6da00d95cffc34bd53ec7164f --- /dev/null +++ b/tensorflow/contrib/checkpoint/python/split_dependency.py @@ -0,0 +1,136 @@ +"""Utility for creating multiple dependencies with synchronized save/restore.""" +# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 + +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.training import checkpointable as checkpointable +from tensorflow.python.training import saver as saver_lib + + +class _CallbackSaveable(saver_lib.BaseSaverBuilder.SaveableObject): + """Wraps save and restore callbacks as a `SaveableObject`.""" + + def __init__(self, name, dtype, save_callback, restore_callback): + self._restore_callback = restore_callback + spec = saver_lib.BaseSaverBuilder.SaveSpec( + tensor=save_callback, + slice_spec="", + name=name, + dtype=dtype) + super(_CallbackSaveable, self).__init__( + save_callback, [spec], name) + + def restore(self, restored_tensors, restored_shapes): + """Restore the same value into both variables.""" + tensor, = restored_tensors + return self._restore_callback(tensor) + + +class _SplitDependency(checkpointable.CheckpointableBase): + """Looks like a regular variable while synchronizing save/restores.""" + + def __init__(self, save_buffer, restore_buffer, name, dtype, num_components, + fill_save_buffer_fn, consume_restore_buffer_fn): + self._save_buffer = save_buffer + self._restore_buffer = restore_buffer + self._name = name + self._dtype = dtype + self._num_components = num_components + self._fill_save_buffer_fn = fill_save_buffer_fn + self._consume_restore_buffer_fn = consume_restore_buffer_fn + + def _save(self): + """Pull from the shared buffer, populating it if necessary.""" + if self._name not in self._save_buffer: + if self._save_buffer: + raise AssertionError( + ("Split dependency %s (%s) unsynchronized. Split dependencies must " + "be saved together.") % (self._name, self)) + self._fill_save_buffer_fn(self._save_buffer) + return self._save_buffer.pop(self._name) + + def _restore(self, tensor): + """Push into the shared buffer, flushing it if necessary.""" + if self._name in self._restore_buffer: + raise AssertionError( + ("Split dependency %s (%s) unsynchronized. Split dependencies must " + "be restored together.") % (self._name, self)) + self._restore_buffer[self._name] = tensor + if len(self._restore_buffer) == self._num_components: + op = self._consume_restore_buffer_fn(self._restore_buffer) + self._restore_buffer.clear() + return op + else: + return control_flow_ops.no_op() + + def _gather_saveables_for_checkpoint(self): + """Looks to Checkpointable like a regular variable.""" + return { + checkpointable.VARIABLE_VALUE_KEY: + functools.partial(_CallbackSaveable, + dtype=self._dtype, + save_callback=self._save, + restore_callback=self._restore) + } + + +def split_dependency(component_names, component_dtypes, + fill_save_buffer_fn, consume_restore_buffer_fn): + """Creates multiple dependencies with a synchronized save/restore. + + Useful when a single op produces `Tensor`s which should each be saved under + different objects, or when `Tensor`s saved with many different objects need to + be restored together as inputs to a single op (i.e. an object which uses a + single fused op may be swapped out for a subgraph of objects, and these two + programs are checkpoint compatible). + + Args: + component_names: A sequence of names for the split + dependencies. `fill_save_buffer_fn` must add these keys to the dictionary + it is passed, and `consume_restore_buffer_fn` will receive a dictionary + with these keys. + component_dtypes: Data types for the `Tensor`s being saved and restored, a + sequence corresponding to `component_names`. + fill_save_buffer_fn: A function which takes an empty dictionary as an + argument and adds `Tensor`s with `component_names` as keys. These + `Tensor`s will be saved as if they were individual variables. + consume_restore_buffer_fn: A function which takes a dictionary with + `component_names` as keys mapping to restored individual `Tensor`s and + returns a restore op (or if executing eagerly, runs the restoration and + may return `None`). + + Returns: + A dictionary mapping from names to Checkpointable objects. If one is + reachable from an object as a dependency, the others should be too; adding + dependencies on some but not all of the objects will result in errors. + """ + save_buffer = {} + restore_buffer = {} + split_dependencies = {} + for name, dtype in zip(component_names, component_dtypes): + split_dependencies[name] = _SplitDependency( + save_buffer=save_buffer, + restore_buffer=restore_buffer, + name=name, + dtype=dtype, + num_components=len(component_names), + fill_save_buffer_fn=fill_save_buffer_fn, + consume_restore_buffer_fn=consume_restore_buffer_fn) + return split_dependencies diff --git a/tensorflow/contrib/checkpoint/python/split_dependency_test.py b/tensorflow/contrib/checkpoint/python/split_dependency_test.py new file mode 100644 index 0000000000000000000000000000000000000000..f1d9d19b047ee69281cf8bdba38a28dc87947e38 --- /dev/null +++ b/tensorflow/contrib/checkpoint/python/split_dependency_test.py @@ -0,0 +1,112 @@ +# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os + +from tensorflow.contrib.checkpoint.python import split_dependency +from tensorflow.python.eager import test +from tensorflow.python.framework import test_util +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import resource_variable_ops +from tensorflow.python.training import checkpointable +from tensorflow.python.training import checkpointable_utils + + +def _split_variable_closure(variable): + def _fill_save_buffer_fn(save_buffer): + save_buffer["first_half"] = variable[:2] + save_buffer["second_half"] = variable[2:] + return _fill_save_buffer_fn + + +def _combine_variable_closure(variable): + def _consume_restore_buffer_fn(restore_buffer): + return variable.assign( + array_ops.concat([restore_buffer["first_half"], + restore_buffer["second_half"]], + axis=0)) + return _consume_restore_buffer_fn + + +class SaveTensorSlicesAsDeps(checkpointable.CheckpointableBase): + + def __init__(self): + self.combined = resource_variable_ops.ResourceVariable([0., 0., 0., 0.]) + split_dependencies = split_dependency.split_dependency( + component_names=("first_half", "second_half"), + component_dtypes=(self.combined.dtype,) * 2, + fill_save_buffer_fn=_split_variable_closure( + self.combined), + consume_restore_buffer_fn=_combine_variable_closure( + self.combined)) + for name, dep in split_dependencies.items(): + self._track_checkpointable(dep, name=name) + + +class HasRegularDeps(checkpointable.Checkpointable): + + def __init__(self): + self.first_half = resource_variable_ops.ResourceVariable([0., 0.]) + self.second_half = resource_variable_ops.ResourceVariable([0., 0.]) + + +class OnlyOneDep(checkpointable.Checkpointable): + + def __init__(self): + self.first_half = resource_variable_ops.ResourceVariable([0., 0.]) + + +class SplitTests(test.TestCase): + + @test_util.run_in_graph_and_eager_modes() + def testSaveRestoreSplitDep(self): + save_checkpoint = checkpointable_utils.Checkpoint( + dep=SaveTensorSlicesAsDeps()) + self.evaluate(save_checkpoint.dep.combined.assign([1., 2., 3., 4.])) + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + save_path = save_checkpoint.save(checkpoint_prefix) + + regular_deps = HasRegularDeps() + regular_restore_checkpoint = checkpointable_utils.Checkpoint( + dep=regular_deps) + regular_restore_checkpoint.restore( + save_path).assert_consumed().run_restore_ops() + self.assertAllEqual([1., 2.], self.evaluate(regular_deps.first_half)) + self.assertAllEqual([3., 4.], self.evaluate(regular_deps.second_half)) + + one_dep = OnlyOneDep() + one_dep_restore_checkpoint = checkpointable_utils.Checkpoint(dep=one_dep) + status = one_dep_restore_checkpoint.restore(save_path) + with self.assertRaises(AssertionError): + # Missing the second dependency. + status.assert_consumed() + status.run_restore_ops() + self.assertAllEqual([1., 2.], self.evaluate(one_dep.first_half)) + + restore_checkpoint = checkpointable_utils.Checkpoint() + status = restore_checkpoint.restore(save_path) + restore_checkpoint.dep = SaveTensorSlicesAsDeps() + status.assert_consumed().run_restore_ops() + self.assertAllEqual( + [1., 2., 3., 4.], + self.evaluate(restore_checkpoint.dep.combined)) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/checkpoint/python/visualize.py b/tensorflow/contrib/checkpoint/python/visualize.py new file mode 100644 index 0000000000000000000000000000000000000000..86fbdb41d2c37803f2bd71b5aa2f72845c87d448 --- /dev/null +++ b/tensorflow/contrib/checkpoint/python/visualize.py @@ -0,0 +1,111 @@ +"""Utilities for visualizing dependency graphs.""" +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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.core.protobuf import checkpointable_object_graph_pb2 +from tensorflow.python import pywrap_tensorflow +from tensorflow.python.framework import errors_impl +from tensorflow.python.training import checkpointable + + +def dot_graph_from_checkpoint(save_path): + r"""Visualizes an object-based checkpoint (from `tf.train.Checkpoint`). + + Useful for inspecting checkpoints and debugging loading issues. + + Example usage from Python (requires pydot): + ```python + import tensorflow as tf + import pydot + + dot_string = tf.contrib.checkpoint.dot_graph_from_checkpoint('/path/to/ckpt') + parsed, = pydot.graph_from_dot_data(dot_string) + parsed.write_svg('/tmp/tensorflow/visualized_checkpoint.svg') + ``` + + Example command line usage: + ```sh + python -c "import tensorflow as tf;\ + print(tf.contrib.checkpoint.dot_graph_from_checkpoint('/path/to/ckpt'))"\ + | dot -Tsvg > /tmp/tensorflow/checkpoint_viz.svg + ``` + + Args: + save_path: The checkpoint prefix, as returned by `tf.train.Checkpoint.save` + or `tf.train.latest_checkpoint`. + Returns: + A graph in DOT format as a string. + """ + reader = pywrap_tensorflow.NewCheckpointReader(save_path) + try: + object_graph_string = reader.get_tensor( + checkpointable.OBJECT_GRAPH_PROTO_KEY) + except errors_impl.NotFoundError: + raise ValueError( + ('The specified checkpoint "%s" does not appear to be object-based (it ' + 'is missing the key "%s"). Likely it was created with a name-based ' + 'saver and does not contain an object dependency graph.') % ( + save_path, checkpointable.OBJECT_GRAPH_PROTO_KEY)) + shape_map = reader.get_variable_to_shape_map() + dtype_map = reader.get_variable_to_dtype_map() + object_graph = ( + checkpointable_object_graph_pb2.CheckpointableObjectGraph()) + object_graph.ParseFromString(object_graph_string) + graph = 'digraph {\n' + def _escape(name): + return name.replace('"', '\\"') + slot_ids = set() + for node in object_graph.nodes: + for slot_reference in node.slot_variables: + slot_ids.add(slot_reference.slot_variable_node_id) + for node_id, node in enumerate(object_graph.nodes): + if (len(node.attributes) == 1 + and node.attributes[0].name == checkpointable.VARIABLE_VALUE_KEY): + if node_id in slot_ids: + color = 'orange' + tooltip_prefix = 'Slot variable' + else: + color = 'blue' + tooltip_prefix = 'Variable' + attribute = node.attributes[0] + graph += ('N_%d [shape=point label="" color=%s width=.25' + ' tooltip="%s %s shape=%s %s"]\n') % ( + node_id, + color, + tooltip_prefix, + _escape(attribute.full_name), + shape_map[attribute.checkpoint_key], + dtype_map[attribute.checkpoint_key].name) + elif node.slot_variables: + graph += ('N_%d [shape=point label="" width=.25 color=red,' + 'tooltip="Optimizer"]\n') % node_id + else: + graph += 'N_%d [shape=point label="" width=.25]\n' % node_id + for reference in node.children: + graph += 'N_%d -> N_%d [label="%s"]\n' % ( + node_id, reference.node_id, _escape(reference.local_name)) + for slot_reference in node.slot_variables: + graph += 'N_%d -> N_%d [label="%s" style=dotted]\n' % ( + node_id, + slot_reference.slot_variable_node_id, + _escape(slot_reference.slot_name)) + graph += 'N_%d -> N_%d [style=dotted]\n' % ( + slot_reference.original_variable_node_id, + slot_reference.slot_variable_node_id) + graph += '}\n' + return graph diff --git a/tensorflow/contrib/checkpoint/python/visualize_test.py b/tensorflow/contrib/checkpoint/python/visualize_test.py new file mode 100644 index 0000000000000000000000000000000000000000..1d9ab789235cb964521315b4864563f89745ae75 --- /dev/null +++ b/tensorflow/contrib/checkpoint/python/visualize_test.py @@ -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. +# ============================================================================== +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import functools +import os + +from tensorflow.contrib.checkpoint.python import visualize + +from tensorflow.python.eager import context +from tensorflow.python.eager import test +from tensorflow.python.framework import constant_op +from tensorflow.python.keras._impl.keras.engine import training +from tensorflow.python.keras._impl.keras.layers import core +from tensorflow.python.ops import resource_variable_ops +from tensorflow.python.training import adam +from tensorflow.python.training import checkpointable_utils + +try: + import pydot # pylint: disable=g-import-not-at-top +except ImportError: + pydot = None + + +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) + + def call(self, values): + ret = self._second(self._named_dense(values)) + return ret + + +class DotGraphTests(test.TestCase): + + def testMakeDotGraph(self): + with context.eager_mode(): + input_value = constant_op.constant([[3.]]) + model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + optimizer_step = resource_variable_ops.ResourceVariable(12) + save_checkpoint = checkpointable_utils.Checkpoint( + optimizer=optimizer, model=model, optimizer_step=optimizer_step) + optimizer.minimize(functools.partial(model, input_value)) + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, 'ckpt') + save_path = save_checkpoint.save(checkpoint_prefix) + prefix = save_checkpoint.save(save_path) + + dot_graph_string = visualize.dot_graph_from_checkpoint(prefix) + + # The remainder of this test is more-or-less optional since it's so + # dependent on pydot/platform/Python versions. + if pydot is None: + self.skipTest('pydot is required for the remainder of this test.') + try: + parsed, = pydot.graph_from_dot_data(dot_graph_string) + except NameError as e: + if "name 'dot_parser' is not defined" in str(e): + self.skipTest("pydot isn't working") + else: + raise + # Check that the graph isn't completely trivial + self.assertEqual( + '"model"', + parsed.obj_dict['edges'][('N_0', 'N_1')][0]['attributes']['label']) + image_path = os.path.join(self.get_temp_dir(), 'saved.svg') + try: + parsed.write_svg(image_path) + except Exception as e: # pylint: disable=broad-except + # For some reason PyDot's "dot not available" error is an Exception, not + # something more specific. + if '"dot" not found in path' in str(e): + self.skipTest("pydot won't save SVGs (dot not available)") + else: + raise + +if __name__ == '__main__': + test.main() 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 5a2771229d9ffe2b5b389d1077fe02a230e9a4c0..1403483d287041b02dfbf538f7e7ddee11662f47 100644 --- a/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py +++ b/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py @@ -245,7 +245,7 @@ class TPUClusterResolver(ClusterResolver): else: if not self._tpu.startswith(compat.as_bytes('grpc://')): # Case 3. - return server_lib.ClusterSpec({}) + return None # Case 2. cluster_spec = {self._job_name: [self._tpu[len( compat.as_bytes('grpc://')):]]} 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 dff7a03b6847fb6e159dc2fa9832fceb3dfe2d54..5b3f9be5a11237f9dceebefa1db294efaf7e482d 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 @@ -356,8 +356,7 @@ class TPUClusterResolverTest(test.TestCase): 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()) + self.assertEqual(None, tpu_cluster_resolver.cluster_spec()) def testGkeEnvironment(self): os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'] = 'grpc://10.120.27.5:8470' diff --git a/tensorflow/contrib/cmake/CMakeLists.txt b/tensorflow/contrib/cmake/CMakeLists.txt index 23b31ae1dcc83d8a7152354ac147de9ada320429..5f38a8e5c755cd11a568e437b3ade665d983461f 100644 --- a/tensorflow/contrib/cmake/CMakeLists.txt +++ b/tensorflow/contrib/cmake/CMakeLists.txt @@ -31,10 +31,14 @@ option(tensorflow_BUILD_PYTHON_TESTS "Build python unit tests " OFF) option(tensorflow_BUILD_MORE_PYTHON_TESTS "Build more python unit tests for contrib packages" OFF) option(tensorflow_BUILD_SHARED_LIB "Build TensorFlow as a shared library" OFF) option(tensorflow_OPTIMIZE_FOR_NATIVE_ARCH "Enable compiler optimizations for the native processor architecture (if available)" ON) -option(tensorflow_WIN_CPU_SIMD_OPTIONS "Enables CPU SIMD instructions") option(tensorflow_ENABLE_SNAPPY_SUPPORT "Enable SNAPPY compression support" ON) option(tensorflow_DISABLE_EIGEN_FORCEINLINE "Disable forceinline, to speed up build on windows." OFF) +# SIMD, MKL and MKLDNN options +option(tensorflow_WIN_CPU_SIMD_OPTIONS "Enables CPU SIMD instructions" OFF) +option(tensorflow_ENABLE_MKL_SUPPORT "Enable Intel MKL support" OFF) +option(tensorflow_ENABLE_MKLDNN_SUPPORT "Enable Intel MKLDNN support, requires MKL enabled" OFF) + # GPU, CUDA and cuDNN options option(tensorflow_ENABLE_GPU "Enable GPU support" OFF) set(tensorflow_CUDA_VERSION "9.0" CACHE STRING "CUDA version to build against") @@ -124,8 +128,16 @@ endif() add_definitions(-DEIGEN_AVOID_STL_ARRAY) if(WIN32) - add_definitions(-DNOMINMAX -D_WIN32_WINNT=0x0A00 -DLANG_CXX11 -DCOMPILER_MSVC) - add_definitions(-DWIN32 -DOS_WIN -D_MBCS -DWIN64 -DWIN32_LEAN_AND_MEAN -DNOGDI -DPLATFORM_WINDOWS) + if(CMAKE_SIZEOF_VOID_P EQUAL 8) + # 64 bits + add_definitions(-DWIN64) + elseif(CMAKE_SIZEOF_VOID_P EQUAL 4) + # 32 bits + # temporary fix for #18241 + add_definitions(-DEIGEN_DEFAULT_DENSE_INDEX_TYPE=std::int64_t) + endif() + add_definitions(-DNOMINMAX -D_WIN32_WINNT=0x0A00 -DLANG_CXX11) + add_definitions(-DWIN32 -DOS_WIN -D_MBCS -DWIN32_LEAN_AND_MEAN -DNOGDI -DPLATFORM_WINDOWS) add_definitions(-DTENSORFLOW_USE_EIGEN_THREADPOOL -DEIGEN_HAS_C99_MATH) add_definitions(-DTF_COMPILE_LIBRARY) add_definitions(/bigobj /nologo /EHsc /GF /MP /Gm-) @@ -162,12 +174,21 @@ endif() # MSVC SIMD instructions if (tensorflow_WIN_CPU_SIMD_OPTIONS) + include(CheckCXXCompilerFlag) + if (tensorflow_ENABLE_MKL_SUPPORT) + add_definitions(-DINTEL_MKL -DEIGEN_USE_VML) + if (NOT tensorflow_ENABLE_MKLDNN_SUPPORT) + add_definitions(-DINTEL_MKL_ML) + endif() + endif() + CHECK_CXX_COMPILER_FLAG("-fopenmp" COMPILER_OPT_OPENMP_SUPPORT) + if (COMPILER_OPT_OPENMP_SUPPORT) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fopenmp") + endif() if (WIN32) - CHECK_CXX_COMPILER_FLAG("${tensorflow_WIN_CPU_SIMD_OPTIONS}" COMPILER_OPT_WIN_CPU_SIMD_SUPPORTED) + CHECK_CXX_COMPILER_FLAG(${tensorflow_WIN_CPU_SIMD_OPTIONS} COMPILER_OPT_WIN_CPU_SIMD_SUPPORTED) if(COMPILER_OPT_WIN_CPU_SIMD_SUPPORTED) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${tensorflow_WIN_CPU_SIMD_OPTIONS}") - else() - message(FATAL_ERROR "${tensorflow_WIN_CPU_SIMD_OPTIONS} not supported") endif() endif() endif() @@ -193,6 +214,7 @@ include(protobuf) include(re2) include(cub) include(sqlite) +include(double_conversion) if (tensorflow_BUILD_CC_TESTS) include(googletest) endif() @@ -213,6 +235,7 @@ set(tensorflow_EXTERNAL_LIBRARIES ${protobuf_STATIC_LIBRARIES} ${re2_STATIC_LIBRARIES} ${sqlite_STATIC_LIBRARIES} + ${double_conversion_STATIC_LIBRARIES} ) if (systemlib_ZLIB) @@ -240,6 +263,7 @@ set(tensorflow_EXTERNAL_DEPENDENCIES fft2d re2 sqlite_copy_headers_to_destination + double_conversion ) include_directories( @@ -262,6 +286,7 @@ include_directories( ${PROTOBUF_INCLUDE_DIRS} ${re2_INCLUDE_DIR} ${sqlite_INCLUDE_DIR} + ${double_conversion_INCLUDE_DIR} ) if(tensorflow_ENABLE_SSL_SUPPORT) @@ -298,6 +323,43 @@ if(HAIKU) list(APPEND tensorflow_EXTERNAL_LIBRARIES network) endif() +if (tensorflow_ENABLE_MKL_SUPPORT) + if (WIN32) + find_path(MKL_HOME_PLATFORM mkl + PATHS ${MKL_HOME} ${MKL_HOME}/../ ${MKL_HOME}/../../ + PATH_SUFFIXES windows) + set(MKL_INCLUDE_DIRS ${MKL_HOME_PLATFORM}/mkl/include) + set(MKL_LINK_DIRS + ${MKL_HOME_PLATFORM}/mkl/lib/intel64 + ${MKL_HOME_PLATFORM}/tbb/lib/intel64/vc_mt + ${MKL_HOME_PLATFORM}/compiler/lib/intel64 + ${MKL_HOME_PLATFORM}/mkl/tools/builder/lib) + set(MKL_REDIST_DLL_DIRS + ${MKL_HOME_PLATFORM}/redist/intel64/mkl + ${MKL_HOME_PLATFORM}/redist/intel64/tbb/vc_mt + ${MKL_HOME_PLATFORM}/redist/intel64/compiler) + list(APPEND tensorflow_EXTERNAL_LIBRARIES + mkl_intel_lp64_dll mkl_sequential_dll mkl_core_dll mkl_rt mkl_cdll_intel64) + endif() + if (UNIX) + # Fix me: complete the path on linux + find_path(MKL_HOME_PLATFORM mkl + HINTS ${MKL_HOME} ${MKL_HOME}/../ ${MKL_HOME}/../../ + PATH_SUFFIXES linux) + set(MKL_INCLUDE_DIRS ${MKL_HOME_PLATFORM}/mkl/include) + set(MKL_LINK_DIRS) # incompleted + set(MKL_REDIST_SO_DIRS) # incompleted + endif() + include_directories(${MKL_INCLUDE_DIRS}) + link_directories(${MKL_LINK_DIRS}) + if (tensorflow_ENABLE_MKLDNN_SUPPORT) + include(mkldnn) + list(APPEND tensorflow_EXTERNAL_LIBRARIES ${mkldnn_STATIC_LIBRARIES}) + list(APPEND tensorflow_EXTERNAL_DEPENDENCIES mkldnn) + include_directories(${mkldnn_INCLUDE_DIRS}) + endif() +endif (tensorflow_ENABLE_MKL_SUPPORT) + if (tensorflow_ENABLE_GPU) if (NOT WIN32) # Default install paths for cuda libraries in Linux diff --git a/tensorflow/contrib/cmake/README.md b/tensorflow/contrib/cmake/README.md index fe83bb32046cd75328c92a74cdb4fdb6ce44560e..0b79f718d4823a987e02804f59a432ee46d0ada3 100644 --- a/tensorflow/contrib/cmake/README.md +++ b/tensorflow/contrib/cmake/README.md @@ -128,6 +128,18 @@ Step-by-step Windows build D:\local\cuda\bin ``` + * When building with MKL support after installing [MKL](https://software.intel.com/en-us/mkl) from INTEL, append its bin directories to your PATH environment variable. + + In case TensorFlow fails to find the MKL dll's during initialization, check your PATH environment variable. + It should contain the directory of the MKL dlls. For example: + + ``` + D:\Tools\IntelSWTools\compilers_and_libraries\windows\redist\intel64\mkl + D:\Tools\IntelSWTools\compilers_and_libraries\windows\redist\intel64\compiler + D:\Tools\IntelSWTools\compilers_and_libraries\windows\redist\intel64\tbb\vc_mt + ``` + + * We assume that `cmake` and `git` are installed and in your `%PATH%`. If for example `cmake` is not in your path and it is installed in `C:\Program Files (x86)\CMake\bin\cmake.exe`, you can add this directory @@ -166,7 +178,15 @@ Step-by-step Windows build More? -Dtensorflow_ENABLE_GPU=ON ^ More? -DCUDNN_HOME="D:\...\cudnn" ``` + To build with MKL support add "^" at the end of the last line above following with: + + ``` + More? -Dtensorflow_ENABLE_MKL_SUPPORT=ON ^ + More? -DMKL_HOME="D:\...\compilers_and_libraries" + ``` + To enable SIMD instructions with MSVC, as AVX and SSE, define it as follows: + ``` More? -Dtensorflow_WIN_CPU_SIMD_OPTIONS=/arch:AVX ``` @@ -226,6 +246,7 @@ Step-by-step Windows build ``` ctest -C RelWithDebInfo ``` + * `-Dtensorflow_BUILD_MORE_PYTHON_TESTS=(ON|OFF)`. Defaults to `OFF`. This enables python tests on serveral major packages. This option is only valid if this and tensorflow_BUILD_PYTHON_TESTS are both set as `ON`. After building the python wheel, you need to install the new wheel before running the tests. @@ -234,6 +255,12 @@ Step-by-step Windows build ctest -C RelWithDebInfo ``` + * `-Dtensorflow_ENABLE_MKL_SUPPORT=(ON|OFF)`. Defaults to `OFF`. Include MKL support. If MKL is enabled you need to install the [Intel Math Kernal Library](https://software.intel.com/en-us/mkl). + CMake will expect the location of MKL in -MKL_HOME=path_you_install_mkl. + + * `-Dtensorflow_ENABLE_MKLDNN_SUPPORT=(ON|OFF)`. Defaults to `OFF`. Include MKL DNN support. MKL DNN is [Intel(R) Math Kernel Library for Deep Neural Networks (Intel(R) MKL-DNN)](https://github.com/intel/mkl-dnn). You have to add `-Dtensorflow_ENABLE_MKL_SUPPORT=ON` before including MKL DNN support. + + 4. Invoke MSBuild to build TensorFlow. To build the C++ example program, which will be created as a `.exe` @@ -251,6 +278,7 @@ Step-by-step Windows build D:\...\build> MSBuild /p:Configuration=Release tf_python_build_pip_package.vcxproj ``` + Linux Continuous Integration build ================================== diff --git a/tensorflow/contrib/cmake/external/double_conversion.cmake b/tensorflow/contrib/cmake/external/double_conversion.cmake new file mode 100644 index 0000000000000000000000000000000000000000..527ccdc8d887cb4c2e7d2412c99a8bc682568472 --- /dev/null +++ b/tensorflow/contrib/cmake/external/double_conversion.cmake @@ -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 (ExternalProject) + +set(double_conversion_INCLUDE_DIR ${CMAKE_CURRENT_BINARY_DIR}/double_conversion/src/double_conversion) +set(double_conversion_URL https://github.com/google/double-conversion.git) +set(double_conversion_TAG 5664746) +set(double_conversion_BUILD ${double_conversion_INCLUDE_DIR}) +set(double_conversion_LIBRARIES ${double_conversion_BUILD}/double-conversion/libdouble-conversion.so) +set(double_conversion_INCLUDES ${double_conversion_BUILD}) + +if(WIN32) + set(double_conversion_STATIC_LIBRARIES ${double_conversion_BUILD}/double-conversion/$(Configuration)/double-conversion.lib) +else() + set(double_conversion_STATIC_LIBRARIES ${double_conversion_BUILD}/double-conversion/libdouble-conversion.a) +endif() + +set(double_conversion_HEADERS + "${double_conversion_INCLUDE_DIR}/double-conversion/bignum-dtoa.h" + "${double_conversion_INCLUDE_DIR}/double-conversion/cached-powers.h" + "${double_conversion_INCLUDE_DIR}/double-conversion/double-conversion.h" + "${double_conversion_INCLUDE_DIR}/double-conversion/fixed-dtoa.h" + "${double_conversion_INCLUDE_DIR}/double-conversion/strtod.h" + "${double_conversion_INCLUDE_DIR}/double-conversion/bignum.h" + "${double_conversion_INCLUDE_DIR}/double-conversion/diy-fp.h" + "${double_conversion_INCLUDE_DIR}/double-conversion/fast-dtoa.h" + "${double_conversion_INCLUDE_DIR}/double-conversion/ieee.h" + "${double_conversion_INCLUDE_DIR}/double-conversion/utils.h" +) + +ExternalProject_Add(double_conversion + PREFIX double_conversion + GIT_REPOSITORY ${double_conversion_URL} + GIT_TAG ${double_conversion_TAG} + DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" + BUILD_IN_SOURCE 1 + INSTALL_COMMAND "" + CMAKE_CACHE_ARGS + -DCMAKE_BUILD_TYPE:STRING=Release + -DCMAKE_VERBOSE_MAKEFILE:BOOL=OFF + -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON +) diff --git a/tensorflow/contrib/cmake/external/gemmlowp.cmake b/tensorflow/contrib/cmake/external/gemmlowp.cmake index a235442dc5c0a07e249653381436eeae81575883..cdaa6b73b93666d272faacb869e8272561a2c74c 100644 --- a/tensorflow/contrib/cmake/external/gemmlowp.cmake +++ b/tensorflow/contrib/cmake/external/gemmlowp.cmake @@ -14,8 +14,8 @@ # ============================================================================== include (ExternalProject) -set(gemmlowp_URL https://github.com/google/gemmlowp/archive/6a2a90822e8546fc2bfa7044de0faf1c1cb4862f.zip) -set(gemmlowp_HASH SHA256=3447948d219f3270383766bbe08942888c0eb4e0ca6663c0e0548502ec5bb77d) +set(gemmlowp_URL https://github.com/google/gemmlowp/archive/38ebac7b059e84692f53e5938f97a9943c120d98.zip) +set(gemmlowp_HASH SHA256=b87faa7294dfcc5d678f22a59d2c01ca94ea1e2a3b488c38a95a67889ed0a658) set(gemmlowp_BUILD ${CMAKE_CURRENT_BINARY_DIR}/gemmlowp/src/gemmlowp) set(gemmlowp_INCLUDE_DIR ${CMAKE_CURRENT_BINARY_DIR}/gemmlowp/src/gemmlowp) diff --git a/tensorflow/contrib/cmake/external/grpc.cmake b/tensorflow/contrib/cmake/external/grpc.cmake index 1fefb731a775d9cd2478cbb654662ec6ba673fed..693dc7cd673233b889b35a3f3170b57581da9a9f 100644 --- a/tensorflow/contrib/cmake/external/grpc.cmake +++ b/tensorflow/contrib/cmake/external/grpc.cmake @@ -17,7 +17,7 @@ include (ExternalProject) set(GRPC_INCLUDE_DIRS ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/include) set(GRPC_URL https://github.com/grpc/grpc.git) set(GRPC_BUILD ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc) -set(GRPC_TAG bd6bdf93279a39a8cd92978fd7c9d14eccd98fc2) +set(GRPC_TAG d184fa229d75d336aedea0041bd59cb93e7e267f) if(WIN32) if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") diff --git a/tensorflow/contrib/cmake/external/mkldnn.cmake b/tensorflow/contrib/cmake/external/mkldnn.cmake new file mode 100644 index 0000000000000000000000000000000000000000..a639fdee367f060d4c8a79267803da6ffe3dc503 --- /dev/null +++ b/tensorflow/contrib/cmake/external/mkldnn.cmake @@ -0,0 +1,44 @@ +# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +include (ExternalProject) + +set(mkldnn_INCLUDE_DIRS ${CMAKE_CURRENT_BINARY_DIR}/mkldnn/src/mkldnn/include) +set(mkldnn_URL https://github.com/01org/mkl-dnn.git) +set(mkldnn_BUILD ${CMAKE_CURRENT_BINARY_DIR}/mkldnn/src/mkldnn/src) +set(mkldnn_TAG 3063b2e4c943983f6bf5f2fb9a490d4a998cd291) + +if(WIN32) + if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") + set(mkldnn_STATIC_LIBRARIES ${CMAKE_CURRENT_BINARY_DIR}/mkldnn/src/mkldnn/src/Release/mkldnn.lib) + else() + set(mkldnn_STATIC_LIBRARIES ${CMAKE_CURRENT_BINARY_DIR}/mkldnn/src/mkldnn/src/mkldnn.lib) + endif() +else() + set(mkldnn_STATIC_LIBRARIES ${CMAKE_CURRENT_BINARY_DIR}/mkldnn/src/mkldnn/src/libmkldnn.a) +endif() + +ExternalProject_Add(mkldnn + PREFIX mkldnn + GIT_REPOSITORY ${mkldnn_URL} + GIT_TAG ${mkldnn_TAG} + DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" + BUILD_IN_SOURCE 1 + BUILD_BYPRODUCTS ${mkldnn_STATIC_LIBRARIES} + INSTALL_COMMAND "" + CMAKE_CACHE_ARGS + -DCMAKE_BUILD_TYPE:STRING=Release + -DCMAKE_VERBOSE_MAKEFILE:BOOL=OFF + -DMKLINC:STRING=${MKL_INCLUDE_DIRS} +) diff --git a/tensorflow/contrib/cmake/external/png.cmake b/tensorflow/contrib/cmake/external/png.cmake index 6cd66a65990e7a2b963b52b310061b551752cd4d..ad2af01bc002555ce48f8b9bfb7d8d724a1a7dc8 100644 --- a/tensorflow/contrib/cmake/external/png.cmake +++ b/tensorflow/contrib/cmake/external/png.cmake @@ -15,32 +15,33 @@ include (ExternalProject) set(png_INCLUDE_DIR ${CMAKE_CURRENT_BINARY_DIR}/external/png_archive) -set(png_URL https://storage.googleapis.com/libpng-public-archive/libpng-1.2.53.tar.gz) -set(png_HASH SHA256=e05c9056d7f323088fd7824d8c6acc03a4a758c4b4916715924edc5dd3223a72) +set(png_URL https://mirror.bazel.build/github.com/glennrp/libpng/archive/v1.6.34.tar.gz) +set(png_HASH SHA256=e45ce5f68b1d80e2cb9a2b601605b374bdf51e1798ef1c2c2bd62131dfcf9eef) set(png_BUILD ${CMAKE_BINARY_DIR}/png/src/png) set(png_INSTALL ${CMAKE_BINARY_DIR}/png/install) if(WIN32) 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) + debug ${CMAKE_BINARY_DIR}/png/install/lib/libpng16_staticd.lib + optimized ${CMAKE_BINARY_DIR}/png/install/lib/libpng16_static.lib) else() if(CMAKE_BUILD_TYPE EQUAL Debug) set(png_STATIC_LIBRARIES - ${CMAKE_BINARY_DIR}/png/install/lib/libpng12_staticd.lib) + ${CMAKE_BINARY_DIR}/png/install/lib/libpng16_staticd.lib) else() set(png_STATIC_LIBRARIES - ${CMAKE_BINARY_DIR}/png/install/lib/libpng12_static.lib) + ${CMAKE_BINARY_DIR}/png/install/lib/libpng16_static.lib) endif() endif() else() - set(png_STATIC_LIBRARIES ${CMAKE_BINARY_DIR}/png/install/lib/libpng12.a) + set(png_STATIC_LIBRARIES ${CMAKE_BINARY_DIR}/png/install/lib/libpng16.a) endif() set(png_HEADERS - "${png_INSTALL}/include/libpng12/png.h" - "${png_INSTALL}/include/libpng12/pngconf.h" + "${png_INSTALL}/include/libpng16/png.h" + "${png_INSTALL}/include/libpng16/pngconf.h" + "${png_INSTALL}/include/libpng16/pnglibconf.h" ) ExternalProject_Add(png diff --git a/tensorflow/contrib/cmake/external/sqlite.cmake b/tensorflow/contrib/cmake/external/sqlite.cmake index 57c4ae76517e4d7247093edd5e5bd95a83258d87..7f835d2d519273a6d52d12f92ed585a4ddbeb973 100644 --- a/tensorflow/contrib/cmake/external/sqlite.cmake +++ b/tensorflow/contrib/cmake/external/sqlite.cmake @@ -15,8 +15,8 @@ include (ExternalProject) set(sqlite_INCLUDE_DIR ${CMAKE_CURRENT_BINARY_DIR}/external/sqlite) -set(sqlite_URL https://mirror.bazel.build/www.sqlite.org/2017/sqlite-amalgamation-3200000.zip) -set(sqlite_HASH SHA256=208780b3616f9de0aeb50822b7a8f5482f6515193859e91ed61637be6ad74fd4) +set(sqlite_URL https://mirror.bazel.build/www.sqlite.org/2018/sqlite-amalgamation-3230100.zip) +set(sqlite_HASH SHA256=4239a1f69e5721d07d9a374eb84d594225229e54be4ee628da2995f4315d8dfc) set(sqlite_BUILD ${CMAKE_CURRENT_BINARY_DIR}/sqlite/src/sqlite) set(sqlite_INSTALL ${CMAKE_CURRENT_BINARY_DIR}/sqlite/install) diff --git a/tensorflow/contrib/cmake/python_modules.txt b/tensorflow/contrib/cmake/python_modules.txt index de84af866ba1cc382004be3f78cba48dc7a71759..fbcdf7e753de45519203c5747b1b2217c36527f0 100644 --- a/tensorflow/contrib/cmake/python_modules.txt +++ b/tensorflow/contrib/cmake/python_modules.txt @@ -130,6 +130,8 @@ tensorflow/contrib/boosted_trees/ops tensorflow/contrib/boosted_trees/proto tensorflow/contrib/boosted_trees/python tensorflow/contrib/boosted_trees/python/ops +tensorflow/contrib/checkpoint +tensorflow/contrib/checkpoint/python tensorflow/contrib/cloud tensorflow/contrib/cloud/kernels tensorflow/contrib/cloud/ops @@ -170,7 +172,6 @@ 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 diff --git a/tensorflow/contrib/cmake/python_protos.txt b/tensorflow/contrib/cmake/python_protos.txt index 0c80d529af5230ed6d36b265e12ee4b749a14ec4..d63c41db844af243f0c6600b1565635ac9b91cac 100644 --- a/tensorflow/contrib/cmake/python_protos.txt +++ b/tensorflow/contrib/cmake/python_protos.txt @@ -5,7 +5,6 @@ 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/tf_core_framework.cmake b/tensorflow/contrib/cmake/tf_core_framework.cmake index 973c191c4762210ae87600f102b1b79ccd516663..b47c32f1c48b3d42fe5b4ba115cc2a511b7ee5f4 100644 --- a/tensorflow/contrib/cmake/tf_core_framework.cmake +++ b/tensorflow/contrib/cmake/tf_core_framework.cmake @@ -341,9 +341,3 @@ add_dependencies(tf_core_framework tf_core_lib proto_text ) - -if(WIN32) - # Cmake > 3.6 will quote this as -D"__VERSION__=\"MSVC\"" which nvcc fails on. - # Instead of defining this global, limit it to tf_core_framework where its used. - target_compile_definitions(tf_core_framework PRIVATE __VERSION__="MSVC") -endif() diff --git a/tensorflow/contrib/cmake/tf_python.cmake b/tensorflow/contrib/cmake/tf_python.cmake index f6aaf41f7351c3870d643bb713e821e668281d57..c4bdb69d828b269e6246777e74c3756ba1c4b96f 100755 --- a/tensorflow/contrib/cmake/tf_python.cmake +++ b/tensorflow/contrib/cmake/tf_python.cmake @@ -554,12 +554,13 @@ if(WIN32) set(pywrap_tensorflow_deffile "${CMAKE_CURRENT_BINARY_DIR}/pywrap_tensorflow.def") endif() set_source_files_properties(${pywrap_tensorflow_deffile} PROPERTIES GENERATED TRUE) - + math(EXPR tensorflow_target_bitness "${CMAKE_SIZEOF_VOID_P}*8") add_custom_command(TARGET pywrap_tensorflow_internal_static POST_BUILD COMMAND ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/tools/create_def_file.py --input "${pywrap_tensorflow_internal_static_dependencies}" --output "${pywrap_tensorflow_deffile}" --target _pywrap_tensorflow_internal.pyd + --bitness "${tensorflow_target_bitness}" BYPRODUCTS ${pywrap_tensorflow_deffile} # Required for Ninja ) endif(WIN32) @@ -589,6 +590,12 @@ add_library(pywrap_tensorflow_internal SHARED ${pywrap_tensorflow_deffile} ) +# There is a bug in GCC 5 resulting in undefined reference to a __cpu_model function when +# linking to the tensorflow library. Adding the following libraries fixes it. +if(CMAKE_COMPILER_IS_GNUCC AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER 5.0) + target_link_libraries(pywrap_tensorflow_internal PRIVATE gcc_s gcc) +endif() + if(WIN32) add_dependencies(pywrap_tensorflow_internal pywrap_tensorflow_internal_static) endif(WIN32) diff --git a/tensorflow/contrib/cmake/tf_shared_lib.cmake b/tensorflow/contrib/cmake/tf_shared_lib.cmake index 9738bbeb9aebaeb67495127528e26634887d392c..38f40452b533fdc0dba6ac686a0ff43a2ef13cb8 100644 --- a/tensorflow/contrib/cmake/tf_shared_lib.cmake +++ b/tensorflow/contrib/cmake/tf_shared_lib.cmake @@ -52,12 +52,13 @@ if(WIN32) set(tensorflow_deffile "${CMAKE_CURRENT_BINARY_DIR}/tensorflow.def") endif() set_source_files_properties(${tensorflow_deffile} PROPERTIES GENERATED TRUE) - + math(EXPR tensorflow_target_bitness "${CMAKE_SIZEOF_VOID_P}*8") add_custom_command(TARGET tensorflow_static POST_BUILD COMMAND ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/tools/create_def_file.py --input "${tensorflow_static_dependencies}" --output "${tensorflow_deffile}" --target tensorflow.dll + --bitness "${tensorflow_target_bitness}" ) endif(WIN32) diff --git a/tensorflow/contrib/cmake/tools/create_def_file.py b/tensorflow/contrib/cmake/tools/create_def_file.py index 53c2285699a6ca94e1e6b147080338b507f4d768..cffe069aa352f8a6f2c436bc70b62f54e2336ac6 100644 --- a/tensorflow/contrib/cmake/tools/create_def_file.py +++ b/tensorflow/contrib/cmake/tools/create_def_file.py @@ -63,7 +63,7 @@ INCLUDE_RE = re.compile(r"^(TF_\w*)$|" r"^(TFE_\w*)$|" r"tensorflow::|" r"functor::|" - r"nsync_|" + r"\?nsync_|" r"perftools::gputools") # We want to identify data members explicitly in the DEF file, so that no one @@ -87,6 +87,7 @@ def get_args(): required=True) parser.add_argument("--output", help="output deffile", required=True) parser.add_argument("--target", help="name of the target", required=True) + parser.add_argument("--bitness", help="build target bitness", required=True) args = parser.parse_args() return args @@ -125,7 +126,10 @@ def main(): # 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") + if args.bitness == "64": + def_fp.write("\t??1OpDef@tensorflow@@UEAA@XZ\n") + else: + def_fp.write("\t??1OpDef@tensorflow@@UAE@XZ\n") # Each symbols returned by undname matches the same position in candidates. # We compare on undname but use the decorated name from candidates. diff --git a/tensorflow/contrib/coder/BUILD b/tensorflow/contrib/coder/BUILD index ce12e38248785987e51befa47d04143e235554fe..9ca4ce8a9c765677865f77ea4982ad8613ce334c 100644 --- a/tensorflow/contrib/coder/BUILD +++ b/tensorflow/contrib/coder/BUILD @@ -92,6 +92,34 @@ tf_cc_test( ], ) +tf_kernel_library( + name = "pmf_to_cdf_op", + srcs = ["kernels/pmf_to_cdf_op.cc"], + visibility = ["//visibility:public"], + deps = [ + ":coder_ops_op_lib", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + ], +) + +tf_cc_test( + name = "pmf_to_cdf_op_test", + size = "small", + srcs = ["kernels/pmf_to_cdf_op_test.cc"], + deps = [ + ":pmf_to_cdf_op", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + "//tensorflow/core:testlib", + "//tensorflow/core/kernels:ops_testutil", + ], +) + cc_library( name = "all_ops", deps = [":coder_ops_op_lib"], @@ -99,12 +127,16 @@ cc_library( cc_library( name = "all_kernels", - deps = [":range_coder_ops"], + deps = [ + ":pmf_to_cdf_op", + ":range_coder_ops", + ], ) tf_custom_op_library( name = "python/ops/_coder_ops.so", srcs = [ + "kernels/pmf_to_cdf_op.cc", "kernels/range_coder.cc", "kernels/range_coder.h", "kernels/range_coder_ops.cc", diff --git a/tensorflow/contrib/coder/kernels/pmf_to_cdf_op.cc b/tensorflow/contrib/coder/kernels/pmf_to_cdf_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..bd5272ee6f20ac3537a2e378225ede5ee90782c5 --- /dev/null +++ b/tensorflow/contrib/coder/kernels/pmf_to_cdf_op.cc @@ -0,0 +1,196 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#define EIGEN_USE_THREADS + +#include +#include +#include +#include +#include + +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/tensor_types.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/core/threadpool.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/platform/types.h" + +namespace tensorflow { +namespace { +using errors::InvalidArgument; + +class PmfToCdfOp : public OpKernel { + public: + explicit PmfToCdfOp(OpKernelConstruction* context) : OpKernel(context) { + OP_REQUIRES_OK(context, context->GetAttr("precision", &precision_)); + OP_REQUIRES( + context, 0 < precision_ && precision_ <= 16, + InvalidArgument("`precision` must be in [1, 16]: ", precision_)); + } + + void Compute(OpKernelContext* context) override { + const Tensor& pmf_tensor = context->input(0); + + TensorShape shape = pmf_tensor.shape(); + OP_REQUIRES(context, TensorShapeUtils::IsVectorOrHigher(shape), + InvalidArgument("`pmf` should be at least 1-D.")); + OP_REQUIRES( + context, shape.dim_size(shape.dims() - 1) > 1, + InvalidArgument("`pmf` size should be at least 2 in the last axis.")); + shape.set_dim(shape.dims() - 1, shape.dim_size(shape.dims() - 1) + 1); + + Tensor* cdf_tensor; + OP_REQUIRES_OK(context, context->allocate_output(0, shape, &cdf_tensor)); + + auto pmf = pmf_tensor.flat_inner_dims(); + auto cdf = cdf_tensor->flat_inner_dims(); + CHECK_EQ(pmf.dimension(0), cdf.dimension(0)); + CHECK_EQ(pmf.dimension(1) + 1, cdf.dimension(1)); + + const double n = pmf.dimension(1); + const int64 cost_per_unit = static_cast(50.0 * n * std::log2(n)); + thread::ThreadPool* thread_pool = + context->device()->tensorflow_cpu_worker_threads()->workers; + thread_pool->ParallelFor( + pmf.dimension(0), cost_per_unit, + [this, pmf, &cdf](int64 start, int64 limit) { + const gtl::ArraySlice::size_type pmf_size = pmf.dimension(1); + for (int64 i = start; i < limit; ++i) { + cdf(i, 0) = 0; + PerShard({&pmf(i, 0), pmf_size}, {&cdf(i, 1), pmf_size}); + } + }); + } + + private: + struct PenaltyItem { + PenaltyItem(int32* p, double mass) : pointer(p), mass(mass) { + penalty = ComputeNextPenalty(); + } + + void Decrease() { + CHECK_GT(*pointer, 1); + --*pointer; + penalty = ComputeNextPenalty(); + } + + friend bool operator<(const PenaltyItem& lhs, const PenaltyItem& rhs) { + return lhs.penalty < rhs.penalty; + } + + double ComputeNextPenalty() { + if (*pointer <= 1) { + return std::numeric_limits::infinity(); + } + return mass * (std::log2(*pointer) - std::log2(*pointer - 1)); + } + + int32* pointer; + double mass; + double penalty; + }; + + struct GainItem { + GainItem(int32* p, double mass) : pointer(p), mass(mass) { + gain = ComputeNextGain(); + } + + void Increase() { + CHECK_GT(*pointer, 0); + ++*pointer; + gain = ComputeNextGain(); + } + + friend bool operator>(const GainItem& lhs, const GainItem& rhs) { + return lhs.gain > rhs.gain; + } + + double ComputeNextGain() { + // Never increment zero value to non-zero value. + if (*pointer < 1) { + return -std::numeric_limits::infinity(); + } + return mass * (std::log2(*pointer + 1) - std::log2(*pointer)); + } + + int32* pointer; + double mass; + double gain; + }; + + void PerShard(gtl::ArraySlice pmf, + gtl::MutableArraySlice cdf) const { + CHECK_EQ(pmf.size(), cdf.size()); + + const int32 normalizer = 1 << precision_; + std::transform(pmf.begin(), pmf.end(), cdf.begin(), + [normalizer](float mass) { + int32 value = std::rint(mass * normalizer); + // NOTE: Consider checking if mass > 0. + value = std::max(value, 1); + return value; + }); + + int32 sum = std::accumulate(cdf.begin(), cdf.end(), 0); + if (sum > normalizer) { + std::vector queue; + queue.reserve(cdf.size()); + for (int i = 0; i < cdf.size(); ++i) { + queue.emplace_back(&cdf[i], pmf[i]); + } + + std::sort(queue.begin(), queue.end()); + while (sum-- > normalizer) { + queue[0].Decrease(); + // Performs a linear search because this find_if is likely to return + // iterator very close to the begin. + auto iter = std::find_if( + std::next(queue.begin()), queue.end(), + [&queue](const PenaltyItem& rhs) { return queue[0] < rhs; }); + std::rotate(queue.begin(), std::next(queue.begin()), iter); + } + } else if (sum < normalizer) { + std::vector queue; + queue.reserve(cdf.size()); + for (int i = 0; i < cdf.size(); ++i) { + queue.emplace_back(&cdf[i], pmf[i]); + } + + std::sort(queue.begin(), queue.end(), std::greater()); + while (sum++ < normalizer) { + queue[0].Increase(); + // Performs a linear search because this find_if is likely to return + // iterator very close to the begin. + auto iter = std::find_if( + std::next(queue.begin()), queue.end(), + [&queue](const GainItem& rhs) { return queue[0] > rhs; }); + std::rotate(queue.begin(), std::next(queue.begin()), iter); + } + } + std::partial_sum(cdf.begin(), cdf.end(), cdf.begin()); + } + + int precision_; +}; + +REGISTER_KERNEL_BUILDER(Name("PmfToQuantizedCdf").Device(DEVICE_CPU), + PmfToCdfOp); +} // namespace +} // namespace tensorflow diff --git a/tensorflow/contrib/coder/kernels/pmf_to_cdf_op_test.cc b/tensorflow/contrib/coder/kernels/pmf_to_cdf_op_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..3408f6b519a33fbb8f23d19c16bc7138fc34c121 --- /dev/null +++ b/tensorflow/contrib/coder/kernels/pmf_to_cdf_op_test.cc @@ -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. +==============================================================================*/ + +#include +#include + +#include "tensorflow/core/framework/fake_input.h" +#include "tensorflow/core/framework/node_def.pb.h" +#include "tensorflow/core/framework/node_def_builder.h" +#include "tensorflow/core/framework/shape_inference_testutil.h" +#include "tensorflow/core/framework/tensor_types.h" +#include "tensorflow/core/framework/types.pb.h" +#include "tensorflow/core/kernels/ops_testutil.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/lib/random/philox_random.h" +#include "tensorflow/core/lib/random/random.h" +#include "tensorflow/core/lib/random/simple_philox.h" +#include "tensorflow/core/platform/test.h" + +namespace tensorflow { +namespace { +class PmfToQuantizedCdfOpTest : public OpsTestBase { + protected: + void SetupOp(int precision, Tensor* input) { + TF_ASSERT_OK(NodeDefBuilder("pmf_to_cdf", "PmfToQuantizedCdf") + .Input(FakeInput(DT_FLOAT)) + .Attr("precision", precision) + .Finalize(node_def())); + TF_ASSERT_OK(InitOp()); + + inputs_.clear(); + inputs_.emplace_back(input); + } + + void GenerateData(random::SimplePhilox* rand, + gtl::MutableArraySlice slice) { + constexpr float minimum = std::numeric_limits::epsilon(); + float sum = 0; + for (float& value : slice) { + value = std::max(rand->RandFloat(), minimum); + sum += value; + } + for (float& value : slice) { + value /= sum; + } + } + + void Verify(int precision, const Tensor& pmf_tensor, + const Tensor& cdf_tensor) { + ASSERT_EQ(pmf_tensor.dims(), cdf_tensor.dims()); + const int n = pmf_tensor.dims(); + + for (int i = 0; i < n - 1; ++i) { + EXPECT_EQ(pmf_tensor.dim_size(i), cdf_tensor.dim_size(i)); + } + + auto pmf = pmf_tensor.flat_inner_dims(); + auto cdf = cdf_tensor.flat_inner_dims(); + EXPECT_EQ(pmf.dimension(1) + 1, cdf.dimension(1)); + + const int normalizer = 1 << precision; + for (int i = 0; i < pmf.dimension(0); ++i) { + EXPECT_EQ(0, cdf(i, 0)); + + TTypes::UnalignedConstVec cdf_slice(&cdf(i, 0), cdf.dimension(1)); + + for (int j = 1; j < cdf_slice.size(); ++j) { + const int32 diff = cdf_slice(j) - cdf_slice(j - 1); + EXPECT_GT(diff, 0); + } + + EXPECT_EQ(cdf_slice(cdf_slice.size() - 1), normalizer); + } + } +}; + +TEST_F(PmfToQuantizedCdfOpTest, UnderSum) { + Tensor pmf(DT_FLOAT, {1, 10, 1, 32}); + auto matrix = pmf.flat_inner_dims(); + const std::size_t n = matrix.dimension(1); + + random::PhiloxRandom gen(random::New64(), random::New64()); + random::SimplePhilox rand(&gen); + for (int64 i = 0; i < matrix.dimension(0); ++i) { + GenerateData(&rand, {&matrix(i, 0), n}); + } + + pmf.flat() = pmf.flat() * 0.85f; + + constexpr int kPrecision = 10; + SetupOp(kPrecision, &pmf); + TF_ASSERT_OK(RunOpKernel()); + + Verify(kPrecision, pmf, *GetOutput(0)); +} + +TEST_F(PmfToQuantizedCdfOpTest, OverSum) { + Tensor pmf(DT_FLOAT, {10, 1, 1, 100}); + auto matrix = pmf.flat_inner_dims(); + + // Half of each PMF is filled with zeros. The op will round up zeros to ones, + // post quantization. These round ups are likely to make the sum over + // normalizer value. + matrix.setZero(); + const std::size_t n = matrix.dimension(1) / 2; + + random::PhiloxRandom gen(random::New64(), random::New64()); + random::SimplePhilox rand(&gen); + for (int64 i = 0; i < matrix.dimension(0); ++i) { + GenerateData(&rand, {&matrix(i, 0), n}); + } + + constexpr int kPrecision = 7; + SetupOp(kPrecision, &pmf); + TF_ASSERT_OK(RunOpKernel()); + + Verify(kPrecision, pmf, *GetOutput(0)); +} + +TEST_F(PmfToQuantizedCdfOpTest, ShapeFn) { + ShapeInferenceTestOp op("PmfToQuantizedCdf"); + + INFER_OK(op, "?", "?"); + INFER_OK(op, "[3]", "[4]"); + INFER_OK(op, "[3,4]", "[d0_0,5]"); + INFER_OK(op, "[3,4,5]", "[d0_0,d0_1,6]"); +} +} // namespace +} // namespace tensorflow diff --git a/tensorflow/contrib/coder/ops/coder_ops.cc b/tensorflow/contrib/coder/ops/coder_ops.cc index 9056d1a6963d7be92f499db31385fb6afe2dc515..a185e07913f84a813d76a8c63741bd22a832c8b9 100644 --- a/tensorflow/contrib/coder/ops/coder_ops.cc +++ b/tensorflow/contrib/coder/ops/coder_ops.cc @@ -19,6 +19,7 @@ limitations under the License. #include "tensorflow/core/lib/core/status.h" namespace tensorflow { +using shape_inference::DimensionHandle; using shape_inference::InferenceContext; using shape_inference::ShapeHandle; @@ -76,7 +77,7 @@ are incorrect. For this reason, the range coder uses integer arithmetics and avoids using any floating point operations internally, and `cdf` should contain integers representing quantized probability mass rather than floating points. -data: An int32 tensor. +data: An int16 tensor. cdf: An int32 tensor representing the CDF's of `data`. Each integer is divided by `2^precision` to represent a fraction. encoded: A range-coded scalar string. @@ -111,9 +112,38 @@ potential performance issues, the decoder does not return error status. encoded: A scalar string tensor from RangeEncode. shape: An int32 1-D tensor representing the shape of the data encoded by RangeEncode. -decoded: An int32 tensor with shape equal to `shape`. +decoded: An int16 tensor with shape equal to `shape`. precision: The number of bits for probability quantization. Must be <= 16, and must match the precision used by RangeEncode that produced `encoded`. )doc"); + +REGISTER_OP("PmfToQuantizedCdf") + .Input("pmf: float") + .Output("cdf: int32") + .Attr("precision: int >= 1") + .SetShapeFn([] (InferenceContext* c) { + ShapeHandle in; + TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 1, &in)); + DimensionHandle last; + TF_RETURN_IF_ERROR(c->Add(c->Dim(in, -1), 1, &last)); + ShapeHandle out; + TF_RETURN_IF_ERROR(c->ReplaceDim(in, -1, last, &out)); + c->set_output(0, out); + return Status::OK(); + }) + .Doc(R"doc( +Converts PMF to quantized CDF. This op uses floating-point operations +internally. Therefore the quantized output may not be consistent across multiple +platforms. For entropy encoders and decoders to have the same quantized CDF on +different platforms, the quantized CDF should be produced once and saved, then +the saved quantized CDF should be used everywhere. + +After quantization, if PMF does not sum to 2^precision, then some values of PMF +are increased or decreased to adjust the sum to equal to 2^precision. + +Note that the input PMF is pre-quantization. The input PMF is not normalized +by this op prior to quantization. Therefore the user is responsible for +normalizing PMF if necessary. +)doc"); // clang-format on } // namespace tensorflow diff --git a/tensorflow/contrib/crf/python/kernel_tests/crf_test.py b/tensorflow/contrib/crf/python/kernel_tests/crf_test.py index 721dc4d0801d1f0e116921888e3851a95e0b72b0..a5e065b93a23c3dd2838d81e7cf537dec226f4f9 100644 --- a/tensorflow/contrib/crf/python/kernel_tests/crf_test.py +++ b/tensorflow/contrib/crf/python/kernel_tests/crf_test.py @@ -281,6 +281,21 @@ class CrfTest(test.TestCase): self.assertEqual(list(tf_actual_max_sequence[:sequence_lengths]), expected_max_sequence[:sequence_lengths]) + def testCrfDecodeZeroSeqLength(self): + """ + Test that crf_decode works when sequence_length contains one or more zeros. + """ + with self.test_session() as sess: + inputs = constant_op.constant(np.ones([2, 10, 5], + dtype=np.float32)) + transition_params = constant_op.constant(np.ones([5, 5], + dtype=np.float32)) + sequence_lengths = constant_op.constant(np.zeros([2], + dtype=np.int32)) + values = crf.crf_decode(inputs, transition_params, sequence_lengths) + tags, scores = sess.run(values) + self.assertEqual(len(tags.shape), 2) + self.assertEqual(len(scores.shape), 1) if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/crf/python/ops/crf.py b/tensorflow/contrib/crf/python/ops/crf.py index 1233c8f251c404c57d9e2b38993e7a386b1e6ceb..e37c029cebf30eba59c560bc00ed73d2eea86213 100644 --- a/tensorflow/contrib/crf/python/ops/crf.py +++ b/tensorflow/contrib/crf/python/ops/crf.py @@ -479,15 +479,17 @@ def crf_decode(potentials, transition_params, sequence_length): initial_state = array_ops.slice(potentials, [0, 0, 0], [-1, 1, -1]) initial_state = array_ops.squeeze(initial_state, axis=[1]) # [B, O] inputs = array_ops.slice(potentials, [0, 1, 0], [-1, -1, -1]) # [B, T-1, O] + # sequence length is not allowed to be less than zero + sequence_length_less_one = math_ops.maximum(0, sequence_length - 1) backpointers, last_score = rnn.dynamic_rnn( # [B, T - 1, O], [B, O] crf_fwd_cell, inputs=inputs, - sequence_length=sequence_length - 1, + sequence_length=sequence_length_less_one, initial_state=initial_state, time_major=False, dtype=dtypes.int32) backpointers = gen_array_ops.reverse_sequence( # [B, T - 1, O] - backpointers, sequence_length - 1, seq_dim=1) + backpointers, sequence_length_less_one, seq_dim=1) # Computes backward decoding. Extract tag indices from backpointers. crf_bwd_cell = CrfDecodeBackwardRnnCell(num_tags) @@ -497,7 +499,7 @@ def crf_decode(potentials, transition_params, sequence_length): decode_tags, _ = rnn.dynamic_rnn( # [B, T - 1, 1] crf_bwd_cell, inputs=backpointers, - sequence_length=sequence_length - 1, + sequence_length=sequence_length_less_one, initial_state=initial_state, time_major=False, dtype=dtypes.int32) diff --git a/tensorflow/contrib/cudnn_rnn/BUILD b/tensorflow/contrib/cudnn_rnn/BUILD index d68015ae1565b778b1ba0744f515d09007175e93..aeefa3cee62281c74388765ea5e2cbc7f16ff927 100644 --- a/tensorflow/contrib/cudnn_rnn/BUILD +++ b/tensorflow/contrib/cudnn_rnn/BUILD @@ -25,7 +25,7 @@ tf_custom_op_py_library( srcs_version = "PY2AND3", visibility = ["//visibility:public"], deps = [ - "//tensorflow/contrib/eager/python:checkpointable_utils", + "//tensorflow/contrib/checkpoint/python:split_dependency", "//tensorflow/contrib/util:util_py", "//tensorflow/python:array_ops", "//tensorflow/python:control_flow_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 9cc6ca09ad6c58a6cdc5909ec755ccdd49424872..6fb56b0858786662546ecab425b1a2564fbd9a64 100644 --- a/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_test.py +++ b/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_test.py @@ -29,7 +29,6 @@ import numpy as np from tensorflow.contrib.cudnn_rnn.python.layers import cudnn_rnn from tensorflow.contrib.cudnn_rnn.python.ops import cudnn_rnn_ops -from tensorflow.contrib.eager.python import checkpointable_utils from tensorflow.contrib.rnn.python.ops import rnn as contrib_rnn_lib from tensorflow.python.eager import backprop from tensorflow.python.eager import context @@ -55,6 +54,7 @@ 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 checkpointable_utils from tensorflow.python.training import gradient_descent from tensorflow.python.training import momentum from tensorflow.python.training import rmsprop diff --git a/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py b/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py index 00d9544602ae4472cd04f04f544396b9814302fe..d58198faf353aab68430d2fa153a18de359112de 100644 --- a/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py +++ b/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py @@ -358,7 +358,8 @@ class _CudnnRNN(base_layer.Layer): "CUDA/CuDNN generations.") # Initialize opaque params with a tensor. self.kernel = vs.get_variable( - "opaque_kernel", initializer=opaque_params_t, validate_shape=False) + "opaque_kernel", dtype=self._plain_dtype, + initializer=opaque_params_t, validate_shape=False) # Create saveable in the outer scope of the cudnn subgraph, such that # alternative subgraph with platform-independent rnn cells can load the # checkpoints directly. diff --git a/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py b/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py index 9796aae4b01e9efc6c9a08ab803bda4ee614abf2..a1ede4471ef24adbfb7db32ee7ea276f33ed20a9 100644 --- a/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py +++ b/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py @@ -17,13 +17,13 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.eager.python import checkpointable_utils +from tensorflow.contrib.checkpoint.python import split_dependency from tensorflow.contrib.rnn.python.ops import lstm_ops from tensorflow.python.framework import common_shapes from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import random_seed -from tensorflow.python.layers import base as base_layer +from tensorflow.python.keras._impl.keras.engine import base_layer from tensorflow.python.ops import array_ops from tensorflow.python.ops import gen_cudnn_rnn_ops from tensorflow.python.ops import init_ops @@ -318,7 +318,7 @@ class CudnnOpaqueParamsSaveable(saver.BaseSaverBuilder.SaveableObject): dependencies too (typically the cuDNN `Layer`). dtype: The dtype for the canonical parameter Tensors. """ - split_dependencies = checkpointable_utils.split_dependency( + split_dependencies = split_dependency.split_dependency( component_names=self._param_names, component_dtypes=(dtype,) * len(self._param_names), fill_save_buffer_fn=self._checkpointable_save, @@ -524,10 +524,7 @@ class CudnnLSTMSaveable(CudnnOpaqueParamsSaveable): _rnn_mode = CUDNN_LSTM _num_params_per_layer = CUDNN_LSTM_PARAMS_PER_LAYER - # pylint:disable=protected-access - _rnn_cell_name = base_layer._to_snake_case(CudnnCompatibleLSTMCell.__name__) - - # pylint:enable=protected-access + _rnn_cell_name = base_layer.to_snake_case(CudnnCompatibleLSTMCell.__name__) def _cudnn_to_tf_gate_params(self, *cu_gate_order): i_g, f_g, c_g, o_g = cu_gate_order @@ -648,10 +645,7 @@ class CudnnGRUSaveable(CudnnOpaqueParamsSaveable): _rnn_mode = CUDNN_GRU _num_params_per_layer = CUDNN_GRU_PARAMS_PER_LAYER - # pylint:disable=protected-access - _rnn_cell_name = base_layer._to_snake_case(CudnnCompatibleGRUCell.__name__) - - # pylint:enable=protected-access + _rnn_cell_name = base_layer.to_snake_case(CudnnCompatibleGRUCell.__name__) def _cudnn_to_tf_weights(self, *cu_weights): r"""Stitching cudnn canonical weights to generate tf canonical weights.""" @@ -730,11 +724,7 @@ class CudnnGRUSaveable(CudnnOpaqueParamsSaveable): class CudnnRNNSimpleSaveable(CudnnLSTMSaveable): """SaveableObject implementation handling Cudnn RNN Tanh opaque params.""" - # pylint:disable=protected-access - _rnn_cell_name = base_layer._to_snake_case( - rnn_cell_impl.BasicRNNCell.__name__) - - # pylint:enable=protected-access + _rnn_cell_name = base_layer.to_snake_case(rnn_cell_impl.BasicRNNCell.__name__) def _cudnn_to_tf_weights(self, *cu_weights): r"""Stitching cudnn canonical weights to generate tf canonical weights.""" @@ -1650,31 +1640,6 @@ class CudnnRNNRelu(_CudnnRNNNoInputC): _NUM_PARAMS_PER_LAYER = CUDNN_RNN_RELU_PARAMS_PER_LAYER -@ops.RegisterGradient("CudnnRNN") -def _cudnn_rnn_backward(op, *grad): - if not op.get_attr("is_training"): - raise ValueError( - "CudnnRNN must set is_training to True to be used in gradients") - return gen_cudnn_rnn_ops.cudnn_rnn_backprop( - input=op.inputs[0], - input_h=op.inputs[1], - input_c=op.inputs[2], - params=op.inputs[3], - output=op.outputs[0], - output_h=op.outputs[1], - output_c=op.outputs[2], - output_backprop=grad[0], - output_h_backprop=grad[1], - output_c_backprop=grad[2], - reserve_space=op.outputs[3], - dropout=op.get_attr("dropout"), - seed=op.get_attr("seed"), - seed2=op.get_attr("seed2"), - rnn_mode=op.get_attr("rnn_mode"), - input_mode=op.get_attr("input_mode"), - direction=op.get_attr("direction")) - - ops.RegisterShape("CudnnRNNParamsSize")(common_shapes.call_cpp_shape_fn) ops.RegisterShape("CudnnRNNParamsToCanonical")(common_shapes.call_cpp_shape_fn) ops.RegisterShape("CudnnRNNCanonicalToParams")(common_shapes.call_cpp_shape_fn) diff --git a/tensorflow/contrib/data/BUILD b/tensorflow/contrib/data/BUILD index 7bb0dc1c0f695f4d1c7739fa11764ded4ff9410a..8bdbba83ef6a8541158d956e36caf6a9be435c5b 100644 --- a/tensorflow/contrib/data/BUILD +++ b/tensorflow/contrib/data/BUILD @@ -22,13 +22,7 @@ py_library( deps = [ "//tensorflow/contrib/data/python/ops:dataset_ops", "//tensorflow/contrib/data/python/ops:iterator_ops", - "//tensorflow/contrib/data/python/ops:prefetching_ops", - "//tensorflow/contrib/data/python/ops:readers", - "//tensorflow/contrib/data/python/ops:shuffle_ops", - "//tensorflow/contrib/data/python/ops:transformation_ops", - "//tensorflow/python:parsing_ops", "//tensorflow/python:util", - "//tensorflow/python/data/ops:iterator_ops", ], ) diff --git a/tensorflow/contrib/data/__init__.py b/tensorflow/contrib/data/__init__.py index 125260b4c1f6b63c8f83f28d1829afe2d9d3ea97..077cbba9d2ae41a83f6c358a63ae27aec5741e2c 100644 --- a/tensorflow/contrib/data/__init__.py +++ b/tensorflow/contrib/data/__init__.py @@ -41,6 +41,7 @@ See the @{$datasets$Importing Data} Programmer's Guide for an overview. @@prefetch_to_device @@read_batch_features @@rejection_resample +@@sample_from_datasets @@scan @@shuffle_and_repeat @@sliding_window_batch @@ -69,6 +70,7 @@ from tensorflow.contrib.data.python.ops.get_single_element import get_single_ele 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 sample_from_datasets 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.prefetching_ops import prefetch_to_device @@ -80,8 +82,6 @@ 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 83ada6fb67dcbff595a38ce9e8609bdd1219b075..c56910c7833d4c54fa8db27cd061b404013f3f54 100644 --- a/tensorflow/contrib/data/kernels/BUILD +++ b/tensorflow/contrib/data/kernels/BUILD @@ -18,6 +18,17 @@ cc_library( alwayslink = 1, ) +cc_library( + name = "directed_interleave_dataset_op", + srcs = ["directed_interleave_dataset_op.cc"], + deps = [ + "//tensorflow/core:framework_headers_lib", + "//third_party/eigen3", + "@protobuf_archive//:protobuf_headers", + ], + alwayslink = 1, +) + cc_library( name = "ignore_errors_dataset_op", srcs = ["ignore_errors_dataset_op.cc"], @@ -52,6 +63,7 @@ cc_library( cc_library( name = "dataset_kernels", deps = [ + ":directed_interleave_dataset_op", ":ignore_errors_dataset_op", ":prefetching_kernels", ":threadpool_dataset_op", diff --git a/tensorflow/contrib/data/kernels/directed_interleave_dataset_op.cc b/tensorflow/contrib/data/kernels/directed_interleave_dataset_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..48d3734162525ffc6ace076e4f0523c1d0cae511 --- /dev/null +++ b/tensorflow/contrib/data/kernels/directed_interleave_dataset_op.cc @@ -0,0 +1,274 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT 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/lib/hash/hash.h" + +namespace tensorflow { + +namespace { + +// See documentation in ../ops/dataset_ops.cc for a high-level +// description of the following op. + +class DirectedInterleaveDatasetOp : public DatasetOpKernel { + public: + explicit DirectedInterleaveDatasetOp(OpKernelConstruction* ctx) + : DatasetOpKernel(ctx) {} + + void MakeDataset(OpKernelContext* ctx, DatasetBase** output) override { + DatasetBase* selector_input; + OP_REQUIRES_OK(ctx, + GetDatasetFromVariantTensor(ctx->input(0), &selector_input)); + + OP_REQUIRES( + ctx, + selector_input->output_dtypes().size() == 1 && + selector_input->output_dtypes()[0] == DT_INT64 && + selector_input->output_shapes().size() == 1 && + selector_input->output_shapes()[0].IsCompatibleWith( + PartialTensorShape({})), + errors::InvalidArgument( + "The selector input must be a dataset of scalar int64 elements.")); + + std::vector data_inputs; + for (size_t i = 1; i < ctx->num_inputs(); ++i) { + DatasetBase* input; + OP_REQUIRES_OK(ctx, GetDatasetFromVariantTensor(ctx->input(i), &input)); + data_inputs.push_back(input); + + OP_REQUIRES( + ctx, data_inputs[0]->output_dtypes() == input->output_dtypes(), + errors::InvalidArgument( + "All inputs must have the same output_dtypes. First input " + "has types ", + DataTypeVectorString(data_inputs[0]->output_dtypes()), + ", and input ", i - 1, " has types ", + DataTypeVectorString(input->output_dtypes()))); + } + *output = new Dataset(ctx, selector_input, std::move(data_inputs)); + } + + private: + class Dataset : public GraphDatasetBase { + public: + Dataset(OpKernelContext* ctx, const DatasetBase* selector_input, + std::vector data_inputs) + : GraphDatasetBase(ctx), + selector_input_(selector_input), + data_inputs_(std::move(data_inputs)) { + selector_input_->Ref(); + + output_shapes_ = data_inputs_[0]->output_shapes(); + data_inputs_[0]->Ref(); + for (size_t i = 1; i < data_inputs_.size(); ++i) { + const DatasetBase* data_input = data_inputs_[i]; + data_input->Ref(); + for (size_t j = 0; j < output_shapes_.size(); ++j) { + output_shapes_[j] = MostSpecificCompatibleShape( + output_shapes_[j], data_input->output_shapes()[j]); + } + } + } + + ~Dataset() override { + selector_input_->Unref(); + for (DatasetBase* data_input : data_inputs_) { + data_input->Unref(); + } + } + + std::unique_ptr MakeIterator( + const string& prefix) const override { + return std::unique_ptr(new Iterator( + {this, strings::StrCat(prefix, "::DirectedInterleave")})); + } + + const DataTypeVector& output_dtypes() const override { + return data_inputs_[0]->output_dtypes(); + } + + const std::vector& output_shapes() const override { + return output_shapes_; + } + + string DebugString() override { + return strings::StrCat("DirectedInterleaveDatasetOp::Dataset"); + } + + protected: + Status AsGraphDefInternal(OpKernelContext* ctx, DatasetGraphDefBuilder* b, + Node** output) const override { + Node* selector_input_node; + TF_RETURN_IF_ERROR( + b->AddParentDataset(ctx, selector_input_, &selector_input_node)); + std::vector data_input_nodes(data_inputs_.size()); + for (size_t i = 0; i < data_inputs_.size(); ++i) { + TF_RETURN_IF_ERROR( + b->AddParentDataset(ctx, data_inputs_[i], &data_input_nodes[i])); + } + TF_RETURN_IF_ERROR(b->AddDataset(this, {{0, selector_input_node}}, + {{1, data_input_nodes}}, {}, output)); + return Status::OK(); + } + + private: + class Iterator : public DatasetIterator { + public: + explicit Iterator(const Params& params) + : DatasetIterator(params), + selector_input_impl_(params.dataset->selector_input_->MakeIterator( + params.prefix + ".selector")), + num_active_inputs_(params.dataset->data_inputs_.size()) { + data_input_impls_.reserve(params.dataset->data_inputs_.size()); + for (size_t i = 0; i < params.dataset->data_inputs_.size(); ++i) { + const DatasetBase* data_input = params.dataset->data_inputs_[i]; + data_input_impls_.push_back(data_input->MakeIterator( + strings::StrCat(params.prefix, "[", i, "]"))); + } + } + + Status GetNextInternal(IteratorContext* ctx, + std::vector* out_tensors, + bool* end_of_sequence) override { + mutex_lock l(mu_); + if (!selector_input_impl_) { + *end_of_sequence = true; + return Status::OK(); + } + + while (true) { + std::vector selector_result; + *end_of_sequence = false; + TF_RETURN_IF_ERROR(selector_input_impl_->GetNext( + ctx, &selector_result, end_of_sequence)); + if (*end_of_sequence) { + selector_input_impl_.reset(); + for (auto& data_input_impl : data_input_impls_) { + data_input_impl.reset(); + } + return Status::OK(); + } + + int64 selected_input = selector_result[0].scalar()(); + if (selected_input < 0 || selected_input > data_input_impls_.size()) { + return errors::InvalidArgument( + "Selector index out of range: ", selected_input, + " >= ", data_input_impls_.size()); + } + + if (data_input_impls_[selected_input]) { + bool end_of_selected_input = false; + TF_RETURN_IF_ERROR(data_input_impls_[selected_input]->GetNext( + ctx, out_tensors, &end_of_selected_input)); + + if (!end_of_selected_input) { + return Status::OK(); + } + + data_input_impls_[selected_input].reset(); + --num_active_inputs_; + + if (num_active_inputs_ == 0) { + selector_input_impl_.reset(); + *end_of_sequence = true; + return Status::OK(); + } + } + + LOG(WARNING) << "DirectedInterleave selected an exhausted input: " + << selected_input; + } + } + + protected: + Status SaveInternal(IteratorStateWriter* writer) override { + mutex_lock l(mu_); + if (selector_input_impl_) { + TF_RETURN_IF_ERROR(SaveParent(writer, selector_input_impl_)); + } else { + TF_RETURN_IF_ERROR( + writer->WriteScalar(full_name("selector_input_impl_empty"), "")); + } + for (size_t i = 0; i < data_input_impls_.size(); ++i) { + const auto& data_input_impl = data_input_impls_[i]; + if (data_input_impl) { + TF_RETURN_IF_ERROR(SaveParent(writer, data_input_impl)); + } else { + TF_RETURN_IF_ERROR(writer->WriteScalar( + full_name(strings::StrCat("data_input_impl_empty[", i, "]")), + "")); + } + } + return Status::OK(); + } + + Status RestoreInternal(IteratorContext* ctx, + IteratorStateReader* reader) override { + mutex_lock l(mu_); + if (!reader->Contains(full_name("selector_input_impl_empty"))) { + TF_RETURN_IF_ERROR(RestoreParent(ctx, reader, selector_input_impl_)); + } else { + selector_input_impl_.reset(); + } + for (size_t i = 0; i < data_input_impls_.size(); ++i) { + if (!reader->Contains(full_name( + strings::StrCat("data_input_impl_empty[", i, "]")))) { + TF_RETURN_IF_ERROR( + RestoreParent(ctx, reader, data_input_impls_[i])); + } else { + data_input_impls_[i].reset(); + } + } + return Status::OK(); + } + + private: + mutex mu_; + std::unique_ptr selector_input_impl_ GUARDED_BY(mu_); + std::vector> data_input_impls_ + GUARDED_BY(mu_); + int64 num_active_inputs_ GUARDED_BY(mu_); + }; + + static PartialTensorShape MostSpecificCompatibleShape( + const PartialTensorShape& ts1, const PartialTensorShape& ts2) { + PartialTensorShape output_tensorshape; + if (ts1.dims() != ts2.dims() || ts1.unknown_rank() || ts2.unknown_rank()) + return output_tensorshape; + auto dims1 = ts1.dim_sizes(); + auto dims2 = ts2.dim_sizes(); + for (int d = 0; d < ts1.dims(); d++) { + if (dims1[d] == dims2[d]) + output_tensorshape.Concatenate(dims1[d]); + else + output_tensorshape.Concatenate(-1); + } + return output_tensorshape; + } + + const DatasetBase* const selector_input_; + const std::vector data_inputs_; + std::vector output_shapes_; + }; +}; + +REGISTER_KERNEL_BUILDER(Name("DirectedInterleaveDataset").Device(DEVICE_CPU), + DirectedInterleaveDatasetOp); + +} // namespace + +} // namespace tensorflow diff --git a/tensorflow/contrib/data/ops/dataset_ops.cc b/tensorflow/contrib/data/ops/dataset_ops.cc index cf0a8bbccb5813c799e7e6db91d73e2ecf4107f8..137deb63527f0bdde7da8d5be83ed038f430e581 100644 --- a/tensorflow/contrib/data/ops/dataset_ops.cc +++ b/tensorflow/contrib/data/ops/dataset_ops.cc @@ -17,6 +17,23 @@ limitations under the License. namespace tensorflow { +REGISTER_OP("DirectedInterleaveDataset") + .Input("selector_input_dataset: variant") + .Input("data_input_datasets: N * variant") + .Output("handle: variant") + .Attr("output_types: list(type) >= 1") + .Attr("output_shapes: list(shape) >= 1") + .Attr("N: int >= 1") + .SetShapeFn(shape_inference::ScalarShape) + .Doc(R"doc( +A substitute for `InterleaveDataset` on a fixed list of `N` datasets. + +selector_input_dataset: A dataset of scalar `DT_INT64` elements that determines + which of the `N` data inputs should produce the next output element. +data_input_datasets: `N` datasets with the same type that will be interleaved + according to the values of `selector_input_dataset`. +)doc"); + REGISTER_OP("IgnoreErrorsDataset") .Input("input_dataset: variant") .Output("handle: variant") diff --git a/tensorflow/contrib/data/python/kernel_tests/BUILD b/tensorflow/contrib/data/python/kernel_tests/BUILD index fa5662ce0bb1741a23a368eda71d6a0d446d1679..d59dd17aea42618075e69516bcfa4ee2b9eafc81 100644 --- a/tensorflow/contrib/data/python/kernel_tests/BUILD +++ b/tensorflow/contrib/data/python/kernel_tests/BUILD @@ -8,14 +8,13 @@ load("//tensorflow:tensorflow.bzl", "cuda_py_test", "py_test", "tf_py_test") py_test( name = "batch_dataset_op_test", - size = "small", + size = "medium", srcs = ["batch_dataset_op_test.py"], srcs_version = "PY2AND3", tags = ["no_pip"], deps = [ ":dataset_serialization_test", - "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/contrib/data/python/ops:batching", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:constant_op", @@ -38,8 +37,7 @@ py_test( srcs_version = "PY2AND3", deps = [ ":dataset_serialization_test", - "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/contrib/data/python/ops:grouping", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:constant_op", @@ -60,10 +58,10 @@ py_test( srcs_version = "PY2AND3", deps = [ ":dataset_serialization_test", - "//tensorflow/contrib/data/python/ops:dataset_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:errors", "//tensorflow/python:tensor_shape", + "//tensorflow/python/data/ops:dataset_ops", "//tensorflow/python/data/util:nest", "//third_party/py/numpy", ], @@ -80,8 +78,7 @@ py_test( ], deps = [ ":dataset_serialization_test", - "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/contrib/data/python/ops:batching", "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", @@ -125,16 +122,19 @@ py_test( size = "small", srcs = ["filter_dataset_op_test.py"], srcs_version = "PY2AND3", - tags = ["no_pip"], + tags = [ + "no_pip", + "optonly", + ], deps = [ ":dataset_serialization_test", - "//tensorflow/contrib/data/python/ops:dataset_ops", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:dtypes", "//tensorflow/python:errors", "//tensorflow/python:functional_ops", "//tensorflow/python:math_ops", + "//tensorflow/python/data/ops:dataset_ops", "//third_party/py/numpy", ], ) @@ -146,7 +146,7 @@ tf_py_test( additional_deps = [ ":dataset_serialization_test", "//third_party/py/numpy", - "//tensorflow/contrib/data/python/ops:dataset_ops", + "//tensorflow/python/data/ops:dataset_ops", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:constant_op", @@ -176,8 +176,7 @@ py_test( ], deps = [ ":dataset_serialization_test", - "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/contrib/data/python/ops:interleave_ops", "//tensorflow/python:array_ops", "//tensorflow/python:client", "//tensorflow/python:client_testlib", @@ -188,6 +187,7 @@ py_test( "//tensorflow/python:sparse_ops", "//tensorflow/python:sparse_tensor", "//tensorflow/python:training", + "//tensorflow/python/data/ops:dataset_ops", "//third_party/py/numpy", ], ) @@ -198,7 +198,8 @@ tf_py_test( srcs = ["get_single_element_test.py"], additional_deps = [ "//third_party/py/numpy", - "//tensorflow/contrib/data/python/ops:dataset_ops", + "//tensorflow/contrib/data/python/ops:get_single_element", + "//tensorflow/python/data/ops:dataset_ops", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:constant_op", @@ -213,11 +214,14 @@ py_test( size = "medium", srcs = ["map_dataset_op_test.py"], srcs_version = "PY2AND3", - tags = ["no_pip"], + tags = [ + "no_pip", + "noasan", # times out + "optonly", + ], deps = [ ":dataset_serialization_test", - "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/contrib/data/python/ops:error_ops", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:constant_op", @@ -262,8 +266,8 @@ py_test( srcs_version = "PY2AND3", deps = [ ":dataset_serialization_test", - "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/contrib/data/python/ops:counter", + "//tensorflow/contrib/data/python/ops:enumerate_ops", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:constant_op", @@ -275,6 +279,7 @@ py_test( "//tensorflow/python:parsing_ops", "//tensorflow/python:tensor_shape", "//tensorflow/python:variables", + "//tensorflow/python/data/ops:dataset_ops", ], ) @@ -308,14 +313,17 @@ py_test( srcs = ["resample_test.py"], shard_count = 2, srcs_version = "PY2AND3", - tags = ["noasan"], + tags = [ + "noasan", + "optonly", + ], deps = [ - "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/contrib/data/python/ops:resampling", "//tensorflow/python:client_testlib", "//tensorflow/python:errors", "//tensorflow/python:string_ops", "//tensorflow/python:util", + "//tensorflow/python/data/ops:dataset_ops", "//third_party/py/numpy", ], ) @@ -328,13 +336,14 @@ py_test( tags = ["no_pip"], deps = [ ":dataset_serialization_test", - "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/contrib/data/python/ops:scan_ops", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:constant_op", "//tensorflow/python:dtypes", "//tensorflow/python:errors", "//tensorflow/python/data/ops:dataset_ops", + "//tensorflow/python/eager:context", "//third_party/py/numpy", ], ) @@ -347,11 +356,11 @@ py_test( tags = ["no_pip"], deps = [ ":dataset_serialization_test", - "//tensorflow/contrib/data/python/ops:dataset_ops", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:dtypes", "//tensorflow/python:errors", + "//tensorflow/python/data/ops:dataset_ops", "//third_party/py/numpy", ], ) @@ -379,7 +388,6 @@ py_test( tags = ["no_pip"], deps = [ ":dataset_serialization_test", - "//tensorflow/contrib/data/python/ops:dataset_ops", "//tensorflow/contrib/data/python/ops:shuffle_ops", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", @@ -416,10 +424,10 @@ py_test( tags = ["no_pip"], deps = [ ":dataset_serialization_test", - "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/contrib/data/python/ops:stats_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:errors", + "//tensorflow/python/data/ops:dataset_ops", ], ) @@ -430,10 +438,11 @@ py_test( srcs_version = "PY2AND3", tags = ["no_pip"], deps = [ - "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/contrib/data/python/ops:threadpool", + "//tensorflow/contrib/data/python/ops:unique", "//tensorflow/python:client_testlib", "//tensorflow/python:errors", + "//tensorflow/python/data/ops:dataset_ops", ], ) @@ -445,13 +454,13 @@ py_test( tags = ["no_pip"], deps = [ ":dataset_serialization_test", - "//tensorflow/contrib/data/python/ops:dataset_ops", - "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/contrib/data/python/ops:unique", "//tensorflow/contrib/stateless", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:dtypes", "//tensorflow/python:errors", + "//tensorflow/python/data/ops:dataset_ops", "//third_party/py/numpy", ], ) @@ -464,11 +473,11 @@ py_test( tags = ["no_pip"], deps = [ ":dataset_serialization_test", - "//tensorflow/contrib/data/python/ops:dataset_ops", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:dtypes", "//tensorflow/python:errors", + "//tensorflow/python/data/ops:dataset_ops", "//third_party/py/numpy", ], ) @@ -497,8 +506,8 @@ tf_py_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/data/ops:dataset_ops", + "//tensorflow/contrib/data/python/ops:sliding", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:dtypes", @@ -508,3 +517,23 @@ tf_py_test( "//third_party/py/numpy", ], ) + +tf_py_test( + name = "writer_ops_test", + size = "small", + srcs = ["writer_ops_test.py"], + additional_deps = [ + "//tensorflow/contrib/data/python/ops:writers", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:dataset_ops_gen", + "//tensorflow/python:dtypes", + "//tensorflow/python:errors", + "//tensorflow/python:framework_ops", + "//tensorflow/python:io_ops", + "//tensorflow/python:lib", + "//tensorflow/python:tensor_shape", + "//tensorflow/python:util", + "//tensorflow/python/data/ops:readers", + ], +) 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 413d8737978b695ac443c92036d6641e5c73f28c..a4a0ce79b6013d8813f2d8d294168ea8189d53ef 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 @@ -18,15 +18,18 @@ from __future__ import division from __future__ import print_function import math +import time 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.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.framework import ops from tensorflow.python.framework import sparse_tensor from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import array_ops @@ -34,6 +37,7 @@ from tensorflow.python.ops import math_ops from tensorflow.python.ops import script_ops from tensorflow.python.ops import string_ops from tensorflow.python.platform import test +from tensorflow.python.util import compat class BatchDatasetTest(test.TestCase): @@ -151,6 +155,69 @@ class BatchDatasetTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(op) + def testUnbatchDatasetWithStrings(self): + data = tuple([math_ops.range(10) for _ in range(3)]) + data = dataset_ops.Dataset.from_tensor_slices(data) + data = data.map(lambda x, y, z: (x, string_ops.as_string(y), z)) + expected_types = (dtypes.int32, dtypes.string, dtypes.int32) + data = data.batch(2) + self.assertEqual(expected_types, data.output_types) + data = data.apply(batching.unbatch()) + self.assertEqual(expected_types, data.output_types) + + iterator = data.make_one_shot_iterator() + op = iterator.get_next() + + with self.test_session() as sess: + for i in range(10): + self.assertEqual((i, compat.as_bytes(str(i)), i), sess.run(op)) + + with self.assertRaises(errors.OutOfRangeError): + sess.run(op) + + def testUnbatchDatasetWithSparseTensor(self): + st = sparse_tensor.SparseTensorValue( + indices=[[i, i] for i in range(10)], + values=list(range(10)), + dense_shape=[10, 10]) + data = dataset_ops.Dataset.from_tensors(st) + data = data.apply(batching.unbatch()) + data = data.batch(5) + data = data.apply(batching.unbatch()) + iterator = data.make_one_shot_iterator() + next_element = iterator.get_next() + + with self.test_session() as sess: + for i in range(10): + st_row = sess.run(next_element) + self.assertEqual([i], st_row.indices) + self.assertEqual([i], st_row.values) + self.assertEqual([10], st_row.dense_shape) + with self.assertRaises(errors.OutOfRangeError): + sess.run(next_element) + + def testUnbatchDatasetWithDenseAndSparseTensor(self): + st = sparse_tensor.SparseTensorValue( + indices=[[i, i] for i in range(10)], + values=list(range(10)), + dense_shape=[10, 10]) + data = dataset_ops.Dataset.from_tensors((list(range(10)), st)) + data = data.apply(batching.unbatch()) + data = data.batch(5) + data = data.apply(batching.unbatch()) + iterator = data.make_one_shot_iterator() + next_element = iterator.get_next() + + with self.test_session() as sess: + for i in range(10): + dense_elem, st_row = sess.run(next_element) + self.assertEqual(i, dense_elem) + self.assertEqual([i], st_row.indices) + self.assertEqual([i], st_row.values) + self.assertEqual([10], st_row.dense_shape) + with self.assertRaises(errors.OutOfRangeError): + sess.run(next_element) + def testUnbatchSingleElementTupleDataset(self): data = tuple([(math_ops.range(10),) for _ in range(3)]) data = dataset_ops.Dataset.from_tensor_slices(data) @@ -191,6 +258,53 @@ class BatchDatasetTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(op) + def testUnbatchEmpty(self): + data = dataset_ops.Dataset.from_tensors( + (constant_op.constant([]), constant_op.constant([], shape=[0, 4]), + constant_op.constant([], shape=[0, 4, 0]))) + data = data.apply(batching.unbatch()) + iterator = data.make_one_shot_iterator() + next_element = iterator.get_next() + + with self.test_session() as sess: + with self.assertRaises(errors.OutOfRangeError): + sess.run(next_element) + + def testUnbatchStaticShapeMismatch(self): + data = dataset_ops.Dataset.from_tensors((np.arange(7), np.arange(8), + np.arange(9))) + with self.assertRaises(ValueError): + data.apply(batching.unbatch()) + + def testUnbatchDynamicShapeMismatch(self): + ph1 = array_ops.placeholder(dtypes.int32, shape=[None]) + ph2 = array_ops.placeholder(dtypes.int32, shape=None) + data = dataset_ops.Dataset.from_tensors((ph1, ph2)) + data = data.apply(batching.unbatch()) + iterator = data.make_initializable_iterator() + next_element = iterator.get_next() + + with self.test_session() as sess: + # Mismatch in the 0th dimension. + sess.run( + iterator.initializer, + feed_dict={ + ph1: np.arange(7).astype(np.int32), + ph2: np.arange(8).astype(np.int32) + }) + with self.assertRaises(errors.InvalidArgumentError): + print(sess.run(next_element)) + + # No 0th dimension (i.e. scalar value) for one component. + sess.run( + iterator.initializer, + feed_dict={ + ph1: np.arange(7).astype(np.int32), + ph2: 7 + }) + with self.assertRaises(errors.InvalidArgumentError): + print(sess.run(next_element)) + def testBatchAndDropRemainder(self): components = (np.arange(7), np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis], @@ -545,6 +659,59 @@ class BatchDatasetSerializationTest( self.run_core_tests(self._build_dataset_nested_sparse, None, 1) +class UnbatchDatasetSerializationTest( + dataset_serialization_test_base.DatasetSerializationTestBase): + + def build_dataset(self, multiplier=15.0, tensor_slice_len=2, batch_size=2): + components = ( + np.arange(tensor_slice_len), + np.array([[1, 2, 3]]) * np.arange(tensor_slice_len)[:, np.newaxis], + np.array(multiplier) * np.arange(tensor_slice_len)) + + return dataset_ops.Dataset.from_tensor_slices(components).batch( + batch_size).apply(batching.unbatch()) + + def testCore(self): + tensor_slice_len = 8 + batch_size = 2 + num_outputs = tensor_slice_len + self.run_core_tests( + lambda: self.build_dataset(15.0, tensor_slice_len, batch_size), + lambda: self.build_dataset(20.0, tensor_slice_len, batch_size), + num_outputs) + + +class MapAndBatchDatasetSerializationTest( + dataset_serialization_test_base.DatasetSerializationTestBase): + + def testSerializationCore(self): + range_size = 11 + num_repeats = 2 + batch_size = 5 + total_outputs = range_size * num_repeats + num_outputs_drop_remainder = total_outputs // batch_size + num_outputs_keep_remainder = int(math.ceil(total_outputs / batch_size)) + num_parallel_batches = 2 + + def build_ds(range_start, drop_remainder=False): + + def _map_fn(x): + return math_ops.square(x) + + return dataset_ops.Dataset.range( + range_start, range_start + range_size).repeat(num_repeats).apply( + batching.map_and_batch( + map_func=_map_fn, + batch_size=batch_size, + num_parallel_batches=num_parallel_batches, + drop_remainder=drop_remainder)) + + self.run_core_tests(lambda: build_ds(10), lambda: build_ds(15), + num_outputs_keep_remainder) + self.run_core_tests(lambda: build_ds(10, True), lambda: build_ds(15, True), + num_outputs_drop_remainder) + + class PaddedBatchDatasetSerializationTest( dataset_serialization_test_base.DatasetSerializationTestBase): @@ -586,10 +753,12 @@ class RestructuredDatasetTest(test.TestCase): def test_assert_element_shape(self): def create_unknown_shape_dataset(x): - return script_ops.py_func(lambda _: (np.ones(2, dtype=np.float32), - np.zeros((3, 4), dtype=np.int32)), - [x], - [dtypes.float32, dtypes.int32]) + return script_ops.py_func( + lambda _: ( # pylint: disable=g-long-lambda + np.ones(2, dtype=np.float32), + np.zeros((3, 4), dtype=np.int32)), + [x], + [dtypes.float32, dtypes.int32]) dataset = dataset_ops.Dataset.range(5).map(create_unknown_shape_dataset) unknown_shapes = (tensor_shape.TensorShape(None), @@ -626,10 +795,12 @@ class RestructuredDatasetTest(test.TestCase): def test_assert_wrong_element_shape_on_unknown_shape_dataset(self): def create_unknown_shape_dataset(x): - return script_ops.py_func(lambda _: (np.ones(2, dtype=np.float32), - np.zeros((3, 4), dtype=np.int32)), - [x], - [dtypes.float32, dtypes.int32]) + return script_ops.py_func( + lambda _: ( # pylint: disable=g-long-lambda + np.ones(2, dtype=np.float32), + np.zeros((3, 4), dtype=np.int32)), + [x], + [dtypes.float32, dtypes.int32]) dataset = dataset_ops.Dataset.range(3).map(create_unknown_shape_dataset) unknown_shapes = (tensor_shape.TensorShape(None), @@ -649,5 +820,77 @@ class RestructuredDatasetTest(test.TestCase): sess.run(get_next) +class UnbatchDatasetBenchmark(test.Benchmark): + + def benchmarkNativeUnbatch(self): + batch_sizes = [1, 2, 5, 10, 20, 50] + elems_per_trial = 10000 + with ops.Graph().as_default(): + dataset = dataset_ops.Dataset.from_tensors("element").repeat(None) + batch_size_placeholder = array_ops.placeholder(dtypes.int64, shape=[]) + dataset = dataset.batch(batch_size_placeholder) + dataset = dataset.apply(batching.unbatch()) + dataset = dataset.skip(elems_per_trial) + iterator = dataset.make_initializable_iterator() + next_element = iterator.get_next() + + with session.Session() as sess: + for batch_size in batch_sizes: + deltas = [] + for _ in range(5): + sess.run( + iterator.initializer, + feed_dict={batch_size_placeholder: batch_size}) + start = time.time() + sess.run(next_element.op) + end = time.time() + deltas.append((end - start) / elems_per_trial) + + median_wall_time = np.median(deltas) + print("Unbatch (native) batch size: %d Median wall time per element:" + " %f microseconds" % (batch_size, median_wall_time * 1e6)) + self.report_benchmark( + iters=10000, + wall_time=median_wall_time, + name="benchmark_unbatch_dataset_native_batch_size_%d" % + batch_size) + + # Include a benchmark of the previous `unbatch()` implementation that uses + # a composition of more primitive ops. Eventually we'd hope to generate code + # that is as good in both cases. + def benchmarkOldUnbatchImplementation(self): + batch_sizes = [1, 2, 5, 10, 20, 50] + elems_per_trial = 10000 + with ops.Graph().as_default(): + dataset = dataset_ops.Dataset.from_tensors("element").repeat(None) + batch_size_placeholder = array_ops.placeholder(dtypes.int64, shape=[]) + dataset = dataset.batch(batch_size_placeholder) + dataset = dataset.flat_map(dataset_ops.Dataset.from_tensor_slices) + dataset = dataset.skip(elems_per_trial) + iterator = dataset.make_initializable_iterator() + next_element = iterator.get_next() + + with session.Session() as sess: + for batch_size in batch_sizes: + deltas = [] + for _ in range(5): + sess.run( + iterator.initializer, + feed_dict={batch_size_placeholder: batch_size}) + start = time.time() + sess.run(next_element.op) + end = time.time() + deltas.append((end - start) / elems_per_trial) + + median_wall_time = np.median(deltas) + print("Unbatch (unfused) batch size: %d Median wall time per element:" + " %f microseconds" % (batch_size, median_wall_time * 1e6)) + self.report_benchmark( + iters=10000, + wall_time=median_wall_time, + name="benchmark_unbatch_dataset_unfused_batch_size_%d" % + batch_size) + + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/data/python/kernel_tests/bucketing_test.py b/tensorflow/contrib/data/python/kernel_tests/bucketing_test.py index 6002cc73c8b41c2f20beaf0158af813807e58c90..55a56b83a8efba899c6b296264d766839a824da5 100644 --- a/tensorflow/contrib/data/python/kernel_tests/bucketing_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/bucketing_test.py @@ -61,7 +61,7 @@ class GroupByWindowTest(test.TestCase): self.assertEqual(len(components), sum(counts)) num_full_batches = len([c for c in counts if c == 4]) - self.assertGreaterEqual(num_full_batches, 23) + self.assertGreaterEqual(num_full_batches, 24) self.assertTrue(all(c == 4 for c in counts[:num_full_batches])) def testImmediateOutput(self): 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 dbc35097ddda9f0375060d43aeb43efa8107f929..78ecce8f7daaf84002ae78d8d77820755b967d89 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 @@ -163,7 +163,7 @@ class DatasetSerializationTestBase(test.TestCase): num_outputs, sparse_tensors=False, verify_exhausted=True): - """Verifies that restoring into an already initilized iterator works. + """Verifies that restoring into an already initialized iterator works. Args: ds_fn: See `run_core_tests`. 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 256ad8d94dc1a7c2b26df3f1ebf8e8e321882c15..43aa4b1bd02791ff304a990c0bbe8e45534c0c77 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 @@ -30,6 +30,7 @@ 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 random_seed from tensorflow.python.framework import sparse_tensor from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops @@ -94,6 +95,76 @@ class InterleaveDatasetSerializationTest( self.run_core_tests(_build_dataset, None, 20) +class ParallelInterleaveDatasetSerializationTest( + dataset_serialization_test_base.DatasetSerializationTestBase): + + def setUp(self): + self.input_values = np.array([4, 5, 6], dtype=np.int64) + self.num_repeats = 2 + self.num_outputs = np.sum(self.input_values) * 2 + + def _build_ds(self, cycle_length, block_length, sloppy=False): + return (dataset_ops.Dataset.from_tensor_slices( + self.input_values).repeat(self.num_repeats).apply( + interleave_ops.parallel_interleave( + lambda x: dataset_ops.Dataset.range(10 * x, 11 * x), + cycle_length, block_length, sloppy))) + + def testSerializationCore(self): + # cycle_length > 1, block_length > 1 + cycle_length = 2 + block_length = 3 + self.run_core_tests( + lambda: self._build_ds(cycle_length, block_length), + lambda: self._build_ds(cycle_length * 2, block_length * 1), + self.num_outputs) + # cycle_length = 1 + cycle_length = 1 + block_length = 3 + self.run_core_tests(lambda: self._build_ds(cycle_length, block_length), + None, self.num_outputs) + # block_length = 1 + cycle_length = 2 + block_length = 1 + self.run_core_tests(lambda: self._build_ds(cycle_length, block_length), + None, self.num_outputs) + + def testSerializationWithSloppy(self): + break_points = self.gen_break_points(self.num_outputs, 10) + expected_outputs = np.repeat( + np.concatenate([np.arange(10 * x, 11 * x) for x in self.input_values]), + self.num_repeats).tolist() + + def run_test(cycle_length, block_length): + actual = self.gen_outputs( + lambda: self._build_ds(cycle_length, block_length, True), + break_points, self.num_outputs) + self.assertSequenceEqual(sorted(actual), expected_outputs) + + # cycle_length > 1, block_length > 1 + run_test(2, 3) + # cycle_length = 1 + run_test(1, 3) + # block_length = 1 + run_test(2, 1) + + 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).apply( + interleave_ops.parallel_interleave(_interleave_fn, 1)) + + self.run_core_tests(_build_dataset, None, 20) + + class ParallelInterleaveDatasetTest(test.TestCase): def setUp(self): @@ -338,7 +409,7 @@ class ParallelInterleaveDatasetTest(test.TestCase): def _testTwoThreadsNoContentionWithRaces(self, sloppy=False): """Tests where all the workers race in producing elements. - Note: this is in contrast with the prevous test which carefully sequences + Note: this is in contrast with the previous test which carefully sequences the execution of the map functions. Args: @@ -424,7 +495,7 @@ class ParallelInterleaveDatasetTest(test.TestCase): def _testTwoThreadsNoContentionWithRacesAndBlocking(self, sloppy=False): """Tests where all the workers race in producing elements. - Note: this is in contrast with the prevous test which carefully sequences + Note: this is in contrast with the previous test which carefully sequences the execution of the map functions. @@ -836,5 +907,114 @@ class ParallelInterleaveDatasetTest(test.TestCase): sess.run(self.next_element) +class DirectedInterleaveDatasetTest(test.TestCase): + + def testBasic(self): + selector_dataset = dataset_ops.Dataset.range(10).repeat(100) + input_datasets = [ + dataset_ops.Dataset.from_tensors(i).repeat(100) for i in range(10) + ] + dataset = interleave_ops.DirectedInterleaveDataset(selector_dataset, + input_datasets) + iterator = dataset.make_initializable_iterator() + next_element = iterator.get_next() + + with self.test_session() as sess: + sess.run(iterator.initializer) + for _ in range(100): + for i in range(10): + self.assertEqual(i, sess.run(next_element)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(next_element) + + def _normalize(self, vec): + return vec / vec.sum() + + def _chi2(self, expected, actual): + actual = np.asarray(actual) + expected = np.asarray(expected) + diff = actual - expected + chi2 = np.sum(diff * diff / expected, axis=0) + return chi2 + + def _testSampleFromDatasetsHelper(self, weights, num_datasets, num_samples): + # Create a dataset that samples each integer in `[0, num_datasets)` + # with probability given by `weights[i]`. + dataset = interleave_ops.sample_from_datasets([ + dataset_ops.Dataset.from_tensors(i).repeat(None) + for i in range(num_datasets) + ], weights) + dataset = dataset.take(num_samples) + iterator = dataset.make_one_shot_iterator() + next_element = iterator.get_next() + + with self.test_session() as sess: + freqs = np.zeros([num_datasets]) + for _ in range(num_samples): + freqs[sess.run(next_element)] += 1 + with self.assertRaises(errors.OutOfRangeError): + sess.run(next_element) + + return freqs + + def testSampleFromDatasets(self): + random_seed.set_random_seed(1619) + num_samples = 10000 + rand_probs = self._normalize(np.random.random_sample((15,))) + + # Use chi-squared test to assert that the observed distribution matches the + # expected distribution. Based on the implementation in + # "tensorflow/python/kernel_tests/multinomial_op_test.py". + for probs in [[.85, .05, .1], rand_probs]: + probs = np.asarray(probs) + classes = len(probs) + freqs = self._testSampleFromDatasetsHelper(probs, classes, num_samples) + self.assertLess(self._chi2(probs, freqs / num_samples), 1e-3) + + # Also check that `weights` as a dataset samples correctly. + probs_ds = dataset_ops.Dataset.from_tensors(probs).repeat() + freqs = self._testSampleFromDatasetsHelper(probs_ds, classes, num_samples) + self.assertLess(self._chi2(probs, freqs / num_samples), 1e-3) + + def testErrors(self): + with self.assertRaisesRegexp(ValueError, + r"vector of length `len\(datasets\)`"): + interleave_ops.sample_from_datasets( + [dataset_ops.Dataset.range(10), + dataset_ops.Dataset.range(20)], + weights=[0.25, 0.25, 0.25, 0.25]) + + with self.assertRaisesRegexp(TypeError, "`tf.float32` or `tf.float64`"): + interleave_ops.sample_from_datasets( + [dataset_ops.Dataset.range(10), + dataset_ops.Dataset.range(20)], + weights=[1, 1]) + + with self.assertRaisesRegexp(TypeError, "must have the same type"): + interleave_ops.sample_from_datasets([ + dataset_ops.Dataset.from_tensors(0), + dataset_ops.Dataset.from_tensors(0.0) + ]) + + +class SampleFromDatasetsSerializationTest( + dataset_serialization_test_base.DatasetSerializationTestBase): + + def _build_dataset(self, probs, num_samples): + dataset = interleave_ops.sample_from_datasets( + [ + dataset_ops.Dataset.from_tensors(i).repeat(None) + for i in range(len(probs)) + ], + probs, + seed=1813) + return dataset.take(num_samples) + + def testSerializationCore(self): + self.run_core_tests( + lambda: self._build_dataset([0.5, 0.5], 100), + lambda: self._build_dataset([0.25, 0.25, 0.25, 0.25], 1000), 100) + + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/data/python/kernel_tests/reader_dataset_ops_test.py b/tensorflow/contrib/data/python/kernel_tests/reader_dataset_ops_test.py index f3e93024092da7c99d4e529f0bd6a6eec06068bf..1075302bae96ca2e0111efbacdf5e919ea76897d 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 @@ -295,18 +295,20 @@ class ReadBatchFeaturesTest(test.TestCase): ).get_next() def _record(self, f, r): - example = example_pb2.Example(features=feature_pb2.Features( - feature={ - "file": - feature_pb2.Feature(int64_list=feature_pb2.Int64List( - value=[f])), - "record": - feature_pb2.Feature(int64_list=feature_pb2.Int64List( - value=[r])), - "keywords": - feature_pb2.Feature(bytes_list=feature_pb2.BytesList( - value=self._get_keywords(f, r))) - })) + example = example_pb2.Example( + features=feature_pb2.Features( + feature={ + "file": + feature_pb2.Feature( + int64_list=feature_pb2.Int64List(value=[f])), + "record": + feature_pb2.Feature( + int64_list=feature_pb2.Int64List(value=[r])), + "keywords": + feature_pb2.Feature( + bytes_list=feature_pb2.BytesList( + value=self._get_keywords(f, r))) + })) return example.SerializeToString() def _get_keywords(self, f, r): @@ -374,8 +376,8 @@ class ReadBatchFeaturesTest(test.TestCase): record_batch.append(r) keywords = self._get_keywords(f, r) keywords_batch_values.extend(keywords) - keywords_batch_indices.extend([[batch_index, i] - for i in range(len(keywords))]) + keywords_batch_indices.extend( + [[batch_index, i] for i in range(len(keywords))]) batch_index += 1 keywords_batch_max_len = max(keywords_batch_max_len, len(keywords)) if len(file_batch) == batch_size: @@ -475,9 +477,10 @@ class ReadBatchFeaturesTest(test.TestCase): "file": parsing_ops.FixedLenFeature([], dtypes.int64), "record": parsing_ops.FixedLenFeature([], dtypes.int64), } - dataset = (core_readers.TFRecordDataset(self.test_filenames) - .map(lambda x: parsing_ops.parse_single_example(x, features)) - .repeat(10).batch(2)) + 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() init_op = iterator.initializer next_element = iterator.get_next() @@ -607,20 +610,25 @@ class MakeCsvDatasetTest(test.TestCase): "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 _write_file(self, filename, rows): + for i in range(len(rows)): + if isinstance(rows[i], list): + rows[i] = ",".join(str(v) if v is not None else "" for v in rows[i]) + fn = os.path.join(self.get_temp_dir(), filename) + f = open(fn, "w") + f.write("\n".join(rows)) + f.close() + return fn def _create_file(self, fileno, header=True, comment=True): - fn = os.path.join(self.get_temp_dir(), "csv_file%d.csv" % fileno) - f = open(fn, "w") + rows = [] if header: - f.write(",".join(self.COLUMNS) + "\n") + rows.append(self.COLUMNS) for recno in range(self._num_records): - f.write(self._csv_record(fileno, recno) + "\n") + rows.append(self._csv_values(fileno, recno)) if comment: - f.write("# Some comment goes here. Should be ignored!\n") - f.close() - return fn + rows.append("# Some comment goes here. Ignore me.") + return self._write_file("csv_file%d.csv" % fileno, rows) def _create_files(self): filenames = [] @@ -634,6 +642,7 @@ class MakeCsvDatasetTest(test.TestCase): defaults, column_names=COLUMNS, label_name=LABEL, + select_cols=None, batch_size=1, num_epochs=1, shuffle=False, @@ -656,6 +665,7 @@ class MakeCsvDatasetTest(test.TestCase): comment=comment, na_value=na_value, default_float_type=default_float_type, + select_columns=select_cols, ) def _next_actual_batch(self, file_indices, batch_size, num_epochs, defaults): @@ -712,7 +722,7 @@ class MakeCsvDatasetTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) - def test_make_csv_dataset(self): + def testMakeCSVDataset(self): defaults = self.DEFAULTS with ops.Graph().as_default() as g: @@ -739,7 +749,7 @@ class MakeCsvDatasetTest(test.TestCase): self._verify_records( sess, dataset, range(self._num_files), batch_size=2, num_epochs=10) - def test_make_csv_dataset_with_bad_columns(self): + def testMakeCSVDataset_withBadColumns(self): """Tests that exception is raised when input is malformed. """ dupe_columns = self.COLUMNS[:-1] + self.COLUMNS[:1] @@ -755,7 +765,7 @@ class MakeCsvDatasetTest(test.TestCase): self._make_csv_dataset( self._test_filenames, defaults, label_name="not_a_real_label") - def test_make_csv_dataset_with_no_label(self): + def testMakeCSVDataset_withNoLabel(self): """Tests that CSV datasets can be created when no label is specified. """ defaults = self.DEFAULTS @@ -776,7 +786,7 @@ class MakeCsvDatasetTest(test.TestCase): num_epochs=10, label_name=None) - def test_make_csv_dataset_with_no_comments(self): + def testMakeCSVDataset_withNoComments(self): """Tests that datasets can be created from CSV files with no header line. """ defaults = self.DEFAULTS @@ -799,7 +809,7 @@ class MakeCsvDatasetTest(test.TestCase): num_epochs=10, ) - def test_make_csv_dataset_with_no_header(self): + def testMakeCSVDataset_withNoHeader(self): """Tests that datasets can be created from CSV files with no header line. """ defaults = self.DEFAULTS @@ -822,7 +832,7 @@ class MakeCsvDatasetTest(test.TestCase): num_epochs=10, ) - def test_make_csv_dataset_with_types(self): + def testMakeCSVDataset_withTypes(self): """Tests that defaults can be a dtype instead of a Tensor for required vals. """ defaults = [d for d in self.COLUMN_TYPES[:-1]] @@ -832,7 +842,7 @@ class MakeCsvDatasetTest(test.TestCase): dataset = self._make_csv_dataset(self._test_filenames, defaults) self._verify_records(sess, dataset, range(self._num_files)) - def test_make_csv_dataset_with_no_col_names(self): + def testMakeCSVDataset_withNoColNames(self): """Tests that datasets can be created when column names are not specified. In that case, we should infer the column names from the header lines. @@ -851,7 +861,17 @@ class MakeCsvDatasetTest(test.TestCase): self._verify_records( sess, dataset, range(self._num_files), batch_size=2, num_epochs=10) - def test_make_csv_dataset_type_inference(self): + def testMakeCSVDataset_withTypeInferenceMismatch(self): + # Test that error is thrown when num fields doesn't match columns + with self.assertRaises(ValueError): + self._make_csv_dataset( + self._test_filenames, + column_names=self.COLUMNS + ["extra_name"], + defaults=None, + batch_size=2, + num_epochs=10) + + def testMakeCSVDataset_withTypeInference(self): """Tests that datasets can be created when no defaults are specified. In that case, we should infer the types from the first N records. @@ -875,19 +895,16 @@ class MakeCsvDatasetTest(test.TestCase): dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float32, dtypes.string, dtypes.string ] - rows = [[0, 0, 0, "NAN", "", "a"], [1, 2**31 + 1, 2**64, 123, "NAN", ""], + col_names = ["col%d" % i for i in range(len(expected_dtypes))] + rows = [[None, None, None, "NAN", "", + "a"], [1, 2**31 + 1, 2**64, 123, "NAN", ""], ['"123"', 2, 2**64, 123.4, "NAN", '"cd,efg"']] expected = [[0, 0, 0, 0, "", "a"], [1, 2**31 + 1, 2**64, 123, "", ""], [123, 2, 2**64, 123.4, "", "cd,efg"]] for row in expected: row[-1] = row[-1].encode("utf-8") # py3 expects byte strings row[-2] = row[-2].encode("utf-8") # py3 expects byte strings - col_names = ["col%d" % i for i in range(len(expected_dtypes))] - with open(fn, "w") as f: - f.write(",".join(col_names)) - f.write("\n") - for row in rows: - f.write(",".join([str(v) if v else "" for v in row]) + "\n") + self._write_file("file.csv", [col_names] + rows) with ops.Graph().as_default() as g: with self.test_session(graph=g) as sess: @@ -895,8 +912,6 @@ class MakeCsvDatasetTest(test.TestCase): fn, defaults=None, column_names=None, - batch_size=1, - num_epochs=1, label_name=None, na_value="NAN", default_float_type=dtypes.float32, @@ -919,8 +934,6 @@ class MakeCsvDatasetTest(test.TestCase): fn, defaults=None, column_names=None, - batch_size=1, - num_epochs=1, label_name=None, na_value="NAN", default_float_type=dtypes.float64, @@ -928,11 +941,99 @@ class MakeCsvDatasetTest(test.TestCase): features = dataset.make_one_shot_iterator().get_next() # Check that types match for i in range(len(expected_dtypes)): - assert features["col%d" % i].dtype == expected_dtypes[i] + self.assertAllEqual(features["col%d" % i].dtype, expected_dtypes[i]) for i in range(len(rows)): - assert sess.run(features) == dict(zip(col_names, expected[i])) + self.assertAllEqual( + sess.run(features), dict(zip(col_names, expected[i]))) + + def testMakeCSVDataset_withSelectColsError(self): + data = [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]] + col_names = ["col%d" % i for i in range(5)] + fn = self._write_file("file.csv", [col_names] + data) + with self.assertRaises(ValueError): + # Mismatch in number of defaults and number of columns selected, + # should raise an error + self._make_csv_dataset( + fn, + defaults=[[0]] * 5, + column_names=col_names, + label_name=None, + select_cols=[1, 3]) + with self.assertRaises(ValueError): + # Invalid column name should raise an error + self._make_csv_dataset( + fn, + defaults=[[0]], + column_names=col_names, + label_name=None, + select_cols=["invalid_col_name"]) + + def testMakeCSVDataset_withSelectCols(self): + data = [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]] + col_names = ["col%d" % i for i in range(5)] + fn = self._write_file("file.csv", [col_names] + data) + # If select_cols is specified, should only yield a subset of columns + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as sess: + dataset = self._make_csv_dataset( + fn, + defaults=[[0], [0]], + column_names=col_names, + label_name=None, + select_cols=[1, 3]) + expected = [[1, 3], [6, 8]] + features = dataset.make_one_shot_iterator().get_next() + for i in range(len(data)): + self.assertAllEqual( + sess.run(features), + dict(zip([col_names[1], col_names[3]], expected[i]))) + # Can still do default inference with select_cols + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as sess: + dataset = self._make_csv_dataset( + fn, + defaults=None, + column_names=col_names, + label_name=None, + select_cols=[1, 3]) + expected = [[1, 3], [6, 8]] + features = dataset.make_one_shot_iterator().get_next() + for i in range(len(data)): + self.assertAllEqual( + sess.run(features), + dict(zip([col_names[1], col_names[3]], expected[i]))) + # Can still do column name inference + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as sess: + dataset = self._make_csv_dataset( + fn, + defaults=None, + column_names=None, + label_name=None, + select_cols=[1, 3]) + expected = [[1, 3], [6, 8]] + features = dataset.make_one_shot_iterator().get_next() + for i in range(len(data)): + self.assertAllEqual( + sess.run(features), + dict(zip([col_names[1], col_names[3]], expected[i]))) + # Can specify column names instead of indices + with ops.Graph().as_default() as g: + with self.test_session(graph=g) as sess: + dataset = self._make_csv_dataset( + fn, + defaults=None, + column_names=None, + label_name=None, + select_cols=[col_names[1], col_names[3]]) + expected = [[1, 3], [6, 8]] + features = dataset.make_one_shot_iterator().get_next() + for i in range(len(data)): + self.assertAllEqual( + sess.run(features), + dict(zip([col_names[1], col_names[3]], expected[i]))) - def test_make_csv_dataset_with_shuffle(self): + def testMakeCSVDataset_withShuffle(self): total_records = self._num_files * self._num_records defaults = self.DEFAULTS for batch_size in [1, 2]: diff --git a/tensorflow/contrib/data/python/kernel_tests/scan_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/scan_dataset_op_test.py index e0494736b72ae52f586cb80d42a5c1e50ac17a61..1a97a84b2cba13e82c8af9c4c8ee413ee8264a5e 100644 --- a/tensorflow/contrib/data/python/kernel_tests/scan_dataset_op_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/scan_dataset_op_test.py @@ -24,9 +24,11 @@ import numpy as np from tensorflow.contrib.data.python.kernel_tests import dataset_serialization_test_base from tensorflow.contrib.data.python.ops import scan_ops from tensorflow.python.data.ops import dataset_ops +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 test_util from tensorflow.python.ops import array_ops from tensorflow.python.platform import test @@ -57,19 +59,24 @@ class ScanDatasetTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(next_element) + @test_util.run_in_graph_and_eager_modes() def testFibonacci(self): iterator = dataset_ops.Dataset.from_tensors(1).repeat(None).apply( scan_ops.scan([0, 1], lambda a, _: ([a[1], a[0] + a[1]], a[1])) ).make_one_shot_iterator() - next_element = iterator.get_next() - with self.test_session() as sess: - self.assertEqual(1, sess.run(next_element)) - self.assertEqual(1, sess.run(next_element)) - self.assertEqual(2, sess.run(next_element)) - self.assertEqual(3, sess.run(next_element)) - self.assertEqual(5, sess.run(next_element)) - self.assertEqual(8, sess.run(next_element)) + if context.executing_eagerly(): + next_element = iterator.get_next + else: + get_next = iterator.get_next() + next_element = lambda: get_next + + self.assertEqual(1, self.evaluate(next_element())) + self.assertEqual(1, self.evaluate(next_element())) + self.assertEqual(2, self.evaluate(next_element())) + self.assertEqual(3, self.evaluate(next_element())) + self.assertEqual(5, self.evaluate(next_element())) + self.assertEqual(8, self.evaluate(next_element())) def testChangingStateShape(self): # Test the fixed-point shape invariant calculations: start with 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 b13ad9ba4e533e1bcef5161d983c8e6578d549b2..d0cb203a3afd2775756c8542a1e86faedc5cee53 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 @@ -48,8 +48,8 @@ class SequenceDatasetSerializationTest( self.run_core_tests(lambda: self._build_skip_dataset(0), None, 10) def testInvalidSkip(self): - with self.assertRaisesRegexp( - ValueError, 'Shape must be rank 0 but is rank 1'): + with self.assertRaisesRegexp(ValueError, + 'Shape must be rank 0 but is rank 1'): self.run_core_tests(lambda: self._build_skip_dataset([1, 2]), None, 0) def _build_take_dataset(self, count): @@ -75,8 +75,8 @@ class SequenceDatasetSerializationTest( self.run_core_tests(lambda: self._build_take_dataset(0), None, 0) def testInvalidTake(self): - with self.assertRaisesRegexp( - ValueError, 'Shape must be rank 0 but is rank 1'): + with self.assertRaisesRegexp(ValueError, + 'Shape must be rank 0 but is rank 1'): self.run_core_tests(lambda: self._build_take_dataset([1, 2]), None, 0) def _build_repeat_dataset(self, count, take_count=3): diff --git a/tensorflow/contrib/data/python/kernel_tests/stats_dataset_ops_test.py b/tensorflow/contrib/data/python/kernel_tests/stats_dataset_ops_test.py index 07bdf920446e953c2a1abaf495d2e9e1256106fd..5c74ed6ae7210e8e22efb6e8fdb773397459ce1e 100644 --- a/tensorflow/contrib/data/python/kernel_tests/stats_dataset_ops_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/stats_dataset_ops_test.py @@ -50,17 +50,17 @@ class StatsDatasetTest(test.TestCase): self.fail("Expected tag %r not found in summary %r" % (tag, summary_proto)) def testBytesProduced(self): + stats_aggregator = stats_ops.StatsAggregator() dataset = dataset_ops.Dataset.range(100).map( lambda x: array_ops.tile([x], ops.convert_to_tensor([x]))).apply( - stats_ops.bytes_produced_stats("bytes_produced")) + stats_ops.bytes_produced_stats("bytes_produced")).apply( + stats_ops.set_stats_aggregator(stats_aggregator)) iterator = dataset.make_initializable_iterator() - stats_aggregator = stats_ops.StatsAggregator() - stats_aggregator_subscriber = stats_aggregator.subscribe(iterator) next_element = iterator.get_next() summary_t = stats_aggregator.get_summary() with self.test_session() as sess: - sess.run([iterator.initializer, stats_aggregator_subscriber]) + sess.run(iterator.initializer) expected_sum = 0.0 for i in range(100): self.assertAllEqual( @@ -76,16 +76,16 @@ class StatsDatasetTest(test.TestCase): self._assertSummaryHasSum(summary_str, "bytes_produced", expected_sum) def testLatencyStats(self): + stats_aggregator = stats_ops.StatsAggregator() dataset = dataset_ops.Dataset.range(100).apply( - stats_ops.latency_stats("record_latency")) + stats_ops.latency_stats("record_latency")).apply( + stats_ops.set_stats_aggregator(stats_aggregator)) iterator = dataset.make_initializable_iterator() - stats_aggregator = stats_ops.StatsAggregator() - stats_aggregator_subscriber = stats_aggregator.subscribe(iterator) next_element = iterator.get_next() summary_t = stats_aggregator.get_summary() with self.test_session() as sess: - sess.run([iterator.initializer, stats_aggregator_subscriber]) + sess.run(iterator.initializer) for i in range(100): self.assertEqual(i, sess.run(next_element)) self._assertSummaryHasCount( @@ -95,16 +95,15 @@ class StatsDatasetTest(test.TestCase): self._assertSummaryHasCount(sess.run(summary_t), "record_latency", 100.0) def testReinitialize(self): + stats_aggregator = stats_ops.StatsAggregator() dataset = dataset_ops.Dataset.range(100).apply( - stats_ops.latency_stats("record_latency")) + stats_ops.latency_stats("record_latency")).apply( + stats_ops.set_stats_aggregator(stats_aggregator)) iterator = dataset.make_initializable_iterator() - stats_aggregator = stats_ops.StatsAggregator() - stats_aggregator_subscriber = stats_aggregator.subscribe(iterator) next_element = iterator.get_next() summary_t = stats_aggregator.get_summary() with self.test_session() as sess: - sess.run(stats_aggregator_subscriber) for j in range(5): sess.run(iterator.initializer) for i in range(100): @@ -130,17 +129,17 @@ class StatsDatasetTest(test.TestCase): sess.run(next_element) def testMultipleTags(self): + stats_aggregator = stats_ops.StatsAggregator() dataset = dataset_ops.Dataset.range(100).apply( stats_ops.latency_stats("record_latency")).apply( - stats_ops.latency_stats("record_latency_2")) + stats_ops.latency_stats("record_latency_2")).apply( + stats_ops.set_stats_aggregator(stats_aggregator)) iterator = dataset.make_initializable_iterator() - stats_aggregator = stats_ops.StatsAggregator() - stats_aggregator_subscriber = stats_aggregator.subscribe(iterator) next_element = iterator.get_next() summary_t = stats_aggregator.get_summary() with self.test_session() as sess: - sess.run([iterator.initializer, stats_aggregator_subscriber]) + sess.run(iterator.initializer) for i in range(100): self.assertEqual(i, sess.run(next_element)) self._assertSummaryHasCount( @@ -154,17 +153,17 @@ class StatsDatasetTest(test.TestCase): sess.run(summary_t), "record_latency_2", 100.0) def testRepeatedTags(self): + stats_aggregator = stats_ops.StatsAggregator() dataset = dataset_ops.Dataset.range(100).apply( stats_ops.latency_stats("record_latency")).apply( - stats_ops.latency_stats("record_latency")) + stats_ops.latency_stats("record_latency")).apply( + stats_ops.set_stats_aggregator(stats_aggregator)) iterator = dataset.make_initializable_iterator() - stats_aggregator = stats_ops.StatsAggregator() - stats_aggregator_subscriber = stats_aggregator.subscribe(iterator) next_element = iterator.get_next() summary_t = stats_aggregator.get_summary() with self.test_session() as sess: - sess.run([iterator.initializer, stats_aggregator_subscriber]) + sess.run(iterator.initializer) for i in range(100): self.assertEqual(i, sess.run(next_element)) self._assertSummaryHasCount( @@ -174,19 +173,17 @@ class StatsDatasetTest(test.TestCase): self._assertSummaryHasCount(sess.run(summary_t), "record_latency", 200.0) def testMultipleIteratorsSameAggregator(self): + stats_aggregator = stats_ops.StatsAggregator() dataset = dataset_ops.Dataset.range(100).apply( - stats_ops.latency_stats("record_latency")) + stats_ops.latency_stats("record_latency")).apply( + stats_ops.set_stats_aggregator(stats_aggregator)) iterator_0 = dataset.make_initializable_iterator() iterator_1 = dataset.make_initializable_iterator() - stats_aggregator = stats_ops.StatsAggregator() - stats_aggregator_subscribers = [stats_aggregator.subscribe(iterator_0), - stats_aggregator.subscribe(iterator_1)] next_element = iterator_0.get_next() + iterator_1.get_next() summary_t = stats_aggregator.get_summary() with self.test_session() as sess: - sess.run([iterator_0.initializer, iterator_1.initializer, - stats_aggregator_subscribers]) + sess.run([iterator_0.initializer, iterator_1.initializer]) for i in range(100): self.assertEqual(i * 2, sess.run(next_element)) self._assertSummaryHasCount( @@ -195,20 +192,6 @@ class StatsDatasetTest(test.TestCase): sess.run(next_element) self._assertSummaryHasCount(sess.run(summary_t), "record_latency", 200.0) - def testMultipleStatsAggregatorsSameIteratorFail(self): - dataset = dataset_ops.Dataset.range(100).apply( - stats_ops.latency_stats("record_latency")) - iterator = dataset.make_initializable_iterator() - stats_aggregator_0 = stats_ops.StatsAggregator() - stats_aggregator_1 = stats_ops.StatsAggregator() - - with self.test_session() as sess: - sess.run(stats_aggregator_0.subscribe(iterator)) - # TODO(mrry): Consider making this allowable (and also allowing - # aggregators to unsubscribe). - with self.assertRaises(errors.FailedPreconditionError): - sess.run(stats_aggregator_1.subscribe(iterator)) - class StatsDatasetSerializationTest( dataset_serialization_test_base.DatasetSerializationTestBase): @@ -218,6 +201,14 @@ class StatsDatasetSerializationTest( lambda x: array_ops.tile([x], ops.convert_to_tensor([x]))).apply( stats_ops.bytes_produced_stats("bytes_produced")) + def test_bytes_produced_stats_invalid_tag_shape(self): + with self.assertRaisesRegexp( + ValueError, 'Shape must be rank 0 but is rank 1'): + self.run_core_tests( + lambda: dataset_ops.Dataset.range(100).apply( + stats_ops.bytes_produced_stats(["bytes_produced"])), + None, 100) + def testBytesStatsDatasetSaveableCore(self): num_outputs = 100 self.run_core_tests( @@ -235,6 +226,14 @@ class StatsDatasetSerializationTest( return dataset_ops.Dataset.range(num_elements).apply( stats_ops.latency_stats(tag1)).apply(stats_ops.latency_stats(tag2)) + def test_latency_stats_invalid_tag_shape(self): + with self.assertRaisesRegexp( + ValueError, 'Shape must be rank 0 but is rank 1'): + self.run_core_tests( + lambda: dataset_ops.Dataset.range(100).apply( + stats_ops.latency_stats(["record_latency", "record_latency_2"])), + None, 100) + def testLatencyStatsDatasetSaveableCore(self): num_outputs = 100 @@ -253,5 +252,9 @@ class StatsDatasetSerializationTest( None, num_outputs) +# TODO(shivaniagrawal): Can not checkpoint input_pipeline with the +# transformation `stats_ops.set_stats_aggregator`, since we don't support +# serializing StatsAggregator yet. + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/data/python/kernel_tests/writer_ops_test.py b/tensorflow/contrib/data/python/kernel_tests/writer_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..c603ecc5ab27a711557376246b093fd5f80f8aec --- /dev/null +++ b/tensorflow/contrib/data/python/kernel_tests/writer_ops_test.py @@ -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. +# ============================================================================== +"""Tests for the experimental input pipeline ops.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os + +from tensorflow.contrib.data.python.ops import writers +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.data.ops import readers +from tensorflow.python.framework import dtypes +from tensorflow.python.lib.io import python_io +from tensorflow.python.lib.io import tf_record +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import test +from tensorflow.python.util import compat + + +class TFRecordWriterTest(test.TestCase): + + def setUp(self): + super(TFRecordWriterTest, self).setUp() + self._num_records = 7 + self.filename = array_ops.placeholder(dtypes.string, shape=[]) + self.compression_type = array_ops.placeholder_with_default("", shape=[]) + + input_dataset = readers.TFRecordDataset([self.filename], + self.compression_type) + self.writer = writers.TFRecordWriter( + self._outputFilename(), self.compression_type).write(input_dataset) + + def _record(self, i): + return compat.as_bytes("Record %d" % (i)) + + def _createFile(self, options=None): + filename = self._inputFilename() + writer = python_io.TFRecordWriter(filename, options) + for i in range(self._num_records): + writer.write(self._record(i)) + writer.close() + return filename + + def _inputFilename(self): + return os.path.join(self.get_temp_dir(), "tf_record.in.txt") + + def _outputFilename(self): + return os.path.join(self.get_temp_dir(), "tf_record.out.txt") + + def testWrite(self): + with self.test_session() as sess: + sess.run( + self.writer, feed_dict={ + self.filename: self._createFile(), + }) + for i, r in enumerate(tf_record.tf_record_iterator(self._outputFilename())): + self.assertAllEqual(self._record(i), r) + + def testWriteZLIB(self): + options = tf_record.TFRecordOptions(tf_record.TFRecordCompressionType.ZLIB) + with self.test_session() as sess: + sess.run( + self.writer, + feed_dict={ + self.filename: self._createFile(options), + self.compression_type: "ZLIB", + }) + for i, r in enumerate( + tf_record.tf_record_iterator(self._outputFilename(), options=options)): + self.assertAllEqual(self._record(i), r) + + def testWriteGZIP(self): + options = tf_record.TFRecordOptions(tf_record.TFRecordCompressionType.GZIP) + with self.test_session() as sess: + sess.run( + self.writer, + feed_dict={ + self.filename: self._createFile(options), + self.compression_type: "GZIP", + }) + for i, r in enumerate( + tf_record.tf_record_iterator(self._outputFilename(), options=options)): + self.assertAllEqual(self._record(i), r) + + def testFailDataset(self): + with self.assertRaises(TypeError): + writers.TFRecordWriter(self._outputFilename(), + self.compression_type).write("whoops") + + def testFailDType(self): + input_dataset = dataset_ops.Dataset.from_tensors(10) + with self.assertRaises(TypeError): + writers.TFRecordWriter(self._outputFilename(), + self.compression_type).write(input_dataset) + + def testFailShape(self): + input_dataset = dataset_ops.Dataset.from_tensors([["hello"], ["world"]]) + with self.assertRaises(TypeError): + writers.TFRecordWriter(self._outputFilename(), + self.compression_type).write(input_dataset) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/data/python/ops/BUILD b/tensorflow/contrib/data/python/ops/BUILD index a1a5c9ed05ff226086885e4e204875d3ca933590..5b04c5316cfbb7577b3f8b3b6d364fc665d14c21 100644 --- a/tensorflow/contrib/data/python/ops/BUILD +++ b/tensorflow/contrib/data/python/ops/BUILD @@ -12,18 +12,26 @@ load( load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") py_library( - name = "dataset_ops", - srcs = [ - "counter.py", - "get_single_element.py", + name = "counter", + srcs = ["counter.py"], + srcs_version = "PY2AND3", + deps = [ + ":scan_ops", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", + "//tensorflow/python/data/ops:dataset_ops", ], +) + +py_library( + name = "get_single_element", + srcs = ["get_single_element.py"], srcs_version = "PY2AND3", deps = [ - ":transformation_ops", "//tensorflow/python:dataset_ops_gen", - "//tensorflow/python:util", "//tensorflow/python/data/ops:dataset_ops", "//tensorflow/python/data/util:nest", + "//tensorflow/python/data/util:sparse", ], ) @@ -66,7 +74,8 @@ py_library( ], srcs_version = "PY2AND3", deps = [ - ":dataset_ops", + ":batching", + ":interleave_ops", ":shuffle_ops", "//tensorflow/python:constant_op", "//tensorflow/python:dataset_ops_gen", @@ -94,51 +103,192 @@ py_library( ], srcs_version = "PY2AND3", deps = [ - ":random_ops", - ":transformation_ops", "//tensorflow/python/data/ops:dataset_ops", ], ) py_library( - name = "transformation_ops", - srcs = [ - "batching.py", - "enumerate_ops.py", - "error_ops.py", - "grouping.py", - "interleave_ops.py", - "resampling.py", - "scan_ops.py", - "sliding.py", - "stats_ops.py", - "threadpool.py", - "unique.py", + name = "batching", + srcs = ["batching.py"], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/contrib/framework:framework_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:dataset_ops_gen", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", + "//tensorflow/python:math_ops", + "//tensorflow/python:tensor_shape", + "//tensorflow/python:tensor_util", + "//tensorflow/python/data/ops:dataset_ops", + "//tensorflow/python/data/util:nest", + "//tensorflow/python/data/util:sparse", ], +) + +py_library( + name = "enumerate_ops", + srcs = ["enumerate_ops.py"], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/python:dtypes", + "//tensorflow/python/data/ops:dataset_ops", + ], +) + +py_library( + name = "error_ops", + srcs = ["error_ops.py"], srcs_version = "PY2AND3", deps = [ ":contrib_op_loader", ":gen_dataset_ops", - "//tensorflow/contrib/framework:framework_py", + "//tensorflow/python/data/ops:dataset_ops", + "//tensorflow/python/data/util:nest", + "//tensorflow/python/data/util:sparse", + ], +) + +py_library( + name = "grouping", + srcs = ["grouping.py"], + srcs_version = "PY2AND3", + deps = [ "//tensorflow/python:array_ops", - "//tensorflow/python:control_flow_ops", - "//tensorflow/python:dataset_ops_gen", + "//tensorflow/python:check_ops", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", "//tensorflow/python:function", + "//tensorflow/python:math_ops", + "//tensorflow/python:tensor_shape", + "//tensorflow/python/data/ops:dataset_ops", + "//tensorflow/python/data/util:nest", + "//tensorflow/python/data/util:sparse", + ], +) + +py_library( + name = "interleave_ops", + srcs = ["interleave_ops.py"], + srcs_version = "PY2AND3", + deps = [ + ":contrib_op_loader", + ":gen_dataset_ops", + ":random_ops", + "//tensorflow/contrib/stateless", + "//tensorflow/python:array_ops", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", + "//tensorflow/python:math_ops", + "//tensorflow/python:util", + "//tensorflow/python/data/ops:readers", + "//tensorflow/python/data/util:nest", + "//tensorflow/python/data/util:sparse", + ], +) + +py_library( + name = "resampling", + srcs = ["resampling.py"], + srcs_version = "PY2AND3", + deps = [ + ":batching", + ":scan_ops", + "//tensorflow/python:array_ops", + "//tensorflow/python:control_flow_ops", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", "//tensorflow/python:logging_ops", "//tensorflow/python:math_ops", "//tensorflow/python:random_ops", + "//tensorflow/python/data/ops:dataset_ops", + ], +) + +py_library( + name = "scan_ops", + srcs = ["scan_ops.py"], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/python:dataset_ops_gen", + "//tensorflow/python:framework_ops", + "//tensorflow/python:function", + "//tensorflow/python/data/ops:dataset_ops", + "//tensorflow/python/data/util:nest", + "//tensorflow/python/data/util:sparse", + ], +) + +py_library( + name = "sliding", + srcs = ["sliding.py"], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/python:dataset_ops_gen", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", + "//tensorflow/python:function", + "//tensorflow/python/data/ops:dataset_ops", + "//tensorflow/python/data/util:nest", + "//tensorflow/python/data/util:sparse", + ], +) + +py_library( + name = "stats_ops", + srcs = ["stats_ops.py"], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/python:dataset_ops_gen", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", + "//tensorflow/python/data/ops:dataset_ops", + "//tensorflow/python/data/ops:iterator_ops", + "//tensorflow/python/data/util:nest", + "//tensorflow/python/data/util:sparse", + ], +) + +py_library( + name = "threadpool", + srcs = ["threadpool.py"], + srcs_version = "PY2AND3", + deps = [ + ":contrib_op_loader", + ":gen_dataset_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", - "//third_party/py/numpy", + "//tensorflow/python/eager:context", + ], +) + +py_library( + name = "unique", + srcs = [ + "unique.py", + ], + srcs_version = "PY2AND3", + deps = [ + ":contrib_op_loader", + ":gen_dataset_ops", + "//tensorflow/python:dtypes", + "//tensorflow/python/data/ops:dataset_ops", + "//tensorflow/python/data/util:nest", + "//tensorflow/python/data/util:sparse", + ], +) + +py_library( + name = "writers", + srcs = [ + "writers.py", + ], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/python:dtypes", + "//tensorflow/python/data/ops:dataset_ops", ], ) @@ -184,3 +334,30 @@ py_library( "//tensorflow/python/data/util:sparse", ], ) + +py_library( + name = "dataset_ops", + deps = [ + ":batching", + ":counter", + ":enumerate_ops", + ":error_ops", + ":get_single_element", + ":grouping", + ":interleave_ops", + ":prefetching_ops", + ":readers", + ":resampling", + ":scan_ops", + ":shuffle_ops", + ":sliding", + ":stats_ops", + ":threadpool", + ":unique", + ":writers", + "//tensorflow/python:dataset_ops_gen", + "//tensorflow/python:util", + "//tensorflow/python/data/ops:dataset_ops", + "//tensorflow/python/data/util:nest", + ], +) diff --git a/tensorflow/contrib/data/python/ops/batching.py b/tensorflow/contrib/data/python/ops/batching.py index 1eba010b562a60ec9469f808fd657ca330a8f5d9..2152bcde84aae6b0c2b368e43750aafab3a04bf2 100644 --- a/tensorflow/contrib/data/python/ops/batching.py +++ b/tensorflow/contrib/data/python/ops/batching.py @@ -80,28 +80,98 @@ def dense_to_sparse_batch(batch_size, row_shape): return _apply_fn +class UnbatchDataset(dataset_ops.Dataset): + """A dataset that splits the elements of its input into multiple elements.""" + + def __init__(self, input_dataset): + """See `unbatch()` for more details.""" + super(UnbatchDataset, self).__init__() + flat_shapes = nest.flatten(input_dataset.output_shapes) + if any(s.ndims == 0 for s in flat_shapes): + raise ValueError("Cannot unbatch an input with scalar components.") + known_batch_dim = tensor_shape.Dimension(None) + for s in flat_shapes: + try: + known_batch_dim = known_batch_dim.merge_with(s[0]) + except ValueError: + raise ValueError("Cannot unbatch an input whose components have " + "different batch sizes.") + self._input_dataset = input_dataset + + def _as_variant_tensor(self): + return gen_dataset_ops.unbatch_dataset( + self._input_dataset._as_variant_tensor(), # 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_classes(self): + return self._input_dataset.output_classes + + @property + def output_shapes(self): + return nest.map_structure(lambda s: s[1:], + self._input_dataset.output_shapes) + + @property + def output_types(self): + return self._input_dataset.output_types + + def unbatch(): - """A Transformation which splits the elements of a dataset. + """Splits elements of a dataset into multiple elements on the batch dimension. For example, if elements of the dataset are shaped `[B, a0, a1, ...]`, - where `B` may vary from element to element, then for each element in - the dataset, the unbatched dataset will contain `B` consecutive elements + where `B` may vary for each input element, then for each element in the + dataset, the unbatched dataset will contain `B` consecutive elements of shape `[a0, a1, ...]`. + ```python + # NOTE: The following example uses `{ ... }` to represent the contents + # of a dataset. + a = { ['a', 'b', 'c'], ['a', 'b'], ['a', 'b', 'c', 'd'] } + + a.apply(tf.contrib.data.unbatch()) == { + 'a', 'b', 'c', 'a', 'b', 'a', 'b', 'c', 'd'} + ``` + Returns: A `Dataset` transformation function, which can be passed to @{tf.data.Dataset.apply}. """ def _apply_fn(dataset): - - def unbatch_map(arg, *rest): + """Function from `Dataset` to `Dataset` that applies the transformation.""" + if not sparse.any_sparse(dataset.output_classes): + return UnbatchDataset(dataset) + + # NOTE(mrry): We must ensure that any SparseTensors in `dataset` + # are normalized to the rank-1 dense representation, so that the + # sparse-oblivious unbatching logic will slice them + # appropriately. This leads to a somewhat inefficient re-encoding step + # for all SparseTensor components. + # TODO(mrry): Consider optimizing this in future + # if it turns out to be a bottleneck. + def normalize(arg, *rest): if rest: - return dataset_ops.Dataset.from_tensor_slices((arg,) + rest) + return sparse.serialize_many_sparse_tensors((arg,) + rest) else: - return dataset_ops.Dataset.from_tensor_slices(arg) + return sparse.serialize_many_sparse_tensors(arg) + + normalized_dataset = dataset.map(normalize) - return dataset.flat_map(map_func=unbatch_map) + # NOTE(mrry): Our `map()` has lost information about the sparseness + # of any SparseTensor components, so re-apply the structure of the + # original dataset. + restructured_dataset = _RestructuredDataset( + normalized_dataset, + dataset.output_types, + dataset.output_shapes, + dataset.output_classes, + allow_unsafe_cast=True) + return UnbatchDataset(restructured_dataset) return _apply_fn @@ -265,7 +335,8 @@ class _RestructuredDataset(dataset_ops.Dataset): dataset, output_types, output_shapes=None, - output_classes=None): + output_classes=None, + allow_unsafe_cast=False): """Creates a new dataset with the given output types and shapes. The given `dataset` must have a structure that is convertible: @@ -283,6 +354,10 @@ class _RestructuredDataset(dataset_ops.Dataset): If omitted, the shapes will be inherited from `dataset`. output_classes: (Optional.) A nested structure of class types. If omitted, the class types will be inherited from `dataset`. + allow_unsafe_cast: (Optional.) If `True`, the caller may switch the + reported output types and shapes of the restructured dataset, e.g. to + switch a sparse tensor represented as `tf.variant` to its user-visible + type and shape. Raises: ValueError: If either `output_types` or `output_shapes` is not compatible @@ -291,14 +366,15 @@ class _RestructuredDataset(dataset_ops.Dataset): super(_RestructuredDataset, self).__init__() self._dataset = dataset - # Validate that the types are compatible. - output_types = nest.map_structure(dtypes.as_dtype, output_types) - flat_original_types = nest.flatten(dataset.output_types) - flat_new_types = nest.flatten(output_types) - if flat_original_types != flat_new_types: - raise ValueError( - "Dataset with output types %r cannot be restructured to have output " - "types %r" % (dataset.output_types, output_types)) + if not allow_unsafe_cast: + # Validate that the types are compatible. + output_types = nest.map_structure(dtypes.as_dtype, output_types) + flat_original_types = nest.flatten(dataset.output_types) + flat_new_types = nest.flatten(output_types) + if flat_original_types != flat_new_types: + raise ValueError( + "Dataset with output types %r cannot be restructured to have " + "output types %r" % (dataset.output_types, output_types)) self._output_types = output_types @@ -308,18 +384,19 @@ class _RestructuredDataset(dataset_ops.Dataset): nest.flatten( dataset.output_shapes)) else: - # Validate that the shapes are compatible. - nest.assert_same_structure(output_types, output_shapes) - flat_original_shapes = nest.flatten(dataset.output_shapes) - flat_new_shapes = nest.flatten_up_to(output_types, output_shapes) - - for original_shape, new_shape in zip(flat_original_shapes, - flat_new_shapes): - if not original_shape.is_compatible_with(new_shape): - raise ValueError( - "Dataset with output shapes %r cannot be restructured to have " - "incompatible output shapes %r" % (dataset.output_shapes, - output_shapes)) + if not allow_unsafe_cast: + # Validate that the shapes are compatible. + nest.assert_same_structure(output_types, output_shapes) + flat_original_shapes = nest.flatten(dataset.output_shapes) + flat_new_shapes = nest.flatten_up_to(output_types, output_shapes) + + for original_shape, new_shape in zip(flat_original_shapes, + flat_new_shapes): + if not original_shape.is_compatible_with(new_shape): + raise ValueError( + "Dataset with output shapes %r cannot be restructured to have " + "incompatible output shapes %r" % (dataset.output_shapes, + output_shapes)) self._output_shapes = nest.map_structure_up_to( output_types, tensor_shape.as_shape, output_shapes) if output_classes is None: @@ -370,9 +447,10 @@ def assert_element_shape(expected_shapes): def _check_shape(*elements): flatten_tensors = nest.flatten(elements) flatten_shapes = nest.flatten(expected_shapes) - checked_tensors = [with_shape(shape, tensor) - for shape, tensor in zip(flatten_shapes, - flatten_tensors)] + checked_tensors = [ + with_shape(shape, tensor) + for shape, tensor in zip(flatten_shapes, flatten_tensors) + ] return nest.pack_sequence_as(elements, checked_tensors) def _apply_fn(dataset): diff --git a/tensorflow/contrib/data/python/ops/grouping.py b/tensorflow/contrib/data/python/ops/grouping.py index 36591c055ae8f2c54981525ffcc3df128a990a61..0531f9cbb9da6e6df85fa46940ab1661ad742eb4 100644 --- a/tensorflow/contrib/data/python/ops/grouping.py +++ b/tensorflow/contrib/data/python/ops/grouping.py @@ -108,7 +108,7 @@ def bucket_by_sequence_length(element_length_func, fraction of padding in a batch which increases training step efficiency. Args: - element_length_func: function from element in `Dataset` to `tf.int64`, + element_length_func: function from element in `Dataset` to `tf.int32`, determines the length of the element, which will determine the bucket it goes into. bucket_boundaries: `list`, upper length boundaries of the buckets. diff --git a/tensorflow/contrib/data/python/ops/interleave_ops.py b/tensorflow/contrib/data/python/ops/interleave_ops.py index 91f19da02d4a479820782822475d9121125fc38e..812a50ecbf105393f7e422edbbdf5c87311d72c1 100644 --- a/tensorflow/contrib/data/python/ops/interleave_ops.py +++ b/tensorflow/contrib/data/python/ops/interleave_ops.py @@ -17,7 +17,18 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.contrib import stateless +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.contrib.data.python.ops import random_ops +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import readers +from tensorflow.python.data.util import nest +from tensorflow.python.data.util import sparse +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.util import deprecation @@ -140,3 +151,92 @@ def sloppy_interleave(map_func, cycle_length, block_length=1): prefetch_input_elements=None) return _apply_fn + + +class DirectedInterleaveDataset(dataset_ops.Dataset): + """A substitute for `Dataset.interleave()` on a fixed list of datasets.""" + + def __init__(self, selector_input, data_inputs): + self._selector_input = selector_input + self._data_inputs = list(data_inputs) + + for data_input in data_inputs[1:]: + if (data_input.output_types != data_inputs[0].output_types or + data_input.output_classes != data_inputs[0].output_classes): + raise TypeError("All datasets must have the same type.") + + def _as_variant_tensor(self): + # pylint: disable=protected-access + return gen_dataset_ops.directed_interleave_dataset( + self._selector_input._as_variant_tensor(), + [data_input._as_variant_tensor() for data_input in self._data_inputs], + 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))) + # pylint: enable=protected-access + + @property + def output_classes(self): + return self._data_inputs[0].output_classes + + @property + def output_shapes(self): + ret = self._data_inputs[0].output_shapes + for data_input in self._data_inputs[1:]: + ret = nest.pack_sequence_as(ret, [ + ts1.most_specific_compatible_shape(ts2) for (ts1, ts2) in zip( + nest.flatten(ret), nest.flatten(data_input.output_shapes)) + ]) + return ret + + @property + def output_types(self): + return self._data_inputs[0].output_types + + +def sample_from_datasets(datasets, weights=None, seed=None): + """Samples elements at random from the datasets in `datasets`. + + Args: + datasets: A list of @{tf.data.Dataset} objects with compatible structure. + weights: (Optional.) A list of `len(datasets)` floating-point values where + `weights[i]` represents the probability with which an element should be + sampled from `datasets[i]`, or a @{tf.data.Dataset} object where each + element is such a list. Defaults to a uniform distribution across + `datasets`. + 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 that interleaves elements from `datasets` at random, according to + `weights` if provided, otherwise with uniform probability. + + Raises: + TypeError: If the `datasets` or `weights` arguments have the wrong type. + ValueError: If the `weights` argument is specified and does not match the + length of the `datasets` element. + """ + num_datasets = len(datasets) + if weights is None: + weights = dataset_ops.Dataset.from_tensors([1.0] * num_datasets).repeat() + elif not isinstance(weights, dataset_ops.Dataset): + weights = ops.convert_to_tensor(weights, name="weights") + if weights.dtype not in (dtypes.float32, dtypes.float64): + raise TypeError("`weights` must be convertible to a tensor of " + "`tf.float32` or `tf.float64` elements.") + if not weights.shape.is_compatible_with([num_datasets]): + raise ValueError("`weights` must be a vector of length `len(datasets)`.") + weights = dataset_ops.Dataset.from_tensors(weights).repeat() + + # The `stateless_multinomial()` op expects log-probabilities, as opposed to + # weights. + logits_ds = weights.map(lambda *p: math_ops.log(p, name="logits")) + def select_dataset(logits, seed): + return array_ops.squeeze( + stateless.stateless_multinomial(logits, 1, seed=seed), axis=[0, 1]) + selector_input = dataset_ops.Dataset.zip( + (logits_ds, random_ops.RandomDataset(seed).batch(2))).map(select_dataset) + + return DirectedInterleaveDataset(selector_input, datasets) diff --git a/tensorflow/contrib/data/python/ops/readers.py b/tensorflow/contrib/data/python/ops/readers.py index b8eb09978e490d706595f18525f424548679ae64..bbb808fbd7730002e48cab47fa8d0fe09e2124d2 100644 --- a/tensorflow/contrib/data/python/ops/readers.py +++ b/tensorflow/contrib/data/python/ops/readers.py @@ -124,18 +124,21 @@ def _next_csv_row(filenames, num_cols, field_delim, use_quote_delim, header, def _infer_column_defaults(filenames, num_cols, field_delim, use_quote_delim, na_value, header, comment, float_dtype, - rows_for_inference): + num_rows_for_inference, select_columns): """Infers column types from the first N valid CSV records of files.""" - inferred_types = [None] * num_cols + if select_columns is None: + select_columns = range(num_cols) + inferred_types = [None] * len(select_columns) - for rows_read, csv_row in enumerate( + for i, csv_row in enumerate( _next_csv_row(filenames, num_cols, field_delim, use_quote_delim, header, comment)): - if rows_for_inference is not None and rows_read >= rows_for_inference: + if num_rows_for_inference is not None and i >= num_rows_for_inference: break - for i, str_val in enumerate(csv_row): - inferred_types[i] = _infer_type(str_val, na_value, inferred_types[i], - float_dtype) + + for j, col_index in enumerate(select_columns): + inferred_types[j] = _infer_type(csv_row[col_index], na_value, + inferred_types[j], float_dtype) # Replace None's with a default type inferred_types = [t or dtypes.string for t in inferred_types] @@ -153,21 +156,55 @@ def _infer_column_names(filenames, field_delim, use_quote_delim): "quoting": csv.QUOTE_MINIMAL if use_quote_delim else csv.QUOTE_NONE } with file_io.FileIO(filenames[0], "r") as f: - column_names = next(csv.reader(f, **csv_kwargs)) + try: + column_names = next(csv.reader(f, **csv_kwargs)) + except StopIteration: + raise ValueError(("Received StopIteration when reading the header line " + "of %s. Empty file?") % filenames[0]) for name in filenames[1:]: with file_io.FileIO(name, "r") as f: - if next(csv.reader(f, **csv_kwargs)) != column_names: - raise ValueError("Files have different column names in the header row.") + try: + if next(csv.reader(f, **csv_kwargs)) != column_names: + raise ValueError( + "Files have different column names in the header row.") + except StopIteration: + raise ValueError(("Received StopIteration when reading the header line " + "of %s. Empty file?") % filenames[0]) return column_names +def _get_sorted_col_indices(select_columns, column_names): + """Transforms select_columns argument into sorted column indices.""" + names_to_indices = {n: i for i, n in enumerate(column_names)} + num_cols = len(column_names) + for i, v in enumerate(select_columns): + if isinstance(v, int): + if v < 0 or v >= num_cols: + raise ValueError( + "Column index %d specified in select_columns out of valid range." % + v) + continue + if v not in names_to_indices: + raise ValueError( + "Value '%s' specified in select_columns not a valid column index or " + "name." % v) + select_columns[i] = names_to_indices[v] + + # Sort and ensure there are no duplicates + result = sorted(set(select_columns)) + if len(result) != len(select_columns): + raise ValueError("select_columns contains duplicate columns") + return result + + def make_csv_dataset( file_pattern, batch_size, column_names=None, column_defaults=None, label_name=None, + select_columns=None, field_delim=",", use_quote_delim=True, na_value="", @@ -201,20 +238,32 @@ def make_csv_dataset( provided, infers the column names from the first row of the records. These names will be the keys of the features dict of each dataset element. column_defaults: A optional list of default values for the CSV fields. One - item per column of the input record. Each item in the list is either a - valid CSV dtype (float32, float64, int32, int64, or string), or a + item per selected column of the input record. Each item in the list is + either a valid CSV dtype (float32, float64, int32, int64, or string), or a `Tensor` with one of the aforementioned types. The tensor can either be a scalar default value (if the column is optional), or an empty tensor (if the column is required). If a dtype is provided instead of a tensor, the column is also treated as required. If this list is not provided, tries to infer types based on reading the first num_rows_for_inference rows of files specified, and assumes all columns are optional, defaulting to `0` - for numeric values and `""` for string values. + for numeric values and `""` for string values. If both this and + `select_columns` are specified, these must have the same lengths, and + `column_defaults` is assumed to be sorted in order of increasing column + index. label_name: 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. + select_columns: An optional list of integer indices or string column + names, that specifies a subset of columns of CSV data to select. If + column names are provided, these must correspond to names provided in + `column_names` or inferred from the file header lines. When this argument + is specified, only a subset of CSV columns will be parsed and returned, + corresponding to the columns specified. Using this results in faster + parsing and lower memory usage. If both this and `column_defaults` are + specified, these must have the same lengths, and `column_defaults` is + assumed to be sorted in order of increasing column index. 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 @@ -279,6 +328,9 @@ def make_csv_dataset( if len(column_names) != len(set(column_names)): raise ValueError("Cannot have duplicate column names.") + if select_columns is not None: + select_columns = _get_sorted_col_indices(select_columns, column_names) + if column_defaults is not None: column_defaults = [ constant_op.constant([], dtype=x) if x in _ACCEPTABLE_CSV_TYPES else x @@ -289,7 +341,17 @@ def make_csv_dataset( # construction time column_defaults = _infer_column_defaults( filenames, len(column_names), field_delim, use_quote_delim, na_value, - header, comment, default_float_type, num_rows_for_inference) + header, comment, default_float_type, num_rows_for_inference, + select_columns) + + if select_columns is not None and len(column_defaults) != len(select_columns): + raise ValueError( + "If specified, column_defaults and select_columns must have same " + "length." + ) + if select_columns is not None and len(column_names) > len(select_columns): + # Pick the relevant subset of column names + column_names = [column_names[i] for i in select_columns] if label_name is not None and label_name not in column_names: raise ValueError("`label_name` provided must be one of the columns.") @@ -322,6 +384,7 @@ def make_csv_dataset( field_delim=field_delim, use_quote_delim=use_quote_delim, na_value=na_value, + select_cols=select_columns, ) features = dict(zip(column_names, columns)) if label_name is not None: diff --git a/tensorflow/contrib/data/python/ops/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py index b465397437adbdfaf865efb8ed2f80e57f48fcab..a182dddd38d23d096979eebb8de29f07573833dd 100644 --- a/tensorflow/contrib/data/python/ops/resampling.py +++ b/tensorflow/contrib/data/python/ops/resampling.py @@ -110,7 +110,6 @@ def rejection_resample(class_func, target_dist, initial_dist=None, seed=None): .filter(lambda _1, p, _2: random_ops.random_uniform([], seed=seed) < p)) return filtered_ds.map(lambda class_value, _, data: (class_value, data)) - return _apply_fn diff --git a/tensorflow/contrib/data/python/ops/scan_ops.py b/tensorflow/contrib/data/python/ops/scan_ops.py index 1c88366273f5d186509454188e02350d4ea9f66b..60ef7efba4bb2bc281bc624ec3f58117ffa9a824 100644 --- a/tensorflow/contrib/data/python/ops/scan_ops.py +++ b/tensorflow/contrib/data/python/ops/scan_ops.py @@ -57,7 +57,7 @@ class _ScanDataset(dataset_ops.Dataset): self._output_shapes = None self._output_types = None - # Iteratively rerun the scan function until reaching a fixed pont on + # Iteratively rerun the scan function until reaching a fixed point on # `self._state_shapes`. need_to_rerun = True while need_to_rerun: @@ -144,6 +144,7 @@ class _ScanDataset(dataset_ops.Dataset): weakened_state_shapes) self._scan_func = tf_scan_func + self._scan_func.add_to_graph(ops.get_default_graph()) def _as_variant_tensor(self): input_t = self._input_dataset._as_variant_tensor() # pylint: disable=protected-access diff --git a/tensorflow/contrib/data/python/ops/stats_ops.py b/tensorflow/contrib/data/python/ops/stats_ops.py index b5cf0fcfe91ebc22444302fca5d488a278ef2994..d39172039683fe8f333572bb1dbf12abfab65113 100644 --- a/tensorflow/contrib/data/python/ops/stats_ops.py +++ b/tensorflow/contrib/data/python/ops/stats_ops.py @@ -18,7 +18,6 @@ from __future__ import division from __future__ import print_function from tensorflow.python.data.ops import dataset_ops -from tensorflow.python.data.ops import iterator_ops from tensorflow.python.data.util import nest from tensorflow.python.data.util import sparse from tensorflow.python.framework import dtypes @@ -85,25 +84,53 @@ class StatsAggregator(object): """ return gen_dataset_ops.stats_aggregator_summary(self._resource) - def subscribe(self, iterator): - """Returns a @{tf.Operation} to associate this aggregator with `iterator`. - Note: Each @{tf.data.Iterator} can be associated with at most one - `StatsAggregator`. After running the operation that this function - returns, all statistics recorded in the iteration of `iterator` - will be stored in `stats_aggregator`. +class _SetStatsAggregatorDataset(dataset_ops.Dataset): + """A `Dataset` that acts as an identity, and sets given stats_aggregator.""" - Args: - iterator: A @{tf.data.Iterator} object. + def __init__(self, input_dataset, stats_aggregator): + super(_SetStatsAggregatorDataset, self).__init__() + self._input_dataset = input_dataset + self._stats_aggregator = stats_aggregator - Returns: - A @{tf.Operation} that, when run, associates this aggregator with - `iterator`. - """ - if not isinstance(iterator, iterator_ops.Iterator): - raise TypeError("`iterator` must be a `tf.data.Iterator` object.") - return gen_dataset_ops.iterator_set_stats_aggregator( - iterator._iterator_resource, self._resource) # pylint: disable=protected-access + def _as_variant_tensor(self): + return gen_dataset_ops.set_stats_aggregator_dataset( + self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access + self._stats_aggregator._resource, # pylint: disable=protected-access + output_types=nest.flatten( + sparse.as_dense_types(self.output_types, self.output_classes)), + output_shapes=nest.flatten( + sparse.as_dense_shapes(self.output_shapes, self.output_classes))) + + @property + def output_shapes(self): + 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 + + +# TODO(shivaniagrawal): Expose these methods in `tf.contrib.data`. +def set_stats_aggregator(stats_aggregator): + """Set the given stats_aggregator for aggregating the input dataset stats. + + Args: + stats_aggregator: A `StatsAggregator` object. + + Returns: + A `Dataset` transformation function, which can be passed to + @{tf.data.Dataset.apply}. + """ + + def _apply_fn(dataset): + return _SetStatsAggregatorDataset(dataset, stats_aggregator) + + return _apply_fn def bytes_produced_stats(tag): diff --git a/tensorflow/contrib/data/python/ops/writers.py b/tensorflow/contrib/data/python/ops/writers.py new file mode 100644 index 0000000000000000000000000000000000000000..f53bd3f7383950d6cfdb35e12811fb1daf24b320 --- /dev/null +++ b/tensorflow/contrib/data/python/ops/writers.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. +# ============================================================================== +"""Python wrappers for tf.data writers.""" +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.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 TFRecordWriter(object): + """Writes data to a TFRecord file.""" + + def __init__(self, filename, compression_type=None): + self._filename = ops.convert_to_tensor( + filename, dtypes.string, name="filename") + self._compression_type = convert.optional_param_to_tensor( + "compression_type", + compression_type, + argument_default="", + argument_dtype=dtypes.string) + + def write(self, dataset): + """Returns a @{tf.Operation} to write a dataset to a file. + + Args: + dataset: a @{tf.data.Dataset} whose elements are to be written to a file + + Returns: + A @{tf.Operation} that, when run, writes contents of `dataset` to a file. + """ + if not isinstance(dataset, dataset_ops.Dataset): + raise TypeError("`dataset` must be a `tf.data.Dataset` object.") + if (dataset.output_types != dtypes.string or + dataset.output_shapes != tensor_shape.scalar()): + raise TypeError( + "`dataset` must produce scalar `DT_STRING` tensors whereas it " + "produces shape {0} and types {1}".format(dataset.output_shapes, + dataset.output_types)) + return gen_dataset_ops.dataset_to_tf_record( + dataset._as_variant_tensor(), self._filename, self._compression_type) # pylint: disable=protected-access diff --git a/tensorflow/contrib/distribute/README.md b/tensorflow/contrib/distribute/README.md index 14de1e8f491634051b6f7c2c5f2bf64f1d73e9a3..44a4481021c380e72b535cf0aca39df2bf04d3b7 100644 --- a/tensorflow/contrib/distribute/README.md +++ b/tensorflow/contrib/distribute/README.md @@ -116,7 +116,8 @@ in the input function gives a solid boost in performance. When using ## Caveats This feature is in early stages and there are a lot of improvements forthcoming: -* Metrics are not yet supported during distributed training. +* Metrics are not yet supported during distributed training. They are still +supported during the evaluation. * Summaries are only computed in the first tower in `MirroredStrategy`. * Evaluation is not yet distributed. * Eager support is in the works; performance can be more challenging with eager diff --git a/tensorflow/contrib/distribute/python/BUILD b/tensorflow/contrib/distribute/python/BUILD index 78b2b0054aa95701ad192b4fb9a0727ce287de4b..c2834d822664b9d60690c5d5dd527bbbd01a106f 100644 --- a/tensorflow/contrib/distribute/python/BUILD +++ b/tensorflow/contrib/distribute/python/BUILD @@ -22,11 +22,13 @@ py_library( visibility = ["//tensorflow:internal"], deps = [ ":prefetching_ops_v2", - "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/contrib/data/python/ops:batching", "//tensorflow/contrib/eager/python:datasets", "//tensorflow/python:array_ops", "//tensorflow/python:checkpointable", "//tensorflow/python:control_flow_ops", + "//tensorflow/python:device_util", + "//tensorflow/python:distribute", "//tensorflow/python:framework_ops", "//tensorflow/python:training", "//tensorflow/python:util", @@ -51,6 +53,7 @@ cuda_py_test( "//tensorflow/python:training", "//tensorflow/python:variable_scope", "//tensorflow/python/eager:context", + "//tensorflow/python:device_util", "//tensorflow/python/eager:test", "//tensorflow/python/estimator:model_fn", ], @@ -66,6 +69,8 @@ py_library( ":values", "//tensorflow/python:array_ops", "//tensorflow/python:device", + "//tensorflow/python:device_util", + "//tensorflow/python:distribute", "//tensorflow/python:framework_ops", "//tensorflow/python:pywrap_tensorflow", "//tensorflow/python:training", @@ -84,9 +89,9 @@ py_library( ":values", "//tensorflow/contrib/eager/python:datasets", "//tensorflow/python:array_ops", + "//tensorflow/python:distribute", "//tensorflow/python:framework_ops", "//tensorflow/python:math_ops", - "//tensorflow/python:training", "//tensorflow/python/eager:context", "@six_archive//:six", ], @@ -104,6 +109,7 @@ py_library( "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", "//tensorflow/python:constant_op", + "//tensorflow/python:distribute", "//tensorflow/python:framework_ops", "//tensorflow/python:layers", "//tensorflow/python:training", @@ -125,6 +131,7 @@ py_library( deps = [ ":mirrored_strategy", ":one_device_strategy", + ":tpu_strategy", "//tensorflow/contrib/optimizer_v2:training", "//tensorflow/python:framework_ops", "//tensorflow/python:training", @@ -156,8 +163,8 @@ py_test( deps = [ ":mirrored_strategy", ":strategy_test_lib", + "//tensorflow/python:distribute", "//tensorflow/python:framework_test_lib", - "//tensorflow/python:training", "//tensorflow/python:variable_scope", "//tensorflow/python/eager:context", "//tensorflow/python/eager:test", @@ -186,10 +193,10 @@ cuda_py_test( ":mirrored_strategy", ":values", ":strategy_test_lib", + "//tensorflow/python:distribute", "//tensorflow/core:protos_all_py", "//tensorflow/python:constant_op", "//tensorflow/python:layers", - "//tensorflow/python:training", "//tensorflow/python:variable_scope", "//tensorflow/python:array_ops", "//tensorflow/python:framework_test_lib", @@ -219,21 +226,50 @@ py_library( ], ) -cuda_py_test( - name = "minimize_loss_test", +py_library( + name = "tpu_strategy", + srcs = ["tpu_strategy.py"], + visibility = ["//tensorflow:internal"], + deps = [ + ":one_device_strategy", + ":values", + "//tensorflow/contrib/tpu", + "//tensorflow/contrib/tpu:tpu_py", + "//tensorflow/python:constant_op", + "//tensorflow/python:control_flow_ops", + "//tensorflow/python:framework_ops", + "//tensorflow/python:util", + ], +) + +py_library( + name = "minimize_loss_test_lib", + testonly = 1, srcs = ["minimize_loss_test.py"], - additional_deps = [ + deps = [ ":combinations", + ":mirrored_strategy", ":single_loss_example", - "@absl_py//absl/testing:parameterized", - "//third_party/py/numpy", + "//tensorflow/contrib/tpu:tpu_lib", "//tensorflow/python:control_flow_ops", + "//tensorflow/python:framework_ops", "//tensorflow/python:math_ops", + "//tensorflow/python:variable_scope", "//tensorflow/python:variables", "//tensorflow/python/data/ops:dataset_ops", "//tensorflow/python/eager:context", "//tensorflow/python/eager:test", "//tensorflow/python/ops/losses", + "//third_party/py/numpy", + "@absl_py//absl/testing:parameterized", + ], +) + +cuda_py_test( + name = "minimize_loss_test", + srcs = ["minimize_loss_test.py"], + additional_deps = [ + ":minimize_loss_test_lib", ], tags = [ "multi_and_single_gpu", @@ -291,6 +327,7 @@ py_library( srcs = ["single_loss_example.py"], deps = [ ":step_fn", + "//tensorflow/contrib/data/python/ops:batching", "//tensorflow/python:array_ops", "//tensorflow/python:constant_op", "//tensorflow/python:layers", diff --git a/tensorflow/contrib/distribute/python/combinations.py b/tensorflow/contrib/distribute/python/combinations.py index 02b1e7ef9fcd4767c59898bd343e712e285e67d5..946310aa6fc2101d75e86d3ff2e9f3284e6c6625 100644 --- a/tensorflow/contrib/distribute/python/combinations.py +++ b/tensorflow/contrib/distribute/python/combinations.py @@ -45,6 +45,7 @@ from absl.testing import parameterized from tensorflow.contrib.distribute.python import mirrored_strategy from tensorflow.contrib.distribute.python import one_device_strategy +from tensorflow.contrib.distribute.python import tpu_strategy from tensorflow.contrib.optimizer_v2 import adam as adam_v2 from tensorflow.contrib.optimizer_v2 import gradient_descent as gradient_descent_v2 from tensorflow.python.eager import context @@ -55,6 +56,7 @@ from tensorflow.python.util import tf_inspect GPU_TEST = "test_gpu" in sys.argv[0] +TPU_TEST = "test_tpu" in sys.argv[0] def generate(combinations): @@ -108,6 +110,11 @@ def generate(combinations): if "distribution" in kwargs: distribution = kwargs["distribution"] kwargs["distribution"] = distribution.strategy + if distribution.required_tpu and not TPU_TEST: + self.skipTest("Test requires a TPU, but it's not available.") + if not distribution.required_tpu and TPU_TEST: + self.skipTest("Test that doesn't require a TPU.") + if not distribution.required_gpus: if GPU_TEST: self.skipTest("Test that doesn't require GPUs.") @@ -232,10 +239,12 @@ class NamedObject(object): class NamedDistribution(object): """Translates DistributionStrategy and its data into a good name.""" - def __init__(self, name, distribution, required_gpus): + def __init__(self, name, distribution, required_gpus=None, + required_tpu=False): self._distribution = distribution self._name = name self._required_gpus = required_gpus + self._required_tpu = required_tpu def __repr__(self): return self._name @@ -248,10 +257,16 @@ class NamedDistribution(object): def required_gpus(self): return self._required_gpus + @property + def required_tpu(self): + return self._required_tpu + one_device_strategy = NamedDistribution( "OneDeviceCPU", one_device_strategy.OneDeviceStrategy("/cpu:0"), None) +tpu_strategy = NamedDistribution( + "TPU", tpu_strategy.TPUStrategy(), required_tpu=True) mirrored_strategy_with_gpu_and_cpu = NamedDistribution( "MirroredCPUAndGPU", mirrored_strategy.MirroredStrategy(["/gpu:0", "/cpu:0"]), 1) diff --git a/tensorflow/contrib/distribute/python/cross_tower_ops.py b/tensorflow/contrib/distribute/python/cross_tower_ops.py index bbe5e877d59518056db3fea251cdae0ed854d0e4..cff717db80f0bdd377b3c9c7e8ca3578ff273930 100644 --- a/tensorflow/contrib/distribute/python/cross_tower_ops.py +++ b/tensorflow/contrib/distribute/python/cross_tower_ops.py @@ -488,7 +488,8 @@ class AllReduceCrossTowerOps(CrossTowerOps): "agg_small_grads_max_group = %d", len(per_device_values), self.all_reduce_alg, self.agg_small_grads_max_bytes, self.agg_small_grads_max_group) - tensor_packer = AggregateSmallTensorPacker(100, 10) + tensor_packer = AggregateSmallTensorPacker( + self.agg_small_grads_max_bytes, self.agg_small_grads_max_group) device_grad_packs = tensor_packer.pack(grouped) else: logging.info( diff --git a/tensorflow/contrib/distribute/python/minimize_loss_test.py b/tensorflow/contrib/distribute/python/minimize_loss_test.py index 0fa90df79bbcd621fe7b7d0da04256b7a59d5bfe..e134fe34e10be402f028db986b8cbf14222db07f 100644 --- a/tensorflow/contrib/distribute/python/minimize_loss_test.py +++ b/tensorflow/contrib/distribute/python/minimize_loss_test.py @@ -25,6 +25,7 @@ from tensorflow.contrib.distribute.python import combinations from tensorflow.contrib.distribute.python import mirrored_strategy from tensorflow.contrib.distribute.python.single_loss_example import batchnorm_example from tensorflow.contrib.distribute.python.single_loss_example import minimize_loss_example +from tensorflow.contrib.tpu.python.tpu import tpu from tensorflow.python.data.ops import dataset_ops from tensorflow.python.eager import context from tensorflow.python.eager import test @@ -42,16 +43,27 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): combinations.times( combinations.distributions_and_v1_optimizers(), combinations.combine(mode=["graph"], use_callable_loss=[True, False]) - + combinations.combine(mode=["eager"], use_callable_loss=[True]))) - def testTrainNetwork(self, distribution, optimizer_fn, - use_callable_loss=True): + + combinations.combine(mode=["eager"], use_callable_loss=[True]), + combinations.combine(is_tpu=[False])) + + combinations.combine( + distribution=[combinations.tpu_strategy], + optimizer_fn=[combinations.adam_optimizer_v1_fn], + mode=["graph"], + use_callable_loss=[False], + is_tpu=[True])) + def testTrainNetwork(self, distribution, optimizer_fn, use_callable_loss, + is_tpu): with distribution.scope(): - model_fn, dataset, layer = minimize_loss_example( - optimizer_fn, - use_bias=True, - use_callable_loss=use_callable_loss) + model_fn, dataset_fn, layer = minimize_loss_example( + optimizer_fn, use_bias=True, use_callable_loss=use_callable_loss) - iterator = distribution.distribute_dataset(dataset) + # TODO(isaprykin): Eliminate `is_tpu`. Probably add a + # `DistributionStrategy.create_monitor` so that each DistributionStrategy + # could influence its training loop. That method would return an instance + # of Monitor. TPUMonitor would execute tpu.initialize_system() and + # tpu.shutdown_system(). + iterator = distribution.distribute_dataset( + dataset_fn).make_one_shot_iterator() def run_step(): return distribution.group( @@ -60,6 +72,8 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): if not context.executing_eagerly(): with self.test_session() as sess: + if is_tpu: + sess.run(tpu.initialize_system()) run_step = sess.make_callable(run_step()) self.evaluate(variables_lib.global_variables_initializer()) @@ -70,6 +84,10 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): weights.append(self.evaluate(distribution.fetch(layer.kernel))) biases.append(self.evaluate(distribution.fetch(layer.bias))) + if is_tpu: + with self.test_session() as sess: + sess.run(tpu.shutdown_system()) + error = abs(numpy.add(numpy.squeeze(weights), numpy.squeeze(biases)) - 1) is_not_increasing = all(y <= x for x, y in zip(error, error[1:])) self.assertTrue(is_not_increasing) @@ -78,8 +96,17 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): combinations.times( combinations.distributions_and_v1_optimizers() + combinations.distributions_and_v2_optimizers(), - combinations.combine(mode=["graph", "eager"]))) - def testOptimizerInsideModelFn(self, distribution, optimizer_fn): + combinations.combine(mode=["graph", "eager"], is_tpu=[False])) + + combinations.combine( + distribution=[combinations.tpu_strategy], + optimizer_fn=[ + combinations.adam_optimizer_v1_fn, + combinations.gradient_descent_optimizer_v1_fn + ], + mode=["graph"], + is_tpu=[True])) + + def testOptimizerInsideModelFn(self, distribution, optimizer_fn, is_tpu): created_variables = [] trainable_variables = [] @@ -94,13 +121,14 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): # `distribution.scope`. with variable_scope.variable_creator_scope( appending_creator), distribution.scope(): - model_fn, dataset, layer = minimize_loss_example( + model_fn, dataset_fn, layer = minimize_loss_example( optimizer_fn, use_bias=True, use_callable_loss=True, create_optimizer_inside_model_fn=True) - iterator = distribution.distribute_dataset(dataset) + iterator = distribution.distribute_dataset( + dataset_fn).make_one_shot_iterator() def run_step(): return distribution.group( @@ -109,11 +137,17 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): if not context.executing_eagerly(): with self.test_session() as sess: + if is_tpu: + sess.run(tpu.initialize_system()) run_step = sess.make_callable(run_step()) self.evaluate(variables_lib.global_variables_initializer()) run_step() + if is_tpu: + with self.test_session() as sess: + sess.run(tpu.shutdown_system()) + def get_expected_variables(optimizer_fn, num_parameter_devices): variables_map = { "GradientDescent": ["dense/kernel", "dense/bias"], @@ -147,7 +181,7 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): """Verifies that moving mean updates are reduced across towers.""" with distribution.scope(): num_towers = len(distribution.worker_devices) - model_fn, dataset, batchnorm = batchnorm_example( + model_fn, dataset_fn, batchnorm = batchnorm_example( optimizer_fn, batch_per_epoch=num_towers, momentum=momentum, @@ -158,7 +192,8 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): # on each device. if isinstance(distribution, mirrored_strategy.MirroredStrategy): distribution._prefetch_on_device = False - iterator = distribution.distribute_dataset(dataset) + iterator = distribution.distribute_dataset( + dataset_fn).make_one_shot_iterator() def run_step(): return control_flow_ops.group( @@ -230,10 +265,13 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): else: return optimizer.minimize(loss_fn()) - features = dataset_ops.Dataset.from_tensors([[2.], [7.]]) - labels = dataset_ops.Dataset.from_tensors([[6.], [21.]]) - dataset = dataset_ops.Dataset.zip((features, labels)).repeat() - iterator = distribution.distribute_dataset(dataset) + def dataset_fn(): + features = dataset_ops.Dataset.from_tensors([[2.], [7.]]) + labels = dataset_ops.Dataset.from_tensors([[6.], [21.]]) + return dataset_ops.Dataset.zip((features, labels)).repeat() + + iterator = distribution.distribute_dataset( + dataset_fn).make_one_shot_iterator() def run_step(): return distribution.group( diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy.py b/tensorflow/contrib/distribute/python/mirrored_strategy.py index eb0edb3a11df7788991ca14f957494d87593a449..6efd578a775da7bf326826289bd5bd50a57be892 100644 --- a/tensorflow/contrib/distribute/python/mirrored_strategy.py +++ b/tensorflow/contrib/distribute/python/mirrored_strategy.py @@ -140,10 +140,10 @@ class MirroredStrategy(distribute_lib.DistributionStrategy): g.add_to_collections(collections, result) return result - def distribute_dataset(self, dataset): - per_device_dataset = values.PerDeviceDataset( - dataset, self._devices, self._prefetch_on_device) - return per_device_dataset.make_one_shot_iterator() + def distribute_dataset(self, dataset_fn): + return values.PerDeviceDataset( + self._call_dataset_fn(dataset_fn), self._devices, + self._prefetch_on_device) def _broadcast(self, tensor, destinations): # TODO(josh11b): In eager mode, use one thread per device, or async mode. diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py b/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py index 9e9f06da8e2ed185c2c32f79a5a4f5407165fb1d..6c5c055070c0fc88ed8f3a459e3f346596f077a6 100644 --- a/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py +++ b/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py @@ -247,8 +247,9 @@ class MirroredStrategyVariableCreationTest(test.TestCase): dist = mirrored_strategy.MirroredStrategy( ["/device:GPU:0", "/device:CPU:0"]) - features = dataset_ops.Dataset.from_tensors([[1.]]).repeat(10) - features = dist.distribute_dataset(features).get_next() + features = dist.distribute_dataset( + lambda: dataset_ops.Dataset.from_tensors([[1.]]).repeat(10) + ).make_one_shot_iterator().get_next() with dist.scope(): result = dist.call_for_each_tower( diff --git a/tensorflow/contrib/distribute/python/one_device_strategy.py b/tensorflow/contrib/distribute/python/one_device_strategy.py index 39c49442b9c3245cfd0b67a51be68773a6fd3ff4..646d2a5c3b3b0bfcce6f89be0e588baacc6b9237 100644 --- a/tensorflow/contrib/distribute/python/one_device_strategy.py +++ b/tensorflow/contrib/distribute/python/one_device_strategy.py @@ -21,8 +21,6 @@ from __future__ import print_function import six from tensorflow.contrib.distribute.python import values -from tensorflow.contrib.eager.python import datasets -from tensorflow.python.eager import context from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops @@ -62,11 +60,8 @@ class OneDeviceStrategy(distribute_lib.DistributionStrategy): with ops.colocate_with(colocate_with): return next_creator(*args, **kwargs) - def distribute_dataset(self, dataset): - if context.executing_eagerly(): - return datasets.Iterator(dataset) - else: - return dataset.make_one_shot_iterator() + def distribute_dataset(self, dataset_fn): + return self._call_dataset_fn(dataset_fn) def _broadcast(self, tensor, destinations): return tensor diff --git a/tensorflow/contrib/distribute/python/optimizer_v2_test.py b/tensorflow/contrib/distribute/python/optimizer_v2_test.py index a0912b625f44342d22acc0ce9bb52a6b632c75a0..abd3a65ac4e19ece6b69b9834f4218fde55b60c2 100644 --- a/tensorflow/contrib/distribute/python/optimizer_v2_test.py +++ b/tensorflow/contrib/distribute/python/optimizer_v2_test.py @@ -39,10 +39,11 @@ class MinimizeLossOptimizerV2Test(test.TestCase, parameterized.TestCase): def testTrainNetwork(self, distribution, optimizer_fn, use_callable_loss=True): with distribution.scope(): - model_fn, dataset, layer = minimize_loss_example( + model_fn, dataset_fn, layer = minimize_loss_example( optimizer_fn, use_bias=True, use_callable_loss=use_callable_loss) - iterator = distribution.distribute_dataset(dataset) + iterator = distribution.distribute_dataset( + dataset_fn).make_one_shot_iterator() def run_step(): return control_flow_ops.group(distribution.unwrap( diff --git a/tensorflow/contrib/distribute/python/prefetching_ops_v2.py b/tensorflow/contrib/distribute/python/prefetching_ops_v2.py index dfcbb8568f92ebabbeeedb45ee677e4ee23d77dc..7b3670b45aba801cf8c18e04bfea03e23eb67184 100644 --- a/tensorflow/contrib/distribute/python/prefetching_ops_v2.py +++ b/tensorflow/contrib/distribute/python/prefetching_ops_v2.py @@ -26,6 +26,7 @@ from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import iterator_ops from tensorflow.python.data.util import nest as data_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 function from tensorflow.python.framework import ops @@ -34,26 +35,55 @@ from tensorflow.python.util import nest # pylint: disable=protected-access class _PrefetchToDeviceIterator(object): - """A replacement for @{tf.data.Iterator} that prefetches to another device.""" + """A replacement for @{tf.data.Iterator} that prefetches to another device. - def __init__(self, input_dataset, devices, buffer_size): + Args: + input_dataset: The input dataset. + one_shot: If true, we make a one shot iterator that's already initialized. + devices: Devices on which to prefetch. + buffer_size: Size of the prefetching buffer. + shared_name: (Optional.) If non-empty, the returned iterator will be + shared under the given name across multiple sessions that share the + same devices (e.g. when using a remote server). Only used if one_shot + is False. + + Returns: + An Iterator type object. + """ + + def __init__(self, + input_dataset, + one_shot, + devices, + buffer_size, + shared_name=None): self._input_dataset = input_dataset self._get_next_call_count = 0 + self._one_shot = one_shot + if shared_name is None: + shared_name = "" self._devices = devices - input_iterator = input_dataset.make_one_shot_iterator() - input_iterator_handle = input_iterator.string_handle() + + if self._one_shot: + self._input_iterator = input_dataset.make_one_shot_iterator() + else: + self._input_iterator = iterator_ops.Iterator.from_structure( + self._input_dataset.output_types, self._input_dataset.output_shapes, + shared_name, self._input_dataset.output_classes) + input_iterator_handle = self._input_iterator.string_handle() @function.Defun(dtypes.string) def _prefetch_fn(handle): """Prefetches one element from `input_iterator`.""" remote_iterator = iterator_ops.Iterator.from_string_handle( - handle, input_iterator.output_types, input_iterator.output_shapes, - input_iterator.output_classes) + handle, self._input_iterator.output_types, + self._input_iterator.output_shapes, + self._input_iterator.output_classes) ret = remote_iterator.get_next() return nest.flatten(sparse.serialize_sparse_tensors(ret)) target_device = gen_dataset_ops.iterator_get_device( - input_iterator._iterator_resource) + self._input_iterator._iterator_resource) self._buffering_resources = [] for device in nest.flatten(self._devices): with ops.device(device): @@ -61,9 +91,19 @@ class _PrefetchToDeviceIterator(object): f=_prefetch_fn, target_device=target_device, string_arg=input_iterator_handle, - buffer_size=buffer_size) + buffer_size=buffer_size, + shared_name=shared_name) self._buffering_resources.append(buffer_resource_handle) + if not self._one_shot: + reset_ops = [] + for buffer_resource in self._buffering_resources: + reset_ops.append( + prefetching_ops.function_buffering_resource_reset(buffer_resource)) + with ops.control_dependencies(reset_ops): + self._initializer = self._input_iterator.make_initializer( + self._input_dataset) + def get_next(self, name=None): """See @{tf.data.Iterator.get_next}.""" self._get_next_call_count += 1 @@ -92,6 +132,12 @@ class _PrefetchToDeviceIterator(object): return nest.pack_sequence_as(self._devices, flat_result) + @property + def initializer(self): + if self._one_shot: + raise NotImplementedError("Can't initialize a one_shot_iterator") + return self._initializer + @property def output_classes(self): return self._input_dataset.output_classes @@ -115,13 +161,24 @@ class _PrefetchToDeviceDataset(dataset_ops.Dataset): self._buffer_size = buffer_size if buffer_size is not None else 1 def make_one_shot_iterator(self): - return _PrefetchToDeviceIterator(self._input_dataset, self._devices, - self._buffer_size) + return _PrefetchToDeviceIterator( + self._input_dataset, + one_shot=True, + devices=self._devices, + buffer_size=self._buffer_size) def make_initializable_iterator(self, shared_name=None): - raise NotImplementedError("`prefetch_to_devices()` is not currently " - "compatible with initializable iterators. Use " - "`make_one_shot_iterator()` instead.") + if context.executing_eagerly(): + raise RuntimeError( + "make_initializable_iterator is not supported when eager " + "execution is enabled.") + + return _PrefetchToDeviceIterator( + self._input_dataset, + one_shot=False, + devices=self._devices, + buffer_size=self._buffer_size, + shared_name=shared_name) def _as_variant_tensor(self): # TODO(mrry): Raise this error earlier (e.g. when one of the Dataset diff --git a/tensorflow/contrib/distribute/python/prefetching_ops_v2_test.py b/tensorflow/contrib/distribute/python/prefetching_ops_v2_test.py index 8ed16f4607881f2864479c04b4c25e95d9fa1850..a68dbce6c7d03f6a1695ebfcd00178e21ac1cda0 100644 --- a/tensorflow/contrib/distribute/python/prefetching_ops_v2_test.py +++ b/tensorflow/contrib/distribute/python/prefetching_ops_v2_test.py @@ -64,5 +64,27 @@ class PrefetchingOpsV2Test(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(next_element) + def testPrefetchToTwoDevicesWithReinit(self): + if not test_util.is_gpu_available(): + self.skipTest("No GPU available") + + host_dataset = dataset_ops.Dataset.range(10) + device_dataset = host_dataset.apply( + prefetching_ops_v2.prefetch_to_devices(["/cpu:0", "/gpu:0"])) + + iterator = device_dataset.make_initializable_iterator() + next_element = iterator.get_next() + + with self.test_session() as sess: + sess.run(iterator.initializer) + for _ in range(5): + sess.run(next_element) + with self.assertRaises(errors.OutOfRangeError): + sess.run(next_element) + sess.run(iterator.initializer) + for _ in range(5): + sess.run(next_element) + + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/distribute/python/single_loss_example.py b/tensorflow/contrib/distribute/python/single_loss_example.py index cef5fd2f8943d348a0721cd72032bf6cb2199ad9..0db0b59fcacee2785eb8191bb84ed5216a79b081 100644 --- a/tensorflow/contrib/distribute/python/single_loss_example.py +++ b/tensorflow/contrib/distribute/python/single_loss_example.py @@ -18,6 +18,7 @@ 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.distribute.python import step_fn from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import constant_op @@ -29,7 +30,10 @@ from tensorflow.python.ops import math_ops def single_loss_example(optimizer_fn, distribution, use_bias=False): """Build a very simple network to use in tests and examples.""" - dataset = dataset_ops.Dataset.from_tensors([[1.]]).repeat() + + def dataset_fn(): + return dataset_ops.Dataset.from_tensors([[1.]]).repeat() + optimizer = optimizer_fn() layer = core.Dense(1, use_bias=use_bias) @@ -37,8 +41,8 @@ def single_loss_example(optimizer_fn, distribution, use_bias=False): y = array_ops.reshape(layer(x), []) - constant_op.constant(1.) return y * y - single_loss_step = step_fn.StandardSingleLossStep(dataset, loss_fn, optimizer, - distribution) + single_loss_step = step_fn.StandardSingleLossStep(dataset_fn, loss_fn, + optimizer, distribution) # Layer is returned for inspecting the kernels in tests. return single_loss_step, layer @@ -49,7 +53,14 @@ def minimize_loss_example(optimizer_fn, use_callable_loss=True, create_optimizer_inside_model_fn=False): """Example of non-distribution-aware legacy code.""" - dataset = dataset_ops.Dataset.from_tensors([[1.]]).repeat() + + def dataset_fn(): + dataset = dataset_ops.Dataset.from_tensors([[1.]]).repeat() + # TODO(isaprykin): map_and_batch with drop_remainder causes shapes to be + # fully defined for TPU. Remove this when XLA supports dynamic shapes. + return dataset.apply( + batching.map_and_batch(lambda x: x, batch_size=2, drop_remainder=True)) + # An Optimizer instance is created either outside or inside model_fn. outer_optimizer = None if not create_optimizer_inside_model_fn: @@ -57,10 +68,11 @@ def minimize_loss_example(optimizer_fn, layer = core.Dense(1, use_bias=use_bias) - def model_fn(x): + def model_fn(xs): """A very simple model written by the user.""" def loss_fn(): + x = math_ops.reduce_mean(xs, keepdims=True) y = array_ops.reshape(layer(x), []) - constant_op.constant(1.) return y * y @@ -71,7 +83,7 @@ def minimize_loss_example(optimizer_fn, else: return optimizer.minimize(loss_fn()) - return model_fn, dataset, layer + return model_fn, dataset_fn, layer def batchnorm_example(optimizer_fn, @@ -79,12 +91,15 @@ def batchnorm_example(optimizer_fn, momentum=0.9, renorm=False): """Example of non-distribution-aware legacy code with batch normalization.""" - # input shape is [16, 8], input values are increasing in both dimensions. - dataset = dataset_ops.Dataset.from_tensor_slices( - [[[float(x * 8 + y + z * 100) - for y in range(8)] - for x in range(16)] - for z in range(batch_per_epoch)]).repeat() + + def dataset_fn(): + # input shape is [16, 8], input values are increasing in both dimensions. + return dataset_ops.Dataset.from_tensor_slices( + [[[float(x * 8 + y + z * 100) + for y in range(8)] + for x in range(16)] + for z in range(batch_per_epoch)]).repeat() + optimizer = optimizer_fn() batchnorm = normalization.BatchNormalization( renorm=renorm, momentum=momentum, fused=False) @@ -99,4 +114,4 @@ def batchnorm_example(optimizer_fn, # Callable loss. return optimizer.minimize(loss_fn) - return model_fn, dataset, batchnorm + return model_fn, dataset_fn, batchnorm diff --git a/tensorflow/contrib/distribute/python/step_fn.py b/tensorflow/contrib/distribute/python/step_fn.py index 82514c64be40b421c4a9887932f2cfb8e1ac4be0..d1910622b38c748fc5a814f9e83c2294850d5d12 100644 --- a/tensorflow/contrib/distribute/python/step_fn.py +++ b/tensorflow/contrib/distribute/python/step_fn.py @@ -49,12 +49,14 @@ class StandardInputStep(Step): """Step with a standard implementation of input handling. Args: - input_dataset: a tf.data Dataset that provides input. + dataset_fn: a function that returns a tf.data Dataset that produces the + input for the model. """ - def __init__(self, input_dataset, distribution): + def __init__(self, dataset_fn, distribution): Step.__init__(self, distribution) - self._distributed_input = distribution.distribute_dataset(input_dataset) + self._distributed_input = distribution.distribute_dataset( + dataset_fn).make_one_shot_iterator() def inputs(self): return self._distributed_input.get_next() @@ -76,14 +78,15 @@ class StandardSingleLossStep(StandardInputStep): ``` Args: - input_dataset: a tf.data Dataset that provides input. + dataset_fn: a function that returns a tf.data Dataset that produces the + input for the model. loss_fn: a function that returns loss. optimizer: an optimizer that implements an update rule. distribution: a `DistributionStrategy` object. """ - def __init__(self, input_dataset, loss_fn, optimizer, distribution): - StandardInputStep.__init__(self, input_dataset, distribution) + def __init__(self, dataset_fn, loss_fn, optimizer, distribution): + StandardInputStep.__init__(self, dataset_fn, distribution) self._loss_fn = loss_fn self._optimizer = optimizer self._is_run_concurrently = False diff --git a/tensorflow/contrib/distribute/python/tpu_strategy.py b/tensorflow/contrib/distribute/python/tpu_strategy.py new file mode 100644 index 0000000000000000000000000000000000000000..a7e4fe80f3e65907fa4b48c5fe0fcfd422ba033f --- /dev/null +++ b/tensorflow/contrib/distribute/python/tpu_strategy.py @@ -0,0 +1,119 @@ +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""TPU Distribution Strategy. + +This is experimental. It's not ready for general use. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import itertools + +from tensorflow.contrib import tpu +from tensorflow.contrib.distribute.python import one_device_strategy +from tensorflow.contrib.distribute.python import values +from tensorflow.contrib.tpu.python.ops import tpu_ops +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import ops +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.util import nest + + +# TODO(isaprykin): Consider whether inheriting is really appropriate. +class TPUStrategy(one_device_strategy.OneDeviceStrategy): + """Experimental TPU distribution strategy implementation.""" + + def __init__(self, + num_cores_per_host=2, + iterations_per_step=2): + # TODO(isaprykin): Generalize the defaults. They are currently tailored for + # the unit test. + super(TPUStrategy, self).__init__('/cpu:0') + # TODO(isaprykin): Auto-detect number of cores and hosts. + self._num_cores_per_host = num_cores_per_host + # TODO(isaprykin): This might have to be per-call. + self._iterations_per_step = iterations_per_step + + def distribute_dataset(self, dataset_fn): + return values.PerIterationDataset( + self._call_dataset_fn(dataset_fn), self._iterations_per_step, + self._num_cores_per_host) + + def _call_for_each_tower(self, fn, *args, **kwargs): + kwargs.pop('run_concurrently', None) + + inputs = {'args': args, 'kwargs': kwargs} + flat_inputs = nest.flatten(inputs) + + feed_mask = [isinstance(f, values.PerIteration) for f in flat_inputs] + + feeds = lambda: itertools.compress(flat_inputs, feed_mask) + shapes = [f.get_shape() for f in feeds()] + if any([not s.is_fully_defined() for s in shapes]): + raise ValueError( + 'TPU currently requires fully defined shapes. Either use ' + 'set_shape() on the input tensors or use ' + 'dataset.apply(map_and_batch(..., drop_remainder=True)).') + types = [f.get_dtype() for f in feeds()] + + def infeed_input(i): + """Get input, split it and then enqueue.""" + iteration_inputs = [f.get(i) for f in feeds()] + + infeed_inputs = [[inputs_per_core[core_id] + for inputs_per_core in iteration_inputs] + for core_id in range(self._num_cores_per_host)] + + infeed_ops = [] + for core_id, infeed_input in enumerate(infeed_inputs): + infeed_ops.append( + tpu_ops.infeed_enqueue_tuple( + inputs=infeed_input, shapes=shapes, device_ordinal=core_id)) + + with ops.control_dependencies(infeed_ops): + return i + 1 + + with ops.device('/task:0/device:CPU:0'): + enqueue_ops = control_flow_ops.while_loop( + lambda i: i < self._iterations_per_step, + infeed_input, [constant_op.constant(0)], + parallel_iterations=1) + + def dequeueing_fn(*args, **kwargs): + """Dequeue input arguments and supply them to `fn`.""" + del args, kwargs + dequeued = tpu.infeed_dequeue_tuple(dtypes=types, shapes=shapes) + dequeued = iter(dequeued) + + fn_inputs = [] + for inp, is_feed in zip(flat_inputs, feed_mask): + if is_feed: + fn_inputs.append(next(dequeued)) + else: + fn_inputs.append(inp) + + fn_inputs = nest.pack_sequence_as(inputs, fn_inputs) + return fn(*fn_inputs['args'], **fn_inputs['kwargs']) + + def iterate_on_tpu(): + return tpu.repeat(self._iterations_per_step, dequeueing_fn, []) + + with one_device_strategy._OneDeviceTowerContext(self): # pylint: disable=protected-access + tpu_result = tpu.batch_parallel( + iterate_on_tpu, [], num_shards=self._num_cores_per_host) + + return control_flow_ops.group(tpu_result, enqueue_ops) diff --git a/tensorflow/contrib/distribute/python/values.py b/tensorflow/contrib/distribute/python/values.py index 87bf0590384cc74ca0f0575bcef4e84599a8b666..8cb5276579f48f9ea5781c5351cbf9bf3db16e6c 100644 --- a/tensorflow/contrib/distribute/python/values.py +++ b/tensorflow/contrib/distribute/python/values.py @@ -28,7 +28,6 @@ import six from tensorflow.contrib.data.python.ops import batching from tensorflow.contrib.distribute.python import prefetching_ops_v2 -from tensorflow.contrib.eager.python import datasets from tensorflow.python.eager import context from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops @@ -510,6 +509,10 @@ class PerDeviceDataIterator(object): self._devices = devices self._prefetch_on_device = prefetch_on_device + @property + def initializer(self): + return self._iterator.initializer + def get_next(self, name=None): """Scatter the input across devices.""" if self._prefetch_on_device: @@ -545,7 +548,8 @@ class PerDeviceDataset(object): "Prefetching is only supported in graph mode currently") if self._prefetch_on_device: - self._dataset = dataset + self._dataset = dataset.apply( + prefetching_ops_v2.prefetch_to_devices(self._devices)) else: # TODO(priyag): If dropping remainder is not appropriate, find another # approach to distributing the dataset when not possible to divide evenly. @@ -555,19 +559,72 @@ class PerDeviceDataset(object): def make_one_shot_iterator(self): """Get a one time use iterator for the distributed PerDeviceDataset.""" - if self._prefetch_on_device: - on_device_dataset = self._dataset.apply( - prefetching_ops_v2.prefetch_to_devices(self._devices)) - dataset_iterator = on_device_dataset.make_one_shot_iterator() - elif context.executing_eagerly(): - dataset_iterator = datasets.Iterator(self._dataset) - else: - dataset_iterator = self._dataset.make_one_shot_iterator() + dataset_iterator = self._dataset.make_one_shot_iterator() + return PerDeviceDataIterator( + dataset_iterator, self._devices, self._prefetch_on_device) + def make_initializable_iterator(self): + """Get an initializable iterator for the distributed PerDeviceDataset.""" + dataset_iterator = self._dataset.make_initializable_iterator() return PerDeviceDataIterator( dataset_iterator, self._devices, self._prefetch_on_device) +class PerIteration(object): + """Holds input for multiple iterations at once.""" + + def __init__(self, index): + self._index = index + + def get(self, iteration): + return array_ops.gather(self._index, iteration) + + def get_shape(self): + return self._index[-1][-1].get_shape() + + def get_dtype(self): + return self._index[-1][-1].dtype + + +class MultiIterator(object): + """Iterator that returns results of multiple get_next()s.""" + + def __init__(self, dataset_iterator, iterations, batches_per_iteration): + self._dataset_iterator = dataset_iterator + self._iterations = iterations + self._batches_per_iteration = batches_per_iteration + + def get_next(self, name=None): + return PerIteration([[ + self._dataset_iterator.get_next(name=name) + for _ in range(self._batches_per_iteration) + ] + for _ in range(self._iterations)]) + + @property + def initializer(self): + return self._dataset_iterator.initializer + + +class PerIterationDataset(object): + """A dataset that returns MultiIterators.""" + + def __init__(self, dataset, iterations, batches_per_iteration): + self._dataset = dataset + self._iterations = iterations + self._batches_per_iteration = batches_per_iteration + + def make_one_shot_iterator(self): + iterator = self._dataset.make_one_shot_iterator() + return MultiIterator(iterator, self._iterations, + self._batches_per_iteration) + + def make_initializable_iterator(self): + iterator = self._dataset.make_initializable_iterator() + return MultiIterator(iterator, self._iterations, + self._batches_per_iteration) + + class MapOutput(object): """Map can result in multiple outputs per device.""" diff --git a/tensorflow/contrib/distribute/python/values_test.py b/tensorflow/contrib/distribute/python/values_test.py index 5c0d4b7d6c78b7cf63c613201d83d4793ecfe76b..e96ce547415fcb2bf3da8b6085ee11f51717db8d 100644 --- a/tensorflow/contrib/distribute/python/values_test.py +++ b/tensorflow/contrib/distribute/python/values_test.py @@ -32,6 +32,7 @@ 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 random_ops from tensorflow.python.ops import variable_scope from tensorflow.python.training import device_util from tensorflow.python.training import saver as saver_lib @@ -408,6 +409,32 @@ class PerDeviceDatasetTest(test.TestCase): expected_values = [[i, i+1] for i in range(0, 10, 2)] self._test_iterator(devices, dataset, expected_values) + def testInitializableIterator(self): + with context.graph_mode(): + devices = ["/device:CPU:0"] + # Using random input since that is only allowed with initializable + # iterator. + dataset = dataset_ops.Dataset.from_tensor_slices( + random_ops.random_uniform((10,))) + + per_device_dataset = values.PerDeviceDataset( + dataset, devices, prefetch_on_device=False) + iterator = per_device_dataset.make_initializable_iterator() + + self.evaluate(iterator.initializer) + next_element = iterator.get_next() + for _ in range(10): + self.evaluate(next_element) + + # Should fail after the input is finished. + with self.assertRaises(errors.OutOfRangeError): + self.evaluate(next_element) + + # After re-initializing the iterator, should be able to iterate again. + self.evaluate(iterator.initializer) + for _ in range(10): + self.evaluate(next_element) + @test_util.with_c_api class MirroredVariableTest(test.TestCase): diff --git a/tensorflow/contrib/distributions/BUILD b/tensorflow/contrib/distributions/BUILD index fec6eafd4a66d38e9c99163b059bfeb81d8ad120..2d99e8172d220a697c3efaf07c6606874bd7dd01 100644 --- a/tensorflow/contrib/distributions/BUILD +++ b/tensorflow/contrib/distributions/BUILD @@ -877,6 +877,7 @@ cuda_py_test( "//tensorflow/python:framework_test_lib", "//tensorflow/python:platform_test", ], + tags = ["optonly"], ) cuda_py_test( @@ -1174,6 +1175,25 @@ cuda_py_test( ], ) +cuda_py_test( + name = "softsign_test", + size = "small", + srcs = ["python/kernel_tests/bijectors/softsign_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 = "square_test", size = "small", diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/absolute_value_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/absolute_value_test.py index e0d65c79b2654c2949de161d6317f218d11cab43..042c8ebd51c47facfc5c942cae56bd56be9df7c5 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/absolute_value_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/absolute_value_test.py @@ -18,11 +18,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import numpy as np - # pylint: disable=g-importing-member from tensorflow.contrib.distributions.python.ops.bijectors.absolute_value import AbsoluteValue -from tensorflow.python.framework import dtypes from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.platform import test @@ -35,50 +32,38 @@ class AbsoluteValueTest(test.TestCase): def testBijectorVersusNumpyRewriteOfBasicFunctionsEventNdims0(self): with self.test_session() as sess: - bijector = AbsoluteValue(event_ndims=0, validate_args=True) + bijector = AbsoluteValue(validate_args=True) self.assertEqual("absolute_value", bijector.name) x = array_ops.constant([[0., 1., -1], [0., -5., 3.]]) # Shape [2, 3] y = math_ops.abs(x) y_ = y.eval() - zeros = np.zeros((2, 3)) self.assertAllClose(y_, bijector.forward(x).eval()) self.assertAllClose((-y_, y_), sess.run(bijector.inverse(y))) - self.assertAllClose((zeros, zeros), - sess.run(bijector.inverse_log_det_jacobian(y))) + self.assertAllClose((0., 0.), + sess.run(bijector.inverse_log_det_jacobian( + y, event_ndims=0))) # Run things twice to make sure there are no issues in caching the tuples # returned by .inverse* self.assertAllClose(y_, bijector.forward(x).eval()) self.assertAllClose((-y_, y_), sess.run(bijector.inverse(y))) - self.assertAllClose((zeros, zeros), - sess.run(bijector.inverse_log_det_jacobian(y))) - - def testEventNdimsMustBeZeroOrRaiseStatic(self): - with self.test_session(): - with self.assertRaisesRegexp(ValueError, "event_ndims.*was not 0"): - AbsoluteValue(event_ndims=1) - - def testEventNdimsMustBeZeroOrRaiseDynamic(self): - with self.test_session() as sess: - event_ndims = array_ops.placeholder(dtypes.int32) - abs_bijector = AbsoluteValue(event_ndims=event_ndims, validate_args=True) - with self.assertRaisesOpError("event_ndims was not 0"): - sess.run(abs_bijector.inverse_log_det_jacobian([1.]), - feed_dict={event_ndims: 1}) + self.assertAllClose((0., 0.), + sess.run(bijector.inverse_log_det_jacobian( + y, event_ndims=0))) def testNegativeYRaisesForInverseIfValidateArgs(self): with self.test_session() as sess: - bijector = AbsoluteValue(event_ndims=0, validate_args=True) + bijector = AbsoluteValue(validate_args=True) with self.assertRaisesOpError("y was negative"): sess.run(bijector.inverse(-1.)) def testNegativeYRaisesForILDJIfValidateArgs(self): with self.test_session() as sess: - bijector = AbsoluteValue(event_ndims=0, validate_args=True) + bijector = AbsoluteValue(validate_args=True) with self.assertRaisesOpError("y was negative"): - sess.run(bijector.inverse_log_det_jacobian(-1.)) + sess.run(bijector.inverse_log_det_jacobian(-1., event_ndims=0)) if __name__ == "__main__": diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/affine_linear_operator_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/affine_linear_operator_test.py index 405ddd292cacd8ace87d6caeebf3e8cfc347c22d..1e4ad724d00f751a55370ef9aa6dde0003a2098c 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/affine_linear_operator_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/affine_linear_operator_test.py @@ -38,9 +38,11 @@ class AffineLinearOperatorTest(test.TestCase): self.assertEqual(affine.name, "affine_linear_operator") self.assertAllClose(y, affine.forward(x).eval()) self.assertAllClose(x, affine.inverse(y).eval()) - self.assertAllClose(ildj, affine.inverse_log_det_jacobian(y).eval()) - self.assertAllClose(-affine.inverse_log_det_jacobian(y).eval(), - affine.forward_log_det_jacobian(x).eval()) + self.assertAllClose(ildj, affine.inverse_log_det_jacobian( + y, event_ndims=2).eval()) + self.assertAllClose( + -affine.inverse_log_det_jacobian(y, event_ndims=2).eval(), + affine.forward_log_det_jacobian(x, event_ndims=2).eval()) def testDiag(self): with self.test_session(): @@ -58,14 +60,16 @@ class AffineLinearOperatorTest(test.TestCase): self.assertEqual(affine.name, "affine_linear_operator") self.assertAllClose(y, affine.forward(x).eval()) self.assertAllClose(x, affine.inverse(y).eval()) - self.assertAllClose(ildj, affine.inverse_log_det_jacobian(y).eval()) - self.assertAllClose(-affine.inverse_log_det_jacobian(y).eval(), - affine.forward_log_det_jacobian(x).eval()) + self.assertAllClose( + ildj, affine.inverse_log_det_jacobian(y, event_ndims=1).eval()) + self.assertAllClose( + -affine.inverse_log_det_jacobian(y, event_ndims=1).eval(), + affine.forward_log_det_jacobian(x, event_ndims=1).eval()) def testTriL(self): with self.test_session(): shift = np.array([-1, 0, 1], dtype=np.float32) - tril = np.array([[[1, 0, 0], + tril = np.array([[[3, 0, 0], [2, -1, 0], [3, 2, 1]], [[2, 0, 0], @@ -85,15 +89,17 @@ class AffineLinearOperatorTest(test.TestCase): # y = np.matmul(x, tril) + shift. y = np.squeeze(np.matmul(tril, np.expand_dims(x, -1)), -1) + shift ildj = -np.sum(np.log(np.abs(np.diagonal( - tril, axis1=-2, axis2=-1))), - axis=-1) + tril, axis1=-2, axis2=-1)))) self.assertEqual(affine.name, "affine_linear_operator") self.assertAllClose(y, affine.forward(x).eval()) self.assertAllClose(x, affine.inverse(y).eval()) - self.assertAllClose(ildj, affine.inverse_log_det_jacobian(y).eval()) - self.assertAllClose(-affine.inverse_log_det_jacobian(y).eval(), - affine.forward_log_det_jacobian(x).eval()) + self.assertAllClose( + ildj, affine.inverse_log_det_jacobian( + y, event_ndims=2).eval()) + self.assertAllClose( + -affine.inverse_log_det_jacobian(y, event_ndims=2).eval(), + affine.forward_log_det_jacobian(x, event_ndims=2).eval()) if __name__ == "__main__": 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 index 16173a166fd943413345036df12245c2a4ab8343..d2533620bebeb0400b6d4a6346e8315c7e37c5c6 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/affine_scalar_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/affine_scalar_test.py @@ -40,13 +40,13 @@ class AffineScalarBijectorTest(test.TestCase): def testNoBatchScalar(self): with self.test_session() as sess: - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): + def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) + return sess.run(fun(x, **kwargs), feed_dict={x: x_value}) for run in (static_run, dynamic_run): mu = -1. @@ -55,19 +55,20 @@ class AffineScalarBijectorTest(test.TestCase): 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)) + self.assertAllClose( + -np.log(2.), + run(bijector.inverse_log_det_jacobian, x, event_ndims=0)) def testOneBatchScalarViaIdentityIn64BitUserProvidesShiftOnly(self): with self.test_session() as sess: - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): + def dynamic_run(fun, x_value, **kwargs): 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}) + return sess.run(fun(x, **kwargs), feed_dict={x: x_value}) for run in (static_run, dynamic_run): mu = np.float64([1.]) @@ -77,18 +78,20 @@ class AffineScalarBijectorTest(test.TestCase): 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)) + self.assertAllClose( + 0., + run(bijector.inverse_log_det_jacobian, x, event_ndims=0)) def testOneBatchScalarViaIdentityIn64BitUserProvidesScaleOnly(self): with self.test_session() as sess: - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): + def dynamic_run(fun, x_value, **kwargs): 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}) + return sess.run(fun(x, **kwargs), feed_dict={x: x_value}) for run in (static_run, dynamic_run): multiplier = np.float64([2.]) @@ -98,19 +101,20 @@ class AffineScalarBijectorTest(test.TestCase): 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)) + self.assertAllClose( + [np.log(0.5)], + run(bijector.inverse_log_det_jacobian, x, event_ndims=0)) def testTwoBatchScalarIdentityViaIdentity(self): with self.test_session() as sess: - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): - x_value = np.array(x_value) + def dynamic_run(fun, x_value, **kwargs): + x_value = np.array(x_value).astype(np.float32) x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) + return sess.run(fun(x, **kwargs), feed_dict={x: x_value}) for run in (static_run, dynamic_run): mu = [1., -1] @@ -120,18 +124,20 @@ class AffineScalarBijectorTest(test.TestCase): 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)) + self.assertAllClose( + 0., + run(bijector.inverse_log_det_jacobian, x, event_ndims=0)) def testTwoBatchScalarIdentityViaScale(self): with self.test_session() as sess: - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): - x_value = np.array(x_value) + def dynamic_run(fun, x_value, **kwargs): + x_value = np.array(x_value).astype(np.float32) x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) + return sess.run(fun(x, **kwargs), feed_dict={x: x_value}) for run in (static_run, dynamic_run): mu = [1., -1] @@ -142,7 +148,8 @@ class AffineScalarBijectorTest(test.TestCase): 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)) + [-np.log(2), 0.], + run(bijector.inverse_log_det_jacobian, x, event_ndims=0)) def testScalarCongruency(self): with self.test_session(): 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 077e6176b4e7aecb28369d49edad6d1367cc7259..9e14b9a53e6c63876478d876030c476c5d77dbbb 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/affine_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/affine_test.py @@ -40,14 +40,15 @@ class AffineBijectorTest(test.TestCase): def testNoBatchMultivariateIdentity(self): with self.test_session() as sess: + placeholder = array_ops.placeholder(dtypes.float32, name="x") - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): + def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) + return sess.run( + fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = [1., -1] @@ -66,18 +67,20 @@ class AffineBijectorTest(test.TestCase): x = [[1., 1], [-1., -1]] self.assertAllClose([[2., 0], [0., -2]], run(bijector.forward, x)) self.assertAllClose([[0., 2], [-2., 0]], run(bijector.inverse, x)) - self.assertAllClose(0., run(bijector.inverse_log_det_jacobian, x)) + self.assertAllClose( + 0., run(bijector.inverse_log_det_jacobian, x, event_ndims=1)) def testNoBatchMultivariateDiag(self): with self.test_session() as sess: + placeholder = array_ops.placeholder(dtypes.float32, name="x") - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): + def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) + return sess.run( + fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = [1., -1] @@ -89,9 +92,12 @@ class AffineBijectorTest(test.TestCase): # = [-1, -1] + [1, -1] self.assertAllClose([3., 0], run(bijector.forward, x)) self.assertAllClose([0., 2], run(bijector.inverse, x)) - self.assertAllClose(-np.log(2.), - run(bijector.inverse_log_det_jacobian, x)) + self.assertAllClose( + -np.log(2.), + run(bijector.inverse_log_det_jacobian, x, event_ndims=1)) + # Reset bijector. + bijector = Affine(shift=mu, scale_diag=[2., 1]) # x is a 2-batch of 2-vectors. # The first vector is [1, 1], the second is [-1, -1]. # Each undergoes matmul(sigma, x) + shift. @@ -103,8 +109,9 @@ class AffineBijectorTest(test.TestCase): self.assertAllClose([[0., 2], [-1., 0]], run(bijector.inverse, x)) - self.assertAllClose(-np.log(2.), - run(bijector.inverse_log_det_jacobian, x)) + self.assertAllClose( + -np.log(2.), + run(bijector.inverse_log_det_jacobian, x, event_ndims=1)) def testNoBatchMultivariateFullDynamic(self): with self.test_session() as sess: @@ -126,18 +133,20 @@ class AffineBijectorTest(test.TestCase): self.assertAllClose([[0., 1]], sess.run(bijector.inverse(x), feed_dict)) self.assertAllClose( -np.log(4), - sess.run(bijector.inverse_log_det_jacobian(x), feed_dict)) + sess.run(bijector.inverse_log_det_jacobian(x, event_ndims=1), + feed_dict)) def testBatchMultivariateIdentity(self): with self.test_session() as sess: + placeholder = array_ops.placeholder(dtypes.float32, name="x") - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): - x_value = np.array(x_value, dtype=np.float32) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) + def dynamic_run(fun, x_value, **kwargs): + x_value = np.array(x_value) + return sess.run( + fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = [[1., -1]] @@ -147,19 +156,21 @@ class AffineBijectorTest(test.TestCase): x = [[[1., 1]]] self.assertAllClose([[[3., 1]]], run(bijector.forward, x)) self.assertAllClose([[[0., 1]]], run(bijector.inverse, x)) - self.assertAllClose(-np.log(4), - run(bijector.inverse_log_det_jacobian, x)) + self.assertAllClose( + -np.log(4), + run(bijector.inverse_log_det_jacobian, x, event_ndims=1)) def testBatchMultivariateDiag(self): with self.test_session() as sess: + placeholder = array_ops.placeholder(dtypes.float32, name="x") - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): - x_value = np.array(x_value, dtype=np.float32) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) + def dynamic_run(fun, x_value, **kwargs): + x_value = np.array(x_value) + return sess.run( + fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = [[1., -1]] @@ -169,8 +180,9 @@ class AffineBijectorTest(test.TestCase): x = [[[1., 1]]] self.assertAllClose([[[3., 1]]], run(bijector.forward, x)) self.assertAllClose([[[0., 1]]], run(bijector.inverse, x)) - self.assertAllClose([-np.log(4)], - run(bijector.inverse_log_det_jacobian, x)) + self.assertAllClose( + [-np.log(4)], + run(bijector.inverse_log_det_jacobian, x, event_ndims=1)) def testBatchMultivariateFullDynamic(self): with self.test_session() as sess: @@ -191,20 +203,22 @@ class AffineBijectorTest(test.TestCase): 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)], - sess.run( - bijector.inverse_log_det_jacobian(x), feed_dict)) + self.assertAllClose( + [-np.log(4)], + sess.run(bijector.inverse_log_det_jacobian( + x, event_ndims=1), feed_dict)) def testIdentityWithDiagUpdate(self): with self.test_session() as sess: + placeholder = array_ops.placeholder(dtypes.float32, name="x") - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): + def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) + return sess.run( + fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = -1. @@ -216,19 +230,21 @@ class AffineBijectorTest(test.TestCase): 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)) + self.assertAllClose( + -np.log(2.**3), + run(bijector.inverse_log_det_jacobian, x, event_ndims=1)) def testIdentityWithTriL(self): with self.test_session() as sess: + placeholder = array_ops.placeholder(dtypes.float32, name="x") - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): + def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) + return sess.run( + fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = -1. @@ -240,19 +256,21 @@ class AffineBijectorTest(test.TestCase): x = [[1., 2]] # One multivariate sample. self.assertAllClose([[1., 5]], run(bijector.forward, x)) self.assertAllClose([[1., 0.5]], run(bijector.inverse, x)) - self.assertAllClose(-np.log(4.), - run(bijector.inverse_log_det_jacobian, x)) + self.assertAllClose( + -np.log(4.), + run(bijector.inverse_log_det_jacobian, x, event_ndims=1)) def testDiagWithTriL(self): with self.test_session() as sess: + placeholder = array_ops.placeholder(dtypes.float32, name="x") - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): + def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) + return sess.run( + fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = -1. @@ -262,19 +280,21 @@ class AffineBijectorTest(test.TestCase): x = [[1., 2]] # One multivariate sample. self.assertAllClose([[1., 7]], run(bijector.forward, x)) self.assertAllClose([[1., 1 / 3.]], run(bijector.inverse, x)) - self.assertAllClose(-np.log(6.), - run(bijector.inverse_log_det_jacobian, x)) + self.assertAllClose( + -np.log(6.), + run(bijector.inverse_log_det_jacobian, x, event_ndims=1)) def testIdentityAndDiagWithTriL(self): with self.test_session() as sess: + placeholder = array_ops.placeholder(dtypes.float32, name="x") - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): + def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) + return sess.run( + fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = -1. @@ -287,19 +307,21 @@ class AffineBijectorTest(test.TestCase): x = [[1., 2]] # One multivariate sample. self.assertAllClose([[2., 9]], run(bijector.forward, x)) self.assertAllClose([[2 / 3., 5 / 12.]], run(bijector.inverse, x)) - self.assertAllClose(-np.log(12.), - run(bijector.inverse_log_det_jacobian, x)) + self.assertAllClose( + -np.log(12.), + run(bijector.inverse_log_det_jacobian, x, event_ndims=1)) def testIdentityWithVDVTUpdate(self): with self.test_session() as sess: + placeholder = array_ops.placeholder(dtypes.float32, name="x") - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): + def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) + return sess.run( + fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = -1. @@ -319,22 +341,24 @@ class AffineBijectorTest(test.TestCase): self.assertAllClose([0.2, 1.5, 4 / 3.], run(bijector.inverse, x)) self.assertAllClose( run(bijector_ref.inverse, x), run(bijector.inverse, x)) - self.assertAllClose(-np.log(60.), - run(bijector.inverse_log_det_jacobian, x)) self.assertAllClose( - run(bijector.inverse_log_det_jacobian, x), - run(bijector_ref.inverse_log_det_jacobian, x)) + -np.log(60.), + run(bijector.inverse_log_det_jacobian, x, event_ndims=1)) + self.assertAllClose( + run(bijector.inverse_log_det_jacobian, x, event_ndims=1), + run(bijector_ref.inverse_log_det_jacobian, x, event_ndims=1)) def testDiagWithVDVTUpdate(self): with self.test_session() as sess: + placeholder = array_ops.placeholder(dtypes.float32, name="x") - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): + def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) + return sess.run( + fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = -1. @@ -353,22 +377,24 @@ class AffineBijectorTest(test.TestCase): self.assertAllClose([0.2, 1., 0.8], run(bijector.inverse, x)) self.assertAllClose( run(bijector_ref.inverse, x), run(bijector.inverse, x)) - self.assertAllClose(-np.log(150.), - run(bijector.inverse_log_det_jacobian, x)) self.assertAllClose( - run(bijector.inverse_log_det_jacobian, x), - run(bijector_ref.inverse_log_det_jacobian, x)) + -np.log(150.), + run(bijector.inverse_log_det_jacobian, x, event_ndims=1)) + self.assertAllClose( + run(bijector.inverse_log_det_jacobian, x, event_ndims=1), + run(bijector_ref.inverse_log_det_jacobian, x, event_ndims=1)) def testTriLWithVDVTUpdate(self): with self.test_session() as sess: + placeholder = array_ops.placeholder(dtypes.float32, name="x") - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): + def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) + return sess.run( + fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = -1. @@ -388,22 +414,24 @@ class AffineBijectorTest(test.TestCase): self.assertAllClose([0.2, 14 / 15., 4 / 25.], run(bijector.inverse, x)) self.assertAllClose( run(bijector_ref.inverse, x), run(bijector.inverse, x)) - self.assertAllClose(-np.log(150.), - run(bijector.inverse_log_det_jacobian, x)) self.assertAllClose( - run(bijector.inverse_log_det_jacobian, x), - run(bijector_ref.inverse_log_det_jacobian, x)) + -np.log(150.), + run(bijector.inverse_log_det_jacobian, x, event_ndims=1)) + self.assertAllClose( + run(bijector.inverse_log_det_jacobian, x, event_ndims=1), + run(bijector_ref.inverse_log_det_jacobian, x, event_ndims=1)) def testTriLWithVDVTUpdateNoDiagonal(self): with self.test_session() as sess: + placeholder = array_ops.placeholder(dtypes.float32, name="x") - def static_run(fun, x): - return fun(x).eval() + def static_run(fun, x, **kwargs): + return fun(x, **kwargs).eval() - def dynamic_run(fun, x_value): + def dynamic_run(fun, x_value, **kwargs): x_value = np.array(x_value) - x = array_ops.placeholder(dtypes.float32, name="x") - return sess.run(fun(x), feed_dict={x: x_value}) + return sess.run( + fun(placeholder, **kwargs), feed_dict={placeholder: x_value}) for run in (static_run, dynamic_run): mu = -1. @@ -423,11 +451,12 @@ class AffineBijectorTest(test.TestCase): self.assertAllClose([1 / 3., 8 / 9., 4 / 30.], run(bijector.inverse, x)) self.assertAllClose( run(bijector_ref.inverse, x), run(bijector.inverse, x)) - self.assertAllClose(-np.log(90.), - run(bijector.inverse_log_det_jacobian, x)) self.assertAllClose( - run(bijector.inverse_log_det_jacobian, x), - run(bijector_ref.inverse_log_det_jacobian, x)) + -np.log(90.), + run(bijector.inverse_log_det_jacobian, x, event_ndims=1)) + self.assertAllClose( + run(bijector.inverse_log_det_jacobian, x, event_ndims=1), + run(bijector_ref.inverse_log_det_jacobian, x, event_ndims=1)) def testNoBatchMultivariateRaisesWhenSingular(self): with self.test_session(): @@ -530,6 +559,7 @@ class AffineBijectorTest(test.TestCase): backward = np.squeeze(backward, axis=-1) self.assertAllClose(backward, bijector.inverse(x).eval()) + scale *= np.ones(shape=x.shape[:-1], dtype=scale.dtype) ildj = -np.log(np.abs(np.linalg.det(scale))) # TODO(jvdillon): We need to make it so the scale_identity_multiplier # case does not deviate in expected shape. Fixing this will get rid of @@ -540,7 +570,8 @@ class AffineBijectorTest(test.TestCase): ildj = np.squeeze(ildj[0]) elif ildj.ndim < scale.ndim - 2: ildj = np.reshape(ildj, scale.shape[0:-2]) - self.assertAllClose(ildj, bijector.inverse_log_det_jacobian(x).eval()) + self.assertAllClose( + ildj, bijector.inverse_log_det_jacobian(x, event_ndims=1).eval()) def testLegalInputs(self): self._testLegalInputs( 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 index a215a4a2b1ffbea7951bdb9b4352ed567e0b1e41..c832fcaa686c92f83810e4f99ca3b23ae694b723 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/batch_normalization_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/batch_normalization_test.py @@ -83,10 +83,11 @@ class BatchNormTest(test_util.VectorDistributionTestHelpers, 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) + fldj = batch_norm.forward_log_det_jacobian( + x, event_ndims=len(event_dims)) # Use identity to invalidate cache. ildj = batch_norm.inverse_log_det_jacobian( - array_ops.identity(denorm_x)) + array_ops.identity(denorm_x), event_ndims=len(event_dims)) variables.global_variables_initializer().run() # Update variables. norm_x_ = sess.run(norm_x) diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/chain_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/chain_test.py index a748acd667e58f9b527bab11d8bc4d086996e9f3..ca20442c3940664feab7526110229872a6cdc41f 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/chain_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/chain_test.py @@ -20,21 +20,33 @@ from __future__ import print_function import numpy as np +from tensorflow.contrib.distributions.python.ops.bijectors.affine import Affine from tensorflow.contrib.distributions.python.ops.bijectors.chain import Chain from tensorflow.contrib.distributions.python.ops.bijectors.exp import Exp from tensorflow.contrib.distributions.python.ops.bijectors.softmax_centered import SoftmaxCentered from tensorflow.contrib.distributions.python.ops.bijectors.softplus import Softplus from tensorflow.python.framework import tensor_shape +from tensorflow.python.ops.distributions import bijector from tensorflow.python.ops.distributions.bijector_test_util import assert_scalar_congruency from tensorflow.python.platform import test +class ShapeChanging(bijector.Bijector): + """Only used for op_ndims manipulation.""" + + def __init__(self, forward_min_event_ndims=0, inverse_min_event_ndims=3): + super(ShapeChanging, self).__init__( + forward_min_event_ndims=forward_min_event_ndims, + inverse_min_event_ndims=inverse_min_event_ndims, + validate_args=False, name="shape_changer") + + class ChainBijectorTest(test.TestCase): """Tests the correctness of the Y = Chain(bij1, bij2, bij3) transformation.""" def testBijector(self): with self.test_session(): - chain = Chain((Exp(event_ndims=1), Softplus(event_ndims=1))) + chain = Chain((Exp(), Softplus())) self.assertEqual("chain_of_exp_of_softplus", chain.name) x = np.asarray([[[1., 2.], [2., 3.]]]) @@ -42,9 +54,10 @@ class ChainBijectorTest(test.TestCase): self.assertAllClose(np.log(x - 1.), chain.inverse(x).eval()) self.assertAllClose( -np.sum(np.log(x - 1.), axis=2), - chain.inverse_log_det_jacobian(x).eval()) + chain.inverse_log_det_jacobian(x, event_ndims=1).eval()) self.assertAllClose( - np.sum(x, axis=2), chain.forward_log_det_jacobian(x).eval()) + np.sum(x, axis=2), + chain.forward_log_det_jacobian(x, event_ndims=1).eval()) def testBijectorIdentity(self): with self.test_session(): @@ -54,31 +67,126 @@ class ChainBijectorTest(test.TestCase): [2., 3.]]]) self.assertAllClose(x, chain.forward(x).eval()) self.assertAllClose(x, chain.inverse(x).eval()) - self.assertAllClose(0., chain.inverse_log_det_jacobian(x).eval()) - self.assertAllClose(0., chain.forward_log_det_jacobian(x).eval()) + self.assertAllClose( + 0., chain.inverse_log_det_jacobian(x, event_ndims=1).eval()) + self.assertAllClose( + 0., chain.forward_log_det_jacobian(x, event_ndims=1).eval()) def testScalarCongruency(self): with self.test_session(): - bijector = Chain((Exp(), Softplus())) + chain = Chain((Exp(), Softplus())) assert_scalar_congruency( - bijector, lower_x=1e-3, upper_x=1.5, rtol=0.05) + chain, lower_x=1e-3, upper_x=1.5, rtol=0.05) def testShapeGetters(self): with self.test_session(): - bijector = Chain([ + chain = Chain([ SoftmaxCentered(validate_args=True), SoftmaxCentered(validate_args=True), ]) x = tensor_shape.TensorShape([1]) y = tensor_shape.TensorShape([2 + 1]) - self.assertAllEqual(y, bijector.forward_event_shape(x)) + self.assertAllEqual(y, chain.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)) + chain.forward_event_shape_tensor(x.as_list()).eval()) + self.assertAllEqual(x, chain.inverse_event_shape(y)) self.assertAllEqual( x.as_list(), - bijector.inverse_event_shape_tensor(y.as_list()).eval()) + chain.inverse_event_shape_tensor(y.as_list()).eval()) + + def testMinEventNdimsChain(self): + chain = Chain([Exp(), Exp(), Exp()]) + self.assertEqual(0, chain.forward_min_event_ndims) + self.assertEqual(0, chain.inverse_min_event_ndims) + + chain = Chain([Affine(), Affine(), Affine()]) + self.assertEqual(1, chain.forward_min_event_ndims) + self.assertEqual(1, chain.inverse_min_event_ndims) + + chain = Chain([Exp(), Affine()]) + self.assertEqual(1, chain.forward_min_event_ndims) + self.assertEqual(1, chain.inverse_min_event_ndims) + + chain = Chain([Affine(), Exp()]) + self.assertEqual(1, chain.forward_min_event_ndims) + self.assertEqual(1, chain.inverse_min_event_ndims) + + chain = Chain([Affine(), Exp(), Softplus(), Affine()]) + self.assertEqual(1, chain.forward_min_event_ndims) + self.assertEqual(1, chain.inverse_min_event_ndims) + + def testMinEventNdimsShapeChangingAddDims(self): + chain = Chain([ShapeChanging()]) + self.assertEqual(0, chain.forward_min_event_ndims) + self.assertEqual(3, chain.inverse_min_event_ndims) + + chain = Chain([ShapeChanging(), Affine()]) + self.assertEqual(1, chain.forward_min_event_ndims) + self.assertEqual(4, chain.inverse_min_event_ndims) + + chain = Chain([Affine(), ShapeChanging()]) + self.assertEqual(0, chain.forward_min_event_ndims) + self.assertEqual(3, chain.inverse_min_event_ndims) + + chain = Chain([ShapeChanging(), ShapeChanging()]) + self.assertEqual(0, chain.forward_min_event_ndims) + self.assertEqual(6, chain.inverse_min_event_ndims) + + def testMinEventNdimsShapeChangingRemoveDims(self): + chain = Chain([ShapeChanging(3, 0)]) + self.assertEqual(3, chain.forward_min_event_ndims) + self.assertEqual(0, chain.inverse_min_event_ndims) + + chain = Chain([ShapeChanging(3, 0), Affine()]) + self.assertEqual(3, chain.forward_min_event_ndims) + self.assertEqual(0, chain.inverse_min_event_ndims) + + chain = Chain([Affine(), ShapeChanging(3, 0)]) + self.assertEqual(4, chain.forward_min_event_ndims) + self.assertEqual(1, chain.inverse_min_event_ndims) + + chain = Chain([ShapeChanging(3, 0), ShapeChanging(3, 0)]) + self.assertEqual(6, chain.forward_min_event_ndims) + self.assertEqual(0, chain.inverse_min_event_ndims) + + def testMinEventNdimsShapeChangingAddRemoveDims(self): + chain = Chain([ + ShapeChanging(2, 1), + ShapeChanging(3, 0), + ShapeChanging(1, 2)]) + self.assertEqual(4, chain.forward_min_event_ndims) + self.assertEqual(1, chain.inverse_min_event_ndims) + + def testChainExpAffine(self): + scale_diag = np.array([1., 2., 3.], dtype=np.float32) + chain = Chain([Exp(), Affine(scale_diag=scale_diag)]) + x = [0., np.log(2., dtype=np.float32), np.log(3., dtype=np.float32)] + y = [1., 4., 27.] + self.assertAllClose(y, self.evaluate(chain.forward(x))) + self.assertAllClose(x, self.evaluate(chain.inverse(y))) + self.assertAllClose( + np.log(6, dtype=np.float32) + np.sum(scale_diag * x), + self.evaluate(chain.forward_log_det_jacobian(x, event_ndims=1))) + + self.assertAllClose( + -np.log(6, dtype=np.float32) - np.sum(scale_diag * x), + self.evaluate(chain.inverse_log_det_jacobian(y, event_ndims=1))) + + def testChainAffineExp(self): + scale_diag = np.array([1., 2., 3.], dtype=np.float32) + chain = Chain([Affine(scale_diag=scale_diag), Exp()]) + x = [0., np.log(2., dtype=np.float32), np.log(3., dtype=np.float32)] + y = [1., 4., 9.] + self.assertAllClose(y, self.evaluate(chain.forward(x))) + self.assertAllClose(x, self.evaluate(chain.inverse(y))) + self.assertAllClose( + np.log(6, dtype=np.float32) + np.sum(x), + self.evaluate(chain.forward_log_det_jacobian(x, event_ndims=1))) + + self.assertAllClose( + -np.log(6, dtype=np.float32) - np.sum(x), + self.evaluate(chain.inverse_log_det_jacobian(y, event_ndims=1))) if __name__ == "__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 f392e83d2c3da9dac43c2e87070e952ae2060b34..e281e81bdf0698c1f7b2f60fb27783dd1351773f 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 @@ -51,10 +51,13 @@ class CholeskyOuterProductBijectorTest(test.TestCase): 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) + ildj, bijector.inverse_log_det_jacobian( + y, event_ndims=2).eval(), atol=0., rtol=1e-7) self.assertAllClose( - -bijector.inverse_log_det_jacobian(y).eval(), - bijector.forward_log_det_jacobian(x).eval(), + -bijector.inverse_log_det_jacobian( + y, event_ndims=2).eval(), + bijector.forward_log_det_jacobian( + x, event_ndims=2).eval(), atol=0., rtol=1e-7) diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/conditional_bijector_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/conditional_bijector_test.py index 26e0d2a539c78540603281ae0f361987a7bf8d90..8b279ebcd908b6f375b35594ac5f3db9228a1e31 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/conditional_bijector_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/conditional_bijector_test.py @@ -27,7 +27,7 @@ class _TestBijector(ConditionalBijector): def __init__(self): super(_TestBijector, self).__init__( - event_ndims=0, + forward_min_event_ndims=0, graph_parents=[], is_constant_jacobian=True, validate_args=False, @@ -51,11 +51,15 @@ class ConditionalBijectorTest(test.TestCase): def testConditionalBijector(self): b = _TestBijector() - for name in ["forward", "inverse", "inverse_log_det_jacobian", - "forward_log_det_jacobian"]: + for name in ["forward", "inverse"]: method = getattr(b, name) with self.assertRaisesRegexp(ValueError, name + ".*b1.*b2"): - method(1.0, arg1="b1", arg2="b2") + method(1., arg1="b1", arg2="b2") + + for name in ["inverse_log_det_jacobian", "forward_log_det_jacobian"]: + method = getattr(b, name) + with self.assertRaisesRegexp(ValueError, name + ".*b1.*b2"): + method(1., event_ndims=0., arg1="b1", arg2="b2") if __name__ == "__main__": diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/exp_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/exp_test.py index 9970c0b4d86afda188d9401ebaf3c98d3fffbfdf..7be939cd274e6f0e33c9b01c82494755db2caa73 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/exp_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/exp_test.py @@ -31,17 +31,21 @@ class ExpBijectorTest(test.TestCase): def testBijector(self): with self.test_session(): - bijector = Exp(event_ndims=1) + bijector = Exp() self.assertEqual("exp", bijector.name) x = [[[1.], [2.]]] y = np.exp(x) self.assertAllClose(y, bijector.forward(x).eval()) self.assertAllClose(x, bijector.inverse(y).eval()) self.assertAllClose( - -np.sum(np.log(y), axis=-1), - bijector.inverse_log_det_jacobian(y).eval()) - self.assertAllClose(-bijector.inverse_log_det_jacobian(np.exp(x)).eval(), - bijector.forward_log_det_jacobian(x).eval()) + -np.squeeze(np.log(y), axis=-1), + bijector.inverse_log_det_jacobian( + y, event_ndims=1).eval()) + self.assertAllClose( + -bijector.inverse_log_det_jacobian( + np.exp(x), event_ndims=1).eval(), + bijector.forward_log_det_jacobian( + x, event_ndims=1).eval()) def testScalarCongruency(self): with self.test_session(): @@ -51,10 +55,10 @@ class ExpBijectorTest(test.TestCase): def testBijectiveAndFinite(self): with self.test_session(): - bijector = Exp(event_ndims=0) + bijector = Exp() x = np.linspace(-10, 10, num=10).astype(np.float32) y = np.logspace(-10, 10, num=10).astype(np.float32) - assert_bijective_and_finite(bijector, x, y) + assert_bijective_and_finite(bijector, x, y, event_ndims=0) if __name__ == "__main__": diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/gumbel_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/gumbel_test.py index 9a905980c7581a86bbcda8c6c726da57c09fe4f8..54e54c3296a89a4fe29a3cce971760502b65e784 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/gumbel_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/gumbel_test.py @@ -34,7 +34,7 @@ class GumbelBijectorTest(test.TestCase): with self.test_session(): loc = 0.3 scale = 5. - bijector = Gumbel(loc=loc, scale=scale, event_ndims=1, validate_args=True) + bijector = Gumbel(loc=loc, scale=scale, validate_args=True) self.assertEqual("gumbel", bijector.name) x = np.array([[[-3.], [0.], [0.5], [4.2], [12.]]], dtype=np.float32) # Gumbel distribution @@ -43,13 +43,11 @@ class GumbelBijectorTest(test.TestCase): self.assertAllClose(y, bijector.forward(x).eval()) self.assertAllClose(x, bijector.inverse(y).eval()) self.assertAllClose( - # We should lose a dimension from calculating the determinant of the - # jacobian. - np.squeeze(gumbel_dist.logpdf(x), axis=2), - bijector.forward_log_det_jacobian(x).eval()) + np.squeeze(gumbel_dist.logpdf(x), axis=-1), + bijector.forward_log_det_jacobian(x, event_ndims=1).eval()) self.assertAllClose( - -bijector.inverse_log_det_jacobian(y).eval(), - bijector.forward_log_det_jacobian(x).eval(), + -bijector.inverse_log_det_jacobian(y, event_ndims=1).eval(), + bijector.forward_log_det_jacobian(x, event_ndims=1).eval(), rtol=1e-4, atol=0.) @@ -60,10 +58,10 @@ class GumbelBijectorTest(test.TestCase): def testBijectiveAndFinite(self): with self.test_session(): - bijector = Gumbel(loc=0., scale=3.0, event_ndims=0, validate_args=True) + bijector = Gumbel(loc=0., scale=3.0, validate_args=True) x = np.linspace(-10., 10., num=10).astype(np.float32) y = np.linspace(0.01, 0.99, num=10).astype(np.float32) - assert_bijective_and_finite(bijector, x, y, rtol=1e-3) + assert_bijective_and_finite(bijector, x, y, event_ndims=0, rtol=1e-3) if __name__ == "__main__": diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/inline_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/inline_test.py index 739fa6d439a8bce993ab1b4601489d9bbcd69bee..7d3bd758cd2db307f95d2d934923ea2133dc1217 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/inline_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/inline_test.py @@ -33,15 +33,13 @@ class InlineBijectorTest(test.TestCase): def testBijector(self): with self.test_session(): - exp = Exp(event_ndims=1) + exp = Exp() inline = Inline( forward_fn=math_ops.exp, inverse_fn=math_ops.log, - inverse_log_det_jacobian_fn=( - lambda y: -math_ops.reduce_sum( # pylint: disable=g-long-lambda - math_ops.log(y), reduction_indices=-1)), - forward_log_det_jacobian_fn=( - lambda x: math_ops.reduce_sum(x, reduction_indices=-1)), + inverse_log_det_jacobian_fn=lambda y: -math_ops.log(y), + forward_log_det_jacobian_fn=lambda x: x, + forward_min_event_ndims=0, name="exp") self.assertEqual(exp.name, inline.name) @@ -51,9 +49,10 @@ class InlineBijectorTest(test.TestCase): self.assertAllClose(x, inline.inverse(y).eval()) self.assertAllClose( -np.sum(np.log(y), axis=-1), - inline.inverse_log_det_jacobian(y).eval()) - self.assertAllClose(-inline.inverse_log_det_jacobian(y).eval(), - inline.forward_log_det_jacobian(x).eval()) + inline.inverse_log_det_jacobian(y, event_ndims=1).eval()) + self.assertAllClose( + -inline.inverse_log_det_jacobian(y, event_ndims=1).eval(), + inline.forward_log_det_jacobian(x, event_ndims=1).eval()) def testShapeGetters(self): with self.test_session(): @@ -62,6 +61,7 @@ class InlineBijectorTest(test.TestCase): forward_event_shape_fn=lambda x: x.as_list() + [1], inverse_event_shape_tensor_fn=lambda x: x[:-1], inverse_event_shape_fn=lambda x: x[:-1], + forward_min_event_ndims=0, name="shape_only") x = tensor_shape.TensorShape([1, 2, 3]) y = tensor_shape.TensorShape([1, 2, 3, 1]) 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 58ba9cedb1437df4e000ce32fe39664afa76c3b5..8b14c8327f08902044f50483f9f8dfe67b58cd70 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/invert_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/invert_test.py @@ -34,9 +34,9 @@ class InvertBijectorTest(test.TestCase): with self.test_session(): for fwd in [ bijectors.Identity(), - bijectors.Exp(event_ndims=1), + bijectors.Exp(), bijectors.Affine(shift=[0., 1.], scale_diag=[2., 3.]), - bijectors.Softplus(event_ndims=1), + bijectors.Softplus(), bijectors.SoftmaxCentered(), ]: rev = bijectors.Invert(fwd) @@ -46,11 +46,11 @@ class InvertBijectorTest(test.TestCase): 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()) + fwd.forward_log_det_jacobian(x, event_ndims=1).eval(), + rev.inverse_log_det_jacobian(x, event_ndims=1).eval()) self.assertAllClose( - fwd.inverse_log_det_jacobian(x).eval(), - rev.forward_log_det_jacobian(x).eval()) + fwd.inverse_log_det_jacobian(x, event_ndims=1).eval(), + rev.forward_log_det_jacobian(x, event_ndims=1).eval()) def testScalarCongruency(self): with self.test_session(): 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 index 074b5f275d107fa49de42df262476bd4aa48ffae..a8089881f684db9f8876d6dd738e52bf2f1f7606 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/kumaraswamy_bijector_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/kumaraswamy_bijector_test.py @@ -34,8 +34,7 @@ class KumaraswamyBijectorTest(test.TestCase): a = 2. b = 0.3 bijector = Kumaraswamy( - concentration1=a, concentration0=b, - event_ndims=0, validate_args=True) + concentration1=a, concentration0=b, 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). @@ -46,13 +45,11 @@ class KumaraswamyBijectorTest(test.TestCase): (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()) + np.squeeze(kumaraswamy_log_pdf, axis=-1), + bijector.inverse_log_det_jacobian(x, event_ndims=1).eval()) self.assertAllClose( - -bijector.inverse_log_det_jacobian(x).eval(), - bijector.forward_log_det_jacobian(y).eval(), + -bijector.inverse_log_det_jacobian(x, event_ndims=1).eval(), + bijector.forward_log_det_jacobian(y, event_ndims=1).eval(), rtol=1e-4, atol=0.) @@ -73,7 +70,7 @@ class KumaraswamyBijectorTest(test.TestCase): # 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) + assert_bijective_and_finite(bijector, x, y, event_ndims=0, rtol=1e-3) if __name__ == "__main__": diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/masked_autoregressive_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/masked_autoregressive_test.py index dcfb0eb05185d36d96947905c2eb91b2201aece1..5ba5a2083bf11791d7d58146dc2e6283b524d241 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/masked_autoregressive_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/masked_autoregressive_test.py @@ -79,9 +79,10 @@ class MaskedAutoregressiveFlowTest(test_util.VectorDistributionTestHelpers, forward_x = ma.forward(x) # Use identity to invalidate cache. inverse_y = ma.inverse(array_ops.identity(forward_x)) - fldj = ma.forward_log_det_jacobian(x) + fldj = ma.forward_log_det_jacobian(x, event_ndims=1) # Use identity to invalidate cache. - ildj = ma.inverse_log_det_jacobian(array_ops.identity(forward_x)) + ildj = ma.inverse_log_det_jacobian( + array_ops.identity(forward_x), event_ndims=1) variables.global_variables_initializer().run() [ forward_x_, diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/permute_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/permute_test.py index 54590de373441c32cc3214cb04d45cfc2d1807ed..7eef4ab599951bbb624652f13a0091363b36b93d 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/permute_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/permute_test.py @@ -53,8 +53,8 @@ class PermuteBijectorTest(test.TestCase): bijector.permutation, bijector.inverse(expected_y), bijector.forward(expected_x), - bijector.forward_log_det_jacobian(expected_x), - bijector.inverse_log_det_jacobian(expected_y), + bijector.forward_log_det_jacobian(expected_x, event_ndims=1), + bijector.inverse_log_det_jacobian(expected_y, event_ndims=1), ], feed_dict={permutation_ph: expected_permutation}) self.assertEqual("permute", bijector.name) self.assertAllEqual(expected_permutation, permutation_) @@ -78,10 +78,9 @@ class PermuteBijectorTest(test.TestCase): x = np.random.randn(4, 2, 3) y = x[..., permutation] with self.test_session(): - bijector = Permute( - permutation=permutation, - validate_args=True) - assert_bijective_and_finite(bijector, x, y, rtol=1e-6, atol=0) + bijector = Permute(permutation=permutation, validate_args=True) + assert_bijective_and_finite( + bijector, x, y, event_ndims=1, rtol=1e-6, atol=0) if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/power_transform_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/power_transform_test.py index de1659aa9f4d0f7d19ec2e8185715573b78eaf2b..85d22830132816cd6c77cd0b07870f3a22ae9798 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/power_transform_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/power_transform_test.py @@ -32,8 +32,7 @@ class PowerTransformBijectorTest(test.TestCase): def testBijector(self): with self.test_session(): c = 0.2 - bijector = PowerTransform( - power=c, event_ndims=1, validate_args=True) + bijector = PowerTransform(power=c, validate_args=True) self.assertEqual("power_transform", bijector.name) x = np.array([[[-1.], [2.], [-5. + 1e-4]]]) y = (1. + x * c)**(1. / c) @@ -41,27 +40,25 @@ class PowerTransformBijectorTest(test.TestCase): self.assertAllClose(x, bijector.inverse(y).eval()) self.assertAllClose( (c - 1.) * np.sum(np.log(y), axis=-1), - bijector.inverse_log_det_jacobian(y).eval()) + bijector.inverse_log_det_jacobian(y, event_ndims=1).eval()) self.assertAllClose( - -bijector.inverse_log_det_jacobian(y).eval(), - bijector.forward_log_det_jacobian(x).eval(), + -bijector.inverse_log_det_jacobian(y, event_ndims=1).eval(), + bijector.forward_log_det_jacobian(x, event_ndims=1).eval(), rtol=1e-4, atol=0.) def testScalarCongruency(self): with self.test_session(): - bijector = PowerTransform( - power=0.2, validate_args=True) + bijector = PowerTransform(power=0.2, validate_args=True) assert_scalar_congruency( bijector, lower_x=-2., upper_x=1.5, rtol=0.05) def testBijectiveAndFinite(self): with self.test_session(): - bijector = PowerTransform( - power=0.2, event_ndims=0, validate_args=True) + bijector = PowerTransform(power=0.2, validate_args=True) x = np.linspace(-4.999, 10, num=10).astype(np.float32) y = np.logspace(0.001, 10, num=10).astype(np.float32) - assert_bijective_and_finite(bijector, x, y, rtol=1e-3) + assert_bijective_and_finite(bijector, x, y, event_ndims=0, rtol=1e-3) if __name__ == "__main__": diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/real_nvp_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/real_nvp_test.py index 46fe7797419a9906ecdad60dd0dfe1e9d7c743ed..2d52895fbe0967cdd2260d6d298a291286858d09 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/real_nvp_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/real_nvp_test.py @@ -52,24 +52,28 @@ class RealNVPTest(test_util.VectorDistributionTestHelpers, test.TestCase): forward_x = nvp.forward(x) # Use identity to invalidate cache. inverse_y = nvp.inverse(array_ops.identity(forward_x)) - fldj = nvp.forward_log_det_jacobian(x) + forward_inverse_y = nvp.forward(inverse_y) + fldj = nvp.forward_log_det_jacobian(x, event_ndims=1) # Use identity to invalidate cache. - ildj = nvp.inverse_log_det_jacobian(array_ops.identity(forward_x)) + ildj = nvp.inverse_log_det_jacobian( + array_ops.identity(forward_x), event_ndims=1) variables.global_variables_initializer().run() [ forward_x_, inverse_y_, + forward_inverse_y_, ildj_, fldj_, ] = sess.run([ forward_x, inverse_y, + forward_inverse_y, ildj, fldj, ]) self.assertEqual("real_nvp", nvp.name) - self.assertAllClose(forward_x_, forward_x_, rtol=1e-6, atol=0.) - self.assertAllClose(x_, inverse_y_, rtol=1e-5, atol=0.) + self.assertAllClose(forward_x_, forward_inverse_y_, rtol=1e-1, atol=0.) + self.assertAllClose(x_, inverse_y_, rtol=1e-1, atol=0.) self.assertAllClose(ildj_, -fldj_, rtol=1e-6, atol=0.) def testMutuallyConsistent(self): diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/reshape_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/reshape_test.py index e216d88cb190dc16fc0056186f80817d6f2d7c67..46f2c63f9b0f78b25bb1948e6ea55ab20c5cfa6e 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/reshape_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/reshape_test.py @@ -65,8 +65,8 @@ class _ReshapeBijectorTest(object): ildj_) = sess.run(( bijector.inverse(expected_y), bijector.forward(expected_x), - bijector.forward_log_det_jacobian(expected_x), - bijector.inverse_log_det_jacobian(expected_y), + bijector.forward_log_det_jacobian(expected_x, event_ndims=2), + bijector.inverse_log_det_jacobian(expected_y, event_ndims=2), ), feed_dict=feed_dict) self.assertEqual("reshape", bijector.name) self.assertAllClose(expected_y, y_, rtol=1e-6, atol=0) @@ -301,7 +301,8 @@ class ReshapeBijectorTestStatic(test.TestCase, _ReshapeBijectorTest): event_shape_in=[2, 3], event_shape_out=[1, 2, 3], validate_args=True) - assert_bijective_and_finite(bijector, x, y, rtol=1e-6, atol=0) + assert_bijective_and_finite( + bijector, x, y, event_ndims=2, rtol=1e-6, atol=0) def testInvalidDimensionsOpError(self): if ops._USE_C_API: diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/sigmoid_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/sigmoid_test.py index e4f9d72785c301284812a48c0a67614ca439ffae..cea4a62c22af5d98d38ee881b29c773e6a27a4b4 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/sigmoid_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/sigmoid_test.py @@ -36,12 +36,13 @@ class SigmoidBijectorTest(test.TestCase): x = np.linspace(-10., 10., 100).reshape([2, 5, 10]).astype(np.float32) y = special.expit(x) ildj = -np.log(y) - np.log1p(-y) - self.assertAllClose(y, Sigmoid().forward(x).eval(), atol=0., rtol=1e-2) - self.assertAllClose(x, Sigmoid().inverse(y).eval(), atol=0., rtol=1e-4) - self.assertAllClose(ildj, Sigmoid().inverse_log_det_jacobian(y).eval(), - atol=0., rtol=1e-6) - self.assertAllClose(-ildj, Sigmoid().forward_log_det_jacobian(x).eval(), - atol=0., rtol=1e-4) + bijector = Sigmoid() + self.assertAllClose(y, bijector.forward(x).eval(), atol=0., rtol=1e-2) + self.assertAllClose(x, bijector.inverse(y).eval(), atol=0., rtol=1e-4) + self.assertAllClose(ildj, bijector.inverse_log_det_jacobian( + y, event_ndims=0).eval(), atol=0., rtol=1e-6) + self.assertAllClose(-ildj, bijector.forward_log_det_jacobian( + x, event_ndims=0).eval(), atol=0., rtol=1e-4) def testScalarCongruency(self): with self.test_session(): @@ -52,7 +53,8 @@ class SigmoidBijectorTest(test.TestCase): x = np.linspace(-7., 7., 100).astype(np.float32) eps = 1e-3 y = np.linspace(eps, 1. - eps, 100).astype(np.float32) - assert_bijective_and_finite(Sigmoid(), x, y, atol=0., rtol=1e-4) + assert_bijective_and_finite( + Sigmoid(), x, y, event_ndims=0, atol=0., rtol=1e-4) if __name__ == "__main__": diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/sinh_arcsinh_bijector_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/sinh_arcsinh_bijector_test.py index 172c180a44229089f06f250a872bc47a89991cf0..45760a29ee42835da69ef63803ccec7ce82a5a8f 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/sinh_arcsinh_bijector_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/sinh_arcsinh_bijector_test.py @@ -39,7 +39,6 @@ class SinhArcsinhBijectorTest(test.TestCase): bijector = SinhArcsinh( skewness=skewness, tailweight=tailweight, - event_ndims=1, validate_args=True) self.assertEqual("SinhArcsinh", bijector.name) x = np.array([[[-2.01], [2.], [1e-4]]]).astype(np.float32) @@ -50,10 +49,11 @@ class SinhArcsinhBijectorTest(test.TestCase): np.sum( np.log(np.cosh(np.arcsinh(y) / tailweight - skewness)) - np.log(tailweight) - np.log(np.sqrt(y**2 + 1)), - axis=-1), bijector.inverse_log_det_jacobian(y).eval()) + axis=-1), + bijector.inverse_log_det_jacobian(y, event_ndims=1).eval()) self.assertAllClose( - -bijector.inverse_log_det_jacobian(y).eval(), - bijector.forward_log_det_jacobian(x).eval(), + -bijector.inverse_log_det_jacobian(y, event_ndims=1).eval(), + bijector.forward_log_det_jacobian(x, event_ndims=1).eval(), rtol=1e-4, atol=0.) @@ -106,14 +106,15 @@ class SinhArcsinhBijectorTest(test.TestCase): bijector = SinhArcsinh(skewness=-1., tailweight=0.5, validate_args=True) x = np.concatenate((-np.logspace(-2, 10, 1000), [0], np.logspace( -2, 10, 1000))).astype(np.float32) - assert_bijective_and_finite(bijector, x, x, rtol=1e-3) + assert_bijective_and_finite(bijector, x, x, event_ndims=0, rtol=1e-3) def testBijectiveAndFiniteSkewness1Tailweight3(self): with self.test_session(): bijector = SinhArcsinh(skewness=1., tailweight=3., validate_args=True) x = np.concatenate((-np.logspace(-2, 5, 1000), [0], np.logspace( -2, 5, 1000))).astype(np.float32) - assert_bijective_and_finite(bijector, x, x, rtol=1e-3) + assert_bijective_and_finite( + bijector, x, x, event_ndims=0, rtol=1e-3) def testBijectorEndpoints(self): with self.test_session(): @@ -124,7 +125,8 @@ class SinhArcsinhBijectorTest(test.TestCase): [np.finfo(dtype).min, np.finfo(dtype).max], dtype=dtype) # Note that the above bijector is the identity bijector. Hence, the # log_det_jacobian will be 0. Because of this we use atol. - assert_bijective_and_finite(bijector, bounds, bounds, atol=2e-6) + assert_bijective_and_finite( + bijector, bounds, bounds, event_ndims=0, atol=2e-6) def testBijectorOverRange(self): with self.test_session(): @@ -156,12 +158,12 @@ class SinhArcsinhBijectorTest(test.TestCase): np.arcsinh(y_float128) / tailweight - skewness) / np.sqrt( y_float128**2 + 1)) - np.log(tailweight), - bijector.inverse_log_det_jacobian(y).eval(), + bijector.inverse_log_det_jacobian(y, event_ndims=0).eval(), rtol=1e-4, atol=0.) self.assertAllClose( - -bijector.inverse_log_det_jacobian(y).eval(), - bijector.forward_log_det_jacobian(x).eval(), + -bijector.inverse_log_det_jacobian(y, event_ndims=0).eval(), + bijector.forward_log_det_jacobian(x, event_ndims=0).eval(), rtol=1e-4, atol=0.) 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 cad4dd1ac8de0da6405aacb9047714b37eec73e3..0f0a2fa531a0585a709df4c2c3e2631e5c275986 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 @@ -44,12 +44,12 @@ class SoftmaxCenteredBijectorTest(test.TestCase): self.assertAllClose(x, softmax.inverse(y).eval()) self.assertAllClose( -np.sum(np.log(y), axis=1), - softmax.inverse_log_det_jacobian(y).eval(), + softmax.inverse_log_det_jacobian(y, event_ndims=1).eval(), atol=0., rtol=1e-7) self.assertAllClose( - -softmax.inverse_log_det_jacobian(y).eval(), - softmax.forward_log_det_jacobian(x).eval(), + -softmax.inverse_log_det_jacobian(y, event_ndims=1).eval(), + softmax.forward_log_det_jacobian(x, event_ndims=1).eval(), atol=0., rtol=1e-7) @@ -67,14 +67,14 @@ class SoftmaxCenteredBijectorTest(test.TestCase): feed_dict={y: real_y})) self.assertAllClose( -np.sum(np.log(real_y), axis=1), - softmax.inverse_log_det_jacobian(y).eval( + softmax.inverse_log_det_jacobian(y, event_ndims=1).eval( feed_dict={y: real_y}), atol=0., rtol=1e-7) self.assertAllClose( - -softmax.inverse_log_det_jacobian(y).eval( + -softmax.inverse_log_det_jacobian(y, event_ndims=1).eval( feed_dict={y: real_y}), - softmax.forward_log_det_jacobian(x).eval( + softmax.forward_log_det_jacobian(x, event_ndims=1).eval( feed_dict={x: real_x}), atol=0., rtol=1e-7) @@ -104,7 +104,7 @@ class SoftmaxCenteredBijectorTest(test.TestCase): y = np.array([y_0, y_1, y_2]) y /= y.sum(axis=0) y = y.T # y.shape = [5, 3] - assert_bijective_and_finite(softmax, x, y) + assert_bijective_and_finite(softmax, x, y, event_ndims=1) if __name__ == "__main__": diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/softplus_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/softplus_test.py index d9af9aec50d3d69bb10f69f2ffd6ca3a24c316f8..3d8a0a32bba3539f732140e8eb7ebeb532d73ff5 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/softplus_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/softplus_test.py @@ -43,13 +43,13 @@ class SoftplusBijectorTest(test.TestCase): def testHingeSoftnessZeroRaises(self): with self.test_session(): - bijector = Softplus(event_ndims=0, hinge_softness=0., validate_args=True) + bijector = Softplus(hinge_softness=0., validate_args=True) with self.assertRaisesOpError("must be non-zero"): bijector.forward([1., 1.]).eval() def testBijectorForwardInverseEventDimsZero(self): with self.test_session(): - bijector = Softplus(event_ndims=0) + bijector = Softplus() self.assertEqual("softplus", bijector.name) x = 2 * rng.randn(2, 10) y = self._softplus(x) @@ -59,7 +59,7 @@ class SoftplusBijectorTest(test.TestCase): def testBijectorForwardInverseWithHingeSoftnessEventDimsZero(self): with self.test_session(): - bijector = Softplus(event_ndims=0, hinge_softness=1.5) + bijector = Softplus(hinge_softness=1.5) x = 2 * rng.randn(2, 10) y = 1.5 * self._softplus(x / 1.5) @@ -68,16 +68,17 @@ class SoftplusBijectorTest(test.TestCase): def testBijectorLogDetJacobianEventDimsZero(self): with self.test_session(): - bijector = Softplus(event_ndims=0) + bijector = Softplus() y = 2 * rng.rand(2, 10) # No reduction needed if event_dims = 0. ildj = self._softplus_ildj_before_reduction(y) - self.assertAllClose(ildj, bijector.inverse_log_det_jacobian(y).eval()) + self.assertAllClose(ildj, bijector.inverse_log_det_jacobian( + y, event_ndims=0).eval()) def testBijectorForwardInverseEventDimsOne(self): with self.test_session(): - bijector = Softplus(event_ndims=1) + bijector = Softplus() self.assertEqual("softplus", bijector.name) x = 2 * rng.randn(2, 10) y = self._softplus(x) @@ -87,58 +88,59 @@ class SoftplusBijectorTest(test.TestCase): def testBijectorLogDetJacobianEventDimsOne(self): with self.test_session(): - bijector = Softplus(event_ndims=1) + bijector = Softplus() y = 2 * rng.rand(2, 10) ildj_before = self._softplus_ildj_before_reduction(y) ildj = np.sum(ildj_before, axis=1) - self.assertAllClose(ildj, bijector.inverse_log_det_jacobian(y).eval()) + self.assertAllClose(ildj, bijector.inverse_log_det_jacobian( + y, event_ndims=1).eval()) def testScalarCongruency(self): with self.test_session(): - bijector = Softplus(event_ndims=0) + bijector = Softplus() assert_scalar_congruency( bijector, lower_x=-2., upper_x=2.) def testScalarCongruencyWithPositiveHingeSoftness(self): with self.test_session(): - bijector = Softplus(event_ndims=0, hinge_softness=1.3) + bijector = Softplus(hinge_softness=1.3) assert_scalar_congruency( bijector, lower_x=-2., upper_x=2.) def testScalarCongruencyWithNegativeHingeSoftness(self): with self.test_session(): - bijector = Softplus(event_ndims=0, hinge_softness=-1.3) + bijector = Softplus(hinge_softness=-1.3) assert_scalar_congruency( bijector, lower_x=-2., upper_x=2.) def testBijectiveAndFinite32bit(self): with self.test_session(): - bijector = Softplus(event_ndims=0) + bijector = Softplus() x = np.linspace(-20., 20., 100).astype(np.float32) y = np.logspace(-10, 10, 100).astype(np.float32) assert_bijective_and_finite( - bijector, x, y, rtol=1e-2, atol=1e-2) + bijector, x, y, event_ndims=0, rtol=1e-2, atol=1e-2) def testBijectiveAndFiniteWithPositiveHingeSoftness32Bit(self): with self.test_session(): - bijector = Softplus(event_ndims=0, hinge_softness=1.23) + bijector = Softplus(hinge_softness=1.23) x = np.linspace(-20., 20., 100).astype(np.float32) y = np.logspace(-10, 10, 100).astype(np.float32) assert_bijective_and_finite( - bijector, x, y, rtol=1e-2, atol=1e-2) + bijector, x, y, event_ndims=0, rtol=1e-2, atol=1e-2) def testBijectiveAndFiniteWithNegativeHingeSoftness32Bit(self): with self.test_session(): - bijector = Softplus(event_ndims=0, hinge_softness=-0.7) + bijector = Softplus(hinge_softness=-0.7) x = np.linspace(-20., 20., 100).astype(np.float32) y = -np.logspace(-10, 10, 100).astype(np.float32) assert_bijective_and_finite( - bijector, x, y, rtol=1e-2, atol=1e-2) + bijector, x, y, event_ndims=0, rtol=1e-2, atol=1e-2) def testBijectiveAndFinite16bit(self): with self.test_session(): - bijector = Softplus(event_ndims=0) + bijector = Softplus() # softplus(-20) is zero, so we can't use such a large range as in 32bit. x = np.linspace(-10., 20., 100).astype(np.float16) # Note that float16 is only in the open set (0, inf) for a smaller @@ -146,7 +148,7 @@ class SoftplusBijectorTest(test.TestCase): # for the test. y = np.logspace(-6, 3, 100).astype(np.float16) assert_bijective_and_finite( - bijector, x, y, rtol=1e-1, atol=1e-3) + bijector, x, y, event_ndims=0, rtol=1e-1, atol=1e-3) if __name__ == "__main__": diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/softsign_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/softsign_test.py new file mode 100644 index 0000000000000000000000000000000000000000..2ac06fce55b448a5f3da7ccb7f8766b5b1404ad7 --- /dev/null +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/softsign_test.py @@ -0,0 +1,111 @@ +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""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.bijectors.softsign import Softsign +from tensorflow.python.framework import test_util +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 SoftsignBijectorTest(test.TestCase): + """Tests the correctness of the Y = g(X) = X / (1 + |X|) transformation.""" + + def _softsign(self, x): + return x / (1. + np.abs(x)) + + def _softsign_ildj_before_reduction(self, y): + """Inverse log det jacobian, before being reduced.""" + return -2. * np.log1p(-np.abs(y)) + + def setUp(self): + self._rng = np.random.RandomState(42) + + @test_util.run_in_graph_and_eager_modes() + def testBijectorBounds(self): + bijector = Softsign(validate_args=True) + with self.test_session(): + with self.assertRaisesOpError("greater than -1"): + bijector.inverse(-3.).eval() + with self.assertRaisesOpError("greater than -1"): + bijector.inverse_log_det_jacobian(-3., event_ndims=0).eval() + + with self.assertRaisesOpError("less than 1"): + bijector.inverse(3.).eval() + with self.assertRaisesOpError("less than 1"): + bijector.inverse_log_det_jacobian(3., event_ndims=0).eval() + + @test_util.run_in_graph_and_eager_modes() + def testBijectorForwardInverse(self): + bijector = Softsign(validate_args=True) + self.assertEqual("softsign", bijector.name) + x = 2. * self._rng.randn(2, 10) + y = self._softsign(x) + + self.assertAllClose(y, self.evaluate(bijector.forward(x))) + self.assertAllClose(x, self.evaluate(bijector.inverse(y))) + + @test_util.run_in_graph_and_eager_modes() + def testBijectorLogDetJacobianEventDimsZero(self): + bijector = Softsign(validate_args=True) + y = self._rng.rand(2, 10) + # No reduction needed if event_dims = 0. + ildj = self._softsign_ildj_before_reduction(y) + + self.assertAllClose(ildj, self.evaluate( + bijector.inverse_log_det_jacobian(y, event_ndims=0))) + + @test_util.run_in_graph_and_eager_modes() + def testBijectorForwardInverseEventDimsOne(self): + bijector = Softsign(validate_args=True) + self.assertEqual("softsign", bijector.name) + x = 2. * self._rng.randn(2, 10) + y = self._softsign(x) + self.assertAllClose(y, self.evaluate(bijector.forward(x))) + self.assertAllClose(x, self.evaluate(bijector.inverse(y))) + + @test_util.run_in_graph_and_eager_modes() + def testBijectorLogDetJacobianEventDimsOne(self): + bijector = Softsign(validate_args=True) + y = self._rng.rand(2, 10) + ildj_before = self._softsign_ildj_before_reduction(y) + ildj = np.sum(ildj_before, axis=1) + self.assertAllClose( + ildj, self.evaluate( + bijector.inverse_log_det_jacobian(y, event_ndims=1))) + + def testScalarCongruency(self): + with self.test_session(): + bijector = Softsign(validate_args=True) + assert_scalar_congruency(bijector, lower_x=-20., upper_x=20.) + + def testBijectiveAndFinite(self): + with self.test_session(): + bijector = Softsign(validate_args=True) + x = np.linspace(-20., 20., 100).astype(np.float32) + y = np.linspace(-0.99, 0.99, 100).astype(np.float32) + assert_bijective_and_finite( + bijector, x, y, event_ndims=0, rtol=1e-3, atol=1e-3) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/square_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/square_test.py index f03d6f1343a11ae4517f9034ceb0c99ca6fe7fa2..30c7a738c320b609ce90685512e6b8344dffc9dc 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/square_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/square_test.py @@ -41,10 +41,11 @@ class SquareBijectorTest(test.TestCase): 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) + ildj, bijector.inverse_log_det_jacobian( + y, event_ndims=0).eval(), atol=0., rtol=1e-7) self.assertAllClose( - -bijector.inverse_log_det_jacobian(y).eval(), - bijector.forward_log_det_jacobian(x).eval(), + -bijector.inverse_log_det_jacobian(y, event_ndims=0).eval(), + bijector.forward_log_det_jacobian(x, event_ndims=0).eval(), atol=0., rtol=1e-7) diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/weibull_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/weibull_test.py index 7a31228d1ade55ce32b511dca073657d3bab53ae..f57adcda898a1fdb18aacbb0804411db1bb4e4c8 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/weibull_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/weibull_test.py @@ -36,7 +36,7 @@ class WeibullBijectorTest(test.TestCase): concentration = 0.3 bijector = Weibull( scale=scale, concentration=concentration, - event_ndims=1, validate_args=True) + validate_args=True) self.assertEqual("weibull", bijector.name) x = np.array([[[0.], [1.], [14.], [20.], [100.]]], dtype=np.float32) # Weibull distribution @@ -45,13 +45,11 @@ class WeibullBijectorTest(test.TestCase): self.assertAllClose(y, bijector.forward(x).eval()) self.assertAllClose(x, bijector.inverse(y).eval()) self.assertAllClose( - # We should lose a dimension from calculating the determinant of the - # jacobian. - np.squeeze(weibull_dist.logpdf(x), axis=2), - bijector.forward_log_det_jacobian(x).eval()) + weibull_dist.logpdf(x), + bijector.forward_log_det_jacobian(x, event_ndims=0).eval()) self.assertAllClose( - -bijector.inverse_log_det_jacobian(y).eval(), - bijector.forward_log_det_jacobian(x).eval(), + -bijector.inverse_log_det_jacobian(y, event_ndims=0).eval(), + bijector.forward_log_det_jacobian(x, event_ndims=0).eval(), rtol=1e-4, atol=0.) @@ -64,12 +62,12 @@ class WeibullBijectorTest(test.TestCase): def testBijectiveAndFinite(self): with self.test_session(): bijector = Weibull( - scale=20., concentration=2., event_ndims=0, validate_args=True) + scale=20., concentration=2., validate_args=True) x = np.linspace(1., 8., num=10).astype(np.float32) y = np.linspace( -np.expm1(-1 / 400.), -np.expm1(-16), num=10).astype(np.float32) - assert_bijective_and_finite(bijector, x, y, rtol=1e-3) + assert_bijective_and_finite(bijector, x, y, event_ndims=0, rtol=1e-3) if __name__ == "__main__": diff --git a/tensorflow/contrib/distributions/python/kernel_tests/conditional_transformed_distribution_test.py b/tensorflow/contrib/distributions/python/kernel_tests/conditional_transformed_distribution_test.py index 545471907f1eabc822b3d28ea9c57e183a09ff50..4e8989b6c2f93560b1fccbc99491d7809f494263 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/conditional_transformed_distribution_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/conditional_transformed_distribution_test.py @@ -44,6 +44,7 @@ class _ChooseLocation(ConditionalBijector): graph_parents=[self._loc], is_constant_jacobian=True, validate_args=False, + forward_min_event_ndims=0, name=name) def _forward(self, x, z): @@ -52,7 +53,7 @@ class _ChooseLocation(ConditionalBijector): def _inverse(self, x, z): return x - self._gather_loc(z) - def _inverse_log_det_jacobian(self, x, z=None): + def _inverse_log_det_jacobian(self, x, event_ndims, z=None): return 0. def _gather_loc(self, z): diff --git a/tensorflow/contrib/distributions/python/kernel_tests/distribution_test.py b/tensorflow/contrib/distributions/python/kernel_tests/distribution_test.py index 68e0d9cb8277f3953039963fec0da499db7a16d1..f42feae25d851eb9ae0bf48649fc3bbe2a221be0 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/distribution_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/distribution_test.py @@ -190,11 +190,30 @@ class DistributionTest(test.TestCase): y = dist._set_sample_static_shape(x, sample_shape) self.assertTrue(y.get_shape().ndims is None) + def testNameScopeWorksCorrectly(self): + x = tfd.Normal(loc=0., scale=1., name="x") + x_duplicate = tfd.Normal(loc=0., scale=1., name="x") + with ops.name_scope("y") as name: + y = tfd.Bernoulli(logits=0., name=name) + x_sample = x.sample(name="custom_sample") + x_sample_duplicate = x.sample(name="custom_sample") + x_log_prob = x.log_prob(0., name="custom_log_prob") + x_duplicate_sample = x_duplicate.sample(name="custom_sample") + + self.assertEqual(x.name, "x/") + self.assertEqual(x_duplicate.name, "x_1/") + self.assertEqual(y.name, "y/") + self.assertTrue(x_sample.name.startswith("x/custom_sample")) + self.assertTrue(x_sample_duplicate.name.startswith("x/custom_sample_1")) + self.assertTrue(x_log_prob.name.startswith("x/custom_log_prob")) + self.assertTrue(x_duplicate_sample.name.startswith( + "x_1/custom_sample")) + def testStrWorksCorrectlyScalar(self): normal = tfd.Normal(loc=np.float16(0), scale=np.float16(1)) self.assertEqual( ("tf.distributions.Normal(" - "\"Normal\", " + "\"Normal/\", " "batch_shape=(), " "event_shape=(), " "dtype=float16)"), # Got the dtype right. @@ -203,7 +222,7 @@ class DistributionTest(test.TestCase): chi2 = tfd.Chi2(df=np.float32([1., 2.]), name="silly") self.assertEqual( ("tf.distributions.Chi2(" - "\"silly\", " # What a silly name that is! + "\"silly/\", " # What a silly name that is! "batch_shape=(2,), " "event_shape=(), " "dtype=float32)"), @@ -211,7 +230,7 @@ class DistributionTest(test.TestCase): exp = tfd.Exponential(rate=array_ops.placeholder(dtype=dtypes.float32)) self.assertEqual( - ("tf.distributions.Exponential(\"Exponential\", " + ("tf.distributions.Exponential(\"Exponential/\", " # No batch shape. "event_shape=(), " "dtype=float32)"), @@ -222,7 +241,7 @@ class DistributionTest(test.TestCase): loc=np.zeros([2, 2]), name="MVN") self.assertEqual( ("tf.distributions.MultivariateNormalDiag(" - "\"MVN\", " + "\"MVN/\", " "batch_shape=(2,), " "event_shape=(2,), " "dtype=float64)"), @@ -233,7 +252,7 @@ class DistributionTest(test.TestCase): name="MVN2") self.assertEqual( ("tf.distributions.MultivariateNormalDiag(" - "\"MVN2\", " + "\"MVN2/\", " "batch_shape=(?,), " # Partially known. "event_shape=(3,), " "dtype=float32)"), @@ -243,7 +262,7 @@ class DistributionTest(test.TestCase): normal = tfd.Normal(loc=np.float16(0), scale=np.float16(1)) self.assertEqual( (""), # Got the dtype right. @@ -252,7 +271,7 @@ class DistributionTest(test.TestCase): chi2 = tfd.Chi2(df=np.float32([1., 2.]), name="silly") self.assertEqual( (""), @@ -261,7 +280,7 @@ class DistributionTest(test.TestCase): exp = tfd.Exponential(rate=array_ops.placeholder(dtype=dtypes.float32)) self.assertEqual( ("" " event_shape=()" " dtype=float32>"), @@ -272,7 +291,7 @@ class DistributionTest(test.TestCase): loc=np.zeros([2, 2]), name="MVN") self.assertEqual( (""), @@ -283,7 +302,7 @@ class DistributionTest(test.TestCase): name="MVN2") self.assertEqual( (""), diff --git a/tensorflow/contrib/distributions/python/kernel_tests/mvn_diag_test.py b/tensorflow/contrib/distributions/python/kernel_tests/mvn_diag_test.py index 933756aa8e12cca4c42eb98d9193512bbf2ad585..9635134b08db47a47a17c869fe813e0376ae6f1e 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/mvn_diag_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/mvn_diag_test.py @@ -68,7 +68,7 @@ class MultivariateNormalDiagTest(test.TestCase): dist = ds.TransformedDistribution( base_dist, validate_args=True, - bijector=bijectors.Softplus(event_ndims=1)) + bijector=bijectors.Softplus()) samps = dist.sample(5) # Shape [5, 1, 3]. self.assertAllEqual([5, 1], dist.log_prob(samps).get_shape()) diff --git a/tensorflow/contrib/distributions/python/kernel_tests/mvn_full_covariance_test.py b/tensorflow/contrib/distributions/python/kernel_tests/mvn_full_covariance_test.py index 1a02fbefb8e88599f5fedeb38fb06f5a09036439..7435bcbc684c1660a648cef4ab30c888723853f8 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/mvn_full_covariance_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/mvn_full_covariance_test.py @@ -52,7 +52,7 @@ class MultivariateNormalFullCovarianceTest(test.TestCase): mu = [1., 2.] sigma = [[1., 0.], [0., 1.]] mvn = ds.MultivariateNormalFullCovariance(mu, sigma, name="Billy") - self.assertEqual(mvn.name, "Billy") + self.assertEqual(mvn.name, "Billy/") def testDoesNotRaiseIfInitializedWithSymmetricMatrix(self): with self.test_session(): diff --git a/tensorflow/contrib/distributions/python/kernel_tests/shape_test.py b/tensorflow/contrib/distributions/python/kernel_tests/shape_test.py index c8d795c3f6afbec5b41755951174439f7703efb9..243b5a034859288b0e2e120f09258cfee77fbdea 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/shape_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/shape_test.py @@ -584,7 +584,6 @@ class DistributionShapeTest(test.TestCase): def testDistributionShapeGetDimsStatic(self): with self.test_session(): - shaper = _DistributionShape(batch_ndims=0, event_ndims=0) shaper = _DistributionShape(batch_ndims=0, event_ndims=0) x = 1 self.assertAllEqual((_empty_shape, _empty_shape, _empty_shape), diff --git a/tensorflow/contrib/distributions/python/kernel_tests/statistical_testing_test.py b/tensorflow/contrib/distributions/python/kernel_tests/statistical_testing_test.py index 0400c80c29cf0c36090168b7a1a6358ad49fde49..ce6cf702d522792f1ad26066a3d9be42003a0e3c 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/statistical_testing_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/statistical_testing_test.py @@ -21,7 +21,7 @@ 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.framework import ops from tensorflow.python.platform import test @@ -129,13 +129,13 @@ class StatisticalTestingTest(test.TestCase): # Test that the test assertion confirms that the mean of the # standard uniform distribution is not 0.4. - with self.assertRaises(errors.InvalidArgumentError): + with self.assertRaisesOpError("Mean confidence interval too high"): 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): + with self.assertRaisesOpError("Mean confidence interval too low"): sess.run(st.assert_true_mean_equal_by_dkwm( samples, 0., 1., 0.6, false_fail_rate=1e-6)) @@ -172,7 +172,7 @@ class StatisticalTestingTest(test.TestCase): # 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): + with self.assertRaisesOpError("samples1 has a smaller mean"): sess.run(st.assert_true_mean_equal_by_dkwm_two_sample( samples1, 0., 1., beta_high_samples, 0., 1., @@ -190,7 +190,7 @@ class StatisticalTestingTest(test.TestCase): # 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): + with self.assertRaisesOpError("samples2 has a smaller mean"): sess.run(st.assert_true_mean_equal_by_dkwm_two_sample( samples1, 0., 1., beta_low_samples, 0., 1., @@ -198,23 +198,46 @@ class StatisticalTestingTest(test.TestCase): def test_dkwm_argument_validity_checking(self): rng = np.random.RandomState(seed=0) - samples = rng.uniform(size=5000).astype(np.float32) + samples = rng.uniform( + low=[0., 1.], high=[1., 2.], size=(2500, 1, 2)).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): + with self.assertRaisesOpError("maximum value exceeds expectations"): sess.run(st.true_mean_confidence_interval_by_dkwm( - samples, 0., 0.5, error_rate=0.5)) - with self.assertRaises(errors.InvalidArgumentError): + samples, [[0., 1.]], [[0.5, 1.5]], error_rate=0.5)) + with self.assertRaisesOpError("minimum value falls below expectations"): sess.run(st.true_mean_confidence_interval_by_dkwm( - samples, 0.5, 1., error_rate=0.5)) + samples, [[0.5, 1.5]], [[1., 2.]], 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) + samples, [[0., 1.]], [[1., 2.]], error_rate=0.5) _ = sess.run(op) + def test_do_maximum_mean(self): + n = 117 + envelope = 0.02 # > 2 / n, but < 3 / n + rng = np.random.RandomState(seed=8) + samples = rng.uniform(size=n).astype(np.float32) + + # Compute the answer in TF using the code under test + with self.test_session() as sess: + envelope_t = ops.convert_to_tensor(envelope) + max_mean = st._do_maximum_mean(samples, envelope_t, 1) + max_mean = sess.run(max_mean) + + # Compute the correct answer for this case in numpy. In this + # example, `n` and `envelope` are such that `samples[2]` is the + # element that should be taken partially, regardless of the + # content of the `samples` array (see algorithm description in + # `../ops/statistical_testing.py`). + samples = sorted(samples) + weight = 1. / n - (envelope - 2. / n) + answer = samples[2] * weight + sum(samples[3:]) / n + envelope * 1. + self.assertAllClose(max_mean, answer, rtol=1e-9) + 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 f0ba1ec3eb57c67c1a0edb15639e91916a4509b7..5fe1331d2c34612e980c7b376367cd63b627533d 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/transformed_distribution_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/transformed_distribution_test.py @@ -28,6 +28,7 @@ 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 linalg_ops from tensorflow.python.ops import math_ops from tensorflow.python.platform import test @@ -36,6 +37,35 @@ ds = distributions la = linalg +class DummyMatrixTransform(bs.Bijector): + """Tractable matrix transformation. + + This is a non-sensical bijector that has forward/inverse_min_event_ndims=2. + The main use is to check that transformed distribution calculations are done + appropriately. + """ + + def __init__(self): + super(DummyMatrixTransform, self).__init__( + forward_min_event_ndims=2, + is_constant_jacobian=False, + validate_args=False, + name="dummy") + + def _forward(self, x): + return x + + def _inverse(self, y): + return y + + # Note: These jacobians don't make sense. + def _forward_log_det_jacobian(self, x): + return -linalg_ops.matrix_determinant(x) + + def _inverse_log_det_jacobian(self, x): + return linalg_ops.matrix_determinant(x) + + class TransformedDistributionTest(test.TestCase): def _cls(self): @@ -55,7 +85,7 @@ class TransformedDistributionTest(test.TestCase): # you may or may not need a reduce_sum. log_normal = self._cls()( distribution=ds.Normal(loc=mu, scale=sigma), - bijector=bs.Exp(event_ndims=0)) + bijector=bs.Exp()) sp_dist = stats.lognorm(s=sigma, scale=np.exp(mu)) # sample @@ -87,7 +117,7 @@ class TransformedDistributionTest(test.TestCase): sigma = 2.0 abs_normal = self._cls()( distribution=ds.Normal(loc=mu, scale=sigma), - bijector=bs.AbsoluteValue(event_ndims=0)) + bijector=bs.AbsoluteValue()) sp_normal = stats.norm(mu, sigma) # sample @@ -129,7 +159,7 @@ class TransformedDistributionTest(test.TestCase): self.assertAllClose(grid, cdf_, rtol=1e-6, atol=0.) def testCachedSamples(self): - exp_forward_only = bs.Exp(event_ndims=0) + exp_forward_only = bs.Exp() exp_forward_only._inverse = self._make_unimplemented( "inverse") exp_forward_only._inverse_event_shape_tensor = self._make_unimplemented( @@ -153,7 +183,7 @@ class TransformedDistributionTest(test.TestCase): self.assertAllClose(expected_log_pdf, log_pdf_val, rtol=1e-4, atol=0.) def testCachedSamplesInvert(self): - exp_inverse_only = bs.Exp(event_ndims=0) + exp_inverse_only = bs.Exp() exp_inverse_only._forward = self._make_unimplemented( "forward") exp_inverse_only._forward_event_shape_tensor = self._make_unimplemented( @@ -210,8 +240,11 @@ class TransformedDistributionTest(test.TestCase): int_identity = bs.Inline( forward_fn=array_ops.identity, inverse_fn=array_ops.identity, - inverse_log_det_jacobian_fn=lambda x: math_ops.cast(0, dtypes.int32), - forward_log_det_jacobian_fn=lambda x: math_ops.cast(0, dtypes.int32), + inverse_log_det_jacobian_fn=( + lambda y: math_ops.cast(0, dtypes.int32)), + forward_log_det_jacobian_fn=( + lambda x: math_ops.cast(0, dtypes.int32)), + forward_min_event_ndims=0, is_constant_jacobian=True) normal = self._cls()( distribution=ds.Normal(loc=0., scale=1.), @@ -435,6 +468,82 @@ class ScalarToMultiTest(test.TestCase): event_shape=[3], validate_args=True) + def testMatrixEvent(self): + with self.test_session() as sess: + batch_shape = [2] + event_shape = [2, 3, 3] + batch_shape_pl = array_ops.placeholder( + dtypes.int32, name="dynamic_batch_shape") + event_shape_pl = array_ops.placeholder( + dtypes.int32, name="dynamic_event_shape") + feed_dict = {batch_shape_pl: np.array(batch_shape, dtype=np.int32), + event_shape_pl: np.array(event_shape, dtype=np.int32)} + + scale = 2. + loc = 0. + fake_mvn_dynamic = self._cls()( + distribution=ds.Normal( + loc=loc, + scale=scale), + bijector=DummyMatrixTransform(), + batch_shape=batch_shape_pl, + event_shape=event_shape_pl, + validate_args=True) + + fake_mvn_static = self._cls()( + distribution=ds.Normal( + loc=loc, + scale=scale), + bijector=DummyMatrixTransform(), + batch_shape=batch_shape, + event_shape=event_shape, + validate_args=True) + + def actual_mvn_log_prob(x): + # This distribution is the normal PDF, reduced over the + # last 3 dimensions + a jacobian term which corresponds + # to the determinant of x. + return (np.sum( + stats.norm(loc, scale).logpdf(x), axis=(-1, -2, -3)) + + np.sum(np.linalg.det(x), axis=-1)) + + self.assertAllEqual([2, 3, 3], fake_mvn_static.event_shape) + self.assertAllEqual([2], fake_mvn_static.batch_shape) + + self.assertAllEqual(tensor_shape.TensorShape(None), + fake_mvn_dynamic.event_shape) + self.assertAllEqual(tensor_shape.TensorShape(None), + fake_mvn_dynamic.batch_shape) + + num_samples = 5e3 + for fake_mvn, feed_dict in ((fake_mvn_static, {}), + (fake_mvn_dynamic, feed_dict)): + # Ensure sample works by checking first, second moments. + y = fake_mvn.sample(int(num_samples), seed=0) + x = y[0:5, ...] + [ + x_, + fake_event_shape_, + fake_batch_shape_, + fake_log_prob_, + fake_prob_, + ] = sess.run([ + x, + fake_mvn.event_shape_tensor(), + fake_mvn.batch_shape_tensor(), + fake_mvn.log_prob(x), + fake_mvn.prob(x), + ], feed_dict=feed_dict) + + # Ensure all other functions work as intended. + self.assertAllEqual([5, 2, 2, 3, 3], x_.shape) + self.assertAllEqual([2, 3, 3], fake_event_shape_) + self.assertAllEqual([2], fake_batch_shape_) + self.assertAllClose(actual_mvn_log_prob(x_), fake_log_prob_, + atol=0., rtol=1e-6) + self.assertAllClose(np.exp(actual_mvn_log_prob(x_)), fake_prob_, + atol=0., rtol=1e-5) + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/distributions/python/kernel_tests/vector_laplace_diag_test.py b/tensorflow/contrib/distributions/python/kernel_tests/vector_laplace_diag_test.py index c355adeedbfff1072281a81de726ddb0ece07882..1226c66113ec4b43f57371abf4983aef1a529ec1 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/vector_laplace_diag_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/vector_laplace_diag_test.py @@ -61,7 +61,7 @@ class VectorLaplaceDiagTest(test.TestCase): dist = ds.TransformedDistribution( base_dist, validate_args=True, - bijector=bijectors.Softplus(event_ndims=1)) + bijector=bijectors.Softplus()) samps = dist.sample(5) # Shape [5, 1, 3]. self.assertAllEqual([5, 1], dist.log_prob(samps).get_shape()) diff --git a/tensorflow/contrib/distributions/python/ops/autoregressive.py b/tensorflow/contrib/distributions/python/ops/autoregressive.py index 69f3d57ff000d6c9acc8aa9e3d0ad8d9cbb6bb3c..88ed0127841093cc1a1168d988f14e7bb0277b12 100644 --- a/tensorflow/contrib/distributions/python/ops/autoregressive.py +++ b/tensorflow/contrib/distributions/python/ops/autoregressive.py @@ -145,7 +145,7 @@ class Autoregressive(distribution_lib.Distribution): ValueError: if `num_steps < 1`. """ parameters = locals() - with ops.name_scope(name): + with ops.name_scope(name) as name: self._distribution_fn = distribution_fn self._sample0 = sample0 self._distribution0 = (distribution_fn() if sample0 is None diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/__init__.py b/tensorflow/contrib/distributions/python/ops/bijectors/__init__.py index bc6b02542ebf3b83d58f888509dafb86351de8a7..babce80396cfc41b53e99f91038d4f077c7efe82 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/__init__.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/__init__.py @@ -38,6 +38,7 @@ @@SinhArcsinh @@SoftmaxCentered @@Softplus +@@Softsign @@Square @@Weibull @@ -74,6 +75,7 @@ from tensorflow.contrib.distributions.python.ops.bijectors.sigmoid import * 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.softsign 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/absolute_value.py b/tensorflow/contrib/distributions/python/ops/bijectors/absolute_value.py index 0fe9f6aa78fbe845b99d0668f075b0162ec2a9f7..c9e31d7712f09f6c4b4cc6ae51a34c42a19c291d 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/absolute_value.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/absolute_value.py @@ -18,9 +18,7 @@ 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.framework import constant_op from tensorflow.python.ops import check_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import math_ops @@ -72,38 +70,22 @@ class AbsoluteValue(bijector.Bijector): """ - def __init__(self, event_ndims=0, validate_args=False, name="absolute_value"): + def __init__(self, validate_args=False, name="absolute_value"): """Instantiates the `AbsoluteValue` bijector. Args: - 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, in particular whether inputs to `inverse` and `inverse_log_det_jacobian` are non-negative. 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 - 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"): super(AbsoluteValue, self).__init__( - event_ndims=event_ndims, + forward_min_event_ndims=0, + is_constant_jacobian=True, validate_args=validate_args, name=name) @@ -121,8 +103,7 @@ class AbsoluteValue(bijector.Bijector): # If event_ndims = 2, # F^{-1}(y) = (-y, y), so DF^{-1}(y) = (-1, 1), # so Log|DF^{-1}(y)| = Log[1, 1] = [0, 0]. - batch_shape = array_ops.shape(y)[:array_ops.rank(y) - self.event_ndims] - zeros = array_ops.zeros(batch_shape, dtype=y.dtype) + zeros = constant_op.constant(0., dtype=y.dtype) if self.validate_args: zeros = control_flow_ops.with_dependencies( [check_ops.assert_non_negative(y, message="Argument y was negative")], diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/affine.py b/tensorflow/contrib/distributions/python/ops/bijectors/affine.py index bef7bbb49b715497695f7513e19ecab4fa56c47e..b4c2939eb914d50475ba6b1c1e979a804090f641 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/affine.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/affine.py @@ -184,6 +184,7 @@ class Affine(bijector.Bijector): with self._name_scope("init", values=[ shift, scale_identity_multiplier, scale_diag, scale_tril, scale_perturb_diag, scale_perturb_factor]): + # In the absence of `loc` and `scale`, we'll assume `dtype` is `float32`. dtype = dtypes.float32 @@ -234,7 +235,7 @@ class Affine(bijector.Bijector): event_ndims=1, validate_args=validate_args) super(Affine, self).__init__( - event_ndims=1, + forward_min_event_ndims=1, graph_parents=( [self._scale] if tensor_util.is_tensor(self._scale) else self._scale.graph_parents + @@ -360,16 +361,17 @@ class Affine(bijector.Bijector): x, sample_shape, expand_batch_dim=False) return x - def _inverse_log_det_jacobian(self, y): - return -self._forward_log_det_jacobian(y) - def _forward_log_det_jacobian(self, x): + # is_constant_jacobian = True for this bijector, hence the + # `log_det_jacobian` need only be specified for a single input, as this will + # be tiled to match `event_ndims`. if self._is_only_identity_multiplier: # We don't pad in this case and instead let the fldj be applied # via broadcast. 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() def _maybe_check_scale(self): diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/affine_linear_operator.py b/tensorflow/contrib/distributions/python/ops/bijectors/affine_linear_operator.py index 89043b1410370074f11f2cfa59b6b6663fa62521..59f9742d576a7804f401d3a47ba31ae61d6c6e54 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/affine_linear_operator.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/affine_linear_operator.py @@ -22,9 +22,6 @@ from tensorflow.contrib.distributions.python.ops.shape import _DistributionShape 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_util -from tensorflow.python.ops import check_ops -from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops.distributions import bijector from tensorflow.python.ops.linalg import linear_operator @@ -94,7 +91,6 @@ class AffineLinearOperator(bijector.Bijector): def __init__(self, shift=None, scale=None, - event_ndims=1, validate_args=False, name="affine_linear_operator"): """Instantiates the `AffineLinearOperator` bijector. @@ -103,14 +99,11 @@ class AffineLinearOperator(bijector.Bijector): shift: Floating-point `Tensor`. scale: Subclass of `LinearOperator`. Represents the (batch) positive definite matrix `M` in `R^{k x k}`. - event_ndims: Scalar `integer` `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. Raises: - ValueError: if `event_ndims` is not 0 or 1. TypeError: if `scale` is not a `LinearOperator`. TypeError: if `shift.dtype` does not match `scale.dtype`. ValueError: if not `scale.is_non_singular`. @@ -120,20 +113,6 @@ class AffineLinearOperator(bijector.Bijector): self._validate_args = validate_args graph_parents = [] with self._name_scope("init", values=[shift]): - event_ndims = ops.convert_to_tensor(event_ndims, name="event_ndims") - if tensor_util.constant_value(event_ndims) is not None: - event_ndims = tensor_util.constant_value(event_ndims) - if event_ndims not in (0, 1): - raise ValueError("event_ndims({}) was not 0 or 1".format(event_ndims)) - 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) - graph_parents += [event_ndims] - # In the absence of `loc` and `scale`, we'll assume `dtype` is `float32`. dtype = dtypes.float32 @@ -166,10 +145,10 @@ class AffineLinearOperator(bijector.Bijector): self._scale = scale self._shaper = _DistributionShape( batch_ndims=batch_ndims, - event_ndims=event_ndims, + event_ndims=1, validate_args=validate_args) super(AffineLinearOperator, self).__init__( - event_ndims=event_ndims, + forward_min_event_ndims=1, graph_parents=graph_parents, is_constant_jacobian=True, dtype=dtype, @@ -213,12 +192,13 @@ class AffineLinearOperator(bijector.Bijector): x, sample_shape, expand_batch_dim=False) return x - def _inverse_log_det_jacobian(self, y): - return -self._forward_log_det_jacobian(y) - - def _forward_log_det_jacobian(self, x): # pylint: disable=unused-argument + def _forward_log_det_jacobian(self, x): + # is_constant_jacobian = True for this bijector, hence the + # `log_det_jacobian` need only be specified for a single input, as this will + # be tiled to match `event_ndims`. if self.scale is None: - return constant_op.constant(0, dtype=x.dtype.base_dtype) + return constant_op.constant(0., dtype=x.dtype.base_dtype) + with ops.control_dependencies(self._maybe_collect_assertions() if self.validate_args else []): 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 index 8adaa54c843d1b243a02967402a37b7c63fabbdf..cd792e2c8cf48602daf9fb5eb56b8c34bac050c7 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/affine_scalar.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/affine_scalar.py @@ -18,6 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.python.framework import constant_op from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import check_ops @@ -99,7 +100,7 @@ class AffineScalar(bijector.Bijector): self._scale) super(AffineScalar, self).__init__( - event_ndims=0, + forward_min_event_ndims=0, is_constant_jacobian=True, validate_args=validate_args, name=name) @@ -131,8 +132,10 @@ class AffineScalar(bijector.Bijector): return x def _forward_log_det_jacobian(self, x): - log_det_jacobian = array_ops.zeros_like(x) + # is_constant_jacobian = True for this bijector, hence the + # `log_det_jacobian` need only be specified for a single input, as this will + # be tiled to match `event_ndims`. if self.scale is None: - return log_det_jacobian - log_det_jacobian += math_ops.log(math_ops.abs(self.scale)) - return log_det_jacobian + return constant_op.constant(0., dtype=x.dtype.base_dtype) + + return math_ops.log(math_ops.abs(self.scale)) diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/batch_normalization.py b/tensorflow/contrib/distributions/python/ops/bijectors/batch_normalization.py index 33fdd32d7a0a01685690e598c69adca2c95972e9..224cec8a63dba53a528490117efac890312fe8d5 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/batch_normalization.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/batch_normalization.py @@ -157,7 +157,12 @@ class BatchNormalization(bijector.Bijector): gamma_constraint=g_constraint) self._validate_bn_layer(self.batchnorm) self._training = training + if isinstance(self.batchnorm.axis, int): + forward_min_event_ndims = 1 + else: + forward_min_event_ndims = len(self.batchnorm.axis) super(BatchNormalization, self).__init__( + forward_min_event_ndims=forward_min_event_ndims, validate_args=validate_args, name=name) def _validate_bn_layer(self, layer): @@ -186,7 +191,6 @@ class BatchNormalization(bijector.Bijector): 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 diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/chain.py b/tensorflow/contrib/distributions/python/ops/bijectors/chain.py index 3ce7c26213034c7345a20faa803c94a1bfa8d579..85ad23e4133ef09051cdc8b45e489caeea90fbb3 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/chain.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/chain.py @@ -21,6 +21,9 @@ from __future__ import print_function import itertools from tensorflow.python.framework import constant_op +from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_util +from tensorflow.python.ops import array_ops from tensorflow.python.ops.distributions import bijector @@ -29,6 +32,91 @@ __all__ = [ ] +def _use_static_shape(input_tensor, ndims): + return input_tensor.shape.is_fully_defined() and isinstance(ndims, int) + + +def _maybe_get_event_ndims_statically(event_ndims): + static_event_ndims = (event_ndims if isinstance(event_ndims, int) + else tensor_util.constant_value(event_ndims)) + if static_event_ndims is not None: + return static_event_ndims + + return event_ndims + + +def _compute_min_event_ndims(bijector_list, compute_forward=True): + """Computes the min_event_ndims associated with the give list of bijectors. + + Given a list `bijector_list` of bijectors, compute the min_event_ndims that is + associated with the composition of bijectors in that list. + + min_event_ndims is the # of right most dimensions for which the bijector has + done necessary computation on (i.e. the non-broadcastable part of the + computation). + + We can derive the min_event_ndims for a chain of bijectors as follows: + + In the case where there are no rank changing bijectors, this will simply be + `max(b.forward_min_event_ndims for b in bijector_list)`. This is because the + bijector with the most forward_min_event_ndims requires the most dimensions, + and hence the chain also requires operating on those dimensions. + + However in the case of rank changing, more care is needed in determining the + exact amount of dimensions. Padding dimensions causes subsequent bijectors to + operate on the padded dimensions, and Removing dimensions causes bijectors to + operate more left. + + Args: + bijector_list: List of bijectors to be composed by chain. + compute_forward: Boolean. If True, computes the min_event_ndims associated + with a forward call to Chain, and otherwise computes the min_event_ndims + associated with an inverse call to Chain. The latter is the same as the + min_event_ndims associated with a forward call to Invert(Chain(....)). + + Returns: + min_event_ndims + """ + min_event_ndims = 0 + # This is a mouthful, but what this encapsulates is that if not for rank + # changing bijectors, we'd only need to compute the largest of the min + # required ndims. Hence "max_min". Due to rank changing bijectors, we need to + # account for synthetic rank growth / synthetic rank decrease from a rank + # changing bijector. + rank_changed_adjusted_max_min_event_ndims = 0 + + if compute_forward: + bijector_list = reversed(bijector_list) + + for b in bijector_list: + if compute_forward: + current_min_event_ndims = b.forward_min_event_ndims + current_inverse_min_event_ndims = b.inverse_min_event_ndims + else: + current_min_event_ndims = b.inverse_min_event_ndims + current_inverse_min_event_ndims = b.forward_min_event_ndims + + # New dimensions were touched. + if rank_changed_adjusted_max_min_event_ndims < current_min_event_ndims: + min_event_ndims += ( + current_min_event_ndims - rank_changed_adjusted_max_min_event_ndims) + rank_changed_adjusted_max_min_event_ndims = max( + current_min_event_ndims, rank_changed_adjusted_max_min_event_ndims) + + # If the number of dimensions has increased via forward, then + # inverse_min_event_ndims > forward_min_event_ndims, and hence the + # dimensions we computed on, have moved left (so we have operated + # on additional dimensions). + # Conversely, if the number of dimensions has decreased via forward, + # then we have inverse_min_event_ndims < forward_min_event_ndims, + # and so we will have operated on fewer right most dimensions. + + number_of_changed_dimensions = ( + current_min_event_ndims - current_inverse_min_event_ndims) + rank_changed_adjusted_max_min_event_ndims -= number_of_changed_dimensions + return min_event_ndims + + class Chain(bijector.Bijector): """Bijector which applies a sequence of bijectors. @@ -93,21 +181,24 @@ class Chain(bijector.Bijector): raise ValueError("incompatible dtypes: %s" % dtype) elif len(dtype) == 2: dtype = dtype[1] if dtype[0] is None else dtype[0] - event_ndims = bijectors[0].event_ndims elif len(dtype) == 1: dtype = dtype[0] - event_ndims = bijectors[0].event_ndims else: dtype = None - event_ndims = None + + inverse_min_event_ndims = _compute_min_event_ndims( + bijectors, compute_forward=False) + forward_min_event_ndims = _compute_min_event_ndims( + bijectors, compute_forward=True) super(Chain, self).__init__( graph_parents=list(itertools.chain.from_iterable( b.graph_parents for b in bijectors)), + forward_min_event_ndims=forward_min_event_ndims, + inverse_min_event_ndims=inverse_min_event_ndims, is_constant_jacobian=all(b.is_constant_jacobian for b in bijectors), validate_args=validate_args, dtype=dtype, - event_ndims=event_ndims, name=name or ("identity" if not bijectors else "_of_".join(["chain"] + [b.name for b in bijectors]))) @@ -147,10 +238,31 @@ class Chain(bijector.Bijector): return y def _inverse_log_det_jacobian(self, y, **kwargs): - ildj = constant_op.constant(0., dtype=y.dtype, - name="inverse_log_det_jacobian") + ildj = constant_op.constant( + 0., dtype=y.dtype.base_dtype, name="inverse_log_det_jacobian") + + if not self.bijectors: + return ildj + + event_ndims = _maybe_get_event_ndims_statically( + self.inverse_min_event_ndims) + + if _use_static_shape(y, event_ndims): + event_shape = y.shape[y.shape.ndims - event_ndims:] + else: + event_shape = array_ops.shape(y)[array_ops.rank(y) - event_ndims:] + for b in self.bijectors: - ildj += b.inverse_log_det_jacobian(y, **kwargs.get(b.name, {})) + ildj += b.inverse_log_det_jacobian( + y, event_ndims=event_ndims, **kwargs.get(b.name, {})) + + if _use_static_shape(y, event_ndims): + event_shape = b.inverse_event_shape(event_shape) + event_ndims = _maybe_get_event_ndims_statically(event_shape.ndims) + else: + event_shape = b.inverse_event_shape_tensor(event_shape) + event_ndims = _maybe_get_event_ndims_statically( + array_ops.rank(event_shape)) y = b.inverse(y, **kwargs.get(b.name, {})) return ildj @@ -160,9 +272,34 @@ class Chain(bijector.Bijector): return x def _forward_log_det_jacobian(self, x, **kwargs): - fldj = constant_op.constant(0., dtype=x.dtype, - name="forward_log_det_jacobian") + x = ops.convert_to_tensor(x, name="x") + + fldj = constant_op.constant( + 0., dtype=x.dtype, name="inverse_log_det_jacobian") + + if not self.bijectors: + return fldj + + event_ndims = _maybe_get_event_ndims_statically( + self.forward_min_event_ndims) + + if _use_static_shape(x, event_ndims): + event_shape = x.shape[x.shape.ndims - event_ndims:] + else: + event_shape = array_ops.shape(x)[array_ops.rank(x) - event_ndims:] + for b in reversed(self.bijectors): - fldj += b.forward_log_det_jacobian(x, **kwargs.get(b.name, {})) + fldj += b.forward_log_det_jacobian( + x, event_ndims=event_ndims, **kwargs.get(b.name, {})) + if _use_static_shape(x, event_ndims): + event_shape = b.forward_event_shape(event_shape) + event_ndims = _maybe_get_event_ndims_statically(event_shape.ndims) + else: + event_shape = b.forward_event_shape_tensor(event_shape) + event_ndims = _maybe_get_event_ndims_statically( + array_ops.rank(event_shape)) + x = b.forward(x, **kwargs.get(b.name, {})) + return fldj + 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 8f09e16058b766c788ab3acced6940fd0026b521..ecdb8967f43e5960b2285de05125d0c3dbafe63c 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py @@ -80,7 +80,7 @@ class CholeskyOuterProduct(bijector.Bijector): self._graph_parents = [] self._name = name super(CholeskyOuterProduct, self).__init__( - event_ndims=2, + forward_min_event_ndims=2, validate_args=validate_args, name=name) @@ -170,7 +170,7 @@ class CholeskyOuterProduct(bijector.Bijector): sum_weighted_log_diag = array_ops.squeeze( math_ops.matmul(math_ops.log(diag), exponents[..., array_ops.newaxis]), - squeeze_dims=-1) + axis=-1) fldj = p_float * np.log(2.) + sum_weighted_log_diag return fldj diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/conditional_bijector.py b/tensorflow/contrib/distributions/python/ops/bijectors/conditional_bijector.py index ccb1f029277bc07011df7be047a075274f2b3a27..e9e994f839ab2fe0a0f52f5f404fb2a0c8f9cd94 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/conditional_bijector.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/conditional_bijector.py @@ -44,12 +44,16 @@ class ConditionalBijector(bijector.Bijector): "**condition_kwargs": "Named arguments forwarded to subclass implementation."}) def inverse_log_det_jacobian( - self, y, name="inverse_log_det_jacobian", **condition_kwargs): - return self._call_inverse_log_det_jacobian(y, name, **condition_kwargs) + self, y, event_ndims, name="inverse_log_det_jacobian", + **condition_kwargs): + return self._call_inverse_log_det_jacobian( + y, event_ndims, name, **condition_kwargs) @distribution_util.AppendDocstring(kwargs_dict={ "**condition_kwargs": "Named arguments forwarded to subclass implementation."}) def forward_log_det_jacobian( - self, x, name="forward_log_det_jacobian", **condition_kwargs): - return self._call_forward_log_det_jacobian(x, name, **condition_kwargs) + self, x, event_ndims, name="forward_log_det_jacobian", + **condition_kwargs): + return self._call_forward_log_det_jacobian( + x, event_ndims, name, **condition_kwargs) diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/exp.py b/tensorflow/contrib/distributions/python/ops/bijectors/exp.py index b1ff840d62a73c941a4d67dec73b5c9f4d5353f9..9fc1bbf052b419d07a9db149b990c2b80190d72b 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/exp.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/exp.py @@ -33,8 +33,8 @@ class Exp(power_transform.PowerTransform): ```python # Create the Y=g(X)=exp(X) transform which works only on Tensors with 1 - # batch ndim and 2 event ndims (i.e., vector of matrices). - exp = Exp(event_ndims=2) + # batch ndim 2. + exp = Exp() x = [[[1., 2], [3, 4]], [[5, 6], @@ -48,19 +48,17 @@ class Exp(power_transform.PowerTransform): """ def __init__(self, - event_ndims=0, validate_args=False, name="exp"): """Instantiates the `Exp` bijector. Args: - event_ndims: Scalar `int32` `Tensor` indicating the number of dimensions - associated with a particular draw from the distribution. validate_args: Python `bool` indicating whether arguments should be checked for correctness. name: Python `str` name given to ops managed by this object. """ + # forward_min_event_ndims = 0. + # No forward_min_event_ndims specified as this is done in PowerTransform. super(Exp, self).__init__( - event_ndims=event_ndims, validate_args=validate_args, name=name) diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/gumbel.py b/tensorflow/contrib/distributions/python/ops/bijectors/gumbel.py index 67f39785563255be0fe154aca3cbcf01c6a01e73..e656a258e56e71898ecb719dd2af876f158cf799 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/gumbel.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/gumbel.py @@ -48,7 +48,6 @@ class Gumbel(bijector.Bijector): def __init__(self, loc=0., scale=1., - event_ndims=0, validate_args=False, name="gumbel"): """Instantiates the `Gumbel` bijector. @@ -60,8 +59,6 @@ class Gumbel(bijector.Bijector): scale: Positive Float-like `Tensor` that is the same dtype and is broadcastable with `loc`. This is `scale` in `Y = g(X) = exp(-exp(-(X - loc) / scale))`. - event_ndims: Python scalar indicating the number of dimensions associated - with a particular draw from the distribution. validate_args: Python `bool` indicating whether arguments should be checked for correctness. name: Python `str` name given to ops managed by this object. @@ -80,7 +77,9 @@ class Gumbel(bijector.Bijector): ], self._scale) super(Gumbel, self).__init__( - event_ndims=event_ndims, validate_args=validate_args, name=name) + validate_args=validate_args, + forward_min_event_ndims=0, + name=name) @property def loc(self): @@ -102,15 +101,11 @@ class Gumbel(bijector.Bijector): def _inverse_log_det_jacobian(self, y): y = self._maybe_assert_valid_y(y) - event_dims = self._event_dims_tensor(y) - return math_ops.reduce_sum( - math_ops.log(self.scale / (-math_ops.log(y) * y)), axis=event_dims) + return math_ops.log(self.scale / (-math_ops.log(y) * y)) def _forward_log_det_jacobian(self, x): - event_dims = self._event_dims_tensor(x) z = (x - self.loc) / self.scale - return math_ops.reduce_sum( - -z - math_ops.exp(-z) - math_ops.log(self.scale), axis=event_dims) + return -z - math_ops.exp(-z) - math_ops.log(self.scale) def _maybe_assert_valid_y(self, y): if not self.validate_args: diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/inline.py b/tensorflow/contrib/distributions/python/ops/bijectors/inline.py index fab1b22fbf92e7b92a5ec86ec62d66bec71a8c94..2bde956d1345129285acae4684256c5ac828b9a1 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/inline.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/inline.py @@ -40,7 +40,7 @@ class Inline(bijector.Bijector): name="exp") ``` - The above example is equivalent to the `Bijector` `Exp(event_ndims=1)`. + The above example is equivalent to the `Bijector` `Exp()`. """ def __init__(self, @@ -54,6 +54,8 @@ class Inline(bijector.Bijector): inverse_event_shape_tensor_fn=None, is_constant_jacobian=False, validate_args=False, + forward_min_event_ndims=None, + inverse_min_event_ndims=None, name="inline"): """Creates a `Bijector` from callables. @@ -76,10 +78,15 @@ class Inline(bijector.Bijector): constant for all input arguments. validate_args: Python `bool` indicating whether arguments should be checked for correctness. + forward_min_event_ndims: Python `int` indicating the minimal + dimensionality this bijector acts on. + inverse_min_event_ndims: Python `int` indicating the minimal + dimensionality this bijector acts on. name: Python `str`, name given to ops managed by this object. """ super(Inline, self).__init__( - event_ndims=0, + forward_min_event_ndims=forward_min_event_ndims, + inverse_min_event_ndims=inverse_min_event_ndims, is_constant_jacobian=is_constant_jacobian, validate_args=validate_args, name=name) @@ -134,8 +141,8 @@ class Inline(bijector.Bijector): "inverse_log_det_jacobian_fn is not a callable function.") return self._inverse_log_det_jacobian_fn(y, **kwargs) - def _forward_log_det_jacobian(self, y, **kwargs): + def _forward_log_det_jacobian(self, x, **kwargs): if not callable(self._forward_log_det_jacobian_fn): raise NotImplementedError( "forward_log_det_jacobian_fn is not a callable function.") - return self._forward_log_det_jacobian_fn(y, **kwargs) + return self._forward_log_det_jacobian_fn(x, **kwargs) diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/invert.py b/tensorflow/contrib/distributions/python/ops/bijectors/invert.py index 2c603fe61f36dd27f4984fe6c13c11f2fb534321..1904239a0e7009c35cc4f3c8876fd749463a2b83 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/invert.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/invert.py @@ -66,8 +66,9 @@ class Invert(bijector_lib.Bijector): self._bijector = bijector super(Invert, self).__init__( - event_ndims=bijector.event_ndims, graph_parents=bijector.graph_parents, + forward_min_event_ndims=bijector.inverse_min_event_ndims, + inverse_min_event_ndims=bijector.forward_min_event_ndims, is_constant_jacobian=bijector.is_constant_jacobian, validate_args=validate_args, dtype=bijector.dtype, diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/kumaraswamy.py b/tensorflow/contrib/distributions/python/ops/bijectors/kumaraswamy.py index f5de052c9ed18b1ebf4c174aeea3a951b1ddcd9d..97000c17262d3efdef10274711364c2bc2083bd4 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/kumaraswamy.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/kumaraswamy.py @@ -19,7 +19,6 @@ 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 @@ -48,7 +47,6 @@ class Kumaraswamy(bijector.Bijector): def __init__(self, concentration1=None, concentration0=None, - event_ndims=0, validate_args=False, name="kumaraswamy"): """Instantiates the `Kumaraswamy` bijector. @@ -60,31 +58,14 @@ class Kumaraswamy(bijector.Bijector): 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"), @@ -96,7 +77,7 @@ class Kumaraswamy(bijector.Bijector): self._concentration1 = concentration1 self._concentration0 = concentration0 super(Kumaraswamy, self).__init__( - event_ndims=0, + forward_min_event_ndims=0, validate_args=validate_args, name=name) @@ -123,12 +104,10 @@ class Kumaraswamy(bijector.Bijector): 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( + return ( 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) + (self.concentration0 - 1) * math_ops.log1p(-y**self.concentration1)) def _maybe_assert_valid_concentration(self, concentration, validate_args): """Checks the validity of a concentration parameter.""" diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/masked_autoregressive.py b/tensorflow/contrib/distributions/python/ops/bijectors/masked_autoregressive.py index 84b2340c75514c3d2c12bf4d775ba74450a0dc26..ef56cf6ddda4dca2b1575e844b2584689e531b81 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/masked_autoregressive.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/masked_autoregressive.py @@ -61,7 +61,7 @@ class MaskedAutoregressiveFlow(bijector_lib.Bijector): this property by zeroing out weights in its `masked_dense` layers. In the `tf.distributions` framework, a "normalizing flow" is implemented as a - `tf.distributions.bijectors.Bijector`. The `forward` "autoregression" + `tf.contrib.distributions.bijectors.Bijector`. The `forward` "autoregression" is implemented using a `tf.while_loop` and a deep neural network (DNN) with masked weights such that the autoregressive property is automatically met in the `inverse`. @@ -220,6 +220,7 @@ class MaskedAutoregressiveFlow(bijector_lib.Bijector): self._shift_and_log_scale_fn = shift_and_log_scale_fn self._unroll_loop = unroll_loop super(MaskedAutoregressiveFlow, self).__init__( + forward_min_event_ndims=1, is_constant_jacobian=is_constant_jacobian, validate_args=validate_args, name=name) diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/permute.py b/tensorflow/contrib/distributions/python/ops/bijectors/permute.py index 8654cc39d0c41ec4f1b85cd5fc4366ceaf4b224d..4978167803fc38b112c95922519c8c296cee2561 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/permute.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/permute.py @@ -114,6 +114,7 @@ class Permute(bijector_lib.Bijector): ], permutation) self._permutation = permutation super(Permute, self).__init__( + forward_min_event_ndims=1, is_constant_jacobian=True, validate_args=validate_args, name=name or "permute") @@ -132,7 +133,10 @@ class Permute(bijector_lib.Bijector): axis=-1) def _inverse_log_det_jacobian(self, y): - return constant_op.constant(0., dtype=y.dtype) + # is_constant_jacobian = True for this bijector, hence the + # `log_det_jacobian` need only be specified for a single input, as this will + # be tiled to match `event_ndims`. + return constant_op.constant(0., dtype=y.dtype.base_dtype) def _forward_log_det_jacobian(self, x): - return constant_op.constant(0., dtype=x.dtype) + return constant_op.constant(0., dtype=x.dtype.base_dtype) diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/power_transform.py b/tensorflow/contrib/distributions/python/ops/bijectors/power_transform.py index c37db61720d10949f294ff7b2e9778ba6efa57f0..71f123f2a998458edaa9c8da07ea2932f62625ca 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/power_transform.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/power_transform.py @@ -43,7 +43,6 @@ class PowerTransform(bijector.Bijector): def __init__(self, power=0., - event_ndims=0, validate_args=False, name="power_transform"): """Instantiates the `PowerTransform` bijector. @@ -51,8 +50,6 @@ class PowerTransform(bijector.Bijector): Args: power: Python `float` scalar indicating the transform power, i.e., `Y = g(X) = (1 + X * c)**(1 / c)` where `c` is the `power`. - event_ndims: Python scalar indicating the number of dimensions associated - with a particular draw from the distribution. validate_args: Python `bool` indicating whether arguments should be checked for correctness. name: Python `str` name given to ops managed by this object. @@ -70,7 +67,7 @@ class PowerTransform(bijector.Bijector): raise ValueError("`power` must be a non-negative TF constant.") self._power = power super(PowerTransform, self).__init__( - event_ndims=event_ndims, + forward_min_event_ndims=0, validate_args=validate_args, name=name) @@ -97,18 +94,13 @@ class PowerTransform(bijector.Bijector): def _inverse_log_det_jacobian(self, y): y = self._maybe_assert_valid_y(y) - event_dims = self._event_dims_tensor(y) - return (self.power - 1.) * math_ops.reduce_sum( - math_ops.log(y), axis=event_dims) + return (self.power - 1.) * math_ops.log(y) def _forward_log_det_jacobian(self, x): x = self._maybe_assert_valid_x(x) - event_dims = self._event_dims_tensor(x) if self.power == 0.: - return math_ops.reduce_sum(x, axis=event_dims) - return (1. / self.power - 1.) * math_ops.reduce_sum( - math_ops.log1p(x * self.power), - axis=event_dims) + return x + return (1. / self.power - 1.) * math_ops.log1p(x * self.power) def _maybe_assert_valid_x(self, x): if not self.validate_args or self.power == 0.: diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/real_nvp.py b/tensorflow/contrib/distributions/python/ops/bijectors/real_nvp.py index 71ab369d01aafc33854a2c2437f96bbb493cc6fb..f09ab21bce100e9dafb77eff1f3999ce4b71c681 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/real_nvp.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/real_nvp.py @@ -166,7 +166,7 @@ class RealNVP(bijector_lib.Bijector): self._input_depth = None self._shift_and_log_scale_fn = shift_and_log_scale_fn super(RealNVP, self).__init__( - event_ndims=1, + forward_min_event_ndims=1, is_constant_jacobian=is_constant_jacobian, validate_args=validate_args, name=name) @@ -224,7 +224,7 @@ class RealNVP(bijector_lib.Bijector): _, log_scale = self._shift_and_log_scale_fn( x0, self._input_depth - self._num_masked) if log_scale is None: - return constant_op.constant(0., dtype=x.dtype, name="ildj") + return constant_op.constant(0., dtype=x.dtype, name="fldj") return math_ops.reduce_sum(log_scale, axis=-1) diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/reshape.py b/tensorflow/contrib/distributions/python/ops/bijectors/reshape.py index 55eca063126797d577653f0d6bcdfddf8192bdb5..f21b982ba664b312c716827c7925767a0b5a037a 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/reshape.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/reshape.py @@ -128,15 +128,17 @@ class Reshape(bijector_lib.Bijector): self._event_shape_in = event_shape_in self._event_shape_out = event_shape_out - super(Reshape, self).__init__(is_constant_jacobian=True, - validate_args=validate_args, - name=name or "reshape") + super(Reshape, self).__init__( + forward_min_event_ndims=0, + is_constant_jacobian=True, + validate_args=validate_args, + name=name or "reshape") def _maybe_check_valid_shape(self, shape, validate_args): """Check that a shape Tensor is int-type and otherwise sane.""" if not shape.dtype.is_integer: raise TypeError("{} dtype ({}) should be `int`-like.".format( - shape.op.name, shape.dtype.name)) + shape, shape.dtype.name)) assertions = [] @@ -144,10 +146,10 @@ class Reshape(bijector_lib.Bijector): ndims_ = tensor_util.constant_value(ndims) if ndims_ is not None and ndims_ > 1: raise ValueError("`{}` rank ({}) should be <= 1.".format( - shape.op.name, ndims_)) + shape, ndims_)) elif validate_args: assertions.append(check_ops.assert_less_equal( - ndims, 1, message="`{}` rank should be <= 1.".format(shape.op.name))) + ndims, 1, message="`{}` rank should be <= 1.".format(shape))) shape_ = tensor_util.constant_value_as_shape(shape) if shape_.is_fully_defined(): @@ -155,12 +157,12 @@ class Reshape(bijector_lib.Bijector): if sum(es == -1) > 1: raise ValueError( "`{}` must have at most one `-1` (given {})" - .format(shape.op.name, es)) + .format(shape, es)) if np.any(es < -1): raise ValueError( "`{}` elements must be either positive integers or `-1`" "(given {})." - .format(shape.op.name, es)) + .format(shape, es)) elif validate_args: assertions.extend([ check_ops.assert_less_equal( @@ -168,11 +170,11 @@ class Reshape(bijector_lib.Bijector): math_ops.cast(math_ops.equal(shape, -1), dtypes.int32)), 1, message="`{}` elements must have at most one `-1`." - .format(shape.op.name)), + .format(shape)), check_ops.assert_greater_equal( shape, -1, message="`{}` elements must be either positive integers or `-1`." - .format(shape.op.name)), + .format(shape)), ]) return assertions diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/sigmoid.py b/tensorflow/contrib/distributions/python/ops/bijectors/sigmoid.py index a640dfe7dfbcce96261589c7fc49107deaefdd54..5df8c886315ff75cdc884e3b9b4665fb64bb109d 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/sigmoid.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/sigmoid.py @@ -33,7 +33,9 @@ class Sigmoid(bijector.Bijector): def __init__(self, validate_args=False, name="sigmoid"): super(Sigmoid, self).__init__( - event_ndims=0, validate_args=validate_args, name=name) + forward_min_event_ndims=0, + validate_args=validate_args, + name=name) def _forward(self, x): return math_ops.sigmoid(x) diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/sinh_arcsinh.py b/tensorflow/contrib/distributions/python/ops/bijectors/sinh_arcsinh.py index 3a75e4ae9495793901b0da91a5aa3982aab35852..2a32e8abcde940b0056b0faf2955ec1b3bd71803 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/sinh_arcsinh.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/sinh_arcsinh.py @@ -91,7 +91,6 @@ class SinhArcsinh(bijector.Bijector): def __init__(self, skewness=None, tailweight=None, - event_ndims=0, validate_args=False, name="SinhArcsinh"): """Instantiates the `SinhArcsinh` bijector. @@ -101,8 +100,6 @@ class SinhArcsinh(bijector.Bijector): of type `float32`. tailweight: Tailweight parameter. Positive `Tensor` of same `dtype` as `skewness` and broadcastable `shape`. Default is `1` of type `float32`. - event_ndims: Python scalar indicating the number of dimensions associated - with a particular draw from the distribution. validate_args: Python `bool` indicating whether arguments should be checked for correctness. name: Python `str` name given to ops managed by this object. @@ -125,7 +122,9 @@ class SinhArcsinh(bijector.Bijector): message="Argument tailweight was not positive") ], self._tailweight) super(SinhArcsinh, self).__init__( - event_ndims=event_ndims, validate_args=validate_args, name=name) + forward_min_event_ndims=0, + validate_args=validate_args, + name=name) @property def skewness(self): @@ -149,31 +148,29 @@ class SinhArcsinh(bijector.Bijector): # dx/dy # = cosh(arcsinh(y) / tailweight - skewness) # / (tailweight * sqrt(y**2 + 1)) - event_dims = self._event_dims_tensor(y) - return math_ops.reduce_sum( - # This is computed inside the log to avoid catastrophic cancellations - # from cosh((arcsinh(y) / tailweight) - skewness) and sqrt(x**2 + 1). + + # This is computed inside the log to avoid catastrophic cancellations + # from cosh((arcsinh(y) / tailweight) - skewness) and sqrt(x**2 + 1). + return ( math_ops.log(math_ops.cosh( math_ops.asinh(y) / self.tailweight - self.skewness) # TODO(srvasude): Consider using cosh(arcsinh(x)) in cases # where (arcsinh(x) / tailweight) - skewness ~= arcsinh(x). / _sqrtx2p1(y)) - - math_ops.log(self.tailweight), - axis=event_dims) + - math_ops.log(self.tailweight)) def _forward_log_det_jacobian(self, x): # y = sinh((arcsinh(x) + skewness) * tailweight) # Using sinh' = cosh, arcsinh'(x) = 1 / sqrt(x**2 + 1), # dy/dx # = cosh((arcsinh(x) + skewness) * tailweight) * tailweight / sqrt(x**2 + 1) - event_dims = self._event_dims_tensor(x) - return math_ops.reduce_sum( - # This is computed inside the log to avoid catastrophic cancellations - # from cosh((arcsinh(x) + skewness) * tailweight) and sqrt(x**2 + 1). + + # This is computed inside the log to avoid catastrophic cancellations + # from cosh((arcsinh(x) + skewness) * tailweight) and sqrt(x**2 + 1). + return ( math_ops.log(math_ops.cosh( (math_ops.asinh(x) + self.skewness) * self.tailweight) # TODO(srvasude): Consider using cosh(arcsinh(x)) in cases # where (arcsinh(x) + skewness) * tailweight ~= arcsinh(x). / _sqrtx2p1(x)) - + math_ops.log(self.tailweight), - axis=event_dims) + + math_ops.log(self.tailweight)) diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/softmax_centered.py b/tensorflow/contrib/distributions/python/ops/bijectors/softmax_centered.py index dc94fd0a38de29f5a7ee6ca826aab0ecf8712966..f52b91550edff7390d8094a4508d862674e85d59 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/softmax_centered.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/softmax_centered.py @@ -66,7 +66,7 @@ class SoftmaxCentered(bijector.Bijector): self._graph_parents = [] self._name = name super(SoftmaxCentered, self).__init__( - event_ndims=1, + forward_min_event_ndims=1, validate_args=validate_args, name=name) @@ -105,8 +105,6 @@ class SoftmaxCentered(bijector.Bijector): y.shape.assert_is_compatible_with(shape) y.set_shape(shape) - # Since we only support event_ndims in [0, 1] and we do padding, we always - # reduce over the last dimension, i.e., dim=-1 (which is the default). return nn_ops.softmax(y) def _inverse(self, y): @@ -162,8 +160,6 @@ class SoftmaxCentered(bijector.Bijector): # -log_normalization + reduce_sum(logits - log_normalization) log_normalization = nn_ops.softplus( math_ops.reduce_logsumexp(x, axis=-1, keep_dims=True)) - fldj = (-log_normalization + - math_ops.reduce_sum(x - log_normalization, - axis=-1, - keep_dims=True)) - return array_ops.squeeze(fldj, squeeze_dims=-1) + return array_ops.squeeze( + (-log_normalization + math_ops.reduce_sum( + x - log_normalization, axis=-1, keepdims=True)), axis=-1) diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/softplus.py b/tensorflow/contrib/distributions/python/ops/bijectors/softplus.py index 81957fcf78922fa15fd20a25d144071f431161ae..96a938c803418ff818f9c531754b47ba1eb8667a 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/softplus.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/softplus.py @@ -62,7 +62,7 @@ class Softplus(bijector.Bijector): ```python # Create the Y=g(X)=softplus(X) transform which works only on Tensors with 1 # batch ndim and 2 event ndims (i.e., vector of matrices). - softplus = Softplus(event_ndims=2) + softplus = Softplus() x = [[[1., 2], [3, 4]], [[5, 6], @@ -81,7 +81,6 @@ class Softplus(bijector.Bijector): "Nonzero floating point `Tensor`. Controls the softness of what " "would otherwise be a kink at the origin. Default is 1.0")}) def __init__(self, - event_ndims=0, hinge_softness=None, validate_args=False, name="softplus"): @@ -101,7 +100,7 @@ class Softplus(bijector.Bijector): [nonzero_check], self.hinge_softness) super(Softplus, self).__init__( - event_ndims=event_ndims, + forward_min_event_ndims=0, validate_args=validate_args, name=name) @@ -130,14 +129,12 @@ class Softplus(bijector.Bijector): # 1 - exp{-Y} approx Y. if self.hinge_softness is not None: y /= math_ops.cast(self.hinge_softness, y.dtype) - return -math_ops.reduce_sum(math_ops.log(-math_ops.expm1(-y)), - axis=self._event_dims_tensor(y)) + return -math_ops.log(-math_ops.expm1(-y)) def _forward_log_det_jacobian(self, x): if self.hinge_softness is not None: x /= math_ops.cast(self.hinge_softness, x.dtype) - return -math_ops.reduce_sum(nn_ops.softplus(-x), - axis=self._event_dims_tensor(x)) + return -nn_ops.softplus(-x) @property def hinge_softness(self): diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/softsign.py b/tensorflow/contrib/distributions/python/ops/bijectors/softsign.py new file mode 100644 index 0000000000000000000000000000000000000000..b4a658c171b8313358754228aabbfa4bf93fd84d --- /dev/null +++ b/tensorflow/contrib/distributions/python/ops/bijectors/softsign.py @@ -0,0 +1,86 @@ +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Softsign bijector.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +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__ = [ + "Softsign", +] + + +class Softsign(bijector.Bijector): + """Bijector which computes `Y = g(X) = X / (1 + |X|)`. + + The softsign `Bijector` has the following two useful properties: + + * The domain is all real numbers + * `softsign(x) approx sgn(x)`, for large `|x|`. + + #### Examples + + ```python + # Create the Y = softsign(X) transform. + softsign = Softsign() + x = [[[1., 2], + [3, 4]], + [[5, 6], + [7, 8]]] + x / (1 + abs(x)) == softsign.forward(x) + x / (1 - abs(x)) == softsign.inverse(x) + ``` + """ + + def __init__(self, validate_args=False, name="softsign"): + super(Softsign, self).__init__( + forward_min_event_ndims=0, + validate_args=validate_args, + name=name) + + def _forward(self, x): + return x / (1. + math_ops.abs(x)) + + def _inverse(self, y): + y = self._maybe_assert_valid_y(y) + return y / (1. - math_ops.abs(y)) + + def _forward_log_det_jacobian(self, x): + return -2. * math_ops.log1p(math_ops.abs(x)) + + def _inverse_log_det_jacobian(self, y): + y = self._maybe_assert_valid_y(y) + return -2. * math_ops.log1p(-math_ops.abs(y)) + + def _maybe_assert_valid_y(self, y): + if not self.validate_args: + return y + is_valid = [ + check_ops.assert_greater( + y, math_ops.cast(-1., dtype=y.dtype.base_dtype), + message="Inverse transformation input must be greater than -1."), + check_ops.assert_less( + y, math_ops.cast(1., dtype=y.dtype.base_dtype), + message="Inverse transformation input must be less than 1.") + ] + + return control_flow_ops.with_dependencies(is_valid, y) diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/square.py b/tensorflow/contrib/distributions/python/ops/bijectors/square.py index 1e9dbf35091fe51f2478dc085c394a77295ca4ee..2ccfdc95970e387e708603e2614ad29fb6a18db3 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/square.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/square.py @@ -59,7 +59,7 @@ class Square(bijector.Bijector): """ self._name = name super(Square, self).__init__( - event_ndims=0, + forward_min_event_ndims=0, validate_args=validate_args, name=name) diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/weibull.py b/tensorflow/contrib/distributions/python/ops/bijectors/weibull.py index 00520bcda85e9527767e6342bf75f10667c264a8..39129cd22cdbf9ca1b4edd7cb5c3571a33837a29 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/weibull.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/weibull.py @@ -50,7 +50,6 @@ class Weibull(bijector.Bijector): def __init__(self, scale=1., concentration=1., - event_ndims=0, validate_args=False, name="weibull"): """Instantiates the `Weibull` bijector. @@ -62,8 +61,6 @@ class Weibull(bijector.Bijector): concentration: Positive Float-type `Tensor` that is the same dtype and is broadcastable with `scale`. This is `k` in `Y = g(X) = 1 - exp((-x / l) ** k)`. - event_ndims: Python scalar indicating the number of dimensions associated - with a particular draw from the distribution. validate_args: Python `bool` indicating whether arguments should be checked for correctness. name: Python `str` name given to ops managed by this object. @@ -89,7 +86,7 @@ class Weibull(bijector.Bijector): ], self._concentration) super(Weibull, self).__init__( - event_ndims=event_ndims, + forward_min_event_ndims=0, validate_args=validate_args, name=name) @@ -113,22 +110,18 @@ class Weibull(bijector.Bijector): def _inverse_log_det_jacobian(self, y): y = self._maybe_assert_valid_y(y) - event_dims = self._event_dims_tensor(y) - return math_ops.reduce_sum( + return ( -math_ops.log1p(-y) + (1 / self.concentration - 1) * math_ops.log(-math_ops.log1p(-y)) + - math_ops.log(self.scale / self.concentration), - axis=event_dims) + math_ops.log(self.scale / self.concentration)) def _forward_log_det_jacobian(self, x): x = self._maybe_assert_valid_x(x) - event_dims = self._event_dims_tensor(x) - return math_ops.reduce_sum( + return ( -(x / self.scale) ** self.concentration + (self.concentration - 1) * math_ops.log(x) + math_ops.log(self.concentration) + - -self.concentration * math_ops.log(self.scale), - axis=event_dims) + -self.concentration * math_ops.log(self.scale)) def _maybe_assert_valid_x(self, x): if not self.validate_args: diff --git a/tensorflow/contrib/distributions/python/ops/binomial.py b/tensorflow/contrib/distributions/python/ops/binomial.py index 6a1bb39ab28218a411bdf4329965186bcf32bf30..12d16031783b78dc3ea6273af77c1eaeb77ca94e 100644 --- a/tensorflow/contrib/distributions/python/ops/binomial.py +++ b/tensorflow/contrib/distributions/python/ops/binomial.py @@ -164,7 +164,7 @@ class Binomial(distribution.Distribution): name: Python `str` name prefixed to Ops created by this class. """ parameters = locals() - with ops.name_scope(name, values=[total_count, logits, probs]): + with ops.name_scope(name, values=[total_count, logits, probs]) as name: self._total_count = self._maybe_assert_valid_total_count( ops.convert_to_tensor(total_count, name="total_count"), validate_args) diff --git a/tensorflow/contrib/distributions/python/ops/cauchy.py b/tensorflow/contrib/distributions/python/ops/cauchy.py index 6f5d724a2a945ed8f9c159d8314327c6f994d1db..daacfe657fe154dce8d0db98894fe8b73546c476 100644 --- a/tensorflow/contrib/distributions/python/ops/cauchy.py +++ b/tensorflow/contrib/distributions/python/ops/cauchy.py @@ -121,7 +121,7 @@ class Cauchy(distribution.Distribution): TypeError: if `loc` and `scale` have different `dtype`. """ parameters = locals() - with ops.name_scope(name, values=[loc, scale]): + with ops.name_scope(name, values=[loc, scale]) as name: with ops.control_dependencies([check_ops.assert_positive(scale)] if validate_args else []): self._loc = array_ops.identity(loc, name="loc") diff --git a/tensorflow/contrib/distributions/python/ops/chi2.py b/tensorflow/contrib/distributions/python/ops/chi2.py index e610f469e5d5f446b75c734cc39811de30a8cb9a..c77c5fd20895a6220604d76a95a152a22cd3d914 100644 --- a/tensorflow/contrib/distributions/python/ops/chi2.py +++ b/tensorflow/contrib/distributions/python/ops/chi2.py @@ -88,7 +88,7 @@ class Chi2(gamma.Gamma): # not true in the parent class "gamma." therefore, passing # allow_nan_stats=True # through to the parent class results in unnecessary asserts. - with ops.name_scope(name, values=[df]): + with ops.name_scope(name, values=[df]) as name: with ops.control_dependencies([ check_ops.assert_positive(df), ] if validate_args else []): @@ -120,7 +120,7 @@ class Chi2WithAbsDf(Chi2): allow_nan_stats=True, name="Chi2WithAbsDf"): parameters = locals() - with ops.name_scope(name, values=[df]): + with ops.name_scope(name, values=[df]) as name: super(Chi2WithAbsDf, self).__init__( df=math_ops.floor( math_ops.abs(df, name="abs_df"), diff --git a/tensorflow/contrib/distributions/python/ops/conditional_transformed_distribution.py b/tensorflow/contrib/distributions/python/ops/conditional_transformed_distribution.py index 1d4c5660d8d73b7b6a7e758fc834ccfddeb5c8ea..10b45361358b40a3c8fd725f27ad84ef9b8a37f5 100644 --- a/tensorflow/contrib/distributions/python/ops/conditional_transformed_distribution.py +++ b/tensorflow/contrib/distributions/python/ops/conditional_transformed_distribution.py @@ -20,6 +20,7 @@ from __future__ import print_function from tensorflow.contrib.distributions.python.ops import conditional_distribution from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops.distributions import transformed_distribution @@ -105,7 +106,9 @@ class ConditionalTransformedDistribution( bijector_kwargs = bijector_kwargs or {} distribution_kwargs = distribution_kwargs or {} x = self.bijector.inverse(y, **bijector_kwargs) - ildj = self.bijector.inverse_log_det_jacobian(y, **bijector_kwargs) + event_ndims = self._maybe_get_event_ndims_statically() + ildj = self.bijector.inverse_log_det_jacobian( + y, event_ndims=event_ndims, **bijector_kwargs) if self.bijector._is_injective: # pylint: disable=protected-access return self._finish_log_prob_for_one_fiber(y, x, ildj, distribution_kwargs) @@ -128,7 +131,9 @@ class ConditionalTransformedDistribution( bijector_kwargs = bijector_kwargs or {} distribution_kwargs = distribution_kwargs or {} x = self.bijector.inverse(y, **bijector_kwargs) - ildj = self.bijector.inverse_log_det_jacobian(y, **bijector_kwargs) + event_ndims = self._maybe_get_event_ndims_statically() + ildj = self.bijector.inverse_log_det_jacobian( + y, event_ndims=event_ndims, **bijector_kwargs) if self.bijector._is_injective: # pylint: disable=protected-access return self._finish_prob_for_one_fiber(y, x, ildj, distribution_kwargs) @@ -214,3 +219,15 @@ class ConditionalTransformedDistribution( # implies the qth quantile of Y is g(x_q). inv_cdf = self.distribution.quantile(value, **distribution_kwargs) return self.bijector.forward(inv_cdf, **bijector_kwargs) + + def _maybe_get_event_ndims_statically(self): + if self.event_shape.ndims is not None: + return self.event_shape.ndims + + event_ndims = array_ops.size(self.event_shape_tensor()) + static_event_ndims = tensor_util.constant_value(event_ndims) + + if static_event_ndims is not None: + return static_event_ndims + + return event_ndims diff --git a/tensorflow/contrib/distributions/python/ops/deterministic.py b/tensorflow/contrib/distributions/python/ops/deterministic.py index 8049522e9f5dc26b244b7e710a9ae8b981efd6b6..a42350430e98515e521ce357bf5a87ff2daefedc 100644 --- a/tensorflow/contrib/distributions/python/ops/deterministic.py +++ b/tensorflow/contrib/distributions/python/ops/deterministic.py @@ -87,7 +87,7 @@ class _BaseDeterministic(distribution.Distribution): ValueError: If `loc` is a scalar. """ parameters = locals() - with ops.name_scope(name, values=[loc, atol, rtol]): + with ops.name_scope(name, values=[loc, atol, rtol]) as name: loc = ops.convert_to_tensor(loc, name="loc") if is_vector and validate_args: msg = "Argument loc must be at least rank 1." diff --git a/tensorflow/contrib/distributions/python/ops/geometric.py b/tensorflow/contrib/distributions/python/ops/geometric.py index 8f190e48a7148d84082d73771cba4660a1a0d221..53dd42f4c83fcea0ec5b1374c8e3109ebe1dd127 100644 --- a/tensorflow/contrib/distributions/python/ops/geometric.py +++ b/tensorflow/contrib/distributions/python/ops/geometric.py @@ -86,7 +86,7 @@ class Geometric(distribution.Distribution): """ parameters = locals() - with ops.name_scope(name, values=[logits, probs]): + with ops.name_scope(name, values=[logits, probs]) as name: self._logits, self._probs = distribution_util.get_logits_and_probs( logits, probs, validate_args=validate_args, name=name) diff --git a/tensorflow/contrib/distributions/python/ops/gumbel.py b/tensorflow/contrib/distributions/python/ops/gumbel.py index 8d05ad6b8032fb8bada99389959091fb1c28beda..2c261073ee16462599740cb241108bfe08c773ec 100644 --- a/tensorflow/contrib/distributions/python/ops/gumbel.py +++ b/tensorflow/contrib/distributions/python/ops/gumbel.py @@ -125,7 +125,7 @@ class _Gumbel(distribution.Distribution): TypeError: if loc and scale are different dtypes. """ parameters = locals() - with ops.name_scope(name, values=[loc, scale]): + with ops.name_scope(name, values=[loc, scale]) as name: with ops.control_dependencies([check_ops.assert_positive(scale)] if validate_args else []): self._loc = array_ops.identity(loc, name="loc") diff --git a/tensorflow/contrib/distributions/python/ops/half_normal.py b/tensorflow/contrib/distributions/python/ops/half_normal.py index fc0751a6e0b78cb3d79bd3478e740bb05cd26428..d0df2befd6e46ca93e5a0b5d1cb5407d6719c7f2 100644 --- a/tensorflow/contrib/distributions/python/ops/half_normal.py +++ b/tensorflow/contrib/distributions/python/ops/half_normal.py @@ -106,7 +106,7 @@ class HalfNormal(distribution.Distribution): name: Python `str` name prefixed to Ops created by this class. """ parameters = locals() - with ops.name_scope(name, values=[scale]): + with ops.name_scope(name, values=[scale]) as name: with ops.control_dependencies([check_ops.assert_positive(scale)] if validate_args else []): self._scale = array_ops.identity(scale, name="scale") diff --git a/tensorflow/contrib/distributions/python/ops/independent.py b/tensorflow/contrib/distributions/python/ops/independent.py index b1bacb91b03093fa93a7e5f7eb855dc944dafb44..fbde55ef310de1d926b8ddd503499fbed4809373 100644 --- a/tensorflow/contrib/distributions/python/ops/independent.py +++ b/tensorflow/contrib/distributions/python/ops/independent.py @@ -119,7 +119,7 @@ class Independent(distribution_lib.Distribution): parameters = locals() name = name or "Independent" + distribution.name self._distribution = distribution - with ops.name_scope(name): + with ops.name_scope(name) as name: if reinterpreted_batch_ndims is None: reinterpreted_batch_ndims = self._get_default_reinterpreted_batch_ndims( distribution) diff --git a/tensorflow/contrib/distributions/python/ops/inverse_gamma.py b/tensorflow/contrib/distributions/python/ops/inverse_gamma.py index 51ac61dcf640ca89f22c47127bda71316a179ca4..502bd4f493337bab180129cd0ddfaf5a76a0ca4e 100644 --- a/tensorflow/contrib/distributions/python/ops/inverse_gamma.py +++ b/tensorflow/contrib/distributions/python/ops/inverse_gamma.py @@ -126,7 +126,7 @@ class InverseGamma(distribution.Distribution): TypeError: if `concentration` and `rate` are different dtypes. """ parameters = locals() - with ops.name_scope(name, values=[concentration, rate]): + with ops.name_scope(name, values=[concentration, rate]) as name: with ops.control_dependencies([ check_ops.assert_positive(concentration), check_ops.assert_positive(rate), @@ -281,7 +281,7 @@ class InverseGammaWithSoftplusConcentrationRate(InverseGamma): allow_nan_stats=True, name="InverseGammaWithSoftplusConcentrationRate"): parameters = locals() - with ops.name_scope(name, values=[concentration, rate]): + with ops.name_scope(name, values=[concentration, rate]) as name: super(InverseGammaWithSoftplusConcentrationRate, self).__init__( concentration=nn.softplus(concentration, name="softplus_concentration"), diff --git a/tensorflow/contrib/distributions/python/ops/kumaraswamy.py b/tensorflow/contrib/distributions/python/ops/kumaraswamy.py index 192dede6ff1d4de8d4be9965c414e7453d7b5d4b..66682b2ff5493f8565410138e770b45ffc6b5d77 100644 --- a/tensorflow/contrib/distributions/python/ops/kumaraswamy.py +++ b/tensorflow/contrib/distributions/python/ops/kumaraswamy.py @@ -151,10 +151,11 @@ class Kumaraswamy(transformed_distribution.TransformedDistribution): 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") + with ops.name_scope(name, values=[concentration1, concentration0]) as name: + 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), diff --git a/tensorflow/contrib/distributions/python/ops/logistic.py b/tensorflow/contrib/distributions/python/ops/logistic.py index 68e6bca5a554b29a450911073eb5c4fe55f313c6..c83b5bc2e3a8c56f5c52d063a7d0d399be1c1870 100644 --- a/tensorflow/contrib/distributions/python/ops/logistic.py +++ b/tensorflow/contrib/distributions/python/ops/logistic.py @@ -120,7 +120,7 @@ class Logistic(distribution.Distribution): TypeError: if loc and scale are different dtypes. """ parameters = locals() - with ops.name_scope(name, values=[loc, scale]): + with ops.name_scope(name, values=[loc, scale]) as name: with ops.control_dependencies([check_ops.assert_positive(scale)] if validate_args else []): self._loc = array_ops.identity(loc, name="loc") diff --git a/tensorflow/contrib/distributions/python/ops/mixture.py b/tensorflow/contrib/distributions/python/ops/mixture.py index cef6a143fc615901315a3780bf4ed53b8c7cd177..2ef294af2e8bc9beff735ec2e0fd6b619ce96176 100644 --- a/tensorflow/contrib/distributions/python/ops/mixture.py +++ b/tensorflow/contrib/distributions/python/ops/mixture.py @@ -145,7 +145,7 @@ class Mixture(distribution.Distribution): "none of the components provide a static number of ndims") # Ensure that all batch and event ndims are consistent. - with ops.name_scope(name, values=[cat.logits]): + with ops.name_scope(name, values=[cat.logits]) as name: num_components = cat.event_size static_num_components = tensor_util.constant_value(num_components) if static_num_components is None: diff --git a/tensorflow/contrib/distributions/python/ops/mixture_same_family.py b/tensorflow/contrib/distributions/python/ops/mixture_same_family.py index b93bdc5ab4010663baddda1410b302644853648b..0b1301e551728f74bb0048d2dcf3c356ae110c75 100644 --- a/tensorflow/contrib/distributions/python/ops/mixture_same_family.py +++ b/tensorflow/contrib/distributions/python/ops/mixture_same_family.py @@ -131,7 +131,7 @@ class MixtureSameFamily(distribution.Distribution): `components_distribution` rightmost batch shape. """ parameters = locals() - with ops.name_scope(name): + with ops.name_scope(name) as name: self._mixture_distribution = mixture_distribution self._components_distribution = components_distribution self._runtime_assertions = [] diff --git a/tensorflow/contrib/distributions/python/ops/mvn_diag.py b/tensorflow/contrib/distributions/python/ops/mvn_diag.py index e862552880f4073c8fa8e90134d0633e7484b0bf..e3236c2db93695a5e007bba9a1414773f3935f2e 100644 --- a/tensorflow/contrib/distributions/python/ops/mvn_diag.py +++ b/tensorflow/contrib/distributions/python/ops/mvn_diag.py @@ -194,7 +194,7 @@ class MultivariateNormalDiag( ValueError: if at most `scale_identity_multiplier` is specified. """ parameters = locals() - with ops.name_scope(name): + with ops.name_scope(name) as name: with ops.name_scope("init", values=[ loc, scale_diag, scale_identity_multiplier]): # No need to validate_args while making diag_scale. The returned @@ -225,7 +225,7 @@ class MultivariateNormalDiagWithSoftplusScale(MultivariateNormalDiag): allow_nan_stats=True, name="MultivariateNormalDiagWithSoftplusScale"): parameters = locals() - with ops.name_scope(name, values=[scale_diag]): + with ops.name_scope(name, values=[scale_diag]) as name: super(MultivariateNormalDiagWithSoftplusScale, self).__init__( loc=loc, scale_diag=nn.softplus(scale_diag), diff --git a/tensorflow/contrib/distributions/python/ops/mvn_diag_plus_low_rank.py b/tensorflow/contrib/distributions/python/ops/mvn_diag_plus_low_rank.py index 413e88f03ae0286c294f3404549a73e1a47dcff7..2f6a6f198cbcfbdcbd0993d3074ddde1c389585f 100644 --- a/tensorflow/contrib/distributions/python/ops/mvn_diag_plus_low_rank.py +++ b/tensorflow/contrib/distributions/python/ops/mvn_diag_plus_low_rank.py @@ -218,7 +218,7 @@ class MultivariateNormalDiagPlusLowRank( parameters = locals() def _convert_to_tensor(x, name): return None if x is None else ops.convert_to_tensor(x, name=name) - with ops.name_scope(name): + with ops.name_scope(name) as name: with ops.name_scope("init", values=[ loc, scale_diag, scale_identity_multiplier, scale_perturb_factor, scale_perturb_diag]): diff --git a/tensorflow/contrib/distributions/python/ops/mvn_full_covariance.py b/tensorflow/contrib/distributions/python/ops/mvn_full_covariance.py index 4bea99fbb75349f97fde473cb5716fe6c426ce90..86fcd4db54ad851af670839a5a93a68dce047675 100644 --- a/tensorflow/contrib/distributions/python/ops/mvn_full_covariance.py +++ b/tensorflow/contrib/distributions/python/ops/mvn_full_covariance.py @@ -159,7 +159,7 @@ class MultivariateNormalFullCovariance(mvn_tril.MultivariateNormalTriL): parameters = locals() # Convert the covariance_matrix up to a scale_tril and call MVNTriL. - with ops.name_scope(name): + with ops.name_scope(name) as name: with ops.name_scope("init", values=[loc, covariance_matrix]): if covariance_matrix is None: scale_tril = None diff --git a/tensorflow/contrib/distributions/python/ops/mvn_linear_operator.py b/tensorflow/contrib/distributions/python/ops/mvn_linear_operator.py index a7399792892f4c179c05168184d76ec95c168b51..44c92312c7dc758500051f89923ec9fafe850c0e 100644 --- a/tensorflow/contrib/distributions/python/ops/mvn_linear_operator.py +++ b/tensorflow/contrib/distributions/python/ops/mvn_linear_operator.py @@ -176,7 +176,7 @@ class MultivariateNormalLinearOperator( if not scale.dtype.is_floating: raise TypeError("`scale` parameter must have floating-point dtype.") - with ops.name_scope(name, values=[loc] + scale.graph_parents): + with ops.name_scope(name, values=[loc] + scale.graph_parents) as name: # Since expand_dims doesn't preserve constant-ness, we obtain the # non-dynamic value if possible. loc = ops.convert_to_tensor(loc, name="loc") if loc is not None else loc diff --git a/tensorflow/contrib/distributions/python/ops/mvn_tril.py b/tensorflow/contrib/distributions/python/ops/mvn_tril.py index 6c7dc4ca7aaf5b3a20b072e9360d15528ad10556..d6f8b731cbeed5fed3b43365e7c668d0434a267e 100644 --- a/tensorflow/contrib/distributions/python/ops/mvn_tril.py +++ b/tensorflow/contrib/distributions/python/ops/mvn_tril.py @@ -184,7 +184,7 @@ class MultivariateNormalTriL( return None if x is None else ops.convert_to_tensor(x, name=name) if loc is None and scale_tril is None: raise ValueError("Must specify one or both of `loc`, `scale_tril`.") - with ops.name_scope(name): + with ops.name_scope(name) as name: with ops.name_scope("init", values=[loc, scale_tril]): loc = _convert_to_tensor(loc, name="loc") scale_tril = _convert_to_tensor(scale_tril, name="scale_tril") diff --git a/tensorflow/contrib/distributions/python/ops/negative_binomial.py b/tensorflow/contrib/distributions/python/ops/negative_binomial.py index 3a58df80da6c02b056f5e5a63bf41de5fc6d44a4..eeaf9c0a5ebc1323e137ff73f82588f6907031c7 100644 --- a/tensorflow/contrib/distributions/python/ops/negative_binomial.py +++ b/tensorflow/contrib/distributions/python/ops/negative_binomial.py @@ -91,7 +91,7 @@ class NegativeBinomial(distribution.Distribution): """ parameters = locals() - with ops.name_scope(name, values=[total_count, logits, probs]): + with ops.name_scope(name, values=[total_count, logits, probs]) as name: self._logits, self._probs = distribution_util.get_logits_and_probs( logits, probs, validate_args=validate_args, name=name) with ops.control_dependencies( diff --git a/tensorflow/contrib/distributions/python/ops/onehot_categorical.py b/tensorflow/contrib/distributions/python/ops/onehot_categorical.py index e3e40b2e9ca232b9970768f21fb95887fdf0df2d..305b138fdc2318523ee078195213caf865d96b4d 100644 --- a/tensorflow/contrib/distributions/python/ops/onehot_categorical.py +++ b/tensorflow/contrib/distributions/python/ops/onehot_categorical.py @@ -116,7 +116,7 @@ class OneHotCategorical(distribution.Distribution): name: Python `str` name prefixed to Ops created by this class. """ parameters = locals() - with ops.name_scope(name, values=[logits, probs]): + with ops.name_scope(name, values=[logits, probs]) as name: self._logits, self._probs = distribution_util.get_logits_and_probs( name=name, logits=logits, probs=probs, validate_args=validate_args, multidimensional=True) diff --git a/tensorflow/contrib/distributions/python/ops/poisson.py b/tensorflow/contrib/distributions/python/ops/poisson.py index 02e97c0a2fd004c4fa9382d5367af9f5b034a869..a84aad6fc9372395ac021fa3aa006ddf9272e6a9 100644 --- a/tensorflow/contrib/distributions/python/ops/poisson.py +++ b/tensorflow/contrib/distributions/python/ops/poisson.py @@ -94,7 +94,7 @@ class Poisson(distribution.Distribution): TypeError: if `log_rate` is not a float-type. """ parameters = locals() - with ops.name_scope(name, values=[rate]): + with ops.name_scope(name, values=[rate]) as name: if (rate is None) == (log_rate is None): raise ValueError("Must specify exactly one of `rate` and `log_rate`.") elif log_rate is None: diff --git a/tensorflow/contrib/distributions/python/ops/poisson_lognormal.py b/tensorflow/contrib/distributions/python/ops/poisson_lognormal.py index 92f2bba1828696248c9d9460566a08ba372c3358..19c99dcee92978e938a73af9be445cd098e5fe90 100644 --- a/tensorflow/contrib/distributions/python/ops/poisson_lognormal.py +++ b/tensorflow/contrib/distributions/python/ops/poisson_lognormal.py @@ -114,7 +114,7 @@ def quadrature_scheme_lognormal_quantiles( # Create a LogNormal distribution. dist = transformed_lib.TransformedDistribution( distribution=normal_lib.Normal(loc=loc, scale=scale), - bijector=Exp(event_ndims=0), + bijector=Exp(), validate_args=validate_args) batch_ndims = dist.batch_shape.ndims if batch_ndims is None: @@ -256,7 +256,7 @@ class PoissonLogNormalQuadratureCompound(distribution_lib.Distribution): `dtype`. """ parameters = locals() - with ops.name_scope(name, values=[loc, scale]): + with ops.name_scope(name, values=[loc, scale]) as name: if loc is not None: loc = ops.convert_to_tensor(loc, name="loc") if scale is not None: diff --git a/tensorflow/contrib/distributions/python/ops/quantized_distribution.py b/tensorflow/contrib/distributions/python/ops/quantized_distribution.py index 8aebb79b9138cce1373e6472d17cf9072d2bc285..1ef7651d03a3388e72618b1d9bb8b819bde17e92 100644 --- a/tensorflow/contrib/distributions/python/ops/quantized_distribution.py +++ b/tensorflow/contrib/distributions/python/ops/quantized_distribution.py @@ -217,7 +217,7 @@ class QuantizedDistribution(distributions.Distribution): values = ( list(distribution.parameters.values()) + [low, high]) - with ops.name_scope(name, values=values): + with ops.name_scope(name, values=values) as name: self._dist = distribution if low is not None: diff --git a/tensorflow/contrib/distributions/python/ops/relaxed_bernoulli.py b/tensorflow/contrib/distributions/python/ops/relaxed_bernoulli.py index e454a53c6275e0c60edd8c87b1c3be670f2b22de..84c8d29072c2f1f3888329638c4695bccf70eab7 100644 --- a/tensorflow/contrib/distributions/python/ops/relaxed_bernoulli.py +++ b/tensorflow/contrib/distributions/python/ops/relaxed_bernoulli.py @@ -166,7 +166,7 @@ class RelaxedBernoulli(transformed_distribution.TransformedDistribution): ValueError: If both `probs` and `logits` are passed, or if neither. """ parameters = locals() - with ops.name_scope(name, values=[logits, probs, temperature]): + with ops.name_scope(name, values=[logits, probs, temperature]) as name: with ops.control_dependencies([check_ops.assert_positive(temperature)] if validate_args else []): self._temperature = array_ops.identity(temperature, name="temperature") diff --git a/tensorflow/contrib/distributions/python/ops/relaxed_onehot_categorical.py b/tensorflow/contrib/distributions/python/ops/relaxed_onehot_categorical.py index f56ba0781604cb5a4fb3070b79aa86e09ceb6766..325f41e37c928ba8e81e45e63a7f7f8126bc80f8 100644 --- a/tensorflow/contrib/distributions/python/ops/relaxed_onehot_categorical.py +++ b/tensorflow/contrib/distributions/python/ops/relaxed_onehot_categorical.py @@ -163,7 +163,7 @@ class ExpRelaxedOneHotCategorical(distribution.Distribution): name: Python `str` name prefixed to Ops created by this class. """ parameters = locals() - with ops.name_scope(name, values=[logits, probs, temperature]): + with ops.name_scope(name, values=[logits, probs, temperature]) as name: self._logits, self._probs = distribution_util.get_logits_and_probs( name=name, logits=logits, probs=probs, validate_args=validate_args, @@ -409,5 +409,5 @@ class RelaxedOneHotCategorical( validate_args=validate_args, allow_nan_stats=allow_nan_stats) super(RelaxedOneHotCategorical, self).__init__(dist, - bijectors.Exp(event_ndims=1), + bijectors.Exp(), name=name) diff --git a/tensorflow/contrib/distributions/python/ops/shape.py b/tensorflow/contrib/distributions/python/ops/shape.py index bac0b79d5908712f4e64259768fb6f3b4558f620..6a7f28713acefd2285b07a212e2e47a6db1ae5e1 100644 --- a/tensorflow/contrib/distributions/python/ops/shape.py +++ b/tensorflow/contrib/distributions/python/ops/shape.py @@ -439,7 +439,7 @@ class _DistributionShape(object): if self._batch_ndims_is_0 and expand_batch_dim: squeeze_dims += [1] if squeeze_dims: - x = array_ops.squeeze(x, squeeze_dims=squeeze_dims) + x = array_ops.squeeze(x, axis=squeeze_dims) # x.shape: [prod(S)]+B+E _, batch_shape, event_shape = self.get_shape(x) else: diff --git a/tensorflow/contrib/distributions/python/ops/sinh_arcsinh.py b/tensorflow/contrib/distributions/python/ops/sinh_arcsinh.py index 0d8a1926913766da374cb65767dccfa28bf75579..03828fa61277eeaf7ce90de8023b4ed91f6cc4dc 100644 --- a/tensorflow/contrib/distributions/python/ops/sinh_arcsinh.py +++ b/tensorflow/contrib/distributions/python/ops/sinh_arcsinh.py @@ -134,7 +134,8 @@ class SinhArcsinh(transformed_distribution.TransformedDistribution): """ parameters = locals() - with ops.name_scope(name, values=[loc, scale, skewness, tailweight]): + with ops.name_scope(name, + values=[loc, scale, skewness, tailweight]) as name: loc = ops.convert_to_tensor(loc, name="loc") dtype = loc.dtype scale = ops.convert_to_tensor(scale, name="scale", dtype=dtype) @@ -166,13 +167,13 @@ class SinhArcsinh(transformed_distribution.TransformedDistribution): # Make the SAS bijector, 'F'. f = bijectors.SinhArcsinh( - skewness=skewness, tailweight=tailweight, event_ndims=0) + skewness=skewness, tailweight=tailweight) if has_default_skewness: f_noskew = f else: f_noskew = bijectors.SinhArcsinh( skewness=skewness.dtype.as_numpy_dtype(0.), - tailweight=tailweight, event_ndims=0) + tailweight=tailweight) # 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)) diff --git a/tensorflow/contrib/distributions/python/ops/statistical_testing.py b/tensorflow/contrib/distributions/python/ops/statistical_testing.py index 5c52015e5fd5d7f3cc49e337520296e58ea08bfb..9c69435fac109914ff29b307dfad105f62849339 100644 --- a/tensorflow/contrib/distributions/python/ops/statistical_testing.py +++ b/tensorflow/contrib/distributions/python/ops/statistical_testing.py @@ -130,7 +130,7 @@ 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 clip_ops from tensorflow.python.ops import gen_math_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn_ops @@ -169,31 +169,27 @@ def _do_maximum_mean(samples, envelope, high, name=None): 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 + + # The maximum mean is given by taking `envelope`-worth of + # probability from the smallest samples and moving it to the + # maximum value. This amounts to: + # - ignoring the smallest k samples, where `k/n < envelope` + # - taking a `1/n - (envelope - k/n)` part of the index k sample + # - taking all the other samples + # - and adding `envelope * high` at the end. + # The following is a vectorized and batched way of computing this. + # `max_mean_contrib` is a mask implementing the previous. + batch_size = array_ops.shape(samples)[-1] + batch_size = math_ops.cast(batch_size, dtype=samples.dtype.base_dtype) + step = 1. / batch_size + cum_steps = step * math_ops.range( + 1, batch_size + 1, dtype=samples.dtype.base_dtype) + max_mean_contrib = clip_ops.clip_by_value( + cum_steps - envelope[..., array_ops.newaxis], + clip_value_min=0., + clip_value_max=step) + return math_ops.reduce_sum( + samples * max_mean_contrib, axis=-1) + envelope * high def _maximum_mean(samples, envelope, high, name=None): @@ -234,7 +230,7 @@ def _maximum_mean(samples, envelope, high, name=None): envelope = ops.convert_to_tensor(envelope, name="envelope") high = ops.convert_to_tensor(high, name="high") - xmax = math_ops.reduce_max(samples, axis=[-1]) + xmax = math_ops.reduce_max(samples, axis=[0]) msg = "Given sample maximum value exceeds expectations" check_op = check_ops.assert_less_equal(xmax, high, message=msg) with ops.control_dependencies([check_op]): @@ -279,7 +275,7 @@ def _minimum_mean(samples, envelope, low, name=None): envelope = ops.convert_to_tensor(envelope, name="envelope") low = ops.convert_to_tensor(low, name="low") - xmin = math_ops.reduce_min(samples, axis=[-1]) + xmin = math_ops.reduce_min(samples, axis=[0]) 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]): @@ -319,8 +315,8 @@ def _dkwm_cdf_envelope(n, error_rate, name=None): 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. +def _check_shape_dominates(samples, parameters): + """Check that broadcasting `samples` against `parameters` 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 @@ -328,24 +324,27 @@ def _check_shape_dominates(tensor, tensors): sample counts end up inflated. Args: - tensor: A Tensor whose shape is to be protected against broadcasting. - tensors: A list of Tensors to check + samples: A Tensor whose shape is to be protected against broadcasting. + parameters: A list of Tensors who are parameters for the statistical test. Returns: - tensor: `tf.identity(tensor)` with control dependencies attached; - be sure to use that downstream. + samples: Return original `samples` with control dependencies attached + to ensure no broadcasting. """ def check(t): - target = array_ops.shape(tensor)[1:] - result = array_ops.broadcast_dynamic_shape(target, array_ops.shape(t)) + samples_batch_shape = array_ops.shape(samples)[1:] + broadcasted_batch_shape = array_ops.broadcast_dynamic_shape( + samples_batch_shape, 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])) + samples_batch_ndims = array_ops.size(samples_batch_shape) + ge = check_ops.assert_greater_equal( + samples_batch_ndims, array_ops.rank(t)) + eq = check_ops.assert_equal(samples_batch_shape, broadcasted_batch_shape) + return ge, eq + checks = list(itertools.chain(*[check(t) for t in parameters])) with ops.control_dependencies(checks): - return array_ops.identity(array_ops.identity(tensor)) + return array_ops.identity(samples) def true_mean_confidence_interval_by_dkwm( @@ -684,9 +683,13 @@ def assert_true_mean_equal_by_dkwm_two_sample( # 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) + check_confidence_intervals_can_intersect = check_ops.assert_greater_equal( + max_mean_1, min_mean_2, message="Confidence intervals do not " + "intersect: samples1 has a smaller mean than samples2") + with ops.control_dependencies([check_confidence_intervals_can_intersect]): + return check_ops.assert_less_equal( + min_mean_1, max_mean_2, message="Confidence intervals do not " + "intersect: samples2 has a smaller mean than samples1") def min_discrepancy_of_true_means_detectable_by_dkwm_two_sample( diff --git a/tensorflow/contrib/distributions/python/ops/vector_diffeomixture.py b/tensorflow/contrib/distributions/python/ops/vector_diffeomixture.py index 971d65c4a69140161461fdac93bb588014dd3e88..af6ff8162b173015dca2d568e13d63127af7853a 100644 --- a/tensorflow/contrib/distributions/python/ops/vector_diffeomixture.py +++ b/tensorflow/contrib/distributions/python/ops/vector_diffeomixture.py @@ -396,7 +396,7 @@ class VectorDiffeomixture(distribution_lib.Distribution): ValueError: if `not distribution.is_scalar_event`. """ parameters = locals() - with ops.name_scope(name, values=[mix_loc, temperature]): + with ops.name_scope(name, values=[mix_loc, temperature]) as name: if not scale or len(scale) < 2: raise ValueError("Must specify list (or list-like object) of scale " "LinearOperators, one for each component with " @@ -427,7 +427,6 @@ class VectorDiffeomixture(distribution_lib.Distribution): self._endpoint_affine = [ AffineLinearOperator(shift=loc_, scale=scale_, - event_ndims=1, validate_args=validate_args, name="endpoint_affine_{}".format(k)) for k, (loc_, scale_) in enumerate(zip(loc, scale))] @@ -467,7 +466,6 @@ class VectorDiffeomixture(distribution_lib.Distribution): self._interpolated_affine = [ AffineLinearOperator(shift=loc_, scale=scale_, - event_ndims=1, validate_args=validate_args, name="interpolated_affine_{}".format(k)) for k, (loc_, scale_) in enumerate(zip( @@ -621,9 +619,11 @@ class VectorDiffeomixture(distribution_lib.Distribution): log_prob = math_ops.reduce_sum(self.distribution.log_prob(y), axis=-2) # Because the affine transformation has a constant Jacobian, it is the case # that `affine.fldj(x) = -affine.ildj(x)`. This is not true in general. - fldj = array_ops.stack( - [aff.forward_log_det_jacobian(x) for aff in self.interpolated_affine], - axis=-1) + fldj = array_ops.stack([ + aff.forward_log_det_jacobian( + x, + event_ndims=array_ops.rank(self.event_shape_tensor()) + ) for aff in self.interpolated_affine], axis=-1) return math_ops.reduce_logsumexp( self.mixture_distribution.logits - fldj + log_prob, axis=-1) diff --git a/tensorflow/contrib/distributions/python/ops/vector_exponential_diag.py b/tensorflow/contrib/distributions/python/ops/vector_exponential_diag.py index 526fe2d39aef9aed833b889de80e849c469435e7..e265b5d0f7c10b2782a1a8924babdca9b986f622 100644 --- a/tensorflow/contrib/distributions/python/ops/vector_exponential_diag.py +++ b/tensorflow/contrib/distributions/python/ops/vector_exponential_diag.py @@ -176,7 +176,7 @@ class VectorExponentialDiag( ValueError: if at most `scale_identity_multiplier` is specified. """ parameters = locals() - with ops.name_scope(name): + with ops.name_scope(name) as name: with ops.name_scope("init", values=[ loc, scale_diag, scale_identity_multiplier]): # No need to validate_args while making diag_scale. The returned diff --git a/tensorflow/contrib/distributions/python/ops/vector_exponential_linear_operator.py b/tensorflow/contrib/distributions/python/ops/vector_exponential_linear_operator.py index 9d5fd9ac4178a1ae29b1ce32f304b22fd3d234dc..89136d6760bb663b5ff86a77c5945ce900f072b9 100644 --- a/tensorflow/contrib/distributions/python/ops/vector_exponential_linear_operator.py +++ b/tensorflow/contrib/distributions/python/ops/vector_exponential_linear_operator.py @@ -181,7 +181,7 @@ class VectorExponentialLinearOperator( if not scale.dtype.is_floating: raise TypeError("`scale` parameter must have floating-point dtype.") - with ops.name_scope(name, values=[loc] + scale.graph_parents): + with ops.name_scope(name, values=[loc] + scale.graph_parents) as name: # Since expand_dims doesn't preserve constant-ness, we obtain the # non-dynamic value if possible. loc = ops.convert_to_tensor(loc, name="loc") if loc is not None else loc 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 003c66b9413fdcad20fbcc8b4bf47259692932e7..1438ede26500bca4541fa9b2020ff22d4c071098 100644 --- a/tensorflow/contrib/distributions/python/ops/vector_sinh_arcsinh_diag.py +++ b/tensorflow/contrib/distributions/python/ops/vector_sinh_arcsinh_diag.py @@ -169,7 +169,7 @@ class VectorSinhArcsinhDiag(transformed_distribution.TransformedDistribution): name, values=[ loc, scale_diag, scale_identity_multiplier, skewness, tailweight - ]): + ]) as name: loc = ops.convert_to_tensor(loc, name="loc") if loc is not None else loc tailweight = 1. if tailweight is None else tailweight has_default_skewness = skewness is None @@ -215,13 +215,13 @@ class VectorSinhArcsinhDiag(transformed_distribution.TransformedDistribution): tailweight = ops.convert_to_tensor( tailweight, dtype=dtype, name="tailweight") f = bijectors.SinhArcsinh( - skewness=skewness, tailweight=tailweight, event_ndims=1) + skewness=skewness, tailweight=tailweight) if has_default_skewness: f_noskew = f else: f_noskew = bijectors.SinhArcsinh( skewness=skewness.dtype.as_numpy_dtype(0.), - tailweight=tailweight, event_ndims=0) + tailweight=tailweight) # Make the Affine bijector, Z --> loc + C * Z. c = 2 * scale_diag_part / f_noskew.forward( diff --git a/tensorflow/contrib/distributions/python/ops/vector_student_t.py b/tensorflow/contrib/distributions/python/ops/vector_student_t.py index 887981d64ef077e2636f8031581c390f177edac8..7e78ded9df07564126b46b6beeeccf95bf1eef94 100644 --- a/tensorflow/contrib/distributions/python/ops/vector_student_t.py +++ b/tensorflow/contrib/distributions/python/ops/vector_student_t.py @@ -178,7 +178,7 @@ class _VectorStudentT(transformed_distribution.TransformedDistribution): parameters = locals() graph_parents = [df, loc, scale_identity_multiplier, scale_diag, scale_tril, scale_perturb_factor, scale_perturb_diag] - with ops.name_scope(name): + with ops.name_scope(name) as name: with ops.name_scope("init", values=graph_parents): # The shape of the _VectorStudentT distribution is governed by the # relationship between df.batch_shape and affine.batch_shape. In diff --git a/tensorflow/contrib/distributions/python/ops/wishart.py b/tensorflow/contrib/distributions/python/ops/wishart.py index 5a8c94dabf4c3c430bee544a48ee7acfe7dd7ed0..91453fed5d279178a0e062b71dad3b0f957b11b4 100644 --- a/tensorflow/contrib/distributions/python/ops/wishart.py +++ b/tensorflow/contrib/distributions/python/ops/wishart.py @@ -109,7 +109,7 @@ class _WishartLinearOperator(distribution.Distribution): """ parameters = locals() self._cholesky_input_output_matrices = cholesky_input_output_matrices - with ops.name_scope(name) as ns: + with ops.name_scope(name) as name: with ops.name_scope("init", values=[df, scale_operator]): if not scale_operator.dtype.is_floating: raise TypeError( @@ -163,7 +163,7 @@ class _WishartLinearOperator(distribution.Distribution): parameters=parameters, graph_parents=([self._df, self._dimension] + self._scale_operator.graph_parents), - name=ns) + name=name) @property def df(self): @@ -531,7 +531,7 @@ class WishartCholesky(_WishartLinearOperator): name: Python `str` name prefixed to Ops created by this class. """ parameters = locals() - with ops.name_scope(name, values=[scale]): + with ops.name_scope(name, values=[scale]) as name: with ops.name_scope("init", values=[scale]): scale = ops.convert_to_tensor(scale) if validate_args: @@ -647,7 +647,7 @@ class WishartFull(_WishartLinearOperator): name: Python `str` name prefixed to Ops created by this class. """ parameters = locals() - with ops.name_scope(name) as ns: + with ops.name_scope(name) as name: with ops.name_scope("init", values=[scale]): scale = ops.convert_to_tensor(scale) if validate_args: @@ -666,5 +666,5 @@ class WishartFull(_WishartLinearOperator): cholesky_input_output_matrices=cholesky_input_output_matrices, validate_args=validate_args, allow_nan_stats=allow_nan_stats, - name=ns) + name=name) self._parameters = parameters diff --git a/tensorflow/contrib/eager/proto/BUILD b/tensorflow/contrib/eager/proto/BUILD deleted file mode 100644 index b016d2dcb504044372c895e1eedf3511751bc13e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/eager/proto/BUILD +++ /dev/null @@ -1,13 +0,0 @@ -licenses(["notice"]) # Apache 2.0 - -exports_files(["LICENSE"]) - -load("//tensorflow/core:platform/default/build_config.bzl", "tf_proto_library") - -tf_proto_library( - name = "checkpointable_object_graph_proto", - srcs = [ - "checkpointable_object_graph.proto", - ], - visibility = ["//tensorflow/contrib/eager/python:__subpackages__"], -) diff --git a/tensorflow/contrib/eager/python/BUILD b/tensorflow/contrib/eager/python/BUILD index 4e088503bf12341fa70b6d119d717068344d9379..99abbae03fc14f241dae27f317902f7335819037 100644 --- a/tensorflow/contrib/eager/python/BUILD +++ b/tensorflow/contrib/eager/python/BUILD @@ -11,7 +11,6 @@ py_library( srcs_version = "PY2AND3", visibility = ["//visibility:public"], deps = [ - ":checkpointable_utils", ":datasets", ":metrics", ":network", @@ -19,15 +18,14 @@ py_library( "//tensorflow/python:framework_ops", "//tensorflow/python:framework_test_lib", "//tensorflow/python:gradients", - "//tensorflow/python:numerics", "//tensorflow/python:resource_variable_ops", "//tensorflow/python:script_ops", "//tensorflow/python:template", + "//tensorflow/python:training", "//tensorflow/python:util", "//tensorflow/python:variable_scope", "//tensorflow/python/eager:backprop", "//tensorflow/python/eager:context", - "//tensorflow/python/eager:core", "//tensorflow/python/eager:execution_callbacks", "//tensorflow/python/eager:function", ], @@ -70,13 +68,15 @@ cuda_py_test( srcs = ["datasets_test.py"], additional_deps = [ ":datasets", - ":checkpointable_utils", - "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/contrib/data/python/ops:prefetching_ops", + "//tensorflow/contrib/data/python/ops:threadpool", + "//tensorflow/contrib/data/python/ops:unique", "//tensorflow/contrib/lookup:lookup_py", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", "//tensorflow/python:math_ops", "//tensorflow/python:script_ops", + "//tensorflow/python:training", "//tensorflow/python/data", "//tensorflow/python/eager:test", ], @@ -119,14 +119,14 @@ 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:checkpointable", "//tensorflow/python:control_flow_ops", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", "//tensorflow/python:init_ops", "//tensorflow/python:math_ops", + "//tensorflow/python:summary_ops_v2", "//tensorflow/python:util", "//tensorflow/python:variable_scope", "//tensorflow/python/eager:context", @@ -140,11 +140,11 @@ py_test( srcs_version = "PY2AND3", deps = [ ":metrics", - "//tensorflow/contrib/summary:summary_ops", "//tensorflow/contrib/summary:summary_test_util", "//tensorflow/python:array_ops", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", + "//tensorflow/python:summary_ops_v2", "//tensorflow/python:training", "//tensorflow/python/eager:context", "//tensorflow/python/eager:test", @@ -161,10 +161,10 @@ py_library( deps = [ ":datasets", ":metrics", - "//tensorflow/contrib/summary:summary_ops", "//tensorflow/python:control_flow_ops", "//tensorflow/python:errors", "//tensorflow/python:framework_ops", + "//tensorflow/python:summary_ops_v2", "//tensorflow/python/eager:context", "//tensorflow/python/eager:function", "@six_archive//:six", @@ -223,56 +223,3 @@ py_test( "//tensorflow/python/eager:test", ], ) - -py_library( - 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:errors", - "//tensorflow/python:framework_ops", - "//tensorflow/python:init_ops", - "//tensorflow/python:pywrap_tensorflow", - "//tensorflow/python:resource_variable_ops", - "//tensorflow/python:session", - "//tensorflow/python:tensor_shape", - "//tensorflow/python:training", - "//tensorflow/python:util", - "//tensorflow/python:variable_scope", - "//tensorflow/python/eager:context", - ], -) - -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", - "//tensorflow/python/keras", - ], - tags = [ - "no_windows", # TODO: needs investigation on Windows - "notsan", # b/74395663 - ], -) diff --git a/tensorflow/contrib/eager/python/checkpointable_utils.py b/tensorflow/contrib/eager/python/checkpointable_utils.py deleted file mode 100644 index 34cb8d0e0887bd5e440873bae117bf27597de11b..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/eager/python/checkpointable_utils.py +++ /dev/null @@ -1,982 +0,0 @@ -"""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 functools -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 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) - - -def _copy_saver_with_new_var_list(old_saver, new_var_list): - """Copy a `tf.train.Saver`'s state to a new Saver with different variables.""" - new_saver = saver_lib.Saver(var_list=new_var_list) - # TODO(allenl): Move to copying functionality to Saver? - # pylint: disable=protected-access - new_saver._last_checkpoints = old_saver._last_checkpoints - new_saver._checkpoints_to_be_deleted = old_saver._checkpoints_to_be_deleted - new_saver._next_checkpoint_time = old_saver._next_checkpoint_time - # pylint: enable=protected-access - return new_saver - - -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._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 - # When executing eagerly, we need to re-create SaveableObjects each time - # save() is called so they pick up new Tensors passed to their - # constructors. That means the Saver needs to be copied with a new - # var_list. - or context.executing_eagerly()): - if self._last_save_object_graph is not None: - self._last_save_saver = _copy_saver_with_new_var_list( - old_saver=self._last_save_saver, new_var_list=named_variables) - else: - self._last_save_saver = saver_lib.Saver(var_list=named_variables) - self._last_save_object_graph = graph_proto - with ops.device("/cpu:0"): - save_path = self._last_save_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): - """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. - - 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. - - 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: - file_prefix_tensor = self._file_prefix_placeholder - file_prefix_feed_dict = {self._file_prefix_placeholder: save_path} - else: - with ops.device("/cpu:0"): - file_prefix_tensor = constant_op.constant(save_path) - file_prefix_feed_dict = None - reader = pywrap_tensorflow.NewCheckpointReader(save_path) - try: - object_graph_string = reader.get_tensor(_OBJECT_GRAPH_PROTO_KEY) - 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: - 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 - - -class _CallbackSaveable(saver_lib.BaseSaverBuilder.SaveableObject): - """Wraps save and restore callbacks as a `SaveableObject`.""" - - def __init__(self, name, dtype, save_callback, restore_callback): - self._restore_callback = restore_callback - spec = saver_lib.BaseSaverBuilder.SaveSpec( - tensor=save_callback, - slice_spec="", - name=name, - dtype=dtype) - super(_CallbackSaveable, self).__init__( - save_callback, [spec], name) - - def restore(self, restored_tensors, restored_shapes): - """Restore the same value into both variables.""" - tensor, = restored_tensors - return self._restore_callback(tensor) - - -class _SplitDependency(core_checkpointable.CheckpointableBase): - """Looks like a regular variable while synchronizing save/restores.""" - - def __init__(self, save_buffer, restore_buffer, name, dtype, num_components, - fill_save_buffer_fn, consume_restore_buffer_fn): - self._save_buffer = save_buffer - self._restore_buffer = restore_buffer - self._name = name - self._dtype = dtype - self._num_components = num_components - self._fill_save_buffer_fn = fill_save_buffer_fn - self._consume_restore_buffer_fn = consume_restore_buffer_fn - - def _save(self): - """Pull from the shared buffer, populating it if necessary.""" - if self._name not in self._save_buffer: - if self._save_buffer: - raise AssertionError( - ("Split dependency %s (%s) unsynchronized. Split dependencies must " - "be saved together.") % (self._name, self)) - self._fill_save_buffer_fn(self._save_buffer) - return self._save_buffer.pop(self._name) - - def _restore(self, tensor): - """Push into the shared buffer, flushing it if necessary.""" - if self._name in self._restore_buffer: - raise AssertionError( - ("Split dependency %s (%s) unsynchronized. Split dependencies must " - "be restored together.") % (self._name, self)) - self._restore_buffer[self._name] = tensor - if len(self._restore_buffer) == self._num_components: - op = self._consume_restore_buffer_fn(self._restore_buffer) - self._restore_buffer.clear() - return op - else: - return control_flow_ops.no_op() - - def _gather_saveables_for_checkpoint(self): - """Looks to Checkpointable like a regular variable.""" - return { - core_checkpointable.VARIABLE_VALUE_KEY: - functools.partial(_CallbackSaveable, - dtype=self._dtype, - save_callback=self._save, - restore_callback=self._restore) - } - - -def split_dependency(component_names, component_dtypes, - fill_save_buffer_fn, consume_restore_buffer_fn): - """Creates multiple dependencies with a synchronized save/restore. - - Useful when a single op produces `Tensor`s which should each be saved under - different objects, or when `Tensor`s saved with many different objects need to - be restored together as inputs to a single op (i.e. an object which uses a - single fused op may be swapped out for a subgraph of objects, and these two - programs are checkpoint compatible). - - Args: - component_names: A sequence of names for the split - dependencies. `fill_save_buffer_fn` must add these keys to the dictionary - it is passed, and `consume_restore_buffer_fn` will receive a dictionary - with these keys. - component_dtypes: Data types for the `Tensor`s being saved and restored, a - sequence corresponding to `component_names`. - fill_save_buffer_fn: A function which takes an empty dictionary as an - argument and adds `Tensor`s with `component_names` as keys. These - `Tensor`s will be saved as if they were individual variables. - consume_restore_buffer_fn: A function which takes a dictionary with - `component_names` as keys mapping to restored individual `Tensor`s and - returns a restore op (or if executing eagerly, runs the restoration and - may return `None`). - - Returns: - A dictionary mapping from names to Checkpointable objects. If one is - reachable from an object as a dependency, the others should be too; adding - dependencies on some but not all of the objects will result in errors. - """ - save_buffer = {} - restore_buffer = {} - split_dependencies = {} - for name, dtype in zip(component_names, component_dtypes): - split_dependencies[name] = _SplitDependency( - save_buffer=save_buffer, - restore_buffer=restore_buffer, - name=name, - dtype=dtype, - num_components=len(component_names), - fill_save_buffer_fn=fill_save_buffer_fn, - consume_restore_buffer_fn=consume_restore_buffer_fn) - return split_dependencies diff --git a/tensorflow/contrib/eager/python/datasets_test.py b/tensorflow/contrib/eager/python/datasets_test.py index f76a896d3d8d795b5a7a0e97b5f688fb0291575a..7b123707cc3a26073088cf2c57c6211e831c19fd 100644 --- a/tensorflow/contrib/eager/python/datasets_test.py +++ b/tensorflow/contrib/eager/python/datasets_test.py @@ -27,7 +27,6 @@ from tensorflow.contrib import lookup from tensorflow.contrib.data.python.ops import prefetching_ops 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 @@ -38,6 +37,7 @@ from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor from tensorflow.python.ops import math_ops from tensorflow.python.ops import script_ops +from tensorflow.python.training import checkpointable_utils class IteratorTest(test.TestCase): diff --git a/tensorflow/contrib/eager/python/evaluator.py b/tensorflow/contrib/eager/python/evaluator.py index 37c8f0d47adbde6932bf409cdcae9a1845d700b5..7949a3f6da293abdd85512209242bae76ab4d816 100644 --- a/tensorflow/contrib/eager/python/evaluator.py +++ b/tensorflow/contrib/eager/python/evaluator.py @@ -22,12 +22,12 @@ import six from tensorflow.contrib.eager.python import datasets from tensorflow.contrib.eager.python import metrics -from tensorflow.contrib.summary import summary_ops from tensorflow.python.eager import context from tensorflow.python.eager import function from tensorflow.python.framework import errors_impl from tensorflow.python.framework import ops from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import summary_ops_v2 as summary_ops class Evaluator(object): diff --git a/tensorflow/contrib/eager/python/examples/spinn/spinn_test.py b/tensorflow/contrib/eager/python/examples/spinn/spinn_test.py index 9adf47d505fc2933d9c009e5863351bd123c3797..f825a2a7363fbe144162eca96398920ead0c4e50 100644 --- a/tensorflow/contrib/eager/python/examples/spinn/spinn_test.py +++ b/tensorflow/contrib/eager/python/examples/spinn/spinn_test.py @@ -33,8 +33,8 @@ import tensorflow as tf import tensorflow.contrib.eager as tfe from tensorflow.contrib.eager.python.examples.spinn import data from third_party.examples.eager.spinn import spinn -from tensorflow.contrib.eager.proto import checkpointable_object_graph_pb2 from tensorflow.contrib.summary import summary_test_util +from tensorflow.core.protobuf import checkpointable_object_graph_pb2 from tensorflow.python.eager import test from tensorflow.python.framework import test_util from tensorflow.python.training import checkpoint_utils diff --git a/tensorflow/contrib/eager/python/metrics_impl.py b/tensorflow/contrib/eager/python/metrics_impl.py index 2f2347736a073c7d9b3fb6685f52f8d58cc40570..907f9204c2d31a652ca2a0539a23db4722b4e154 100644 --- a/tensorflow/contrib/eager/python/metrics_impl.py +++ b/tensorflow/contrib/eager/python/metrics_impl.py @@ -20,7 +20,6 @@ from __future__ import print_function import re -from tensorflow.contrib.summary import summary_ops from tensorflow.python.eager import context from tensorflow.python.eager import function from tensorflow.python.framework import dtypes @@ -29,6 +28,7 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops +from tensorflow.python.ops import summary_ops_v2 as summary_ops from tensorflow.python.ops import variable_scope from tensorflow.python.training import checkpointable diff --git a/tensorflow/contrib/eager/python/metrics_test.py b/tensorflow/contrib/eager/python/metrics_test.py index 15ac889191e0fe51269bc5740d5e0ab1bc0e2b72..f0fe4ce8c53bb80c03a3f0de37078bcdb975a0b4 100644 --- a/tensorflow/contrib/eager/python/metrics_test.py +++ b/tensorflow/contrib/eager/python/metrics_test.py @@ -21,9 +21,7 @@ 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 from tensorflow.python.eager import context from tensorflow.python.eager import test @@ -31,6 +29,8 @@ 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 summary_ops_v2 as summary_ops +from tensorflow.python.training import checkpointable_utils from tensorflow.python.training import training_util diff --git a/tensorflow/contrib/eager/python/network.py b/tensorflow/contrib/eager/python/network.py index e55a9276ab53f44f76dc5e537b3bdde7c975f463..2f8721324f5fc12565d047a64af22b8df215a92b 100644 --- a/tensorflow/contrib/eager/python/network.py +++ b/tensorflow/contrib/eager/python/network.py @@ -25,6 +25,7 @@ import weakref from tensorflow.python.eager import context from tensorflow.python.estimator import util as estimator_util from tensorflow.python.framework import ops +from tensorflow.python.keras._impl.keras.engine import base_layer as keras_base_layer from tensorflow.python.layers import base from tensorflow.python.ops import variable_scope from tensorflow.python.training import checkpoint_utils @@ -176,7 +177,7 @@ class Network(base.Layer): avoid_names = parent_network._owned_layers name_uid_map = parent_network._sub_layer_name_uids else: - name_uid_map = base._get_default_graph_uid_map() + name_uid_map = keras_base_layer.get_default_graph_uid_map() # Figure out which names we have to avoid based on which variable scope # we're nested in. strip_name = self._default_parent_variable_scope.name @@ -326,6 +327,8 @@ class Network(base.Layer): raise TypeError( "Network.track_layer() passed type %s, not a tf.layers.Layer" % (type(layer),)) + # Always use `ResourceVariable` with legacy layers. + layer._use_resource_variables = True if isinstance(layer, Network): layer._finalize_name(parent_network=self) else: diff --git a/tensorflow/contrib/eager/python/network_test.py b/tensorflow/contrib/eager/python/network_test.py index 3329fc6c513265deff41a368f5688dd605209c14..f43376d5d777a7f17d975e07b746f7b1c731e8ea 100644 --- a/tensorflow/contrib/eager/python/network_test.py +++ b/tensorflow/contrib/eager/python/network_test.py @@ -20,12 +20,10 @@ import gc from tensorflow.contrib.eager.python import network from tensorflow.contrib.layers.python.layers import regularizers -from tensorflow.python.eager import context from tensorflow.python.eager import function from tensorflow.python.eager import test from tensorflow.python.framework import constant_op from tensorflow.python.framework import errors_impl -from tensorflow.python.framework import ops from tensorflow.python.framework import test_util from tensorflow.python.layers import core from tensorflow.python.ops import math_ops @@ -469,36 +467,6 @@ class NetworkTest(test.TestCase): self.assertIsInstance(net.trainable_weights[0], resource_variable_ops.ResourceVariable) - def testGraphOpNames(self): - """Network operation names should match variable naming.""" - - def _check_op_prefixes(expected_prefix, checked_ops): - for operation in ops.get_default_graph().get_operations(): - if operation.name == "ignore": - continue - if operation.name in checked_ops: - continue - checked_ops.add(operation.name) - self.assertStartsWith(expected_start=expected_prefix, - actual=operation.name) - self.assertNotIn("my_network", operation.name[len(expected_prefix):]) - self.assertNotIn("dense", operation.name[len(expected_prefix):]) - - with context.graph_mode(): - net = MyNetwork() - zero = constant_op.constant([[0.]], name="ignore") - net(zero) - checked_ops = set() - _check_op_prefixes(expected_prefix="my_network/dense/", - checked_ops=checked_ops) - net.net2 = net.track_layer(MyNetwork()) - net.net2(zero) - _check_op_prefixes(expected_prefix="my_network/my_network/dense/", - checked_ops=checked_ops) - MyNetwork()(zero) - _check_op_prefixes(expected_prefix="my_network_1/dense/", - checked_ops=checked_ops) - @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) def testVariableRegularizers(self): net = RegularizedNetwork() diff --git a/tensorflow/contrib/eager/python/saver_test.py b/tensorflow/contrib/eager/python/saver_test.py index 1a7f7b85e688e80e3cf482f2754462888187d311..4032e755f6e7dea9dcb42587f14e8386e5db2338 100644 --- a/tensorflow/contrib/eager/python/saver_test.py +++ b/tensorflow/contrib/eager/python/saver_test.py @@ -102,7 +102,6 @@ class SaverTest(test.TestCase): # Can still restore it. saver.restore(ckpt_prefix) self.assertEqual(v1.read_value().numpy(), 1.0) - self.assertEqual(v1.read_value().numpy(), 1.0) # However, cannot restore it with default name. with self.assertRaisesOpError('not found in checkpoint'): saver = _saver.Saver([v1, v2]).restore(ckpt_prefix) diff --git a/tensorflow/contrib/eager/python/tfe.py b/tensorflow/contrib/eager/python/tfe.py index c6f3f20e781147140f2c4b339ed465ab7e919d37..79dd117854e5fe9f066f671d8ce62e08579e0ed9 100644 --- a/tensorflow/contrib/eager/python/tfe.py +++ b/tensorflow/contrib/eager/python/tfe.py @@ -84,8 +84,6 @@ 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 @@ -123,6 +121,8 @@ 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.training.checkpointable_utils import CheckpointableSaver +from tensorflow.python.training.checkpointable_utils import Checkpoint from tensorflow.python.util.all_util import remove_undocumented py_func = script_ops.eager_py_func diff --git a/tensorflow/contrib/estimator/BUILD b/tensorflow/contrib/estimator/BUILD index 9f4cd44afbede286966ba0e7357c5dac92a2b729..62ddb3d290e6c6455d79af50b2cfd894ff9f449a 100644 --- a/tensorflow/contrib/estimator/BUILD +++ b/tensorflow/contrib/estimator/BUILD @@ -210,7 +210,7 @@ py_library( py_test( name = "head_test", - size = "small", + size = "medium", srcs = ["python/estimator/head_test.py"], srcs_version = "PY2AND3", deps = [ @@ -250,7 +250,7 @@ py_library( py_test( name = "linear_test", - size = "small", + size = "medium", srcs = ["python/estimator/linear_test.py"], srcs_version = "PY2AND3", tags = [ @@ -447,6 +447,7 @@ py_test( srcs_version = "PY2AND3", tags = [ "no_pip", + "noasan", # times out "notsan", ], deps = [ diff --git a/tensorflow/contrib/estimator/__init__.py b/tensorflow/contrib/estimator/__init__.py index 9a87fa915d28788c5abd57b2ab7189b47fb10a4c..be20d1b7770d3f3df21ac9c0f811d924bf4152ee 100644 --- a/tensorflow/contrib/estimator/__init__.py +++ b/tensorflow/contrib/estimator/__init__.py @@ -28,6 +28,7 @@ from tensorflow.contrib.estimator.python.estimator.linear import * from tensorflow.contrib.estimator.python.estimator.logit_fns import * from tensorflow.contrib.estimator.python.estimator.multi_head import * from tensorflow.contrib.estimator.python.estimator.replicate_model_fn import * +from tensorflow.contrib.estimator.python.estimator.rnn import * from tensorflow.python.util.all_util import remove_undocumented # pylint: enable=unused-import,line-too-long,wildcard-import diff --git a/tensorflow/contrib/estimator/python/estimator/head.py b/tensorflow/contrib/estimator/python/estimator/head.py index ae2fd8b4902e850292bd6672ec34f084a230dfbe..3dcf0374c8a12b5907fbaf20d1ad72211a45ab5c 100644 --- a/tensorflow/contrib/estimator/python/estimator/head.py +++ b/tensorflow/contrib/estimator/python/estimator/head.py @@ -485,7 +485,7 @@ class _MultiLabelHead(head_lib._Head): # pylint:disable=protected-access reduction=losses.Reduction.NONE) # Averages loss over classes. unweighted_loss = math_ops.reduce_mean( - unweighted_loss, axis=-1, keep_dims=True) + unweighted_loss, axis=-1, keepdims=True) weights = head_lib._get_weights_and_check_match_logits( # pylint:disable=protected-access, features=features, weight_column=self._weight_column, logits=logits) training_loss = losses.compute_weighted_loss( diff --git a/tensorflow/contrib/estimator/python/estimator/replicate_model_fn.py b/tensorflow/contrib/estimator/python/estimator/replicate_model_fn.py index fa2697800ec1a44f215f3d5fc9be2197a9e58219..a8774d6dab9205439e6e312827f9cd1306e3f1ea 100644 --- a/tensorflow/contrib/estimator/python/estimator/replicate_model_fn.py +++ b/tensorflow/contrib/estimator/python/estimator/replicate_model_fn.py @@ -456,7 +456,7 @@ def _get_local_devices(device_type): def _split_batch(features, labels, number_of_shards, device): - """Split input features and labes into batches.""" + """Split input features and labels into batches.""" def ensure_divisible_by_shards(sequence): batch_size = ops_lib.convert_to_tensor(sequence).get_shape()[0] @@ -602,7 +602,7 @@ def _local_device_setter(worker_device, ps_devices, ps_strategy): def _scale_tower_loss(tower_spec, loss_reduction, number_of_towers): - """Produce an EstimatorSpec with approproriately scaled loss.""" + """Produce an EstimatorSpec with appropriately scaled loss.""" if tower_spec.loss is None: return tower_spec diff --git a/tensorflow/contrib/factorization/python/ops/gmm_ops.py b/tensorflow/contrib/factorization/python/ops/gmm_ops.py index 5d77bc77e124378e13667673e4e841c0a1135b31..e076631bc16fd379a2ad31af9055a7388d98c7ca 100644 --- a/tensorflow/contrib/factorization/python/ops/gmm_ops.py +++ b/tensorflow/contrib/factorization/python/ops/gmm_ops.py @@ -54,10 +54,10 @@ def _covariance(x, diag): diagonal matrix just the diagonal is returned. """ num_points = math_ops.to_float(array_ops.shape(x)[0]) - x -= math_ops.reduce_mean(x, 0, keep_dims=True) + x -= math_ops.reduce_mean(x, 0, keepdims=True) if diag: cov = math_ops.reduce_sum( - math_ops.square(x), 0, keep_dims=True) / (num_points - 1) + math_ops.square(x), 0, keepdims=True) / (num_points - 1) else: cov = math_ops.matmul(x, x, transpose_a=True) / (num_points - 1) return cov @@ -313,7 +313,7 @@ class GmmAlgorithm(object): # TODO(xavigonzalvo): look into alternatives to log for # reparametrization of variance parameters. det_expanded = math_ops.reduce_sum( - math_ops.log(self._covs + 1e-3), 1, keep_dims=True) + math_ops.log(self._covs + 1e-3), 1, keepdims=True) diff = shard - self._means x2 = math_ops.square(diff) cov_expanded = array_ops.expand_dims(1.0 / (self._covs + 1e-3), 2) @@ -351,7 +351,7 @@ class GmmAlgorithm(object): shard_id: id of current shard_id. """ self._prior_probs[shard_id] = math_ops.reduce_logsumexp( - self._probs[shard_id], axis=1, keep_dims=True) + self._probs[shard_id], axis=1, keepdims=True) def _define_expectation_operation(self, shard_id): # Shape broadcasting. @@ -375,7 +375,7 @@ class GmmAlgorithm(object): """ # Soft assignment of each data point to each of the two clusters. self._points_in_k[shard_id] = math_ops.reduce_sum( - self._w[shard_id], 0, keep_dims=True) + self._w[shard_id], 0, keepdims=True) # Partial means. w_mul_x = array_ops.expand_dims( math_ops.matmul( @@ -397,7 +397,7 @@ class GmmAlgorithm(object): # Compute the effective number of data points assigned to component k. with ops.control_dependencies(self._w): points_in_k = array_ops.squeeze( - math_ops.add_n(self._points_in_k), squeeze_dims=[0]) + math_ops.add_n(self._points_in_k), axis=[0]) # Update alpha. if 'w' in self._params: final_points_in_k = points_in_k / num_batches @@ -454,7 +454,7 @@ class GmmAlgorithm(object): 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) + math_ops.reduce_logsumexp(op, axis=2, keepdims=True), axis=0) def gmm(inp, diff --git a/tensorflow/contrib/factorization/python/ops/kmeans.py b/tensorflow/contrib/factorization/python/ops/kmeans.py index bfe338c9f9a7b761cfcd627b92f1682af97630c9..9ffdd3ba5e8ac496533d0207f2b6846dbc92bc89 100644 --- a/tensorflow/contrib/factorization/python/ops/kmeans.py +++ b/tensorflow/contrib/factorization/python/ops/kmeans.py @@ -374,11 +374,11 @@ class KMeansClustering(estimator.Estimator): than `num_clusters`, a TensorFlow runtime error occurs. distance_metric: The distance metric used for clustering. One of: * `KMeansClustering.SQUARED_EUCLIDEAN_DISTANCE`: Euclidean distance - between vectors `u` and `v` is defined as `\\(||u - v||_2\\)` + between vectors `u` and `v` is defined as \\(||u - v||_2\\) which is the square root of the sum of the absolute squares of the elements' difference. * `KMeansClustering.COSINE_DISTANCE`: Cosine distance between vectors - `u` and `v` is defined as `\\(1 - (u . v) / (||u||_2 ||v||_2)\\)`. + `u` and `v` is defined as \\(1 - (u . v) / (||u||_2 ||v||_2)\\). random_seed: Python integer. Seed for PRNG used to initialize centers. use_mini_batch: A boolean specifying whether to use the mini-batch k-means algorithm. See explanation above. diff --git a/tensorflow/contrib/framework/__init__.py b/tensorflow/contrib/framework/__init__.py index cbb68bd3eb257f9472515e5c29ce4f02057be321..11397e86bd811e376ff4b60f86109d9ee5dd1a62 100644 --- a/tensorflow/contrib/framework/__init__.py +++ b/tensorflow/contrib/framework/__init__.py @@ -72,6 +72,9 @@ See the @{$python/contrib.framework} guide. @@variable @@VariableDeviceChooser @@convolutional_delta_orthogonal +@@convolutional_orthogonal_1d +@@convolutional_orthogonal_2d +@@convolutional_orthogonal_3d @@zero_initializer @@load_checkpoint @@ -115,9 +118,13 @@ 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.array_ops import broadcast_to from tensorflow.python.ops.init_ops import convolutional_delta_orthogonal +from tensorflow.python.ops.init_ops import convolutional_orthogonal_1d +from tensorflow.python.ops.init_ops import convolutional_orthogonal_2d +from tensorflow.python.ops.init_ops import convolutional_orthogonal_3d from tensorflow.python.util.all_util import remove_undocumented -_allowed_symbols = ['nest'] +_allowed_symbols = ['nest', 'broadcast_to'] remove_undocumented(__name__, allowed_exception_list=_allowed_symbols) diff --git a/tensorflow/contrib/framework/python/framework/tensor_util_test.py b/tensorflow/contrib/framework/python/framework/tensor_util_test.py index a2834b648933772cab53002462c3edbe9a553e94..8fc4f60492b0bfb22ea78cb7b5906e452bb6da58 100644 --- a/tensorflow/contrib/framework/python/framework/tensor_util_test.py +++ b/tensorflow/contrib/framework/python/framework/tensor_util_test.py @@ -48,7 +48,7 @@ class LocalVariabletest(test.TestCase): variables = variables_lib.local_variables() self.assertEquals(2, len(variables)) self.assertRaises(errors_impl.OpError, sess.run, variables) - variables_lib.initialize_variables(variables).run() + variables_lib.variables_initializer(variables).run() self.assertAllEqual(set([value0, value1]), set(sess.run(variables))) diff --git a/tensorflow/contrib/fused_conv/kernels/fused_conv2d_bias_activation_op.cc b/tensorflow/contrib/fused_conv/kernels/fused_conv2d_bias_activation_op.cc index 0e06575d96f9b9538f0245b12d48cfd7c0e8d981..1e8f011b5d8ee6f9dff8fd0c7a650a2085ffce56 100644 --- a/tensorflow/contrib/fused_conv/kernels/fused_conv2d_bias_activation_op.cc +++ b/tensorflow/contrib/fused_conv/kernels/fused_conv2d_bias_activation_op.cc @@ -543,7 +543,8 @@ void LaunchFusedConv2DBiasActivationOp:: fused_conv_parameters, &algorithm_config)) { std::vector algorithms; CHECK(stream->parent()->GetConvolveAlgorithms( - fused_conv_parameters.ShouldIncludeWinogradNonfusedAlgo(), + fused_conv_parameters.ShouldIncludeWinogradNonfusedAlgo( + stream->parent()), &algorithms)); dnn::ProfileResult best_result; dnn::ProfileResult best_result_no_scratch; diff --git a/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_op.py b/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_op.py index a97adf622e6e576f8b5ce2babe004cb3a46d80a5..983b6dc8e5a1512ba81ecbc8d5ca5adaea09afe4 100644 --- a/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_op.py +++ b/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_op.py @@ -65,7 +65,7 @@ def fused_conv2d_bias_activation(conv_input, side_input_scale: A scalar `float32` that will be multiplied by side_input. This is optional and defaults to 0. side_input: A `Tensor` of the format specified by `data_format`. - This is useful for imlementing ResNet blocks. + This is useful for implementing ResNet blocks. activation_mode: (optional) currently must be the default "Relu". Note that in qint8 mode, it also clips to 127, so acts like ReluX. data_format: Specifies the data format. diff --git a/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_op_test.py b/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_op_test.py index bb155aa2496cbafd9f0630d3dffb2ba69395186c..3d0ed899322c26bf4ae428930899d7a5885e9f21 100644 --- a/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_op_test.py +++ b/tensorflow/contrib/fused_conv/python/ops/fused_conv2d_bias_activation_op_test.py @@ -566,7 +566,7 @@ def GetInceptionFwdTest(input_size, filter_size, stride, padding, return Test -def CalculateCovolvedOutputDim(input_dim, filter_dim, stride, padding_type): +def CalculateConvolvedOutputDim(input_dim, filter_dim, stride, padding_type): """Calculates the size of an output dimension of a strided convolution. Given the sizes of the corresponding dimension of the input and filter shapes, @@ -827,10 +827,10 @@ class FusedConvInt8Tests(test.TestCase): maxval=1.0, dtype=dtypes.float32), -1.0, 1.0, dtypes.qint8) - output_height = CalculateCovolvedOutputDim(input_height, filter_height, - vertical_stride, padding_type) - output_width = CalculateCovolvedOutputDim(input_width, filter_width, - horizontal_stride, padding_type) + output_height = CalculateConvolvedOutputDim(input_height, filter_height, + vertical_stride, padding_type) + output_width = CalculateConvolvedOutputDim(input_width, filter_width, + horizontal_stride, padding_type) print("output_height=", output_height, ", output_width=", output_width) side_input, _, _ = gen_array_ops.quantize_v2( diff --git a/tensorflow/contrib/gan/BUILD b/tensorflow/contrib/gan/BUILD index 461066bbb493932b342cee8f8842e899a2d84fff..b305f37791d71f5a6edeada2bb710a2e5f23087d 100644 --- a/tensorflow/contrib/gan/BUILD +++ b/tensorflow/contrib/gan/BUILD @@ -364,6 +364,7 @@ py_test( "//tensorflow/python:framework_ops", "//tensorflow/python:variables", "//third_party/py/numpy", + "@absl_py//absl/testing:parameterized", ], ) 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 663e49bdca3cb2dd9257da326488c877fcc4256d..4fb8d58bc9125664d42260de72b83b2362eff9ba 100644 --- a/tensorflow/contrib/gan/python/eval/python/classifier_metrics_test.py +++ b/tensorflow/contrib/gan/python/eval/python/classifier_metrics_test.py @@ -22,6 +22,7 @@ import os import tarfile import tempfile +from absl.testing import parameterized import numpy as np from scipy import linalg as scp_linalg @@ -182,13 +183,20 @@ def _run_with_mock(function, *args, **kwargs): return function(*args, **kwargs) -class ClassifierMetricsTest(test.TestCase): +class ClassifierMetricsTest(test.TestCase, parameterized.TestCase): - def test_run_inception_graph(self): + @parameterized.named_parameters( + ('GraphDef', False), + ('DefaultGraphDefFn', True)) + def test_run_inception_graph(self, use_default_graph_def): """Test `run_inception` graph construction.""" batch_size = 7 img = array_ops.ones([batch_size, 299, 299, 3]) - logits = _run_with_mock(classifier_metrics.run_inception, img) + + if use_default_graph_def: + logits = _run_with_mock(classifier_metrics.run_inception, img) + else: + logits = classifier_metrics.run_inception(img, _get_dummy_graphdef()) self.assertTrue(isinstance(logits, ops.Tensor)) logits.shape.assert_is_compatible_with([batch_size, 1001]) @@ -196,14 +204,23 @@ class ClassifierMetricsTest(test.TestCase): # Check that none of the model variables are trainable. self.assertListEqual([], variables.trainable_variables()) - def test_run_inception_graph_pool_output(self): + @parameterized.named_parameters( + ('GraphDef', False), + ('DefaultGraphDefFn', True)) + def test_run_inception_graph_pool_output(self, use_default_graph_def): """Test `run_inception` graph construction with pool output.""" batch_size = 3 img = array_ops.ones([batch_size, 299, 299, 3]) - pool = _run_with_mock( - classifier_metrics.run_inception, - img, - output_tensor=classifier_metrics.INCEPTION_FINAL_POOL) + + if use_default_graph_def: + pool = _run_with_mock( + classifier_metrics.run_inception, + img, + output_tensor=classifier_metrics.INCEPTION_FINAL_POOL) + else: + pool = classifier_metrics.run_inception( + img, _get_dummy_graphdef(), + output_tensor=classifier_metrics.INCEPTION_FINAL_POOL) self.assertTrue(isinstance(pool, ops.Tensor)) pool.shape.assert_is_compatible_with([batch_size, 2048]) 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 4b10bc0f8e607c02763d8ea622d6f8f2572c586d..4b1105f6bd4f21a0da02338b0fc9db87a41b145f 100644 --- a/tensorflow/contrib/gan/python/eval/python/sliced_wasserstein_impl.py +++ b/tensorflow/contrib/gan/python/eval/python/sliced_wasserstein_impl.py @@ -161,7 +161,7 @@ def _sliced_wasserstein(a, b, random_sampling_count, random_projection_dim): proj = random_ops.random_normal( [array_ops.shape(a)[1], random_projection_dim]) proj *= math_ops.rsqrt( - math_ops.reduce_sum(math_ops.square(proj), 0, keep_dims=True)) + math_ops.reduce_sum(math_ops.square(proj), 0, keepdims=True)) # Project both distributions and sort them. proj_a = math_ops.matmul(a, proj) proj_b = math_ops.matmul(b, proj) diff --git a/tensorflow/contrib/gan/python/features/python/virtual_batchnorm_impl.py b/tensorflow/contrib/gan/python/features/python/virtual_batchnorm_impl.py index f8b372546b60ec8fa5fd1d72b57adaf67596c059..650eab97a3952e9aec2b489fffcc83c3bc49f2dd 100644 --- a/tensorflow/contrib/gan/python/features/python/virtual_batchnorm_impl.py +++ b/tensorflow/contrib/gan/python/features/python/virtual_batchnorm_impl.py @@ -64,11 +64,11 @@ def _statistics(x, axes): y = math_ops.cast(x, dtypes.float32) if x.dtype == dtypes.float16 else x # Compute true mean while keeping the dims for proper broadcasting. - shift = array_ops.stop_gradient(math_ops.reduce_mean(y, axes, keep_dims=True)) + shift = array_ops.stop_gradient(math_ops.reduce_mean(y, axes, keepdims=True)) - shifted_mean = math_ops.reduce_mean(y - shift, axes, keep_dims=True) + shifted_mean = math_ops.reduce_mean(y - shift, axes, keepdims=True) mean = shifted_mean + shift - mean_squared = math_ops.reduce_mean(math_ops.square(y), axes, keep_dims=True) + mean_squared = math_ops.reduce_mean(math_ops.square(y), axes, keepdims=True) mean = array_ops.squeeze(mean, axes) mean_squared = array_ops.squeeze(mean_squared, axes) diff --git a/tensorflow/contrib/gan/python/train.py b/tensorflow/contrib/gan/python/train.py index 73acd05b60a5fb02601423fd9234a56a34f75276..6fa43059f3125daea080f780210223363d0a89f9 100644 --- a/tensorflow/contrib/gan/python/train.py +++ b/tensorflow/contrib/gan/python/train.py @@ -710,7 +710,10 @@ def gan_train_ops( be used to train a generator/discriminator pair. """ if isinstance(model, namedtuples.CycleGANModel): - saved_params = locals() + # Get and store all arguments other than model and loss from locals. + # Contents of locals should not be modified, may not affect values. So make + # a copy. https://docs.python.org/2/library/functions.html#locals. + saved_params = dict(locals()) saved_params.pop('model', None) saved_params.pop('loss', None) kwargs = saved_params.pop('kwargs', {}) diff --git a/tensorflow/contrib/gdr/gdr_rendezvous_mgr.cc b/tensorflow/contrib/gdr/gdr_rendezvous_mgr.cc index 28f68cec8cce126f1b177a73e197ccd7ab749f4a..94f522c04e5a09ed2d9355fa675125c340407923 100644 --- a/tensorflow/contrib/gdr/gdr_rendezvous_mgr.cc +++ b/tensorflow/contrib/gdr/gdr_rendezvous_mgr.cc @@ -155,7 +155,7 @@ class GdrRemoteRendezvous : public BaseRemoteRendezvous { } Device* dst_device; - Status s = sess->device_mgr->LookupDevice(parsed.dst_device, &dst_device); + Status s = sess->device_mgr()->LookupDevice(parsed.dst_device, &dst_device); if (!s.ok()) { sess->worker_cache->ReleaseWorker(src_worker, rwi); done(s, Args(), recv_args, Tensor{}, false); diff --git a/tensorflow/contrib/graph_editor/tests/transform_test.py b/tensorflow/contrib/graph_editor/tests/transform_test.py index 2603de640735a612cbd883cc6227fe3cd9f11fca..97f38c923f4a19cedf3e16203ca1e66b7e5e45d2 100644 --- a/tensorflow/contrib/graph_editor/tests/transform_test.py +++ b/tensorflow/contrib/graph_editor/tests/transform_test.py @@ -18,9 +18,11 @@ from __future__ import division from __future__ import print_function import collections +import functools import numpy as np from tensorflow.contrib import graph_editor as ge from tensorflow.contrib.graph_editor.tests import match +from tensorflow.core.framework import attr_value_pb2 from tensorflow.python.client import session from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes @@ -42,6 +44,7 @@ class TransformTest(test.TestCase): self.graph = ops.Graph() with self.graph.as_default(): c0 = constant_op.constant(1.0, shape=[10], name="Const") + c0.op._set_attr("_foo", attr_value_pb2.AttrValue(s=b"foo")) c1 = constant_op.constant(1.0, shape=[10], name="Const") c2 = constant_op.constant(1.0, shape=[10], name="Const") i = constant_op.constant(1.0, shape=[10], name="Input") @@ -112,6 +115,32 @@ class TransformTest(test.TestCase): top = ge.select_ops("^AddNoise_2$", graph=graph)[0] self.assertTrue(matcher2(top)) + def test_transform_nodedef_fn(self): + transformer = ge.Transformer() + + def nodedef_fn(node_def): + if "_foo" in node_def.attr: + del node_def.attr["_foo"] + node_def.attr["_bar"].s = b"bar" + return node_def + + my_copy_op_handler = functools.partial( + ge.transform.copy_op_handler, nodedef_fn=nodedef_fn) + transformer.transform_op_handler = my_copy_op_handler + + graph = ops.Graph() + transformer(self.graph, graph, "", "") + + c0_before = self.graph.get_operation_by_name("Const") + c0_after = graph.get_operation_by_name("Const") + self.assertEquals(c0_before.get_attr("_foo"), b"foo") + with self.assertRaises(ValueError): + c0_after.get_attr("_foo") + + all_ops = graph.get_operations() + for op in all_ops: + self.assertEquals(op.get_attr("_bar"), b"bar") + def test_copy_with_input_replacements(self): with self.graph.as_default(): ten = constant_op.constant(10.0, shape=[10], name="Input") diff --git a/tensorflow/contrib/graph_editor/transform.py b/tensorflow/contrib/graph_editor/transform.py index d8a48387a745e7d88cc6a74c96cb21a2ba1cfa1f..a320a3f232fc1dc8c9ccfd1d0f2a9a40225db5cb 100644 --- a/tensorflow/contrib/graph_editor/transform.py +++ b/tensorflow/contrib/graph_editor/transform.py @@ -129,7 +129,7 @@ def transform_op_if_inside_handler(info, op, keep_if_possible=True): return None -def copy_op_handler(info, op, new_inputs, copy_shape=True): +def copy_op_handler(info, op, new_inputs, copy_shape=True, nodedef_fn=None): """Copy a `tf.Operation`. Args: @@ -137,6 +137,11 @@ def copy_op_handler(info, op, new_inputs, copy_shape=True): op: the `tf.Operation` to be copied. new_inputs: The new inputs for this op. copy_shape: also copy the shape of the tensor + nodedef_fn: If provided, a function that will be run on the NodeDef + and should return a mutated NodeDef before a new Operation is created. + This is useful as certain features cannot be set on the Operation and + must be modified in NodeDef. + Returns: A `(op, op_outputs)` tuple containing the transformed op and its outputs. """ @@ -155,6 +160,10 @@ def copy_op_handler(info, op, new_inputs, copy_shape=True): name_ = info.graph_.unique_name(name_) node_def_.name = name_ + # Mutate NodeDef if requested: + if nodedef_fn is not None: + node_def_ = nodedef_fn(node_def_) + # Copy the other inputs needed for initialization output_types_ = op._output_types[:] input_types_ = op._input_types[:] diff --git a/tensorflow/contrib/hvx/README.md b/tensorflow/contrib/hvx/README.md index 163993a3f6bb1bedcdffb32944a98c7cc846878e..68e34f3b0938f795c8ad4c8c75226f6b0afe188d 100644 --- a/tensorflow/contrib/hvx/README.md +++ b/tensorflow/contrib/hvx/README.md @@ -42,11 +42,12 @@ If you've finished walking through the quick start guide, you may want to try bu ### Build libhexagon\_nn\_skel.so -Download Hexagon NN library from codeaurora.org and build it. +Download Hexagon NN library from codeaurora.org and build it. For Hexagon SDK 3.0, we need use the compatible version([721b2d58f](https://source.codeaurora.org/quic/hexagon_nn/nnlib/commit/?id=721b2d58f0f4e2d5b182f41e6b7c4db5356bf0fb)) of nnlib. ```shell git clone https://source.codeaurora.org/quic/hexagon_nn/nnlib cd nnlib +git reset 721b2d58f --hard ``` Just follow the instructions in `README.HOW_TO_BUILD`. You can find the file `libhexagon_nn_skel.so` in `hexagon_Release_dynamic_toolv72_v60/ship`. diff --git a/tensorflow/contrib/image/__init__.py b/tensorflow/contrib/image/__init__.py index e982030bc8959309e72d0f4e02b9755c48535a10..8f406ace1d5dcc13a018e56cc98c621a511da29b 100755 --- a/tensorflow/contrib/image/__init__.py +++ b/tensorflow/contrib/image/__init__.py @@ -25,6 +25,8 @@ projective transforms (including rotation) are supported. @@angles_to_projective_transforms @@compose_transforms @@adjust_yiq_hsv +@@flat_transforms_to_matrices +@@matrices_to_flat_transforms @@random_yiq_hsv @@rotate @@transform @@ -58,6 +60,8 @@ from tensorflow.contrib.image.python.ops.distort_image_ops import random_hsv_in_ from tensorflow.contrib.image.python.ops.image_ops import angles_to_projective_transforms from tensorflow.contrib.image.python.ops.image_ops import compose_transforms from tensorflow.contrib.image.python.ops.image_ops import connected_components +from tensorflow.contrib.image.python.ops.image_ops import flat_transforms_to_matrices +from tensorflow.contrib.image.python.ops.image_ops import matrices_to_flat_transforms from tensorflow.contrib.image.python.ops.image_ops import rotate from tensorflow.contrib.image.python.ops.image_ops import transform from tensorflow.contrib.image.python.ops.image_ops import translate diff --git a/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op_gpu.cu.cc b/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op_gpu.cu.cc index b71ff9cd507faac66b3a33d3c02ec9b5901d814a..645abbf0b0ea5465dadf55d065e997e16940c18d 100644 --- a/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op_gpu.cu.cc +++ b/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op_gpu.cu.cc @@ -53,7 +53,7 @@ void AdjustHsvInYiqGPU::operator()(OpKernelContext* ctx, int channel_count, OP_REQUIRES_OK(ctx, ctx->allocate_temp( DT_FLOAT, TensorShape({kChannelSize * kChannelSize}), &tranformation_matrix)); - // TODO(huangyp): It takes about 3.5 us to comute tranformation_matrix + // TODO(huangyp): It takes about 3.5 us to compute tranformation_matrix // with one thread. Improve its performance if necessary. internal::compute_tranformation_matrix_cuda<<<1, 1, 0, cu_stream>>>( delta_h, scale_s, scale_v, tranformation_matrix.flat().data(), diff --git a/tensorflow/contrib/image/ops/distort_image_ops.cc b/tensorflow/contrib/image/ops/distort_image_ops.cc index b169b0b2b22ad6432baed2cc96711da5ca995875..ca49635d5d0bc7bb84b19508a74be74362d96ddf 100644 --- a/tensorflow/contrib/image/ops/distort_image_ops.cc +++ b/tensorflow/contrib/image/ops/distort_image_ops.cc @@ -36,9 +36,9 @@ REGISTER_OP("AdjustHsvInYiq") Adjust the YIQ hue of one or more images. `images` is a tensor of at least 3 dimensions. The last dimension is -interpretted as channels, and must be three. +interpreted as channels, and must be three. -We used linear transfomation described in: +We used linear transformation described in: beesbuzz.biz/code/hsv_color_transforms.php The input image is considered in the RGB colorspace. Conceptually, the RGB colors are first mapped into YIQ space, rotated around the Y channel by diff --git a/tensorflow/contrib/image/ops/image_ops.cc b/tensorflow/contrib/image/ops/image_ops.cc index 68771b3d054a64ba94141c092e20df1ed6b2339b..ebdcaea7abae2a967786831b62b331897aa3f6a3 100644 --- a/tensorflow/contrib/image/ops/image_ops.cc +++ b/tensorflow/contrib/image/ops/image_ops.cc @@ -93,7 +93,7 @@ row_to_col_match_indices: A vector of length num_rows, which is the number of If `row_to_col_match_indices[i]` is not -1, row i is matched to column `row_to_col_match_indices[i]`. col_to_row_match_indices: A vector of length num_columns, which is the number - of columns of the input ditance matrix. + of columns of the input distance matrix. If `col_to_row_match_indices[j]` is not -1, column j is matched to row `col_to_row_match_indices[j]`. )doc"); diff --git a/tensorflow/contrib/image/ops/single_image_random_dot_stereograms_ops.cc b/tensorflow/contrib/image/ops/single_image_random_dot_stereograms_ops.cc index 8139d4272d6950815bd39a64e86e0f7422e6f799..bd784c6bda0344c092c1ae0af2c60be50fdff102 100755 --- a/tensorflow/contrib/image/ops/single_image_random_dot_stereograms_ops.cc +++ b/tensorflow/contrib/image/ops/single_image_random_dot_stereograms_ops.cc @@ -69,7 +69,7 @@ Outputs a single image random dot stereogram for export via encode_PNG/JPG OP. Given the 2-D tensor 'depth_values' with encoded Z values, this operation will encode 3-D data into a 2-D image. The output of this Op is suitable for the encode_PNG/JPG ops. Be careful with image compression as this may corrupt the -encode 3-D data witin the image. +encode 3-D data within the image. This Op is based upon: 'http://www.learningace.com/doc/4331582/b6ab058d1e206d68ab60e4e1ead2fe6e/sirds-paper' @@ -111,7 +111,7 @@ output_image_shape: Output size of returned image in X,Y, Channels 1-grayscale, output_data_window: Size of "DATA" window, must be equal to or smaller than 'output_image_shape', will be centered and use 'convergence_dots_size' for best fit to avoid overlap if possible -image:= A tensor of size 'output_image_shape' with the encloded 'depth_values' +image:= A tensor of size 'output_image_shape' with the encoded 'depth_values' )doc"); } // namespace tensorflow diff --git a/tensorflow/contrib/image/python/ops/image_ops.py b/tensorflow/contrib/image/python/ops/image_ops.py index c139ae89d8d682d6b87813c3a21703ffa762f28e..cd984c80543886be1f682933e2e003bd3374e425 100644 --- a/tensorflow/contrib/image/python/ops/image_ops.py +++ b/tensorflow/contrib/image/python/ops/image_ops.py @@ -433,7 +433,7 @@ def bipartite_match(distance_mat, of rows of the input `distance_matrix`. If `row_to_col_match_indices[i]` is not -1, row i is matched to column `row_to_col_match_indices[i]`. col_to_row_match_indices: A vector of length num_columns, which is the - number of columns of the input ditance matrix. + number of columns of the input distance matrix. If `col_to_row_match_indices[j]` is not -1, column j is matched to row `col_to_row_match_indices[j]`. """ diff --git a/tensorflow/contrib/image/python/ops/single_image_random_dot_stereograms.py b/tensorflow/contrib/image/python/ops/single_image_random_dot_stereograms.py index d4a6a5bcbb52511d4093587814100b2a0e8b2420..0ceb683ff4c6965a5ee4bcb04846a69d4d8ea0a5 100755 --- a/tensorflow/contrib/image/python/ops/single_image_random_dot_stereograms.py +++ b/tensorflow/contrib/image/python/ops/single_image_random_dot_stereograms.py @@ -45,7 +45,7 @@ def single_image_random_dot_stereograms(depth_values, Given the 2-D tensor 'depth_values' with encoded Z values, this operation will encode 3-D data into a 2-D image. The output of this Op is suitable for the encode_PNG/JPG ops. Be careful with image compression as this may - corrupt the encode 3-D data witin the image. + corrupt the encode 3-D data within the image. Based upon [this paper](http://www.learningace.com/doc/4331582/b6ab058d1e206d68ab60e4e1ead2fe6e/sirds-paper). diff --git a/tensorflow/contrib/kernel_methods/python/mappers/random_fourier_features.py b/tensorflow/contrib/kernel_methods/python/mappers/random_fourier_features.py index 091f0a109801065f06110e2a313c24486d38109f..9a721a9d440e66eb30bb94daf2b6878318f1e75f 100644 --- a/tensorflow/contrib/kernel_methods/python/mappers/random_fourier_features.py +++ b/tensorflow/contrib/kernel_methods/python/mappers/random_fourier_features.py @@ -34,33 +34,31 @@ class RandomFourierFeatureMapper(dkm.DenseKernelMapper): r"""Class that implements Random Fourier Feature Mapping (RFFM) in TensorFlow. The RFFM mapping is used to approximate the Gaussian (RBF) kernel: - ``` $$(exp(-||x-y||_2^2 / (2 * \sigma^2))$$ - ``` The implementation of RFFM is based on the following paper: "Random Features for Large-Scale Kernel Machines" by Ali Rahimi and Ben Recht. (link: https://people.eecs.berkeley.edu/~brecht/papers/07.rah.rec.nips.pdf) - The mapping uses a matrix `\\(Omega \in R^{d x D}\\)` and a bias vector - `\\(b \in R^D\\)` where `d` is the input dimension (number of dense input - features) and `D` is the output dimension (i.e., dimension of the feature - space the input is mapped to). Each entry of `Omega` is sampled i.i.d. from a - (scaled) Gaussian distribution and each entry of `b` is sampled independently - and uniformly from [0, \\(2 * pi\\)]. - - For a single input feature vector x in R^d, its RFFM is defined as: - ``` - $$sqrt(2/D) * cos(x * Omega + b)$$ - ``` - where `cos` is the element-wise cosine function and `x, b` are represented as - row vectors. The aforementioned paper shows that the linear kernel of - RFFM-mapped vectors approximates the Gaussian kernel of the initial vectors. + The mapping uses a matrix \\(\Omega \in R^{d x D}\\) and a bias vector + \\(b \in R^D\\) where \\(d\\) is the input dimension (number of dense input + features) and \\(D\\) is the output dimension (i.e., dimension of the feature + space the input is mapped to). Each entry of \\(\Omega\\) is sampled i.i.d. + from a (scaled) Gaussian distribution and each entry of \\(b\\) is sampled + independently and uniformly from [0, \\(2 * \pi\\)]. + + For a single input feature vector \\(x \in R^d\\), its RFFM is defined as: + $$\sqrt(2/D) * cos(x * \Omega + b)$$ + + where \\(cos\\) is the element-wise cosine function and \\(x, b\\) are + represented as row vectors. The aforementioned paper shows that the linear + kernel of RFFM-mapped vectors approximates the Gaussian kernel of the initial + vectors. """ def __init__(self, input_dim, output_dim, stddev=1.0, seed=1, name=None): - """Constructs a RandomFourierFeatureMapper instance. + r"""Constructs a RandomFourierFeatureMapper instance. Args: input_dim: The dimension (number of features) of the tensors to be mapped. @@ -68,11 +66,11 @@ class RandomFourierFeatureMapper(dkm.DenseKernelMapper): stddev: The standard deviation of the Gaussian kernel to be approximated. The error of the classifier trained using this approximation is very sensitive to this parameter. - seed: An integer used to initialize the parameters (`Omega` and `b`) of - the mapper. For repeatable sequences across different invocations of the - mapper object (for instance, to ensure consistent mapping both at - training and eval/inference if these happen in different invocations), - set this to the same integer. + seed: An integer used to initialize the parameters (\\(\Omega\\) and + \\(b\\)) of the mapper. For repeatable sequences across different + invocations of the mapper object (for instance, to ensure consistent + mapping both at training and eval/inference if these happen in + different invocations), set this to the same integer. name: name for the mapper object. """ # TODO(sibyl-vie3Poto): Maybe infer input_dim and/or output_dim (if not explicitly diff --git a/tensorflow/contrib/kfac/examples/convnet_mnist_single_main.py b/tensorflow/contrib/kfac/examples/convnet_mnist_single_main.py index 3aa52aff196fd2699559f80b0c226f470c94b2a3..2c1f09936073a34816da61d771f59e848b8787af 100644 --- a/tensorflow/contrib/kfac/examples/convnet_mnist_single_main.py +++ b/tensorflow/contrib/kfac/examples/convnet_mnist_single_main.py @@ -32,7 +32,7 @@ flags.DEFINE_string("data_dir", "/tmp/mnist", "local mnist dir") def main(unused_argv): - convnet.train_mnist_single_gpu(FLAGS.data_dir, num_epochs=200) + convnet.train_mnist_single_machine(FLAGS.data_dir, num_epochs=200) if __name__ == "__main__": diff --git a/tensorflow/contrib/kfac/python/ops/layer_collection.py b/tensorflow/contrib/kfac/python/ops/layer_collection.py index 411da033c3a0d5e2148c02207f6e04efcd2a0efc..366e2a82d56602de0df706cbd382c21aba5540af 100644 --- a/tensorflow/contrib/kfac/python/ops/layer_collection.py +++ b/tensorflow/contrib/kfac/python/ops/layer_collection.py @@ -28,6 +28,7 @@ from collections import defaultdict from collections import OrderedDict from contextlib import contextmanager from functools import partial +import warnings import math import six @@ -171,6 +172,9 @@ class LayerCollection(object): def __init__(self, graph=None, name="LayerCollection"): + warnings.warn( + "tf.contrib.kfac is deprecated and will be removed by 2018-11-01. " + "Use https://pypi.python.org/pypi/kfac instead.") self.fisher_blocks = LayerParametersDict() self.fisher_factors = OrderedDict() self._linked_parameters = dict( diff --git a/tensorflow/contrib/kfac/python/ops/loss_functions.py b/tensorflow/contrib/kfac/python/ops/loss_functions.py index e7d4243fc3d1c2d860693f2f62447b1c9aeeee03..42d525c2c21f5ba3457cba041261dc3b225dc11e 100644 --- a/tensorflow/contrib/kfac/python/ops/loss_functions.py +++ b/tensorflow/contrib/kfac/python/ops/loss_functions.py @@ -613,19 +613,19 @@ class CategoricalLogitsNegativeLogProbLoss(DistributionNegativeLogProbLoss, def multiply_fisher(self, vector): probs = self._probs return vector * probs - probs * math_ops.reduce_sum( - vector * probs, axis=-1, keep_dims=True) + vector * probs, axis=-1, keepdims=True) def multiply_fisher_factor(self, vector): probs = self._probs sqrt_probs = self._sqrt_probs return sqrt_probs * vector - probs * math_ops.reduce_sum( - sqrt_probs * vector, axis=-1, keep_dims=True) + sqrt_probs * vector, axis=-1, keepdims=True) def multiply_fisher_factor_transpose(self, vector): probs = self._probs sqrt_probs = self._sqrt_probs return sqrt_probs * vector - sqrt_probs * math_ops.reduce_sum( - probs * vector, axis=-1, keep_dims=True) + probs * vector, axis=-1, keepdims=True) def multiply_fisher_factor_replicated_one_hot(self, index): assert len(index) == 1, "Length of index was {}".format(len(index)) diff --git a/tensorflow/contrib/kfac/python/ops/loss_functions_lib.py b/tensorflow/contrib/kfac/python/ops/loss_functions_lib.py index 705a871d482565897e7ac850327729a6186f1746..4279cb2792854249e3e076d200e2656bc615779d 100644 --- a/tensorflow/contrib/kfac/python/ops/loss_functions_lib.py +++ b/tensorflow/contrib/kfac/python/ops/loss_functions_lib.py @@ -33,7 +33,6 @@ _allowed_symbols = [ "CategoricalLogitsNegativeLogProbLoss", "OnehotCategoricalLogitsNegativeLogProbLoss", "MultiBernoulliNegativeLogProbLoss", - "MultiBernoulliNegativeLogProbLoss", "insert_slice_in_zeros", ] diff --git a/tensorflow/contrib/labeled_tensor/python/ops/ops_test.py b/tensorflow/contrib/labeled_tensor/python/ops/ops_test.py index 0727f4cf88728dc3d919e662d65c93a658ac730b..39e9d65407f3b1e79804317023ea03dd81484ff5 100644 --- a/tensorflow/contrib/labeled_tensor/python/ops/ops_test.py +++ b/tensorflow/contrib/labeled_tensor/python/ops/ops_test.py @@ -660,7 +660,7 @@ class ReduceSumTest(Base): sum_lt = ops.reduce_sum(self.original_lt, {('channel', 'hihowareyou')}) golden_lt = core.LabeledTensor( math_ops.reduce_sum( - self.original_lt.tensor, 1, keep_dims=True), + self.original_lt.tensor, 1, keepdims=True), [self.a0, ('channel', ['hihowareyou']), self.a2, self.a3]) self.assertLabeledTensorsEqual(sum_lt, golden_lt) @@ -668,7 +668,7 @@ class ReduceSumTest(Base): sum_lt = ops.reduce_sum(self.original_lt, ('channel', 'hihowareyou')) golden_lt = core.LabeledTensor( math_ops.reduce_sum( - self.original_lt.tensor, 1, keep_dims=True), + self.original_lt.tensor, 1, keepdims=True), [self.a0, ('channel', ['hihowareyou']), self.a2, self.a3]) self.assertLabeledTensorsEqual(sum_lt, golden_lt) diff --git a/tensorflow/contrib/layers/python/kernel_tests/sparse_feature_cross_op_test.py b/tensorflow/contrib/layers/python/kernel_tests/sparse_feature_cross_op_test.py index f701647c2b297015f025eb53bd191a1a8c54ec62..28ddaa69a14776e0c157c2e68105ee9e17bc3cbb 100644 --- a/tensorflow/contrib/layers/python/kernel_tests/sparse_feature_cross_op_test.py +++ b/tensorflow/contrib/layers/python/kernel_tests/sparse_feature_cross_op_test.py @@ -200,7 +200,7 @@ class SparseCrossOpTest(test.TestCase): self._assert_sparse_tensor_equals(expected_out, sess.run(op)) def test_large_batch(self): - """Tests with large batch size to force multithreding. + """Tests with large batch size to force multithreading. """ batch_size = 5000 col1 = [] diff --git a/tensorflow/contrib/layers/python/layers/embedding_ops.py b/tensorflow/contrib/layers/python/layers/embedding_ops.py index ffa208540dae975cb139ad6d76dcf392678ba0ee..49c3faf3b7f5eaa3b1542a1fdddcfaff99737a24 100644 --- a/tensorflow/contrib/layers/python/layers/embedding_ops.py +++ b/tensorflow/contrib/layers/python/layers/embedding_ops.py @@ -140,6 +140,9 @@ def safe_embedding_lookup_sparse(embedding_weights, # Prune invalid ids and weights. sparse_ids, sparse_weights = _prune_invalid_ids(sparse_ids, sparse_weights) + if combiner != "sum": + sparse_ids, sparse_weights = _prune_invalid_weights( + sparse_ids, sparse_weights) # Fill in dummy values for empty features, if necessary. sparse_ids, is_row_empty = sparse_ops.sparse_fill_empty_rows(sparse_ids, @@ -188,13 +191,23 @@ def _prune_invalid_ids(sparse_ids, sparse_weights): is_id_valid = math_ops.greater_equal(sparse_ids.values, 0) if sparse_weights is not None: is_id_valid = math_ops.logical_and( - is_id_valid, math_ops.greater(sparse_weights.values, 0)) + is_id_valid, + array_ops.ones_like(sparse_weights.values, dtype=dtypes.bool)) sparse_ids = sparse_ops.sparse_retain(sparse_ids, is_id_valid) if sparse_weights is not None: sparse_weights = sparse_ops.sparse_retain(sparse_weights, is_id_valid) return sparse_ids, sparse_weights +def _prune_invalid_weights(sparse_ids, sparse_weights): + """Prune invalid weights (< 0) from the input ids and weights.""" + if sparse_weights is not None: + is_weights_valid = math_ops.greater(sparse_weights.values, 0) + sparse_ids = sparse_ops.sparse_retain(sparse_ids, is_weights_valid) + sparse_weights = sparse_ops.sparse_retain(sparse_weights, is_weights_valid) + return sparse_ids, sparse_weights + + def scattered_embedding_lookup(params, values, dimension, diff --git a/tensorflow/contrib/layers/python/layers/feature_column.py b/tensorflow/contrib/layers/python/layers/feature_column.py index 9ccb589d698ad83c9654f5523ccdcb35b031b3da..3ae07cedab0be2da8ec633cfd84e07cfdfb11457 100644 --- a/tensorflow/contrib/layers/python/layers/feature_column.py +++ b/tensorflow/contrib/layers/python/layers/feature_column.py @@ -48,7 +48,7 @@ you should choose depends on (1) the feature type and (2) the model type. recommended. embedded_dept_column = embedding_column( - sparse_column_with_keys("department", ["math", "philosphy", ...]), + sparse_column_with_keys("department", ["math", "philosophy", ...]), dimension=10) * Wide (aka linear) models (`LinearClassifier`, `LinearRegressor`). diff --git a/tensorflow/contrib/layers/python/layers/feature_column_ops.py b/tensorflow/contrib/layers/python/layers/feature_column_ops.py index 78affea44cbfb92523063968dbc1be98841854db..06060b99e7e58787994f20f037ffa451abbc7459 100644 --- a/tensorflow/contrib/layers/python/layers/feature_column_ops.py +++ b/tensorflow/contrib/layers/python/layers/feature_column_ops.py @@ -815,7 +815,7 @@ class _Transformer(object): """ def __init__(self, columns_to_tensors): - """Initializes transfomer. + """Initializes transformer. Args: columns_to_tensors: A mapping from feature columns to tensors. 'string' @@ -908,7 +908,7 @@ def _gather_feature_columns(feature_columns): def _check_forbidden_sequence_columns(feature_columns): - """Recursively cecks `feature_columns` for `_FORBIDDEN_SEQUENCE_COLUMNS`.""" + """Recursively checks `feature_columns` for `_FORBIDDEN_SEQUENCE_COLUMNS`.""" all_feature_columns = _gather_feature_columns(feature_columns) for feature_column in all_feature_columns: if isinstance(feature_column, _FORBIDDEN_SEQUENCE_COLUMNS): diff --git a/tensorflow/contrib/layers/python/layers/layers.py b/tensorflow/contrib/layers/python/layers/layers.py index 10d7f6d076b4b4c6578d7adcffc4e9cc44d77ac6..2f3e57653c5d6d949c4dcc91635690322b7f90c4 100644 --- a/tensorflow/contrib/layers/python/layers/layers.py +++ b/tensorflow/contrib/layers/python/layers/layers.py @@ -932,7 +932,8 @@ def convolution(inputs, variables_collections=None, outputs_collections=None, trainable=True, - scope=None): + scope=None, + conv_dims=None): """Adds an N-D convolution followed by an optional batch_norm layer. It is required that 1 <= N <= 3. @@ -993,6 +994,10 @@ def convolution(inputs, trainable: If `True` also add variables to the graph collection `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable). scope: Optional scope for `variable_scope`. + conv_dims: Optional convolution dimensionality, when set it would use the + corresponding convolution (e.g. 2 for Conv 2D, 3 for Conv 3D, ..). When + leaved to None it would select the convolution dimensionality based on + the input rank (i.e. Conv ND, with N = input_rank - 2). Returns: A tensor representing the output of the operation. @@ -1015,6 +1020,9 @@ def convolution(inputs, inputs = ops.convert_to_tensor(inputs) input_rank = inputs.get_shape().ndims + if conv_dims is not None and conv_dims + 2 != input_rank: + raise ValueError('Convolution expects input with rank %d, got %d' % + (conv_dims + 2, input_rank)) if input_rank == 3: layer_class = convolutional_layers.Convolution1D elif input_rank == 4: @@ -1061,10 +1069,134 @@ def convolution(inputs, outputs = activation_fn(outputs) return utils.collect_named_outputs(outputs_collections, sc.name, outputs) +@add_arg_scope +def convolution1d(inputs, + num_outputs, + kernel_size, + stride=1, + padding='SAME', + data_format=None, + rate=1, + activation_fn=nn.relu, + normalizer_fn=None, + normalizer_params=None, + weights_initializer=initializers.xavier_initializer(), + weights_regularizer=None, + biases_initializer=init_ops.zeros_initializer(), + biases_regularizer=None, + reuse=None, + variables_collections=None, + outputs_collections=None, + trainable=True, + scope=None): + return convolution(inputs, + num_outputs, + kernel_size, + stride, + padding, + data_format, + rate, + activation_fn, + normalizer_fn, + normalizer_params, + weights_initializer, + weights_regularizer, + biases_initializer, + biases_regularizer, + reuse, + variables_collections, + outputs_collections, + trainable, + scope, + conv_dims=1) + +convolution1d.__doc__ = convolution.__doc__ -convolution2d = convolution -convolution3d = convolution +@add_arg_scope +def convolution2d(inputs, + num_outputs, + kernel_size, + stride=1, + padding='SAME', + data_format=None, + rate=1, + activation_fn=nn.relu, + normalizer_fn=None, + normalizer_params=None, + weights_initializer=initializers.xavier_initializer(), + weights_regularizer=None, + biases_initializer=init_ops.zeros_initializer(), + biases_regularizer=None, + reuse=None, + variables_collections=None, + outputs_collections=None, + trainable=True, + scope=None): + return convolution(inputs, + num_outputs, + kernel_size, + stride, + padding, + data_format, + rate, + activation_fn, + normalizer_fn, + normalizer_params, + weights_initializer, + weights_regularizer, + biases_initializer, + biases_regularizer, + reuse, + variables_collections, + outputs_collections, + trainable, + scope, + conv_dims=2) + +convolution2d.__doc__ = convolution.__doc__ +@add_arg_scope +def convolution3d(inputs, + num_outputs, + kernel_size, + stride=1, + padding='SAME', + data_format=None, + rate=1, + activation_fn=nn.relu, + normalizer_fn=None, + normalizer_params=None, + weights_initializer=initializers.xavier_initializer(), + weights_regularizer=None, + biases_initializer=init_ops.zeros_initializer(), + biases_regularizer=None, + reuse=None, + variables_collections=None, + outputs_collections=None, + trainable=True, + scope=None): + return convolution(inputs, + num_outputs, + kernel_size, + stride, + padding, + data_format, + rate, + activation_fn, + normalizer_fn, + normalizer_params, + weights_initializer, + weights_regularizer, + biases_initializer, + biases_regularizer, + reuse, + variables_collections, + outputs_collections, + trainable, + scope, + conv_dims=3) + +convolution3d.__doc__ = convolution.__doc__ @add_arg_scope def convolution2d_in_plane( @@ -1404,13 +1536,14 @@ def convolution3d_transpose( @add_arg_scope def dense_to_sparse(tensor, eos_token=0, outputs_collections=None, scope=None): """Converts a dense tensor into a sparse tensor. + An example use would be to convert dense labels to sparse ones so that they can be fed to the ctc_loss. Args: tensor: An `int` `Tensor` to be converted to a `Sparse`. eos_token: An integer. - It is part of the target label that signfies the end of a sentence. + It is part of the target label that signifies the end of a sentence. outputs_collections: Collection to add the outputs. scope: Optional scope for name_scope. """ @@ -1554,7 +1687,7 @@ def _inner_flatten(inputs, new_rank, output_collections=None, scope=None): output_collections: Collection to which the outputs will be added. scope: Optional scope for `name_scope`. Returns: - A `Tensor` or `SparseTensor` conataining the same values as `inputs`, but + A `Tensor` or `SparseTensor` containing the same values as `inputs`, but with innermost dimensions flattened to obtain rank `new_rank`. Raises: @@ -2191,11 +2324,16 @@ def images_to_sequence(inputs, 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. + + Raises: + ValueError: If `data_format` is not either NCHW or NHWC. + Returns: (width, num_images*height, depth) sequence tensor """ @@ -2701,6 +2839,7 @@ def sequence_to_images(inputs, 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 @@ -2708,6 +2847,7 @@ def sequence_to_images(inputs, 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. """ @@ -2717,7 +2857,7 @@ def sequence_to_images(inputs, if num_batches is None: num_batches = -1 else: - num_batches = num_batches // height + num_batches //= height reshaped = array_ops.reshape(inputs, [width, num_batches, height, depth]) if output_data_format == 'channels_first': diff --git a/tensorflow/contrib/layers/python/layers/layers_test.py b/tensorflow/contrib/layers/python/layers/layers_test.py index 997f910a2a97567adbd7ffa3e81a31d2ae0bad7e..b01fd5d5c95ac15c76f9dbe7c77f7e76f12149a9 100644 --- a/tensorflow/contrib/layers/python/layers/layers_test.py +++ b/tensorflow/contrib/layers/python/layers/layers_test.py @@ -310,6 +310,17 @@ class BiasAddTest(test.TestCase): class ConvolutionTest(test.TestCase): + def testInvalidShape(self): + with self.test_session(): + images_2d = random_ops.random_uniform((5, 7, 9, 3), seed=1) + with self.assertRaisesRegexp( + ValueError, 'Convolution expects input with rank 5, got 4'): + layers_lib.convolution3d(images_2d, 32, 3) + images_3d = random_ops.random_uniform((5, 6, 7, 9, 3), seed=1) + with self.assertRaisesRegexp( + ValueError, 'Convolution expects input with rank 4, got 5'): + layers_lib.convolution2d(images_3d, 32, 3) + def testInvalidDataFormat(self): height, width = 7, 9 with self.test_session(): @@ -3155,7 +3166,7 @@ class RepeatTests(test.TestCase): with self.test_session(): images = np.random.uniform(size=(5, height, width, 3)).astype(np.float32) output = _layers.repeat(images, 3, layers_lib.conv2d, 32, [3, 3]) - self.assertEqual(output.op.name, 'Repeat/convolution_3/Relu') + self.assertEqual(output.op.name, 'Repeat/convolution2d_3/Relu') self.assertListEqual(output.get_shape().as_list(), [5, 3, 3, 32]) def testRepeatWithScope(self): @@ -3749,7 +3760,7 @@ class StackTests(test.TestCase): layers_lib.convolution2d, [10, 20, 30], kernel_size=[3, 3], padding='SAME') - self.assertEqual(output.op.name, 'Stack/convolution_3/Relu') + self.assertEqual(output.op.name, 'Stack/convolution2d_3/Relu') self.assertListEqual(output.get_shape().as_list(), [5, 3, 3, 30]) def testStackWithScope(self): diff --git a/tensorflow/contrib/layers/python/layers/target_column.py b/tensorflow/contrib/layers/python/layers/target_column.py index 3e639a180ef11af5f7f498c647eb25417f918eb9..69bb6be81453f5f5487f25547f017dc5f87c2f2c 100644 --- a/tensorflow/contrib/layers/python/layers/target_column.py +++ b/tensorflow/contrib/layers/python/layers/target_column.py @@ -270,7 +270,7 @@ class _RegressionTargetColumn(_TargetColumn): def logits_to_predictions(self, logits, proba=False): if self.num_label_columns == 1: - return array_ops.squeeze(logits, squeeze_dims=[1]) + return array_ops.squeeze(logits, axis=[1]) return logits def get_eval_ops(self, features, logits, labels, metrics=None): @@ -418,7 +418,7 @@ def _softmax_cross_entropy_loss(logits, target): "Instead got %s." % target.dtype) # sparse_softmax_cross_entropy_with_logits requires [batch_size] target. if len(target.get_shape()) == 2: - target = array_ops.squeeze(target, squeeze_dims=[1]) + target = array_ops.squeeze(target, axis=[1]) loss_vec = nn.sparse_softmax_cross_entropy_with_logits( labels=target, logits=logits) return loss_vec diff --git a/tensorflow/contrib/layers/python/layers/utils_test.py b/tensorflow/contrib/layers/python/layers/utils_test.py index 3409860add8f8c393ffd342633e7023931867dd9..645dc1291eb6370a5e504306fc00a5454dde77ed 100644 --- a/tensorflow/contrib/layers/python/layers/utils_test.py +++ b/tensorflow/contrib/layers/python/layers/utils_test.py @@ -294,7 +294,6 @@ class NPositiveIntegersTest(test.TestCase): self.assertEqual(utils.n_positive_integers(2, 2), (2, 2)) self.assertEqual(utils.n_positive_integers(2, (2, 3)), (2, 3)) self.assertEqual(utils.n_positive_integers(3, (2, 3, 1)), (2, 3, 1)) - self.assertEqual(utils.n_positive_integers(3, (2, 3, 1)), (2, 3, 1)) self.assertEqual( utils.n_positive_integers(3, tensor_shape.TensorShape([2, 3, 1])), (2, 3, 1)) diff --git a/tensorflow/contrib/learn/BUILD b/tensorflow/contrib/learn/BUILD index d665fc9335cf22cdfa1e7330ab67003042502515..3b053cd4c66952cf6c494186b16c17f38801bcaf 100644 --- a/tensorflow/contrib/learn/BUILD +++ b/tensorflow/contrib/learn/BUILD @@ -281,7 +281,10 @@ py_test( size = "medium", srcs = ["python/learn/estimators/estimator_test.py"], srcs_version = "PY2AND3", - tags = ["manual"], + tags = [ + "manual", + "noasan", # times out + ], deps = [ ":learn", "//tensorflow/contrib/framework:framework_py", diff --git a/tensorflow/contrib/learn/python/learn/estimators/head.py b/tensorflow/contrib/learn/python/learn/estimators/head.py index 2b4b6eff39f4fc8a20a149edfc07d2f4f27a9bae..e28e6854a5097d66cb486be3e82f3726f5cc70fd 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/head.py +++ b/tensorflow/contrib/learn/python/learn/estimators/head.py @@ -777,7 +777,7 @@ class _RegressionHead(_SingleHead): key = prediction_key.PredictionKey.SCORES with ops.name_scope(None, "predictions", (logits,)): if self.logits_dimension == 1: - logits = array_ops.squeeze(logits, squeeze_dims=(1,), name=key) + logits = array_ops.squeeze(logits, axis=(1,), name=key) return {key: self._link_fn(logits)} def _metrics(self, eval_loss, predictions, labels, weights): @@ -974,7 +974,7 @@ def _softmax_cross_entropy_loss(labels, logits, weights=None): is_squeezed_labels = False # TODO(ptucker): This will break for dynamic shapes. if len(labels.get_shape()) == 2: - labels = array_ops.squeeze(labels, squeeze_dims=(1,)) + labels = array_ops.squeeze(labels, axis=(1,)) is_squeezed_labels = True loss = nn.sparse_softmax_cross_entropy_with_logits( diff --git a/tensorflow/contrib/learn/python/learn/estimators/kmeans_test.py b/tensorflow/contrib/learn/python/learn/estimators/kmeans_test.py index b28835a809736a099ad2f08d127dc68d7977a3c1..584556992a0db2345e182e92c4a7f7582d3cd8dc 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/kmeans_test.py +++ b/tensorflow/contrib/learn/python/learn/estimators/kmeans_test.py @@ -36,7 +36,6 @@ from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import data_flow_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops -from tensorflow.python.ops import random_ops from tensorflow.python.platform import benchmark from tensorflow.python.platform import flags from tensorflow.python.platform import test diff --git a/tensorflow/contrib/learn/python/learn/estimators/run_config.py b/tensorflow/contrib/learn/python/learn/estimators/run_config.py index 8c85c431be69caaca6872111896b9487faf9e679..14ee2ba6094760d52180d6de7763ea88b8ee98c8 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/run_config.py +++ b/tensorflow/contrib/learn/python/learn/estimators/run_config.py @@ -299,6 +299,7 @@ class RunConfig(ClusterConfig, core_run_config.RunConfig): # so instead of breaking compatibility with that assumption, we # just manually initialize this field: self._train_distribute = None + self._device_fn = None gpu_options = config_pb2.GPUOptions( per_process_gpu_memory_fraction=gpu_memory_fraction) diff --git a/tensorflow/contrib/learn/python/learn/learn_io/data_feeder_test.py b/tensorflow/contrib/learn/python/learn/learn_io/data_feeder_test.py index 82848be7df653dd60219317d28f233767746f544..1f439965daf956665bbedc919281df0ee07b5d62 100644 --- a/tensorflow/contrib/learn/python/learn/learn_io/data_feeder_test.py +++ b/tensorflow/contrib/learn/python/learn/learn_io/data_feeder_test.py @@ -18,6 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import os.path import numpy as np import six from six.moves import xrange # pylint: disable=redefined-builtin @@ -26,6 +27,7 @@ from six.moves import xrange # pylint: disable=redefined-builtin from tensorflow.contrib.learn.python.learn.learn_io import * from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops +from tensorflow.python.lib.io import file_io from tensorflow.python.platform import test # pylint: enable=wildcard-import @@ -35,6 +37,13 @@ class DataFeederTest(test.TestCase): # pylint: disable=undefined-variable """Tests for `DataFeeder`.""" + def setUp(self): + self._base_dir = os.path.join(self.get_temp_dir(), 'base_dir') + file_io.create_dir(self._base_dir) + + def tearDown(self): + file_io.delete_recursively(self._base_dir) + def _wrap_dict(self, data, prepend=''): return {prepend + '1': data, prepend + '2': data} @@ -45,14 +54,14 @@ class DataFeederTest(test.TestCase): def _assert_dtype(self, expected_np_dtype, expected_tf_dtype, input_data): feeder = data_feeder.DataFeeder(input_data, None, n_classes=0, batch_size=1) if isinstance(input_data, dict): - for k, v in list(feeder.input_dtype.items()): + for v in list(feeder.input_dtype.values()): self.assertEqual(expected_np_dtype, v) else: self.assertEqual(expected_np_dtype, feeder.input_dtype) with ops.Graph().as_default() as g, self.test_session(g): inp, _ = feeder.input_builder() if isinstance(inp, dict): - for k, v in list(inp.items()): + for v in list(inp.values()): self.assertEqual(expected_tf_dtype, v.dtype) else: self.assertEqual(expected_tf_dtype, inp.dtype) @@ -301,7 +310,10 @@ class DataFeederTest(test.TestCase): [0.60000002, 0.2]]) self.assertAllClose(feed_dict[out.name], [[0., 0., 1.], [0., 1., 0.]]) - def test_hdf5_data_feeder(self): + # TODO(rohanj): Fix this test by fixing data_feeder. Currently, h5py doesn't + # support permutation based indexing lookups (More documentation at + # http://docs.h5py.org/en/latest/high/dataset.html#fancy-indexing) + def DISABLED_test_hdf5_data_feeder(self): def func(df): inp, out = df.input_builder() @@ -314,11 +326,12 @@ class DataFeederTest(test.TestCase): import h5py # pylint: disable=g-import-not-at-top x = np.matrix([[1, 2], [3, 4]]) y = np.array([1, 2]) - h5f = h5py.File('test_hdf5.h5', 'w') + file_path = os.path.join(self._base_dir, 'test_hdf5.h5') + h5f = h5py.File(file_path, 'w') h5f.create_dataset('x', data=x) h5f.create_dataset('y', data=y) h5f.close() - h5f = h5py.File('test_hdf5.h5', 'r') + h5f = h5py.File(file_path, 'r') x = h5f['x'] y = h5f['y'] func(data_feeder.DataFeeder(x, y, n_classes=0, batch_size=3)) diff --git a/tensorflow/contrib/learn/python/learn/ops/losses_ops.py b/tensorflow/contrib/learn/python/learn/ops/losses_ops.py index 92976d1539c7ddc226b81f903beee82b798ec8db..9f2cadb01747c5a8e4ee75ac38f423f85e11bbba 100644 --- a/tensorflow/contrib/learn/python/learn/ops/losses_ops.py +++ b/tensorflow/contrib/learn/python/learn/ops/losses_ops.py @@ -40,7 +40,7 @@ def mean_squared_error_regressor(tensor_in, labels, weights, biases, name=None): [tensor_in, labels]): predictions = nn.xw_plus_b(tensor_in, weights, biases) if len(labels.get_shape()) == 1 and len(predictions.get_shape()) == 2: - predictions = array_ops_.squeeze(predictions, squeeze_dims=[1]) + predictions = array_ops_.squeeze(predictions, axis=[1]) return predictions, losses.mean_squared_error(labels, predictions) diff --git a/tensorflow/contrib/legacy_seq2seq/BUILD b/tensorflow/contrib/legacy_seq2seq/BUILD index 8c2c4fd29c0502d4199f27a65e4827b2db973c3d..4ce91a140f816ddc8bdc60287e4cbc807172ec6d 100644 --- a/tensorflow/contrib/legacy_seq2seq/BUILD +++ b/tensorflow/contrib/legacy_seq2seq/BUILD @@ -58,5 +58,8 @@ cuda_py_tests( "//tensorflow/python:variable_scope", "//tensorflow/python:variables", ], - tags = ["noasan"], # times out b/63678675 + tags = [ + "noasan", # times out b/63678675 + "optonly", # times out (flaky) + ], ) diff --git a/tensorflow/contrib/linalg/BUILD b/tensorflow/contrib/linalg/BUILD index a7812f74d1e69276a4bba597b41e442bc4dbbc4a..2e92ad6eb39d8aa8876a34572f50d5b6aff0511a 100644 --- a/tensorflow/contrib/linalg/BUILD +++ b/tensorflow/contrib/linalg/BUILD @@ -58,6 +58,31 @@ cuda_py_test( "//tensorflow/python:math_ops", "//tensorflow/python:platform_test", ], - shard_count = 4, - tags = ["noasan"], + shard_count = 5, + tags = [ + "noasan", + "optonly", + ], +) + +cuda_py_test( + name = "linear_operator_kronecker_test", + size = "medium", + srcs = ["python/kernel_tests/linear_operator_kronecker_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 = 8, + tags = [ + "noasan", + "optonly", + ], ) diff --git a/tensorflow/contrib/linalg/__init__.py b/tensorflow/contrib/linalg/__init__.py index 14cc3b2b4971de1a31960ee33c2f304154b1f411..38bd66b13f79ee918acbdc2d33a6367e4e34d4b1 100644 --- a/tensorflow/contrib/linalg/__init__.py +++ b/tensorflow/contrib/linalg/__init__.py @@ -22,6 +22,7 @@ See the @{$python/contrib.linalg} guide. @@LinearOperatorIdentity @@LinearOperatorScaledIdentity @@LinearOperatorFullMatrix +@@LinearOperatorKronecker @@LinearOperatorLowerTriangular @@LinearOperatorLowRankUpdate @@LinearOperatorComposition @@ -36,6 +37,7 @@ from __future__ import print_function from tensorflow.contrib.linalg.python.ops.linear_operator_addition import * from tensorflow.contrib.linalg.python.ops.linear_operator_block_diag import * +from tensorflow.contrib.linalg.python.ops.linear_operator_kronecker 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 * 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 index cc1a047d6a2b6029080fad3f240aa00f50504f07..e7407ede11409a47f4d9db96ad5b5d801ef1625d 100644 --- 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 @@ -76,6 +76,8 @@ class SquareLinearOperatorBlockDiagTest( build_info((1, 1)), build_info((1, 3, 3)), build_info((5, 5), blocks=[(2, 2), (3, 3)]), + build_info((3, 7, 7), blocks=[(1, 2, 2), (3, 2, 2), (1, 3, 3)]), + build_info((2, 1, 5, 5), blocks=[(2, 1, 2, 2), (1, 3, 3)]), ] def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder): @@ -184,70 +186,5 @@ class SquareLinearOperatorBlockDiagTest( 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/kernel_tests/linear_operator_kronecker_test.py b/tensorflow/contrib/linalg/python/kernel_tests/linear_operator_kronecker_test.py new file mode 100644 index 0000000000000000000000000000000000000000..6574da22a188c7aa25ad8426522d0b446af8f5f3 --- /dev/null +++ b/tensorflow/contrib/linalg/python/kernel_tests/linear_operator_kronecker_test.py @@ -0,0 +1,194 @@ +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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_kronecker as kronecker +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.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 _kronecker_dense(factors): + """Convert a list of factors, into a dense Kronecker product.""" + product = factors[0] + for factor in factors[1:]: + product = product[..., array_ops.newaxis, :, array_ops.newaxis] + factor_to_mul = factor[..., array_ops.newaxis, :, array_ops.newaxis, :] + product *= factor_to_mul + product = array_ops.reshape( + product, + shape=array_ops.concat( + [array_ops.shape(product)[:-4], + [array_ops.shape(product)[-4] * array_ops.shape(product)[-3], + array_ops.shape(product)[-2] * array_ops.shape(product)[-1]] + ], axis=0)) + + return product + + +class KroneckerDenseTest(test.TestCase): + + def testKroneckerDenseMatrix(self): + x = ops.convert_to_tensor([[2., 3.], [1., 2.]], dtype=dtypes.float32) + y = ops.convert_to_tensor([[1., 2.], [5., -1.]], dtype=dtypes.float32) + # From explicitly writing out the kronecker product of x and y. + z = ops.convert_to_tensor([ + [2., 4., 3., 6.], + [10., -2., 15., -3.], + [1., 2., 2., 4.], + [5., -1., 10., -2.]], dtype=dtypes.float32) + # From explicitly writing out the kronecker product of y and x. + w = ops.convert_to_tensor([ + [2., 3., 4., 6.], + [1., 2., 2., 4.], + [10., 15., -2., -3.], + [5., 10., -1., -2.]], dtype=dtypes.float32) + + with self.test_session(): + self.assertAllClose(_kronecker_dense([x, y]).eval(), z.eval()) + self.assertAllClose(_kronecker_dense([y, x]).eval(), w.eval()) + + +class SquareLinearOperatorKroneckerTest( + 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((1, 1), factors=[(1, 1), (1, 1)]), + build_info((8, 8), factors=[(2, 2), (2, 2), (2, 2)]), + build_info((12, 12), factors=[(2, 2), (3, 3), (2, 2)]), + build_info((1, 3, 3), factors=[(1, 1), (1, 3, 3)]), + build_info((3, 6, 6), factors=[(3, 1, 1), (1, 2, 2), (1, 3, 3)]), + ] + + def _operator_and_mat_and_feed_dict(self, build_info, dtype, use_placeholder): + shape = list(build_info.shape) + expected_factors = build_info.__dict__["factors"] + matrices = [ + linear_operator_test_util.random_positive_definite_matrix( + block_shape, dtype, force_well_conditioned=True) + for block_shape in expected_factors + ] + + if use_placeholder: + matrices_ph = [ + array_ops.placeholder(dtype=dtype) for _ in expected_factors + ] + # 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 = kronecker.LinearOperatorKronecker( + [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 = kronecker.LinearOperatorKronecker( + [linalg.LinearOperatorFullMatrix( + m, is_square=True) for m in matrices]) + feed_dict = None + # Should be auto-set. + self.assertTrue(operator.is_square) + + matrices = linear_operator_util.broadcast_matrix_batch_dims(matrices) + + kronecker_dense = _kronecker_dense(matrices) + + if not use_placeholder: + kronecker_dense.set_shape(shape) + + return operator, kronecker_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 = kronecker.LinearOperatorKronecker( + [linalg.LinearOperatorFullMatrix(matrix), + 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 = kronecker.LinearOperatorKronecker( + [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"): + kronecker.LinearOperatorKronecker( + [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 = kronecker.LinearOperatorKronecker([operator_1, operator_2]) + + self.assertEqual("left_x_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"): + kronecker.LinearOperatorKronecker(operators) + + def test_empty_or_one_operators_raises(self): + with self.assertRaisesRegexp(ValueError, ">=1 operators"): + kronecker.LinearOperatorKronecker([]) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/linalg/python/ops/linear_operator_kronecker.py b/tensorflow/contrib/linalg/python/ops/linear_operator_kronecker.py new file mode 100644 index 0000000000000000000000000000000000000000..79080d194f59b7ebce045ab3e3d262ca948d9391 --- /dev/null +++ b/tensorflow/contrib/linalg/python/ops/linear_operator_kronecker.py @@ -0,0 +1,560 @@ +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Construct the Kronecker product of 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 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 check_ops +from tensorflow.python.ops import control_flow_ops +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 + + +def _vec(x): + """Stacks column of matrix to form a single column.""" + return array_ops.reshape( + array_ops.matrix_transpose(x), + array_ops.concat( + [array_ops.shape(x)[:-2], [-1]], axis=0)) + + +def _unvec_by(y, num_col): + """Unstack vector to form a matrix, with a specified amount of columns.""" + return array_ops.matrix_transpose( + array_ops.reshape( + y, + array_ops.concat( + [array_ops.shape(y)[:-1], [num_col, -1]], axis=0))) + + +def _rotate_last_dim(x, rotate_right=False): + """Rotate the last dimension either left or right.""" + ndims = array_ops.rank(x) + if rotate_right: + transpose_perm = array_ops.concat( + [[ndims - 1], math_ops.range(0, ndims - 1)], axis=0) + else: + transpose_perm = array_ops.concat( + [math_ops.range(1, ndims), [0]], axis=0) + return array_ops.transpose(x, transpose_perm) + + +class LinearOperatorKronecker(linear_operator.LinearOperator): + """Kronecker product between two `LinearOperators`. + + This operator composes one or more linear operators `[op1,...,opJ]`, + building a new `LinearOperator` representing the Kronecker product: + `op1 x op2 x .. opJ` (we omit parentheses as the Kronecker product is + associative). + + If `opj` has shape `batch_shape_j` + [M_j, N_j`, then the composed operator + will have shape equal to `broadcast_batch_shape + [prod M_j, prod N_j]`, + where the product is over all operators. + + ```python + # Create a 4 x 4 linear operator composed of two 2 x 2 operators. + operator_1 = LinearOperatorFullMatrix([[1., 2.], [3., 4.]]) + operator_2 = LinearOperatorFullMatrix([[1., 0.], [2., 1.]]) + operator = LinearOperatorKronecker([operator_1, operator_2]) + + operator.to_dense() + ==> [[1., 2., 0., 0.], + [3., 4., 0., 0.], + [2., 4., 1., 2.], + [6., 8., 3., 4.]] + + operator.shape + ==> [4, 4] + + operator.log_abs_determinant() + ==> scalar Tensor + + x = ... Shape [4, 2] Tensor + operator.matmul(x) + ==> Shape [4, 2] Tensor + + # Create a [2, 3] batch of 4 x 5 linear operators. + matrix_45 = tf.random_normal(shape=[2, 3, 4, 5]) + operator_45 = LinearOperatorFullMatrix(matrix) + + # Create a [2, 3] batch of 5 x 6 linear operators. + matrix_56 = tf.random_normal(shape=[2, 3, 5, 6]) + operator_56 = LinearOperatorFullMatrix(matrix_56) + + # Compose to create a [2, 3] batch of 20 x 30 operators. + operator_large = LinearOperatorKronecker([operator_45, operator_56]) + + # Create a shape [2, 3, 20, 2] vector. + x = tf.random_normal(shape=[2, 3, 6, 2]) + operator_large.matmul(x) + ==> Shape [2, 3, 30, 2] Tensor + ``` + + #### Performance + + The performance of `LinearOperatorKronecker` 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=None, + name=None): + r"""Initialize a `LinearOperatorKronecker`. + + `LinearOperatorKronecker` 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, representing the Kronecker + factors. + 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. + name: A name for this `LinearOperator`. Default is the individual + operators names joined with `_x_`. + + Raises: + TypeError: If all operators do not have the same `dtype`. + ValueError: If `operators` is empty. + """ + # Validate operators. + check_ops.assert_proper_iterable(operators) + operators = list(operators) + if not operators: + raise ValueError( + "Expected a list of >=1 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. + # A Kronecker product is invertible, if and only if all factors are + # invertible. + if all(operator.is_non_singular for operator in operators): + if is_non_singular is False: + raise ValueError( + "The Kronecker product 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 Kronecker product of self-adjoint operators is always " + "self-adjoint.") + is_self_adjoint = True + + # The eigenvalues of a Kronecker product are equal to the products of eigen + # values of the corresponding factors. + if all(operator.is_positive_definite for operator in operators): + if is_positive_definite is False: + raise ValueError("The Kronecker product of positive-definite operators " + "is always positive-definite.") + is_positive_definite = True + + # Initialization. + graph_parents = [] + for operator in operators: + graph_parents.extend(operator.graph_parents) + + if name is None: + name = operators[0].name + for operator in operators[1:]: + name += "_x_" + operator.name + with ops.name_scope(name, values=graph_parents): + super(LinearOperatorKronecker, 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=is_square, + name=name) + + @property + def operators(self): + return self._operators + + def _shape(self): + # Get final matrix shape. + domain_dimension = self.operators[0].domain_dimension + for operator in self.operators[1:]: + domain_dimension *= operator.domain_dimension + + range_dimension = self.operators[0].range_dimension + for operator in self.operators[1:]: + range_dimension *= operator.range_dimension + + matrix_shape = tensor_shape.TensorShape([ + range_dimension, domain_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): + domain_dimension = self.operators[0].domain_dimension_tensor() + for operator in self.operators[1:]: + domain_dimension *= operator.domain_dimension_tensor() + + range_dimension = self.operators[0].range_dimension_tensor() + for operator in self.operators[1:]: + range_dimension *= operator.range_dimension_tensor() + + matrix_shape = [range_dimension, domain_dimension] + + # Get broadcast batch shape. + # broadcast_shape checks for compatibility. + batch_shape = self.operators[0].batch_shape_tensor() + for operator in self.operators[1:]: + batch_shape = array_ops.broadcast_dynamic_shape( + batch_shape, operator.batch_shape_tensor()) + + return array_ops.concat((batch_shape, matrix_shape), 0) + + def _matmul(self, x, adjoint=False, adjoint_arg=False): + # Here we heavily rely on Roth's column Lemma [1]: + # (A x B) * vec X = vec BXA^T, + # where vec stacks all the columns of the matrix under each other. In our + # case, x represents a batch of vec X (i.e. we think of x as a batch of + # column vectors, rather than a matrix). Each member of the batch can be + # reshaped to a matrix (hence we get a batch of matrices). + # We can iteratively apply this lemma by noting that if B is a Kronecker + # product, then we can apply the lemma again. + + # [1] W. E. Roth, "On direct product matrices," + # Bulletin of the American Mathematical Society, vol. 40, pp. 461-468, + # 1934 + + # Efficiency + + # Naively doing the Kronecker product, by calculating the dense matrix and + # applying it will can take cubic time in the size of domain_dimension + # (assuming a square matrix). The other issue is that calculating the dense + # matrix can be prohibitively expensive, in that it can take a large amount + # of memory. + # + # This implementation avoids this memory blow up by only computing matmuls + # with the factors. In this way, we don't have to realize the dense matrix. + # In terms of complexity, if we have Kronecker Factors of size: + # (n1, n1), (n2, n2), (n3, n3), ... (nJ, nJ), with N = \prod n_i, and we + # have as input a [N, M] matrix, the naive approach would take O(N^2 M). + # With this approach (ignoring reshaping of tensors and transposes for now), + # the time complexity can be O(M * (\sum n_i) * N). There is also the + # benefit of batched multiplication (In this example, the batch size is + # roughly M * N) so this can be much faster. However, not factored in are + # the costs of the several transposing of tensors, which can affect cache + # behavior. + + # Below we document the shape manipulation for adjoint=False, + # adjoint_arg=False, but the general case of different adjoints is still + # handled. + + if adjoint_arg: + x = linalg.adjoint(x) + + # Always add a batch dimension to enable broadcasting to work. + batch_shape = array_ops.concat( + [array_ops.ones_like(self.batch_shape_tensor()), [1, 1]], 0) + x += array_ops.zeros(batch_shape, dtype=x.dtype.base_dtype) + + # x has shape [B, R, C], where B represent some number of batch dimensions, + # R represents the number of rows, and C represents the number of columns. + # In order to apply Roth's column lemma, we need to operate on a batch of + # column vectors, so we reshape into a batch of column vectors. We put it + # at the front to ensure that broadcasting between operators to the batch + # dimensions B still works. + output = _rotate_last_dim(x, rotate_right=True) + + # Also expand the shape to be [A, C, B, R]. The first dimension will be + # used to accumulate dimensions from each operator matmul. + output = output[array_ops.newaxis, ...] + + # In this loop, A is going to refer to the value of the accumulated + # dimension. A = 1 at the start, and will end up being self.range_dimension. + # V will refer to the last dimension. V = R at the start, and will end up + # being 1 in the end. + for operator in self.operators[:-1]: + # Reshape output from [A, C, B, V] to be + # [A, C, B, V / op.domain_dimension, op.domain_dimension] + if adjoint: + operator_dimension = operator.range_dimension_tensor() + else: + operator_dimension = operator.domain_dimension_tensor() + + output = _unvec_by(output, operator_dimension) + + # We are computing (XA^T) = (AX^T)^T. + # output has [A, C, B, V / op.domain_dimension, op.domain_dimension], + # which is being converted to: + # [A, C, B, V / op.domain_dimension, op.range_dimension] + output = array_ops.matrix_transpose(output) + output = operator.matmul(output, adjoint=adjoint, adjoint_arg=False) + output = array_ops.matrix_transpose(output) + # Rearrange it to [A * op.range_dimension, C, B, V / op.domain_dimension] + output = _rotate_last_dim(output, rotate_right=False) + output = _vec(output) + output = _rotate_last_dim(output, rotate_right=True) + + # After the loop, we will have + # A = self.range_dimension / op[-1].range_dimension + # V = op[-1].domain_dimension + + # We convert that using matvec to get: + # [A, C, B, op[-1].range_dimension] + output = self.operators[-1].matvec(output, adjoint=adjoint) + # Rearrange shape to be [B1, ... Bn, self.range_dimension, C] + output = _rotate_last_dim(output, rotate_right=False) + output = _vec(output) + output = _rotate_last_dim(output, rotate_right=False) + + if x.shape.is_fully_defined(): + column_dim = x.shape[-1] + broadcast_batch_shape = common_shapes.broadcast_shape( + x.shape[:-2], self.batch_shape) + if adjoint: + matrix_dimensions = [self.domain_dimension, column_dim] + else: + matrix_dimensions = [self.range_dimension, column_dim] + + print("x: ", x) + print("bathc_shape:", self.batch_shape) + print("self.shape:", self.shape) + print("output: ", output) + output.set_shape(broadcast_batch_shape.concatenate( + matrix_dimensions)) + + return output + + def _determinant(self): + # Note that we have |X1 x X2| = |X1| ** n * |X2| ** m, where X1 is an m x m + # matrix, and X2 is an n x n matrix. We can iteratively apply this property + # to get the determinant of |X1 x X2 x X3 ...|. If T is the product of the + # domain dimension of all operators, then we have: + # |X1 x X2 x X3 ...| = + # |X1| ** (T / m) * |X2 x X3 ... | ** m = + # |X1| ** (T / m) * |X2| ** (m * (T / m) / n) * ... = + # |X1| ** (T / m) * |X2| ** (T / n) * | X3 x X4... | ** (m * n) + # And by doing induction we have product(|X_i| ** (T / dim(X_i))). + total = self.domain_dimension_tensor() + determinant = 1. + for operator in self.operators: + determinant *= operator.determinant() ** math_ops.cast( + total / operator.domain_dimension_tensor(), + dtype=operator.dtype) + return determinant + + def _log_abs_determinant(self): + # This will be sum((total / dim(x_i)) * log |X_i|) + total = self.domain_dimension_tensor() + log_abs_det = 0. + for operator in self.operators: + log_abs_det += operator.log_abs_determinant() * math_ops.cast( + total / operator.domain_dimension_tensor(), + dtype=operator.dtype) + return log_abs_det + + def _trace(self): + # tr(A x B) = tr(A) * tr(B) + trace = 1. + for operator in self.operators: + trace *= operator.trace() + return trace + + def _solve(self, rhs, adjoint=False, adjoint_arg=False): + # Here we follow the same use of Roth's column lemma as in `matmul`, with + # the key difference that we replace all `matmul` instances with `solve`. + # This follows from the property that inv(A x B) = inv(A) x inv(B). + + # Below we document the shape manipulation for adjoint=False, + # adjoint_arg=False, but the general case of different adjoints is still + # handled. + + if adjoint_arg: + rhs = linalg.adjoint(rhs) + + # Always add a batch dimension to enable broadcasting to work. + batch_shape = array_ops.concat( + [array_ops.ones_like(self.batch_shape_tensor()), [1, 1]], 0) + rhs += array_ops.zeros(batch_shape, dtype=rhs.dtype.base_dtype) + + # rhs has shape [B, R, C], where B represent some number of batch + # dimensions, + # R represents the number of rows, and C represents the number of columns. + # In order to apply Roth's column lemma, we need to operate on a batch of + # column vectors, so we reshape into a batch of column vectors. We put it + # at the front to ensure that broadcasting between operators to the batch + # dimensions B still works. + output = _rotate_last_dim(rhs, rotate_right=True) + + # Also expand the shape to be [A, C, B, R]. The first dimension will be + # used to accumulate dimensions from each operator matmul. + output = output[array_ops.newaxis, ...] + + # In this loop, A is going to refer to the value of the accumulated + # dimension. A = 1 at the start, and will end up being self.range_dimension. + # V will refer to the last dimension. V = R at the start, and will end up + # being 1 in the end. + for operator in self.operators[:-1]: + # Reshape output from [A, C, B, V] to be + # [A, C, B, V / op.domain_dimension, op.domain_dimension] + if adjoint: + operator_dimension = operator.range_dimension_tensor() + else: + operator_dimension = operator.domain_dimension_tensor() + + output = _unvec_by(output, operator_dimension) + + # We are computing (XA^-1^T) = (A^-1 X^T)^T. + # output has [A, C, B, V / op.domain_dimension, op.domain_dimension], + # which is being converted to: + # [A, C, B, V / op.domain_dimension, op.range_dimension] + output = array_ops.matrix_transpose(output) + output = operator.solve(output, adjoint=adjoint, adjoint_arg=False) + output = array_ops.matrix_transpose(output) + # Rearrange it to [A * op.range_dimension, C, B, V / op.domain_dimension] + output = _rotate_last_dim(output, rotate_right=False) + output = _vec(output) + output = _rotate_last_dim(output, rotate_right=True) + + # After the loop, we will have + # A = self.range_dimension / op[-1].range_dimension + # V = op[-1].domain_dimension + + # We convert that using matvec to get: + # [A, C, B, op[-1].range_dimension] + output = self.operators[-1].solvevec(output, adjoint=adjoint) + # Rearrange shape to be [B1, ... Bn, self.range_dimension, C] + output = _rotate_last_dim(output, rotate_right=False) + output = _vec(output) + output = _rotate_last_dim(output, rotate_right=False) + + if rhs.shape.is_fully_defined(): + column_dim = rhs.shape[-1] + broadcast_batch_shape = common_shapes.broadcast_shape( + rhs.shape[:-2], self.batch_shape) + if adjoint: + matrix_dimensions = [self.domain_dimension, column_dim] + else: + matrix_dimensions = [self.range_dimension, column_dim] + + output.set_shape(broadcast_batch_shape.concatenate( + matrix_dimensions)) + + return output + + def _diag_part(self): + diag_part = self.operators[0].diag_part() + for operator in self.operators[1:]: + diag_part = diag_part[..., :, array_ops.newaxis] + op_diag_part = operator.diag_part()[..., array_ops.newaxis, :] + diag_part *= op_diag_part + diag_part = array_ops.reshape( + diag_part, + shape=array_ops.concat( + [array_ops.shape(diag_part)[:-2], [-1]], axis=0)) + if self.range_dimension > self.domain_dimension: + diag_dimension = self.domain_dimension + else: + diag_dimension = self.range_dimension + diag_part.set_shape( + self.batch_shape.concatenate(diag_dimension)) + return diag_part + + def _to_dense(self): + product = self.operators[0].to_dense() + for operator in self.operators[1:]: + # Product has shape [B, R1, 1, C1]. + product = product[ + ..., :, array_ops.newaxis, :, array_ops.newaxis] + # Operator has shape [B, 1, R2, 1, C2]. + op_to_mul = operator.to_dense()[ + ..., array_ops.newaxis, :, array_ops.newaxis, :] + # This is now [B, R1, R2, C1, C2]. + product *= op_to_mul + # Now merge together dimensions to get [B, R1 * R2, C1 * C2]. + product = array_ops.reshape( + product, + shape=array_ops.concat( + [array_ops.shape(product)[:-4], + [array_ops.shape(product)[-4] * array_ops.shape(product)[-3], + array_ops.shape(product)[-2] * array_ops.shape(product)[-1]] + ], axis=0)) + product.set_shape(self.shape) + return product + + def _assert_non_singular(self): + if all(operator.is_square for operator in self.operators): + asserts = [operator.assert_non_singular() for operator in self.operators] + return control_flow_ops.group(asserts) + else: + raise errors.InvalidArgumentError( + node_def=None, op=None, message="All Kronecker factors must be " + "square for the product to be invertible.") + + def _assert_self_adjoint(self): + if all(operator.is_square for operator in self.operators): + asserts = [operator.assert_self_adjoint() for operator in self.operators] + return control_flow_ops.group(asserts) + else: + raise errors.InvalidArgumentError( + node_def=None, op=None, message="All Kronecker factors must be " + "square for the product to be self adjoint.") 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 ac50699f5984d04dc895205464cd6199dd4342f7..6e6c812adcbaf7e90d7c10c05fdfc0e150829329 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 @@ -105,11 +105,13 @@ def make_example_dict(example_protos, example_weights): def make_random_examples_and_variables_dicts(num_examples, dim, num_non_zero): random.seed(1) + sparse_features = [ SparseFeatureColumn( - [int(i / num_non_zero) for i in range(num_examples * num_non_zero)], - [int(random.random() * dim) for _ in range( - num_examples * num_non_zero)], + [i for i in range(num_examples) for _ in range(num_non_zero)], [ + i for _ in range(num_examples) + for i in random.sample(range(dim), num_non_zero) + ], [num_non_zero**(-0.5) for _ in range(num_examples * num_non_zero)]) ] examples_dict = dict( @@ -289,6 +291,34 @@ class SdcaWithLogisticLossTest(SdcaModelTest): # It would be 0.01 without shuffling and 0.02 with adaptive sampling. self.assertNear(0.0, lr.approximate_duality_gap().eval(), err=1e-3) + def testSparseDuplicate(self): + # Setup test data + example_protos = [ + make_example_proto({ + 'age': [0] * 5, + 'gender': [0] * 5 + }, 0), + make_example_proto({ + 'age': [1] * 5, + 'gender': [1] * 5 + }, 1), + ] + example_weights = [1.0, 1.0] + with self._single_threaded_test_session(): + examples = make_example_dict(example_protos, example_weights) + variables = make_variable_dict(1, 1) + options = dict( + symmetric_l2_regularization=1, + symmetric_l1_regularization=0, + loss_type='logistic_loss') + + lr = SdcaModel(examples, variables, options) + variables_lib.global_variables_initializer().run() + train_op = lr.minimize() + with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, + 'Duplicate'): + train_op.run() + def testDistributedSimple(self): # Setup test data example_protos = [ diff --git a/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable.py b/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable.py index ec726bbed41a86eb314e3591ecaedaa6bf0e5e9b..5015fb0848107950dd27eb81431dd308f22858bc 100644 --- a/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable.py +++ b/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable.py @@ -49,6 +49,7 @@ class ShardedMutableDenseHashTable(lookup.LookupInterface): default_value, empty_key, num_shards=1, + checkpoint=True, name='ShardedMutableHashTable'): with ops.name_scope(name, 'sharded_mutable_hash_table') as scope: super(ShardedMutableDenseHashTable, self).__init__(key_dtype, @@ -61,6 +62,7 @@ class ShardedMutableDenseHashTable(lookup.LookupInterface): value_dtype=value_dtype, default_value=default_value, empty_key=empty_key, + checkpoint=checkpoint, name='%s-%d-of-%d' % (name, i + 1, num_shards))) self._table_shards = table_shards # TODO(andreasst): add a value_shape() method to LookupInterface diff --git a/tensorflow/contrib/linear_optimizer/python/sdca_estimator_test.py b/tensorflow/contrib/linear_optimizer/python/sdca_estimator_test.py index 79a5928a21cb9a2633b2aac178f185ba333790d6..bed3d5139fcbf9d9e8b85605c752736f26af6793 100644 --- a/tensorflow/contrib/linear_optimizer/python/sdca_estimator_test.py +++ b/tensorflow/contrib/linear_optimizer/python/sdca_estimator_test.py @@ -30,6 +30,13 @@ from tensorflow.python.platform import test class SDCALogisticClassifierTest(test.TestCase): + def _single_threaded_test_session(self): + # TODO(andreasst): figure out why SDCALinearRegressor needs a single + # threaded session to pass in tsan mode but SDCALogisticClassifier does not. + config = config_pb2.ConfigProto( + inter_op_parallelism_threads=1, intra_op_parallelism_threads=1) + return self.test_session(config=config) + def testRealValuedFeatures(self): """Tests SDCALogisticClassifier works with real valued features.""" @@ -41,7 +48,7 @@ class SDCALogisticClassifierTest(test.TestCase): 'weights': constant_op.constant([[1.0], [1.0]]) }, constant_op.constant([[0], [1]]) - with self.test_session(): + with self._single_threaded_test_session(): maintenance_cost = feature_column_lib.real_valued_column( 'maintenance_cost') sq_footage = feature_column_lib.real_valued_column('sq_footage') @@ -66,7 +73,7 @@ class SDCALogisticClassifierTest(test.TestCase): constant_op.constant([[500.0, 800.0], [200.0, 600.0]]) }, constant_op.constant([[0], [1]]) - with self.test_session(): + with self._single_threaded_test_session(): dense_feature = feature_column_lib.real_valued_column( 'dense_feature', dimension=2) classifier = sdca_estimator.SDCALogisticClassifier( @@ -86,7 +93,7 @@ class SDCALogisticClassifierTest(test.TestCase): 'weights': constant_op.constant([[1.0], [1.0], [1.0]]) }, constant_op.constant([[1], [0], [1]]) - with self.test_session(): + with self._single_threaded_test_session(): price_bucket = feature_column_lib.bucketized_column( feature_column_lib.real_valued_column('price'), boundaries=[500.0, 700.0]) @@ -120,7 +127,7 @@ class SDCALogisticClassifierTest(test.TestCase): constant_op.constant([[1.0], [1.0], [1.0]]) }, constant_op.constant([[1], [0], [1]]) - with self.test_session(): + with self._single_threaded_test_session(): price = feature_column_lib.real_valued_column('price') country = feature_column_lib.sparse_column_with_hash_bucket( 'country', hash_bucket_size=5) @@ -151,7 +158,7 @@ class SDCALogisticClassifierTest(test.TestCase): dense_shape=[3, 5]) }, constant_op.constant([[1], [0], [1]]) - with self.test_session(): + with self._single_threaded_test_session(): country = feature_column_lib.sparse_column_with_hash_bucket( 'country', hash_bucket_size=5) country_weighted_by_price = feature_column_lib.weighted_sparse_column( @@ -163,6 +170,38 @@ class SDCALogisticClassifierTest(test.TestCase): metrics = classifier.evaluate(input_fn=input_fn, steps=1) self.assertGreater(metrics['accuracy'], 0.9) + def testSparseFeaturesWithDuplicates(self): + """Tests SDCALogisticClassifier with duplicated sparse features.""" + + def input_fn(): + return { + 'example_id': + constant_op.constant(['1', '2']), + 'age': + sparse_tensor.SparseTensor( + values=['20-29'] * 5 + ['31-40'] * 5, + indices=[[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [1, 0], + [1, 0], [1, 0], [1, 0], [1, 0]], + dense_shape=[2, 1]), + 'gender': + sparse_tensor.SparseTensor( + values=['m'] * 5 + ['f'] * 5, + indices=[[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [1, 0], + [1, 0], [1, 0], [1, 0], [1, 0]], + dense_shape=[2, 1]), + }, constant_op.constant([[1], [0]]) + + with self._single_threaded_test_session(): + age = feature_column_lib.sparse_column_with_hash_bucket( + 'age', hash_bucket_size=10) + gender = feature_column_lib.sparse_column_with_hash_bucket( + 'gender', hash_bucket_size=10) + classifier = sdca_estimator.SDCALogisticClassifier( + example_id_column='example_id', feature_columns=[age, gender]) + classifier.fit(input_fn=input_fn, steps=50) + metrics = classifier.evaluate(input_fn=input_fn, steps=1) + self.assertLess(metrics['loss'], 0.060) + def testCrossedFeatures(self): """Tests SDCALogisticClassifier with crossed features.""" @@ -182,7 +221,7 @@ class SDCALogisticClassifierTest(test.TestCase): dense_shape=[3, 1]) }, constant_op.constant([[0], [0], [1]]) - with self.test_session(): + with self._single_threaded_test_session(): language = feature_column_lib.sparse_column_with_hash_bucket( 'language', hash_bucket_size=5) country = feature_column_lib.sparse_column_with_hash_bucket( @@ -215,7 +254,7 @@ class SDCALogisticClassifierTest(test.TestCase): constant_op.constant([[3.0], [1.0], [1.0]]) }, constant_op.constant([[1], [0], [1]]) - with self.test_session(): + with self._single_threaded_test_session(): price = feature_column_lib.real_valued_column('price') sq_footage_bucket = feature_column_lib.bucketized_column( feature_column_lib.real_valued_column('sq_footage'), diff --git a/tensorflow/contrib/linear_optimizer/python/sdca_optimizer.py b/tensorflow/contrib/linear_optimizer/python/sdca_optimizer.py index dffdddacfb575636ea0b17797caf8502b8b8098c..5d4572bf6c761e0de2c9e6d7e17193abf0ebb170 100644 --- a/tensorflow/contrib/linear_optimizer/python/sdca_optimizer.py +++ b/tensorflow/contrib/linear_optimizer/python/sdca_optimizer.py @@ -19,6 +19,7 @@ from __future__ import print_function from tensorflow.contrib import layers from tensorflow.contrib.linear_optimizer.python.ops import sdca_ops from tensorflow.contrib.linear_optimizer.python.ops.sparse_feature_column import SparseFeatureColumn +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 @@ -181,28 +182,42 @@ class SDCAOptimizer(object): elif isinstance( column, ( + layers.feature_column._WeightedSparseColumn, # pylint: disable=protected-access layers.feature_column._CrossedColumn, # pylint: disable=protected-access layers.feature_column._SparseColumn)): # pylint: disable=protected-access - sparse_features.append( - SparseFeatureColumn( - array_ops.reshape( - array_ops.split( - value=transformed_tensor.indices, - num_or_size_splits=2, - axis=1)[0], [-1]), - array_ops.reshape(transformed_tensor.values, [-1]), None)) - sparse_feature_weights.append(columns_to_variables[column][0]) - elif isinstance(column, layers.feature_column._WeightedSparseColumn): # pylint: disable=protected-access - id_tensor = column.id_tensor(transformed_tensor) - weight_tensor = column.weight_tensor(transformed_tensor) + + if isinstance(column, layers.feature_column._WeightedSparseColumn): # pylint: disable=protected-access + id_tensor = column.id_tensor(transformed_tensor) + weight_tensor = array_ops.reshape( + column.weight_tensor(transformed_tensor).values, [-1]) + else: + id_tensor = transformed_tensor + weight_tensor = array_ops.ones( + [array_ops.shape(id_tensor.indices)[0]], dtypes.float32) + + example_ids = array_ops.reshape(id_tensor.indices[:, 0], [-1]) + + flat_ids = array_ops.reshape(id_tensor.values, [-1]) + projection_length = math_ops.reduce_max(flat_ids) + 1 + # project ids based on example ids so that we can dedup ids that + # occur multiple times for a single example. + projected_ids = projection_length * example_ids + flat_ids + + # Remove any redudant ids. + ids, idx = array_ops.unique(projected_ids) + # Keep only one example id per duplicated ids. + example_ids_filtered = math_ops.unsorted_segment_min( + example_ids, idx, + array_ops.shape(ids)[0]) + + # reproject ids back feature id space. + reproject_ids = (ids - projection_length * example_ids_filtered) + + weights = array_ops.reshape( + math_ops.unsorted_segment_sum(weight_tensor, idx, + array_ops.shape(ids)[0]), [-1]) sparse_feature_with_values.append( - SparseFeatureColumn( - array_ops.reshape( - array_ops.split( - value=id_tensor.indices, num_or_size_splits=2, axis=1) - [0], [-1]), - array_ops.reshape(id_tensor.values, [-1]), - array_ops.reshape(weight_tensor.values, [-1]))) + SparseFeatureColumn(example_ids_filtered, reproject_ids, weights)) sparse_feature_with_values_weights.append( columns_to_variables[column][0]) else: diff --git a/tensorflow/contrib/lite/BUILD b/tensorflow/contrib/lite/BUILD index 9c4533079c72f5ed68c6f45582fb1cecaa3a3679..1534f97d7600151e78c7fa7e8509d9e871240421 100644 --- a/tensorflow/contrib/lite/BUILD +++ b/tensorflow/contrib/lite/BUILD @@ -137,6 +137,7 @@ cc_library( "//tensorflow/contrib/lite/kernels:eigen_support", "//tensorflow/contrib/lite/kernels:gemm_support", "//tensorflow/contrib/lite/nnapi:nnapi_lib", + "//tensorflow/contrib/lite/profiling:profiler", "//tensorflow/contrib/lite/schema:schema_fbs", ], ) diff --git a/tensorflow/contrib/lite/Makefile b/tensorflow/contrib/lite/Makefile index b4504f246a0f806d35d8c3d659717a86d2f2a4f5..65fba52d461461f4594e2222ef6df3849b741f99 100644 --- a/tensorflow/contrib/lite/Makefile +++ b/tensorflow/contrib/lite/Makefile @@ -90,7 +90,8 @@ $(wildcard tensorflow/contrib/lite/kernels/*.c) \ $(wildcard tensorflow/contrib/lite/kernels/internal/*.c) \ $(wildcard tensorflow/contrib/lite/kernels/internal/optimized/*.c) \ $(wildcard tensorflow/contrib/lite/kernels/internal/reference/*.c) \ -$(wildcard tensorflow/contrib/lite/downloads/farmhash/src/farmhash.cc) +$(wildcard tensorflow/contrib/lite/downloads/farmhash/src/farmhash.cc) \ +$(wildcard tensorflow/contrib/lite/downloads/fft2d/fftsg.c) # Remove any duplicates. CORE_CC_ALL_SRCS := $(sort $(CORE_CC_ALL_SRCS)) CORE_CC_EXCLUDE_SRCS := \ diff --git a/tensorflow/contrib/lite/allocation.cc b/tensorflow/contrib/lite/allocation.cc index 4b322e027d48f4bf9f90d5b873c449d1ec31cc49..a4772731ecda92431c412672610a39c188dabf27 100644 --- a/tensorflow/contrib/lite/allocation.cc +++ b/tensorflow/contrib/lite/allocation.cc @@ -22,6 +22,7 @@ limitations under the License. #include #include #include +#include #include "tensorflow/contrib/lite/allocation.h" #include "tensorflow/contrib/lite/context.h" diff --git a/tensorflow/contrib/lite/arena_planner.cc b/tensorflow/contrib/lite/arena_planner.cc index 8e47e2375e2e306c345a2b6caa2411abd9b3ceb0..4f836d367747e06de682b5764206d33f6e2fb983 100644 --- a/tensorflow/contrib/lite/arena_planner.cc +++ b/tensorflow/contrib/lite/arena_planner.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ #include "tensorflow/contrib/lite/arena_planner.h" +#include namespace tflite { diff --git a/tensorflow/contrib/lite/build_def.bzl b/tensorflow/contrib/lite/build_def.bzl index b8f6b7fd59af9834edb4aa7aefa524c25ede66d2..85216776823eab2ab3ac2a3bc666f21e312acc6c 100644 --- a/tensorflow/contrib/lite/build_def.bzl +++ b/tensorflow/contrib/lite/build_def.bzl @@ -124,19 +124,19 @@ def tf_to_tflite(name, src, options, out): out: name of the output flatbuffer file. """ - toco = "//tensorflow/contrib/lite/toco:toco" + toco_cmdline = " ".join([ + "//tensorflow/contrib/lite/toco:toco", + "--input_format=TENSORFLOW_GRAPHDEF", + "--output_format=TFLITE", + ("--input_file=$(location %s)" % src), + ("--output_file=$(location %s)" % out), + ] + options ) native.genrule( name = name, - srcs=[src, options], + srcs=[src], outs=[out], - cmd = ("$(location %s) " + - " --input_file=$(location %s) " + - " --output_file=$(location %s) " + - " --input_format=TENSORFLOW_GRAPHDEF" + - " --output_format=TFLITE" + - " `cat $(location %s)`") - % (toco, src, out, options), - tools= [toco], + cmd = toco_cmdline, + tools= ["//tensorflow/contrib/lite/toco:toco"], ) def tflite_to_json(name, src, out): diff --git a/tensorflow/contrib/lite/builtin_op_data.h b/tensorflow/contrib/lite/builtin_op_data.h index 2b6c24768c0f35b91d0dabf8a5723e73f040cc3b..4910c89eaebabb7bd9a4e003b75fa6de4d5af69d 100644 --- a/tensorflow/contrib/lite/builtin_op_data.h +++ b/tensorflow/contrib/lite/builtin_op_data.h @@ -53,6 +53,8 @@ typedef struct { TfLitePadding padding; int stride_width; int stride_height; + int dilation_width_factor; + int dilation_height_factor; TfLiteFusedActivation activation; } TfLiteConvParams; @@ -221,6 +223,10 @@ typedef struct { int shrink_axis_mask; } TfLiteStridedSliceParams; +typedef struct { + TfLiteType output_type; +} TfLiteArgMaxParams; + #ifdef __cplusplus } // extern "C" #endif // __cplusplus diff --git a/tensorflow/contrib/lite/builtin_ops.h b/tensorflow/contrib/lite/builtin_ops.h index 17b791e4e2f38d9a1108d35d1298445a1c370727..859bc7ab70dc363e08800ca5c40eb0da6ca426b0 100644 --- a/tensorflow/contrib/lite/builtin_ops.h +++ b/tensorflow/contrib/lite/builtin_ops.h @@ -80,6 +80,9 @@ typedef enum { kTfLiteBuiltinCast = 53, kTfLiteBuiltinPrelu = 54, kTfLiteBuiltinMaximum = 55, + kTfLiteBuiltinArgMax = 56, + kTfLiteBuiltinMinimum = 57, + kTfLiteBuiltinLess = 58, } TfLiteBuiltinOperator; #ifdef __cplusplus diff --git a/tensorflow/contrib/lite/context.h b/tensorflow/contrib/lite/context.h index 45184b05ecefb504c75815ae900f3b605359a443..0b38f43cd32fbdfa0296eec7ef81aab76ebe5461 100644 --- a/tensorflow/contrib/lite/context.h +++ b/tensorflow/contrib/lite/context.h @@ -137,6 +137,7 @@ typedef enum { kTfLiteUInt8 = 3, kTfLiteInt64 = 4, kTfLiteString = 5, + kTfLiteBool = 6, } TfLiteType; // Parameters for asymmetric quantization. Quantized values can be converted @@ -155,6 +156,7 @@ typedef union { char* raw; const char* raw_const; uint8_t* uint8; + bool* b; } TfLitePtrUnion; // Memory allocation strategies. kTfLiteMmapRo is for read-only memory-mapped diff --git a/tensorflow/contrib/lite/download_dependencies.sh b/tensorflow/contrib/lite/download_dependencies.sh index a93ed201d647ddf2359a57254a959871c13fb94f..436c3e1d4cad5e6ee355d7e9cf8ee7da1a8385ce 100755 --- a/tensorflow/contrib/lite/download_dependencies.sh +++ b/tensorflow/contrib/lite/download_dependencies.sh @@ -30,12 +30,15 @@ if [ ! -f $BZL_FILE_PATH ]; then fi EIGEN_URL="$(grep -o 'http.*bitbucket.org/eigen/eigen/get/.*tar\.gz' "${BZL_FILE_PATH}" | grep -v mirror.bazel | head -n1)" -GEMMLOWP_URL="$(grep -o 'https://mirror.bazel.build/github.com/google/gemmlowp/.*zip' "${BZL_FILE_PATH}" | head -n1)" +# TODO (yongtang): Replace the following with 'https://mirror.bazel.build/github.com/google/gemmlowp/.*zip' once +# the archive has been propagated in mirror.bazel.build. +GEMMLOWP_URL="$(grep -o 'https://github.com/google/gemmlowp/.*zip' "${BZL_FILE_PATH}" | head -n1)" GOOGLETEST_URL="https://github.com/google/googletest/archive/release-1.8.0.tar.gz" ABSL_URL="$(grep -o 'https://github.com/abseil/abseil-cpp/.*tar.gz' "${BZL_FILE_PATH}" | head -n1)" NEON_2_SSE_URL="https://github.com/intel/ARM_NEON_2_x86_SSE/archive/master.zip" FARMHASH_URL="https://mirror.bazel.build/github.com/google/farmhash/archive/816a4ae622e964763ca0862d9dbd19324a1eaf45.tar.gz" FLATBUFFERS_URL="https://github.com/google/flatbuffers/archive/master.zip" +FFT2D_URL="https://mirror.bazel.build/www.kurims.kyoto-u.ac.jp/~ooura/fft.tgz" # TODO(petewarden): Some new code in Eigen triggers a clang bug with iOS arm64, # so work around it by patching the source. @@ -91,6 +94,7 @@ 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 "${FFT2D_URL}" "${DOWNLOADS_DIR}/fft2d" replace_by_sed 's#static uint32x4_t p4ui_CONJ_XOR = vld1q_u32( conj_XOR_DATA );#static uint32x4_t p4ui_CONJ_XOR; // = vld1q_u32( conj_XOR_DATA ); - Removed by script#' \ "${DOWNLOADS_DIR}/eigen/Eigen/src/Core/arch/NEON/Complex.h" diff --git a/tensorflow/contrib/lite/examples/ios/camera/tflite_camera_example.xcodeproj/project.pbxproj b/tensorflow/contrib/lite/examples/ios/camera/tflite_camera_example.xcodeproj/project.pbxproj index b0236e9c608ec35437bcfe79c51149a76f9f416e..98d3b5bb8ad45bf34f6996b3361291896a451a6f 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 @@ -326,10 +326,6 @@ GCC_WARN_UNUSED_VARIABLE = YES; HEADER_SEARCH_PATHS = ( "$(inherited)", - ../../../../../../, - ../../../downloads/flatbuffers/include/, - ../../../downloads/eigen/, - ../../../downloads/, ); IPHONEOS_DEPLOYMENT_TARGET = 8.0; MTL_ENABLE_DEBUG_INFO = YES; @@ -373,10 +369,6 @@ GCC_WARN_UNUSED_VARIABLE = YES; HEADER_SEARCH_PATHS = ( "$(inherited)", - ../../../../../../, - ../../../downloads/flatbuffers/include/, - ../../../downloads/eigen/, - ../../../downloads/, ); IPHONEOS_DEPLOYMENT_TARGET = 8.0; MTL_ENABLE_DEBUG_INFO = NO; diff --git a/tensorflow/contrib/lite/g3doc/apis.md b/tensorflow/contrib/lite/g3doc/apis.md index fe208e47d1ac10995881e55c8596ae14ff4242df..50cc146a87ee9ab94aea6a92fb2fb5c531f83369 100644 --- a/tensorflow/contrib/lite/g3doc/apis.md +++ b/tensorflow/contrib/lite/g3doc/apis.md @@ -29,7 +29,7 @@ interpreter->AllocateTensors(); float* input = interpreter->typed_input_tensor(0); // Fill `input`. interpreter->Invoke(); -float* output = interpreter->type_output_tensor(0); +float* output = interpreter->typed_output_tensor(0); ``` ### Data Alignment diff --git a/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md b/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md index 61ea5231e352f5e014f9200eccae69548574c034..203924f03d3101130049b9679328fac1e2da02bd 100644 --- a/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md +++ b/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md @@ -302,6 +302,19 @@ Options { } ``` +**LESS** + +``` +Inputs { + 0: a tensor + 1: a tensor +} +Outputs { + 0: a tensor of type bool, true whenever an element of the first tensor is less + than the corresponding element of the second tensor. +} +``` + **LOCAL_RESPONSE_NORMALIZATION** ``` diff --git a/tensorflow/contrib/lite/interpreter.cc b/tensorflow/contrib/lite/interpreter.cc index 4575fe884dc07963df5f0a26c5fe6680d92e409c..91b6c414bf036fbf57f53fc75f570b05449fa89e 100644 --- a/tensorflow/contrib/lite/interpreter.cc +++ b/tensorflow/contrib/lite/interpreter.cc @@ -14,10 +14,12 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/contrib/lite/interpreter.h" + #include #include #include #include + #include "tensorflow/contrib/lite/arena_planner.h" #include "tensorflow/contrib/lite/context.h" #include "tensorflow/contrib/lite/error_reporter.h" @@ -26,6 +28,7 @@ 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/profiling/profiler.h" #include "tensorflow/contrib/lite/schema/schema_generated.h" #include "tensorflow/contrib/lite/util.h" @@ -245,11 +248,8 @@ TfLiteStatus Interpreter::ReplaceSubgraphsWithDelegateKernels( // 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. + TF_LITE_ENSURE(&context_, tensor->delegate == nullptr || + tensor->delegate == delegate); tensor->delegate = delegate; } @@ -337,9 +337,13 @@ TfLiteStatus Interpreter::BytesRequired(TfLiteType type, const int* dims, case kTfLiteInt64: *bytes = sizeof(int64_t) * count; break; + case kTfLiteBool: + *bytes = sizeof(bool) * count; + break; default: - ReportError(&context_, - "Only float32, int32, int64, uint8 supported currently."); + ReportError( + &context_, + "Only float32, int32, int64, uint8, bool supported currently."); return kTfLiteError; } return kTfLiteOk; @@ -543,6 +547,7 @@ TfLiteStatus Interpreter::Invoke() { TfLiteNode& node = nodes_and_registration_[node_index].first; const TfLiteRegistration& registration = nodes_and_registration_[node_index].second; + SCOPED_OPERATOR_PROFILE(profiler_, node_index); // 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 @@ -566,6 +571,12 @@ TfLiteStatus Interpreter::Invoke() { } } + if (!allow_buffer_handle_output_) { + for (int tensor_index : outputs_) { + EnsureTensorDataIsReadable(tensor_index); + } + } + return status; } diff --git a/tensorflow/contrib/lite/interpreter.h b/tensorflow/contrib/lite/interpreter.h index a6d582a813e0dcaacb2f196cc71dc111638171b2..a49134b95ee47b97ad1e56d07f737fa44f89badc 100644 --- a/tensorflow/contrib/lite/interpreter.h +++ b/tensorflow/contrib/lite/interpreter.h @@ -20,10 +20,12 @@ limitations under the License. #include #include #include + #include "tensorflow/contrib/lite/allocation.h" #include "tensorflow/contrib/lite/context.h" #include "tensorflow/contrib/lite/error_reporter.h" #include "tensorflow/contrib/lite/memory_planner.h" +#include "tensorflow/contrib/lite/profiling/profiler.h" namespace tflite { @@ -48,6 +50,10 @@ template <> constexpr TfLiteType typeToTfLiteType() { return kTfLiteUInt8; } +template <> +constexpr TfLiteType typeToTfLiteType() { + return kTfLiteBool; +} // Forward declare since NNAPIDelegate uses Interpreter. class NNAPIDelegate; @@ -278,6 +284,7 @@ class Interpreter { // 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. + // WARNING: This is an experimental API and subject to change. TfLiteStatus EnsureTensorDataIsReadable(int tensor_index) { TF_LITE_ENSURE(&context_, tensor_index < tensors_size()); TfLiteTensor* tensor = &tensors_[tensor_index]; @@ -316,6 +323,12 @@ class Interpreter { TfLiteBufferHandle* buffer_handle, TfLiteDelegate** delegate); + void SetProfiler(profiling::Profiler* profiler) { profiler_ = profiler; } + + profiling::Profiler* GetProfiler(profiling::Profiler* profiler) { + return profiler_; + } + // The default capacity of `tensors_` vector. static constexpr int kTensorsReservedCapacity = 128; // The capacity headroom of `tensors_` vector before calling ops' @@ -324,6 +337,18 @@ class Interpreter { // pointers to existing tensors. static constexpr int kTensorsCapacityHeadroom = 16; + // Set if buffer handle output is allowed. + // + // When using hardware delegation, Interpreter will make the data of output + // tensors available in `tensor->data` by default. If the application can + // consume the buffer handle directly (e.g. reading output from OpenGL + // texture), it can set this flag to false, so Interpreter won't copy the data + // from buffer handle to CPU memory. + // WARNING: This is an experimental API and subject to change. + void SetAllowBufferHandleOutput(bool allow_buffer_handle_output) { + allow_buffer_handle_output_ = allow_buffer_handle_output; + } + private: // Give 'op_reg' a chance to initialize itself using the contents of // 'buffer'. @@ -514,6 +539,11 @@ class Interpreter { std::unique_ptr nnapi_delegate_; std::unique_ptr memory_planner_; + + bool allow_buffer_handle_output_ = false; + + // Profiler for this interpreter instance. + profiling::Profiler* profiler_; }; } // namespace tflite 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 300786c3ca01b12a46f7f9a6fe8fd720f97a79f4..18f64651889d7eeb4be961afc47554cbcc51a410 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 @@ -54,6 +54,9 @@ import android.view.Surface; import android.view.TextureView; import android.view.View; import android.view.ViewGroup; +import android.widget.CompoundButton; +import android.widget.NumberPicker; +import android.widget.ToggleButton; import android.widget.TextView; import android.widget.Toast; import java.io.IOException; @@ -82,6 +85,8 @@ public class Camera2BasicFragment extends Fragment private boolean runClassifier = false; private boolean checkedPermissions = false; private TextView textView; + private ToggleButton toggle; + private NumberPicker np; private ImageClassifier classifier; /** Max preview width that is guaranteed by Camera2 API */ @@ -289,6 +294,24 @@ public class Camera2BasicFragment extends Fragment public void onViewCreated(final View view, Bundle savedInstanceState) { textureView = (AutoFitTextureView) view.findViewById(R.id.texture); textView = (TextView) view.findViewById(R.id.text); + toggle = (ToggleButton) view.findViewById(R.id.button); + + toggle.setOnCheckedChangeListener(new CompoundButton.OnCheckedChangeListener() { + public void onCheckedChanged(CompoundButton buttonView, boolean isChecked) { + classifier.setUseNNAPI(isChecked); + } + }); + + np = (NumberPicker) view.findViewById(R.id.np); + np.setMinValue(1); + np.setMaxValue(10); + np.setWrapSelectorWheel(true); + np.setOnValueChangedListener(new NumberPicker.OnValueChangeListener() { + @Override + public void onValueChange(NumberPicker picker, int oldVal, int newVal){ + classifier.setNumThreads(newVal); + } + }); } /** Load the model and labels. */ 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 c57bb348c5b386a59327c7b1bc769717ca755269..d32c0779101cf8e795ee9d7e970401c2c03bb33a 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 @@ -142,6 +142,16 @@ public abstract class ImageClassifier { } } + public void setUseNNAPI(Boolean nnapi) { + if (tflite != null) + tflite.setUseNNAPI(nnapi); + } + + public void setNumThreads(int num_threads) { + if (tflite != null) + tflite.setNumThreads(num_threads); + } + /** Closes tflite to release resources. */ public void close() { tflite.close(); diff --git a/tensorflow/contrib/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml b/tensorflow/contrib/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml index 15305c436e0d997af15a326ab4027ea713ed8098..db557ad62f619e88f72426a48a74bffb0f57b818 100644 --- a/tensorflow/contrib/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml +++ b/tensorflow/contrib/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml @@ -22,24 +22,59 @@ android:layout_width="wrap_content" android:layout_height="wrap_content" android:layout_alignParentStart="true" + android:layout_alignParentLeft="true" android:layout_alignParentTop="true" /> - + + + + + + + diff --git a/tensorflow/contrib/lite/java/demo/app/src/main/res/values/strings.xml b/tensorflow/contrib/lite/java/demo/app/src/main/res/values/strings.xml index a08ec3eb629250a727cec49a822375fe5569f455..29a033bcd437c951ef6e8ba78f4fc3a0fcafac96 100644 --- a/tensorflow/contrib/lite/java/demo/app/src/main/res/values/strings.xml +++ b/tensorflow/contrib/lite/java/demo/app/src/main/res/values/strings.xml @@ -21,4 +21,6 @@ NN:On NN:Off Use NNAPI + tflite + NNAPI 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 fc16488a6459eb227fde712055d3e8ccfcce0070..75334cd96e8daadc356dadea063eee30ef6d5245 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 @@ -51,7 +51,11 @@ enum DataType { } } throw new IllegalArgumentException( - "DataType " + c + " is not recognized in Java (version " + TensorFlowLite.version() + ")"); + "DataType error: DataType " + + c + + " is not recognized in Java (version " + + TensorFlowLite.version() + + ")"); } /** Returns byte size of the type. */ @@ -68,7 +72,8 @@ enum DataType { case BYTEBUFFER: return 1; } - throw new IllegalArgumentException("DataType " + this + " is not supported yet"); + throw new IllegalArgumentException( + "DataType error: DataType " + this + " is not supported yet"); } /** Gets string names of the data type. */ @@ -85,7 +90,8 @@ enum DataType { case BYTEBUFFER: return "ByteBuffer"; } - throw new IllegalArgumentException("DataType " + this + " is not supported yet"); + throw new IllegalArgumentException( + "DataType error: DataType " + this + " is not supported yet"); } // Cached to avoid copying it 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 a33959dca4954e3c2aaed987839bdec1ba079b5e..e84ee7112983ec584308b7cbcd919f119eccbcc9 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 @@ -137,17 +137,19 @@ public final class Interpreter implements AutoCloseable { public void runForMultipleInputsOutputs( @NonNull Object[] inputs, @NonNull Map outputs) { if (wrapper == null) { - throw new IllegalStateException("The Interpreter has already been closed."); + throw new IllegalStateException("Internal error: The Interpreter has already been closed."); } Tensor[] tensors = wrapper.run(inputs); if (outputs == null || tensors == null || outputs.size() > tensors.length) { - throw new IllegalArgumentException("Outputs do not match with model outputs."); + throw new IllegalArgumentException("Output error: Outputs do not match with model outputs."); } final int size = tensors.length; for (Integer idx : outputs.keySet()) { if (idx == null || idx < 0 || idx >= size) { throw new IllegalArgumentException( - String.format("Invalid index of output %d (should be in range [0, %d))", idx, size)); + String.format( + "Output error: Invalid index of output %d (should be in range [0, %d))", + idx, size)); } tensors[idx].copyTo(outputs.get(idx)); } @@ -160,7 +162,7 @@ public final class Interpreter implements AutoCloseable { */ public void resizeInput(int idx, @NonNull int[] dims) { if (wrapper == null) { - throw new IllegalStateException("The Interpreter has already been closed."); + throw new IllegalStateException("Internal error: The Interpreter has already been closed."); } wrapper.resizeInput(idx, dims); } @@ -173,7 +175,7 @@ public final class Interpreter implements AutoCloseable { */ public int getInputIndex(String opName) { if (wrapper == null) { - throw new IllegalStateException("The Interpreter has already been closed."); + throw new IllegalStateException("Internal error: The Interpreter has already been closed."); } return wrapper.getInputIndex(opName); } @@ -186,7 +188,7 @@ public final class Interpreter implements AutoCloseable { */ public int getOutputIndex(String opName) { if (wrapper == null) { - throw new IllegalStateException("The Interpreter has already been closed."); + throw new IllegalStateException("Internal error: The Interpreter has already been closed."); } return wrapper.getOutputIndex(opName); } @@ -198,7 +200,7 @@ public final class Interpreter implements AutoCloseable { */ public Long getLastNativeInferenceDurationNanoseconds() { if (wrapper == null) { - throw new IllegalStateException("The interpreter has already been closed."); + throw new IllegalStateException("Internal error: The interpreter has already been closed."); } return wrapper.getLastNativeInferenceDurationNanoseconds(); } @@ -208,8 +210,16 @@ public final class Interpreter implements AutoCloseable { if (wrapper != null) { wrapper.setUseNNAPI(useNNAPI); } else { - throw new IllegalStateException("NativeInterpreterWrapper has already been closed."); + throw new IllegalStateException( + "Internal error: NativeInterpreterWrapper has already been closed."); + } + } + + public void setNumThreads(int num_threads) { + if (wrapper == null) { + throw new IllegalStateException("The interpreter has already been closed."); } + wrapper.setNumThreads(num_threads); } /** Release resources associated with the {@code Interpreter}. */ 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 fc8187acfebf272a72ceb7844333bd589359cc2e..2fc803715be5e5afbc5d548d46c1665785f6055d 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 @@ -80,7 +80,7 @@ final class NativeInterpreterWrapper implements AutoCloseable { /** Sets inputs, runs model inference and returns outputs. */ Tensor[] run(Object[] inputs) { if (inputs == null || inputs.length == 0) { - throw new IllegalArgumentException("Invalid inputs. Inputs should not be null or empty."); + throw new IllegalArgumentException("Input error: Inputs should not be null or empty."); } int[] dataTypes = new int[inputs.length]; Object[] sizes = new Object[inputs.length]; @@ -92,7 +92,7 @@ final class NativeInterpreterWrapper implements AutoCloseable { ByteBuffer buffer = (ByteBuffer) inputs[i]; if (buffer.order() != ByteOrder.nativeOrder()) { throw new IllegalArgumentException( - "Invalid ByteBuffer. It shoud use ByteOrder.nativeOrder()."); + "Input error: ByteBuffer shoud use ByteOrder.nativeOrder()."); } numsOfBytes[i] = buffer.limit(); sizes[i] = getInputDims(interpreterHandle, i, numsOfBytes[i]); @@ -103,7 +103,7 @@ final class NativeInterpreterWrapper implements AutoCloseable { } else { throw new IllegalArgumentException( String.format( - "%d-th element of the %d inputs is not an array or a ByteBuffer.", + "Input error: %d-th element of the %d inputs is not an array or a ByteBuffer.", i, inputs.length)); } } @@ -119,7 +119,7 @@ final class NativeInterpreterWrapper implements AutoCloseable { this, isMemoryAllocated); if (outputsHandles == null || outputsHandles.length == 0) { - throw new IllegalStateException("Interpreter has no outputs."); + throw new IllegalStateException("Internal error: Interpreter has no outputs."); } isMemoryAllocated = true; Tensor[] outputs = new Tensor[outputsHandles.length]; @@ -153,6 +153,10 @@ final class NativeInterpreterWrapper implements AutoCloseable { useNNAPI(interpreterHandle, useNNAPI); } + void setNumThreads(int num_threads) { + numThreads(interpreterHandle, num_threads); + } + /** Gets index of an input given its name. */ int getInputIndex(String name) { if (inputsIndexes == null) { @@ -169,7 +173,8 @@ final class NativeInterpreterWrapper implements AutoCloseable { } else { throw new IllegalArgumentException( String.format( - "%s is not a valid name for any input. The indexes of the inputs are %s", + "Input error: %s is not a valid name for any input. " + + "The indexes of the inputs are %s", name, inputsIndexes.toString())); } } @@ -190,7 +195,8 @@ final class NativeInterpreterWrapper implements AutoCloseable { } else { throw new IllegalArgumentException( String.format( - "%s is not a valid name for any output. The indexes of the outputs are %s", + "Input error: %s is not a valid name for any output. " + + "The indexes of the outputs are %s", name, outputsIndexes.toString())); } } @@ -229,7 +235,8 @@ final class NativeInterpreterWrapper implements AutoCloseable { return DataType.BYTEBUFFER; } } - throw new IllegalArgumentException("cannot resolve DataType of " + o.getClass().getName()); + throw new IllegalArgumentException( + "DataType error: cannot resolve DataType of " + o.getClass().getName()); } /** Returns the shape of an object as an int array. */ @@ -245,7 +252,7 @@ final class NativeInterpreterWrapper implements AutoCloseable { return 0; } if (Array.getLength(o) == 0) { - throw new IllegalArgumentException("array lengths cannot be 0."); + throw new IllegalArgumentException("Array lengths cannot be 0."); } return 1 + numDimensions(Array.get(o, 0)); } @@ -259,7 +266,7 @@ final class NativeInterpreterWrapper implements AutoCloseable { shape[dim] = len; } else if (shape[dim] != len) { throw new IllegalArgumentException( - String.format("mismatched lengths (%d and %d) in dimension %d", shape[dim], len, dim)); + String.format("Mismatched lengths (%d and %d) in dimension %d", shape[dim], len, dim)); } for (int i = 0; i < len; ++i) { fillShape(Array.get(o, i), dim + 1, shape); @@ -321,6 +328,8 @@ final class NativeInterpreterWrapper implements AutoCloseable { private static native void useNNAPI(long interpreterHandle, boolean state); + private static native void numThreads(long interpreterHandle, int num_threads); + private static native long createErrorReporter(int size); private static native long createModel(String modelPathOrBuffer, long errorHandle); diff --git a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Tensor.java b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Tensor.java index 54ace6c63ce5bd1b38be744176d0378e3cc8a1d3..09e887aae3339e9f114c07d689c0d7b5e2fc384b 100644 --- a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Tensor.java +++ b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Tensor.java @@ -34,15 +34,16 @@ final class Tensor { if (NativeInterpreterWrapper.dataTypeOf(dst) != dtype) { throw new IllegalArgumentException( String.format( - "Cannot convert an TensorFlowLite tensor with type %s to a Java object of " - + "type %s (which is compatible with the TensorFlowLite type %s)", + "Output error: Cannot convert an TensorFlowLite tensor with type %s to a Java " + + "object of type %s (which is compatible with the TensorFlowLite type %s)", dtype, dst.getClass().getName(), NativeInterpreterWrapper.dataTypeOf(dst))); } int[] dstShape = NativeInterpreterWrapper.shapeOf(dst); if (!Arrays.equals(dstShape, shapeCopy)) { throw new IllegalArgumentException( String.format( - "Shape of output target %s does not match with the shape of the Tensor %s.", + "Output error: Shape of output target %s does not match with the shape of the " + + "Tensor %s.", Arrays.toString(dstShape), Arrays.toString(shapeCopy))); } readMultiDimensionalArray(nativeHandle, dst); diff --git a/tensorflow/contrib/lite/java/src/main/native/exception_jni.cc b/tensorflow/contrib/lite/java/src/main/native/exception_jni.cc index 1578c9e3ddd034ad9ce17c8c3ae6c942258e2a55..34d91be04cd6c855a2068510ca810c0b93637584 100644 --- a/tensorflow/contrib/lite/java/src/main/native/exception_jni.cc +++ b/tensorflow/contrib/lite/java/src/main/native/exception_jni.cc @@ -44,7 +44,8 @@ BufferErrorReporter::BufferErrorReporter(JNIEnv* env, int limit) { buffer_ = new char[limit]; if (!buffer_) { throwException(env, kNullPointerException, - "Malloc of BufferErrorReporter to hold %d char failed.", + "Internal error: Malloc of BufferErrorReporter to hold %d " + "char failed.", limit); return; } 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 844226203bb02f4017b2f04da34ac81ac2b7a191..45f510da1d940a288e2794cb3e08f66451956b64 100644 --- a/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc +++ b/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc @@ -22,7 +22,7 @@ const int kBufferSize = 256; tflite::Interpreter* convertLongToInterpreter(JNIEnv* env, jlong handle) { if (handle == 0) { throwException(env, kIllegalArgumentException, - "Invalid handle to Interpreter."); + "Internal error: Invalid handle to Interpreter."); return nullptr; } return reinterpret_cast(handle); @@ -30,7 +30,8 @@ tflite::Interpreter* convertLongToInterpreter(JNIEnv* env, jlong handle) { tflite::FlatBufferModel* convertLongToModel(JNIEnv* env, jlong handle) { if (handle == 0) { - throwException(env, kIllegalArgumentException, "Invalid handle to model."); + throwException(env, kIllegalArgumentException, + "Internal error: Invalid handle to model."); return nullptr; } return reinterpret_cast(handle); @@ -39,7 +40,7 @@ tflite::FlatBufferModel* convertLongToModel(JNIEnv* env, jlong handle) { BufferErrorReporter* convertLongToErrorReporter(JNIEnv* env, jlong handle) { if (handle == 0) { throwException(env, kIllegalArgumentException, - "Invalid handle to ErrorReporter."); + "Internal error: Invalid handle to ErrorReporter."); return nullptr; } return reinterpret_cast(handle); @@ -51,7 +52,7 @@ std::vector convertJIntArrayToVector(JNIEnv* env, jintArray inputs) { jint* ptr = env->GetIntArrayElements(inputs, nullptr); if (ptr == nullptr) { throwException(env, kIllegalArgumentException, - "Empty dimensions of input array."); + "Array has empty dimensions."); return {}; } for (int i = 0; i < size; ++i) { @@ -113,7 +114,7 @@ TfLiteStatus checkInputs(JNIEnv* env, tflite::Interpreter* interpreter, jobjectArray sizes) { if (input_size != interpreter->inputs().size()) { throwException(env, kIllegalArgumentException, - "Expected num of inputs is %d but got %d", + "Input error: Expected num of inputs is %d but got %d", interpreter->inputs().size(), input_size); return kTfLiteError; } @@ -121,8 +122,9 @@ TfLiteStatus checkInputs(JNIEnv* env, tflite::Interpreter* interpreter, input_size != env->GetArrayLength(nums_of_bytes) || input_size != env->GetArrayLength(values)) { throwException(env, kIllegalArgumentException, - "Arrays in arguments should be of the same length, but got " - "%d sizes, %d data_types, %d nums_of_bytes, and %d values", + "Internal error: Arrays in arguments should be of the same " + "length, but got %d sizes, %d data_types, %d nums_of_bytes, " + "and %d values", input_size, env->GetArrayLength(data_types), env->GetArrayLength(nums_of_bytes), env->GetArrayLength(values)); @@ -136,8 +138,8 @@ TfLiteStatus checkInputs(JNIEnv* env, tflite::Interpreter* interpreter, int num_dims = static_cast(env->GetArrayLength(dims)); if (target->dims->size != num_dims) { throwException(env, kIllegalArgumentException, - "%d-th input should have %d dimensions, but found %d " - "dimensions", + "Input error: %d-th input should have %d dimensions, but " + "found %d dimensions", i, target->dims->size, num_dims); return kTfLiteError; } @@ -150,7 +152,8 @@ TfLiteStatus checkInputs(JNIEnv* env, tflite::Interpreter* interpreter, num_dims); printDims(obtained_dims.get(), kBufferSize, ptr, num_dims); throwException(env, kIllegalArgumentException, - "%d-th input dimension should be [%s], but found [%s]", + "Input error: %d-th input dimension should be [%s], but " + "found [%s]", i, expected_dims.get(), obtained_dims.get()); env->ReleaseIntArrayElements(dims, ptr, JNI_ABORT); return kTfLiteError; @@ -236,8 +239,8 @@ TfLiteStatus setInputs(JNIEnv* env, tflite::Interpreter* interpreter, TfLiteType type = resolveDataType(data_type[i]); if (type != target->type) { throwException(env, kIllegalArgumentException, - "DataType (%d) of input data does not match with the " - "DataType (%d) of model inputs.", + "Input error: DataType (%d) of input data does not " + "match with the DataType (%d) of model inputs.", type, target->type); return kTfLiteError; } @@ -270,7 +273,8 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_getInputNames(JNIEnv* env, jclass string_class = env->FindClass("java/lang/String"); if (string_class == nullptr) { throwException(env, kUnsupportedOperationException, - "Can not find java/lang/String class to get input names."); + "Internal error: Can not find java/lang/String class to get " + "input names."); return nullptr; } size_t size = interpreter->inputs().size(); @@ -292,7 +296,8 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_getOutputNames(JNIEnv* env, jclass string_class = env->FindClass("java/lang/String"); if (string_class == nullptr) { throwException(env, kUnsupportedOperationException, - "Can not find java/lang/String class to get output names."); + "Internal error: Can not find java/lang/String class to get " + "output names."); return nullptr; } size_t size = interpreter->outputs().size(); @@ -315,6 +320,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) { @@ -351,8 +366,8 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_createModel( path, verifier.get(), error_reporter); if (!model) { throwException(env, kIllegalArgumentException, - "Contents of %s does not encode a valid TensorFlowLite " - "model: %s", + "Contents of %s does not encode a valid " + "TensorFlowLite model: %s", path, error_reporter->CachedErrorMessage()); env->ReleaseStringUTFChars(model_file, path); return 0; @@ -380,8 +395,8 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_createModelWithBuffer( buf, static_cast(capacity), error_reporter); if (!model) { throwException(env, kIllegalArgumentException, - "MappedByteBuffer does not encode a valid TensorFlowLite " - "model: %s", + "MappedByteBuffer does not encode a valid " + "TensorFlowLite model: %s", error_reporter->CachedErrorMessage()); return 0; } @@ -403,7 +418,7 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_createInterpreter( &interpreter, static_cast(num_threads)); if (status != kTfLiteOk) { throwException(env, kIllegalArgumentException, - "Cannot create interpreter: %s", + "Internal error: Cannot create interpreter: %s", error_reporter->CachedErrorMessage()); return 0; } @@ -411,7 +426,7 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_createInterpreter( status = interpreter->AllocateTensors(); if (status != kTfLiteOk) { throwException(env, kNullPointerException, - "Can not allocate memory for the interpreter", + "Internal error: Cannot allocate memory for the interpreter", error_reporter->CachedErrorMessage()); return 0; } @@ -440,7 +455,8 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_run( // resizes inputs status = resizeInputs(env, interpreter, input_size, sizes); if (status != kTfLiteOk) { - throwException(env, kNullPointerException, "Can not resize the input: %s", + throwException(env, kNullPointerException, + "Internal error: Can not resize the input: %s", error_reporter->CachedErrorMessage()); return nullptr; } @@ -448,7 +464,8 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_run( status = interpreter->AllocateTensors(); if (status != kTfLiteOk) { throwException(env, kNullPointerException, - "Can not allocate memory for the given inputs: %s", + "Internal error: Can not allocate memory for the given " + "inputs: %s", error_reporter->CachedErrorMessage()); return nullptr; } @@ -461,7 +478,7 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_run( // runs inference if (interpreter->Invoke() != kTfLiteOk) { throwException(env, kIllegalArgumentException, - "Failed to run on the given Interpreter: %s", + "Internal error: Failed to run on the given Interpreter: %s", error_reporter->CachedErrorMessage()); return nullptr; } @@ -479,8 +496,9 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_run( // returns outputs const std::vector& results = interpreter->outputs(); if (results.empty()) { - throwException(env, kIllegalArgumentException, - "The Interpreter does not have any outputs."); + throwException( + env, kIllegalArgumentException, + "Internal error: The Interpreter does not have any outputs."); return nullptr; } jlongArray outputs = env->NewLongArray(results.size()); @@ -501,7 +519,8 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_getInputDims( const int idx = static_cast(input_idx); 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 error: Out of range: Failed to get %d-th input out of" + " %d inputs", input_idx, interpreter->inputs().size()); return nullptr; } @@ -514,8 +533,8 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_getInputDims( } 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.", + "Input error: Failed to get input dimensions. %d-th input " + "should have %d bytes, but found %d bytes.", idx, expected_num_bytes, num_bytes); return nullptr; } @@ -533,8 +552,8 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_getOutputDataType( 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()); + "Failed to get %d-th output out of %d outputs", output_idx, + interpreter->outputs().size()); return -1; } TfLiteTensor* target = interpreter->tensor(interpreter->outputs()[idx]); @@ -555,7 +574,8 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_resizeInput( 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.", + "Input error: Can not resize %d-th input for a model having " + "%d inputs.", idx, interpreter->inputs().size()); return JNI_FALSE; } @@ -567,7 +587,7 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_resizeInput( interpreter->inputs()[idx], convertJIntArrayToVector(env, dims)); if (status != kTfLiteOk) { throwException(env, kIllegalArgumentException, - "Failed to resize %d-th input: %s", idx, + "Internal error: Failed to resize %d-th input: %s", idx, error_reporter->CachedErrorMessage()); return JNI_FALSE; } 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 0e28a77feea41d72be126d6e60fffbe7ce374a76..eaa765cb343e9764bd0ef018d636a76f4b8a13e4 100644 --- a/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.h +++ b/tensorflow/contrib/lite/java/src/main/native/nativeinterpreterwrapper_jni.h @@ -61,7 +61,7 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_getOutputNames(JNIEnv* env, /* * Class: org_tensorflow_lite_NativeInterpreterWrapper * Method: - * Signature: (JZ) + * Signature: (JZ)V */ JNIEXPORT void JNICALL Java_org_tensorflow_lite_NativeInterpreterWrapper_useNNAPI(JNIEnv* env, @@ -69,6 +69,16 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_useNNAPI(JNIEnv* env, jlong handle, jboolean state); +/* + * Class: org_tensorflow_lite_NativeInterpreterWrapper + * Method: + * Signature: (JI)V + */ +JNIEXPORT void JNICALL +Java_org_tensorflow_lite_NativeInterpreterWrapper_numThreads(JNIEnv* env, + jclass clazz, + jlong handle, + jint num_threads); /* * Class: org_tensorflow_lite_NativeInterpreterWrapper * Method: diff --git a/tensorflow/contrib/lite/java/src/main/native/tensor_jni.cc b/tensorflow/contrib/lite/java/src/main/native/tensor_jni.cc index 65126e78a3003f8a69c69326124d613e878c0f9d..17f4be09c63a9e80feda9d0324d49cc0418fe66a 100644 --- a/tensorflow/contrib/lite/java/src/main/native/tensor_jni.cc +++ b/tensorflow/contrib/lite/java/src/main/native/tensor_jni.cc @@ -23,7 +23,7 @@ namespace { TfLiteTensor* convertLongToTensor(JNIEnv* env, jlong handle) { if (handle == 0) { throwException(env, kIllegalArgumentException, - "Invalid handle to TfLiteTensor."); + "Internal error: Invalid handle to TfLiteTensor."); return nullptr; } return reinterpret_cast(handle); @@ -36,7 +36,8 @@ size_t writeOneDimensionalArray(JNIEnv* env, jobject object, TfLiteType type, size_t to_copy = num_elements * elementByteSize(type); if (to_copy > dst_size) { throwException(env, kIllegalStateException, - "cannot write Java array of %d bytes to Tensor of %d bytes", + "Internal error: cannot write Java array of %d bytes to " + "Tensor of %d bytes", to_copy, dst_size); return 0; } @@ -71,10 +72,10 @@ size_t writeOneDimensionalArray(JNIEnv* env, jobject object, TfLiteType type, } default: { throwException(env, kUnsupportedOperationException, - "TensorFlowLite currently supports float (32 bits), " - "int (32 bits), byte (8 bits), and long (64 bits), " - "support for other types (DataType %d in this case) will " - "be added in the future", + "DataType error: TensorFlowLite currently supports float " + "(32 bits), int (32 bits), byte (8 bits), and long " + "(64 bits), support for other types (DataType %d in this " + "case) will be added in the future", kTfLiteFloat32, type); return 0; } @@ -88,8 +89,9 @@ size_t readOneDimensionalArray(JNIEnv* env, TfLiteType data_type, if (size > src_size) { throwException( env, kIllegalStateException, - "cannot fill a Java array of %d bytes with a Tensor of %d bytes", size, - src_size); + "Internal error: cannot fill a Java array of %d bytes with a Tensor of " + "%d bytes", + size, src_size); return 0; } switch (data_type) { @@ -117,8 +119,8 @@ size_t readOneDimensionalArray(JNIEnv* env, TfLiteType data_type, return size; } default: { - throwException(env, kIllegalStateException, "invalid DataType(%d)", - data_type); + throwException(env, kIllegalStateException, + "DataType error: invalid DataType(%d)", data_type); } } return 0; @@ -152,19 +154,22 @@ size_t elementByteSize(TfLiteType data_type) { switch (data_type) { case kTfLiteFloat32: static_assert(sizeof(jfloat) == 4, - "Java float not compatible with kTfLiteFloat"); + "Interal error: Java float not compatible with " + "kTfLiteFloat"); return 4; case kTfLiteInt32: static_assert(sizeof(jint) == 4, - "Java int not compatible with kTfLiteInt"); + "Interal error: Java int not compatible with kTfLiteInt"); return 4; case kTfLiteUInt8: static_assert(sizeof(jbyte) == 1, - "Java byte not compatible with kTfLiteUInt8"); + "Interal error: Java byte not compatible with " + "kTfLiteUInt8"); return 1; case kTfLiteInt64: static_assert(sizeof(jlong) == 8, - "Java long not compatible with kTfLiteInt64"); + "Interal error: Java long not compatible with " + "kTfLiteInt64"); return 8; default: return 0; @@ -212,7 +217,7 @@ Java_org_tensorflow_lite_Tensor_readMultiDimensionalArray(JNIEnv* env, int num_dims = tensor->dims->size; if (num_dims == 0) { throwException(env, kIllegalArgumentException, - "copyTo() is not meant for scalar Tensors."); + "Internal error: Cannot copy empty/scalar Tensors."); return; } readMultiDimensionalArray(env, tensor->type, tensor->data.raw, tensor->bytes, 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 dbe45e5a05b8227b441de7ca6747f61d010ae210..7c00d3196fd001a288d77d4e01f0b30978d72afe 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 @@ -321,9 +321,7 @@ public final class NativeInterpreterWrapperTest { wrapper.run(inputs); fail(); } catch (IllegalArgumentException e) { - assertThat(e) - .hasMessageThat() - .contains("Invalid inputs. Inputs should not be null or empty."); + assertThat(e).hasMessageThat().contains("Inputs should not be null or empty."); } wrapper.close(); } @@ -440,7 +438,7 @@ public final class NativeInterpreterWrapperTest { NativeInterpreterWrapper.numDimensions(emptyArray); fail(); } catch (IllegalArgumentException e) { - assertThat(e).hasMessageThat().contains("array lengths cannot be 0."); + assertThat(e).hasMessageThat().contains("Array lengths cannot be 0."); } } diff --git a/tensorflow/contrib/lite/kernels/BUILD b/tensorflow/contrib/lite/kernels/BUILD index df0f3cbeb0e99c4a7cb6a9c610ce660f06454744..80cefe83b29192d3ed8f37fc38c4b6d8a9a039a9 100644 --- a/tensorflow/contrib/lite/kernels/BUILD +++ b/tensorflow/contrib/lite/kernels/BUILD @@ -12,10 +12,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -108,10 +105,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":kernel_util", "//tensorflow/contrib/lite/testing:util", @@ -135,12 +129,14 @@ cc_library( srcs = [ "activations.cc", "add.cc", + "arg_max.cc", "audio_spectrogram.cc", "basic_rnn.cc", "batch_to_space_nd.cc", "bidirectional_sequence_lstm.cc", "bidirectional_sequence_rnn.cc", "cast.cc", + "comparisons.cc", "concatenation.cc", "conv.cc", "depthwise_conv.cc", @@ -156,7 +152,7 @@ cc_library( "local_response_norm.cc", "lsh_projection.cc", "lstm.cc", - "maximum.cc", + "maximum_minimum.cc", "mean.cc", "mfcc.cc", "mul.cc", @@ -216,6 +212,7 @@ tf_cc_test( name = "audio_spectrogram_test", size = "small", srcs = ["audio_spectrogram_test.cc"], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -229,6 +226,7 @@ tf_cc_test( name = "mfcc_test", size = "small", srcs = ["mfcc_test.cc"], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -242,10 +240,7 @@ tf_cc_test( name = "activations_test", size = "small", srcs = ["activations_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -258,9 +253,21 @@ tf_cc_test( name = "add_test", size = "small", srcs = ["add_test.cc"], + tags = ["tflite_not_portable_ios"], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + +tf_cc_test( + name = "arg_max_test", + size = "small", + srcs = ["arg_max_test.cc"], tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", + "tflite_not_portable_ios", ], deps = [ ":builtin_ops", @@ -274,10 +281,7 @@ tf_cc_test( name = "div_test", size = "small", srcs = ["div_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -290,10 +294,7 @@ tf_cc_test( name = "sub_test", size = "small", srcs = ["sub_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -306,10 +307,7 @@ tf_cc_test( name = "transpose_test", size = "small", srcs = ["transpose_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -324,10 +322,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -340,10 +335,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -356,6 +348,7 @@ tf_cc_test( name = "cast_test", size = "small", srcs = ["cast_test.cc"], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -368,10 +361,7 @@ tf_cc_test( name = "concatenation_test", size = "small", srcs = ["concatenation_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -384,10 +374,7 @@ tf_cc_test( name = "conv_test", size = "small", srcs = ["conv_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -401,10 +388,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -417,6 +401,7 @@ tf_cc_test( name = "dequantize_test", size = "small", srcs = ["dequantize_test.cc"], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -430,10 +415,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -446,10 +428,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -462,10 +441,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -493,10 +469,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -509,10 +482,7 @@ tf_cc_test( name = "l2norm_test", size = "small", srcs = ["l2norm_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -525,10 +495,7 @@ tf_cc_test( name = "exp_test", size = "small", srcs = ["exp_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -538,9 +505,10 @@ tf_cc_test( ) tf_cc_test( - name = "maximum_test", + name = "maximum_minimum_test", size = "small", - srcs = ["maximum_test.cc"], + srcs = ["maximum_minimum_test.cc"], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -553,10 +521,7 @@ tf_cc_test( name = "mean_test", size = "small", srcs = ["mean_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -569,10 +534,7 @@ tf_cc_test( name = "mul_test", size = "small", srcs = ["mul_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -585,10 +547,7 @@ tf_cc_test( name = "pad_test", size = "small", srcs = ["pad_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -601,10 +560,7 @@ tf_cc_test( name = "reshape_test", size = "small", srcs = ["reshape_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -617,10 +573,7 @@ tf_cc_test( name = "gather_test", size = "small", srcs = ["gather_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:builtin_op_data", @@ -634,10 +587,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:builtin_op_data", @@ -651,10 +601,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -667,10 +614,7 @@ tf_cc_test( name = "svdf_test", size = "small", srcs = ["svdf_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -683,10 +627,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -699,10 +640,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -715,10 +653,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -732,10 +667,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -748,10 +680,7 @@ tf_cc_test( name = "pooling_test", size = "small", srcs = ["pooling_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -764,10 +693,7 @@ tf_cc_test( name = "softmax_test", size = "small", srcs = ["softmax_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -781,10 +707,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -798,10 +721,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -814,10 +734,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -831,10 +748,7 @@ tf_cc_test( name = "lstm_test", size = "small", srcs = ["lstm_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -847,10 +761,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -864,10 +775,7 @@ 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", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -880,10 +788,7 @@ tf_cc_test( name = "split_test", size = "small", srcs = ["split_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -896,10 +801,7 @@ tf_cc_test( name = "squeeze_test", size = "small", srcs = ["squeeze_test.cc"], - tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", - ], + tags = ["tflite_not_portable_ios"], deps = [ ":builtin_ops", "//tensorflow/contrib/lite:framework", @@ -912,9 +814,23 @@ tf_cc_test( name = "strided_slice_test", size = "small", srcs = ["strided_slice_test.cc"], + tags = ["tflite_not_portable_ios"], + deps = [ + ":builtin_ops", + "//tensorflow/contrib/lite:framework", + "//tensorflow/contrib/lite/kernels:test_util", + "@com_google_googletest//:gtest", + ], +) + +tf_cc_test( + name = "comparisons_test", + size = "small", + srcs = [ + "comparisons_test.cc", + ], tags = [ - "tflite_not_portable_ios_arm64", - "tflite_not_portable_ios_x86_64", + "tflite_not_portable_ios", ], deps = [ ":builtin_ops", @@ -924,4 +840,16 @@ tf_cc_test( ], ) +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) + tflite_portable_test_suite() diff --git a/tensorflow/contrib/lite/kernels/add.cc b/tensorflow/contrib/lite/kernels/add.cc index 63ea89df56bafa995950afec3a58267681af304f..e0aa070e2d02cecb9c6ff500ab32b8ad2159db6e 100644 --- a/tensorflow/contrib/lite/kernels/add.cc +++ b/tensorflow/contrib/lite/kernels/add.cc @@ -176,7 +176,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { output); } else { context->ReportError(context, - "Inputs and outputs not all float|unit8 types."); + "Inputs and outputs not all float|uint8 types."); return kTfLiteError; } diff --git a/tensorflow/contrib/lite/kernels/arg_max.cc b/tensorflow/contrib/lite/kernels/arg_max.cc new file mode 100644 index 0000000000000000000000000000000000000000..a2c5e4ceadbc905d22eb02b450c88745a351f58f --- /dev/null +++ b/tensorflow/contrib/lite/kernels/arg_max.cc @@ -0,0 +1,178 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +#include "tensorflow/contrib/lite/builtin_op_data.h" +#include "tensorflow/contrib/lite/context.h" +#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" +#include "tensorflow/contrib/lite/kernels/internal/quantization_util.h" +#include "tensorflow/contrib/lite/kernels/internal/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 arg_max { + +constexpr int kInputTensor = 0; +constexpr int kAxis = 1; +constexpr int kOutputTensor = 0; + +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); + + TfLiteTensor* input = GetInput(context, node, kInputTensor); + TfLiteTensor* axis = GetInput(context, node, kAxis); + // Make sure the axis is only 1 dimension. + TF_LITE_ENSURE_EQ(context, NumElements(axis), 1); + + // Make sure the axis is only either int32 or int64. + TF_LITE_ENSURE(context, + axis->type == kTfLiteInt32 || axis->type == kTfLiteInt64); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + + auto* params = reinterpret_cast(node->builtin_data); + switch (params->output_type) { + case kTfLiteInt32: + output->type = kTfLiteInt32; + break; + case kTfLiteInt64: + output->type = kTfLiteInt64; + break; + default: + context->ReportError(context, "Unknown index output data type"); + return kTfLiteError; + } + + // Check conditions for different types. + switch (input->type) { + case kTfLiteFloat32: + case kTfLiteUInt8: + case kTfLiteInt32: + break; + + default: + context->ReportError(context, "Only float32 and int types are supported"); + return kTfLiteError; + } + + // Copy the input dimensions to output except make the last dimension 1. + TF_LITE_ENSURE(context, NumDimensions(input) >= 1); + TfLiteIntArray* output_size = TfLiteIntArrayCopy(input->dims); + output_size->data[NumDimensions(input) - 1] = 1; + + return context->ResizeTensor(context, output, output_size); +} + +// The current impl actually ignores the axis argument. +// Only determine the index of the maximum value in the last dimension. +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + TfLiteTensor* input = GetInput(context, node, kInputTensor); + TfLiteTensor* axis = GetInput(context, node, kAxis); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + +#define TF_LITE_ARG_MAX(data_type, axis_type, output_type) \ + TF_LITE_ENSURE_EQ(context, GetTensorData(axis)[0], 3); \ + optimized_ops::ArgMax(GetTensorData(axis), \ + GetTensorData(input), GetTensorDims(input), \ + GetTensorData(output), \ + GetTensorDims(output)) + if (axis->type == kTfLiteInt32) { + switch (output->type) { + case kTfLiteInt32: { + switch (input->type) { + case kTfLiteFloat32: + TF_LITE_ARG_MAX(float, int32_t, int32_t); + break; + case kTfLiteUInt8: + TF_LITE_ARG_MAX(uint8_t, int32_t, int32_t); + break; + case kTfLiteInt32: + TF_LITE_ARG_MAX(int32_t, int32_t, int32_t); + break; + default: + return kTfLiteError; + } + } break; + case kTfLiteInt64: { + switch (input->type) { + case kTfLiteFloat32: + TF_LITE_ARG_MAX(float, int32_t, int64_t); + break; + case kTfLiteUInt8: + TF_LITE_ARG_MAX(uint8_t, int32_t, int64_t); + break; + case kTfLiteInt32: + TF_LITE_ARG_MAX(int32_t, int32_t, int64_t); + break; + default: + return kTfLiteError; + } + } break; + default: + return kTfLiteError; + } + } else { + switch (output->type) { + case kTfLiteInt32: { + switch (input->type) { + case kTfLiteFloat32: + TF_LITE_ARG_MAX(float, int64_t, int32_t); + break; + case kTfLiteUInt8: + TF_LITE_ARG_MAX(uint8_t, int64_t, int32_t); + break; + case kTfLiteInt32: + TF_LITE_ARG_MAX(int32_t, int64_t, int32_t); + break; + default: + return kTfLiteError; + } + } break; + case kTfLiteInt64: { + switch (input->type) { + case kTfLiteFloat32: + TF_LITE_ARG_MAX(float, int64_t, int64_t); + break; + case kTfLiteUInt8: + TF_LITE_ARG_MAX(uint8_t, int64_t, int64_t); + break; + case kTfLiteInt32: + TF_LITE_ARG_MAX(int32_t, int64_t, int64_t); + break; + default: + return kTfLiteError; + } + } break; + default: + return kTfLiteError; + } + } +#undef TF_LITE_ARG_MAX + + return kTfLiteOk; +} + +} // namespace arg_max + +TfLiteRegistration* Register_ARG_MAX() { + static TfLiteRegistration r = {nullptr, nullptr, arg_max::Prepare, + arg_max::Eval}; + return &r; +} + +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/arg_max_test.cc b/tensorflow/contrib/lite/kernels/arg_max_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..31b15fe19ab87027c28bde9eaff7d88d03b2c213 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/arg_max_test.cc @@ -0,0 +1,106 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT 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; + +template +class ArgMaxOpModel : public SingleOpModel { + public: + ArgMaxOpModel(std::initializer_list input_shape, TensorType input_type, + TensorType output_type, TensorType index_output_type) { + input_ = AddInput(input_type); + axis_ = AddInput(TensorType_INT32); + output_ = AddOutput(output_type); + SetBuiltinOp(BuiltinOperator_ARG_MAX, BuiltinOptions_ArgMaxOptions, + CreateArgMaxOptions(builder_, index_output_type).Union()); + BuildInterpreter({input_shape, {1, 1, 1, 1}}); + } + + int input() { return input_; } + int axis() { return axis_; } + + std::vector GetOutput() { return ExtractVector(output_); } + std::vector GetOutputShape() { return GetTensorShape(output_); } + + private: + int input_; + int axis_; + int output_; +}; + +TEST(ArgMaxOpTest, GetMaxArgFloat) { + ArgMaxOpModel model({1, 1, 1, 4}, TensorType_FLOAT32, + TensorType_INT32, TensorType_INT32); + model.PopulateTensor(model.input(), {0.1, 0.9, 0.7, 0.3}); + // Currently only support the last dimension. + model.PopulateTensor(model.axis(), {3}); + model.Invoke(); + + EXPECT_THAT(model.GetOutput(), ElementsAreArray({1})); + EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 1})); +} + +TEST(ArgMaxOpTest, GetMaxArgInt) { + ArgMaxOpModel model({1, 1, 1, 4}, TensorType_INT32, TensorType_INT32, + TensorType_INT32); + model.PopulateTensor(model.input(), {1, 9, 7, 3}); + // Currently only support the last dimension. + model.PopulateTensor(model.axis(), {3}); + model.Invoke(); + + EXPECT_THAT(model.GetOutput(), ElementsAreArray({1})); + EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 1})); +} + +TEST(ArgMaxOpTest, GetMaxArgMulDimensions) { + ArgMaxOpModel model({1, 1, 2, 4}, TensorType_INT32, TensorType_INT32, + TensorType_INT32); + model.PopulateTensor(model.input(), {1, 2, 7, 8, 1, 9, 7, 3}); + // Currently only support the last dimension. + model.PopulateTensor(model.axis(), {3}); + model.Invoke(); + + EXPECT_THAT(model.GetOutput(), ElementsAreArray({3, 1})); + EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 2, 1})); +} + +TEST(ArgMaxOpTest, GetMaxArgOutput64) { + ArgMaxOpModel model({1, 1, 2, 4}, TensorType_INT32, TensorType_INT64, + TensorType_INT64); + model.PopulateTensor(model.input(), {10, 2, 7, 8, 1, 9, 7, 3}); + // Currently only support the last dimension. + model.PopulateTensor(model.axis(), {3}); + model.Invoke(); + + EXPECT_THAT(model.GetOutput(), ElementsAreArray({0, 1})); + EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 2, 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/batch_to_space_nd.cc b/tensorflow/contrib/lite/kernels/batch_to_space_nd.cc index bc438f99c6a72fdbc2794dee03524db6a7523834..90edf4f9e3683f2987dd8299a1cd5233caa24479 100644 --- a/tensorflow/contrib/lite/kernels/batch_to_space_nd.cc +++ b/tensorflow/contrib/lite/kernels/batch_to_space_nd.cc @@ -123,6 +123,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { GetTensorDims(op_context.input), \ GetTensorData(op_context.block_shape), \ GetTensorDims(op_context.block_shape), \ + GetTensorData(op_context.crops), \ + GetTensorDims(op_context.crops), \ GetTensorData(op_context.output), \ GetTensorDims(op_context.output)) switch (op_context.input->type) { // Already know in/out types are same. diff --git a/tensorflow/contrib/lite/kernels/comparisons.cc b/tensorflow/contrib/lite/kernels/comparisons.cc new file mode 100644 index 0000000000000000000000000000000000000000..87c413cb982dafd239818040d067738e786d43ff --- /dev/null +++ b/tensorflow/contrib/lite/kernels/comparisons.cc @@ -0,0 +1,119 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +#include "tensorflow/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" +#include "tensorflow/contrib/lite/string_util.h" + +namespace tflite { +namespace ops { +namespace builtin { +namespace comparisons { + +constexpr int kInputTensor1 = 0; +constexpr int kInputTensor2 = 1; +constexpr int kOutputTensor = 0; + +TfLiteStatus LessPrepare(TfLiteContext* context, TfLiteNode* node) { + TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); + TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); + + TfLiteTensor* input1 = GetInput(context, node, kInputTensor1); + TfLiteTensor* input2 = GetInput(context, node, kInputTensor2); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + + // Don't support string and bool. + TF_LITE_ENSURE(context, + input1->type != kTfLiteString || input1->type != kTfLiteBool); + // Currently only support tensors have the same type. + TF_LITE_ENSURE_EQ(context, input1->type, input2->type); + output->type = kTfLiteBool; + + bool requires_broadcast = !HaveSameShapes(input1, input2); + + TfLiteIntArray* output_size = nullptr; + if (requires_broadcast) { + TF_LITE_ENSURE_OK(context, CalculateShapeForBroadcast( + context, input1, input2, &output_size)); + } else { + output_size = TfLiteIntArrayCopy(input1->dims); + } + + return context->ResizeTensor(context, output, output_size); +} + +TfLiteStatus LessEval(TfLiteContext* context, TfLiteNode* node) { + TfLiteTensor* input1 = GetInput(context, node, kInputTensor1); + TfLiteTensor* input2 = GetInput(context, node, kInputTensor2); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + + bool requires_broadcast = !HaveSameShapes(input1, input2); + +#define TF_LITE_LESS(type, opname) \ + reference_ops::opname(GetTensorData(input1), GetTensorDims(input1), \ + GetTensorData(input2), GetTensorDims(input2), \ + GetTensorData(output), GetTensorDims(output)); + + // TODO(renjieliu): Support quantized data. + if (requires_broadcast) { + switch (input1->type) { + case kTfLiteFloat32: + TF_LITE_LESS(float, BroadcastLess); + break; + case kTfLiteInt32: + TF_LITE_LESS(int32_t, BroadcastLess); + break; + case kTfLiteInt64: + TF_LITE_LESS(int64_t, BroadcastLess); + break; + default: + context->ReportError(context, + "Does not support type other than float|int"); + return kTfLiteError; + } + } else { + switch (input1->type) { + case kTfLiteFloat32: + TF_LITE_LESS(float, Less); + break; + case kTfLiteInt32: + TF_LITE_LESS(int32_t, Less); + break; + case kTfLiteInt64: + TF_LITE_LESS(int64_t, Less); + break; + default: + context->ReportError(context, + "Does not support type other than float|int"); + return kTfLiteError; + } + } +#undef TF_LITE_LESS + return kTfLiteOk; +} + +} // namespace comparisons + +TfLiteRegistration* Register_LESS() { + static TfLiteRegistration r = {nullptr, nullptr, comparisons::LessPrepare, + comparisons::LessEval}; + return &r; +} + +} // namespace builtin +} // namespace ops +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/comparisons_test.cc b/tensorflow/contrib/lite/kernels/comparisons_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..da2d7f858984a4d3bb09ca8e485fe1599bea7ded --- /dev/null +++ b/tensorflow/contrib/lite/kernels/comparisons_test.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 +#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 LessOpModel : public SingleOpModel { + public: + LessOpModel(std::initializer_list input1_shape, + std::initializer_list input2_shape, TensorType input_type) { + input1_ = AddInput(input_type); + input2_ = AddInput(input_type); + output_ = AddOutput(TensorType_BOOL); + SetBuiltinOp(BuiltinOperator_LESS, BuiltinOptions_LessOptions, + CreateLessOptions(builder_).Union()); + BuildInterpreter({input1_shape, input2_shape}); + } + + int input1() { return input1_; } + int input2() { return input2_; } + + std::vector GetOutput() { return ExtractVector(output_); } + std::vector GetOutputShape() { return GetTensorShape(output_); } + + private: + int input1_; + int input2_; + int output_; +}; + +TEST(ArgMaxOpTest, LessFloat) { + LessOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_FLOAT32); + model.PopulateTensor(model.input1(), {0.1, 0.9, 0.7, 0.3}); + model.PopulateTensor(model.input2(), {0.1, 0.2, 0.6, 0.5}); + model.Invoke(); + + EXPECT_THAT(model.GetOutput(), ElementsAreArray({false, false, false, true})); + EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4})); +} + +TEST(ArgMaxOpTest, LessInt) { + LessOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT32); + model.PopulateTensor(model.input1(), {-1, 9, 7, 3}); + model.PopulateTensor(model.input2(), {1, 2, 6, 5}); + model.Invoke(); + + EXPECT_THAT(model.GetOutput(), ElementsAreArray({true, false, false, true})); + EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4})); +} + +TEST(ArgMaxOpTest, LessBroadcast) { + LessOpModel model({1, 1, 1, 4}, {1, 1, 1, 1}, TensorType_INT32); + model.PopulateTensor(model.input1(), {-1, 9, 7, 3}); + model.PopulateTensor(model.input2(), {7}); + model.Invoke(); + + EXPECT_THAT(model.GetOutput(), ElementsAreArray({true, false, false, true})); + EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4})); +} + +TEST(ArgMaxOpTest, LessBroadcastTwoD) { + LessOpModel model({1, 1, 2, 4}, {1, 1, 1, 4}, TensorType_INT32); + model.PopulateTensor(model.input1(), {-1, 9, 7, 3, 2, 4, 6, 8}); + model.PopulateTensor(model.input2(), {7, 1, 2, 4}); + model.Invoke(); + + EXPECT_THAT(model.GetOutput(), ElementsAreArray({true, false, false, true, + true, false, false, false})); + EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 2, 4})); +} + +} // 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/conv.cc b/tensorflow/contrib/lite/kernels/conv.cc index 18ff33bf9f55ac1d25bb3392e714686c5305c2b8..3b467b3aa284586ab8e67ede55583adffbe06cc7 100644 --- a/tensorflow/contrib/lite/kernels/conv.cc +++ b/tensorflow/contrib/lite/kernels/conv.cc @@ -225,22 +225,27 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { // Matching GetWindowedOutputSize in TensorFlow. auto padding = params->padding; - auto computeOutSize = [padding](int imageSize, int filterSize, - int stride) -> int { + auto computeOutSize = [padding](int imageSize, int filterSize, int stride, + int dilationRate) -> int { + int effectiveFilterSize = (filterSize - 1) * dilationRate + 1; return padding == kTfLitePaddingSame ? (imageSize + stride - 1) / stride : padding == kTfLitePaddingValid - ? (imageSize - filterSize + stride) / stride + ? (imageSize - effectiveFilterSize + stride) / stride : 0; }; - int outWidth = computeOutSize(width, filter_width, params->stride_width); - int outHeight = computeOutSize(height, filter_height, params->stride_height); + int outWidth = computeOutSize(width, filter_width, params->stride_width, + params->dilation_width_factor); + int outHeight = computeOutSize(height, filter_height, params->stride_height, + params->dilation_height_factor); data->padding.height = - ComputePadding(params->stride_height, height, filter_height, outHeight); + ComputePadding(params->stride_height, params->dilation_height_factor, + height, filter_height, outHeight); data->padding.width = - ComputePadding(params->stride_width, width, filter_width, outWidth); + ComputePadding(params->stride_width, params->dilation_width_factor, width, + filter_width, outWidth); TF_LITE_ENSURE(context, hasBias); @@ -375,28 +380,40 @@ void EvalFloat(TfLiteContext* context, TfLiteNode* node, float output_activation_min, output_activation_max; CalculateActivationRangeFloat(params->activation, &output_activation_min, &output_activation_max); - - switch (kernel_type) { + KernelType effective_kernel_type; + if (((kernel_type == kMultithreadOptimized) || + (kernel_type == kCblasOptimized)) && + ((params->dilation_width_factor != 1) || + (params->dilation_height_factor != 1))) { + // kMultithreadOptimized and kCblasOptimized do not support dilation. + // Therefore, fallback to optimized. + effective_kernel_type = kGenericOptimized; + } else { + effective_kernel_type = kernel_type; + } + switch (effective_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)); + reference_ops::Conv( + GetTensorData(input), GetTensorDims(input), + GetTensorData(filter), GetTensorDims(filter), + GetTensorData(bias), GetTensorDims(bias), params->stride_width, + params->stride_height, params->dilation_width_factor, + params->dilation_height_factor, 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)); + optimized_ops::Conv( + GetTensorData(input), GetTensorDims(input), + GetTensorData(filter), GetTensorDims(filter), + GetTensorData(bias), GetTensorDims(bias), params->stride_width, + params->stride_height, params->dilation_width_factor, + params->dilation_height_factor, data->padding.width, + data->padding.height, output_activation_min, output_activation_max, + GetTensorData(output), GetTensorDims(output), + GetTensorData(im2col), GetTensorDims(im2col)); break; } case kMultithreadOptimized: { diff --git a/tensorflow/contrib/lite/kernels/conv_test.cc b/tensorflow/contrib/lite/kernels/conv_test.cc index d2393c3c97bb9516e2b8a6c8ae037dc0dfdfe64b..0dcfc826fd218d2d2dfbf89201d2c13fbfe6f0e1 100644 --- a/tensorflow/contrib/lite/kernels/conv_test.cc +++ b/tensorflow/contrib/lite/kernels/conv_test.cc @@ -46,7 +46,8 @@ class BaseConvolutionOpModel : public SingleOpModel { 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) { + enum ActivationFunctionType activation = ActivationFunctionType_NONE, + int dilation_width_factor = 1, int dilation_height_factor = 1) { input_ = AddInput(input); filter_ = AddInput(filter); @@ -71,8 +72,9 @@ class BaseConvolutionOpModel : public SingleOpModel { } SetBuiltinOp(BuiltinOperator_CONV_2D, BuiltinOptions_Conv2DOptions, - CreateConv2DOptions(builder_, padding, stride_width, - stride_height, activation) + CreateConv2DOptions( + builder_, padding, stride_width, stride_height, activation, + dilation_width_factor, dilation_height_factor) .Union()); resolver_ = absl::make_unique(BuiltinOperator_CONV_2D, diff --git a/tensorflow/contrib/lite/kernels/depthwise_conv.cc b/tensorflow/contrib/lite/kernels/depthwise_conv.cc index cad9ce114c8387047af2b63bee704035fd329330..eeda1bc3c5ba2da5b6546ce36925a6f20fc9cbae 100644 --- a/tensorflow/contrib/lite/kernels/depthwise_conv.cc +++ b/tensorflow/contrib/lite/kernels/depthwise_conv.cc @@ -140,10 +140,10 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { int out_height = compute_out_size(height, filter_height, params->stride_height); - data->padding.height = - ComputePadding(params->stride_height, height, filter_height, out_height); + data->padding.height = ComputePadding(params->stride_height, 1, height, + filter_height, out_height); data->padding.width = - ComputePadding(params->stride_width, width, filter_width, out_width); + ComputePadding(params->stride_width, 1, width, filter_width, out_width); // Note that quantized inference requires that all tensors have their // parameters set. This is usually done during quantized training. diff --git a/tensorflow/contrib/lite/kernels/div.cc b/tensorflow/contrib/lite/kernels/div.cc index 6dd243ad62ece3e094529d923ce80d1d4a0c19ca..ec380c8e4956e5bcd0d7559bfd8f89a52d9d233c 100644 --- a/tensorflow/contrib/lite/kernels/div.cc +++ b/tensorflow/contrib/lite/kernels/div.cc @@ -106,6 +106,8 @@ void EvalFloat(TfLiteContext* context, TfLiteNode* node, #undef TF_LITE_DIV } + + template TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { auto* params = reinterpret_cast(node->builtin_data); @@ -118,7 +120,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { if (output->type == kTfLiteFloat32) { EvalFloat(context, node, params, data, input1, input2, output); } else { - context->ReportError(context, "Inputs and outputs not all float types."); + context->ReportError(context, + "Div only supports FLOAT32 and quantized UINT8 now."); return kTfLiteError; } diff --git a/tensorflow/contrib/lite/kernels/internal/BUILD b/tensorflow/contrib/lite/kernels/internal/BUILD index 167c0f1fde9202452a915cea69cbb935fa1af7b6..67dd1884966d8addafa54e80f9923aa66354dff9 100644 --- a/tensorflow/contrib/lite/kernels/internal/BUILD +++ b/tensorflow/contrib/lite/kernels/internal/BUILD @@ -154,6 +154,7 @@ cc_library( ], copts = tflite_copts(), deps = [ + ":quantization_util", ":types", ":round", "//third_party/eigen3", @@ -238,6 +239,7 @@ cc_library( "reference/reference_ops.h", ], deps = [ + ":quantization_util", ":round", ":types", "//third_party/eigen3", @@ -430,4 +432,13 @@ cc_library( ), ) +cc_test( + name = "batch_to_space_nd_test", + srcs = ["batch_to_space_nd_test.cc"], + deps = [ + ":optimized_base", + "@com_google_googletest//:gtest_main", + ], +) + exports_files(["optimized/eigen_tensor_reduced_instantiations_oss.h"]) diff --git a/tensorflow/contrib/lite/kernels/internal/batch_to_space_nd_test.cc b/tensorflow/contrib/lite/kernels/internal/batch_to_space_nd_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..5a2901ac8c297265e542cc30d3127fe774c19e78 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/internal/batch_to_space_nd_test.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/contrib/lite/kernels/internal/optimized/optimized_ops.h" + +#include + +namespace tflite { +namespace { + +// A light wrapper of GetIndexRange which returns a pair of start / end +// indices. +std::pair GetIndexRange(int spatial_index_dim, int block_shape_dim, + int input_dim, int output_dim) { + int index_start = 0; + int index_end = 0; + optimized_ops::GetIndexRange(spatial_index_dim, block_shape_dim, input_dim, + output_dim, &index_start, &index_end); + return {index_start, index_end}; +} + +TEST(BatchToSpaceNDTest, TestIndexRange) { + // Simple test case, no cropping. + EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/3, /*block_shape_dim=*/6, + /*input_dim=*/1, /*output_dim=*/6), + std::make_pair(0, 1)); + + // No cropping and input_dim > 1. + EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/2, /*block_shape_dim=*/6, + /*input_dim=*/5, /*output_dim=*/30), + std::make_pair(0, 5)); + + // With small cropping values (can be either at the beginning or at the end). + EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/0, /*block_shape_dim=*/2, + /*input_dim=*/3, /*output_dim=*/4), + std::make_pair(0, 2)); + + // With positive cropping values at the beginning. + EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/-2, /*block_shape_dim=*/2, + /*input_dim=*/3, /*output_dim=*/4), + std::make_pair(1, 3)); + + // Large crop at the beginning. + EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/-30, /*block_shape_dim=*/5, + /*input_dim=*/7, /*output_dim=*/5), + std::make_pair(6, 7)); + + EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/-26, /*block_shape_dim=*/5, + /*input_dim=*/7, /*output_dim=*/5), + std::make_pair(6, 7)); + + // Large crop at the end. + EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/0, /*block_shape_dim=*/5, + /*input_dim=*/7, /*output_dim=*/5), + std::make_pair(0, 1)); + + EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/4, /*block_shape_dim=*/5, + /*input_dim=*/7, /*output_dim=*/5), + std::make_pair(0, 1)); + + // Rounding up incorrectly will fail this test. + EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/3, /*block_shape_dim=*/5, + /*input_dim=*/7, /*output_dim=*/5), + std::make_pair(0, 1)); + + // Extreme cropping with output of a single spatial location. + // Valid position 1, when large crop at the end. + EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/0, /*block_shape_dim=*/5, + /*input_dim=*/7, /*output_dim=*/1), + std::make_pair(0, 1)); + + // Valid position 2, when large crop at the beginning. + EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/-30, /*block_shape_dim=*/5, + /*input_dim=*/7, /*output_dim=*/1), + std::make_pair(6, 7)); + + // Invalid positions. + EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/1, /*block_shape_dim=*/5, + /*input_dim=*/7, /*output_dim=*/1), + std::make_pair(0, 0)); + EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/-29, /*block_shape_dim=*/5, + /*input_dim=*/7, /*output_dim=*/1), + std::make_pair(6, 6)); +} + +} // namespace +} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/compatibility.h b/tensorflow/contrib/lite/kernels/internal/compatibility.h index 51426bb1c584b82af7b1a2ffaf5a675a1dd9a6fd..93fc6b6a76f67e2e75ba3a766e5ea6fb6bada77a 100644 --- a/tensorflow/contrib/lite/kernels/internal/compatibility.h +++ b/tensorflow/contrib/lite/kernels/internal/compatibility.h @@ -77,6 +77,7 @@ limitations under the License. #endif // TODO(ahentz): Clean up. +using int8 = std::int8_t; using uint8 = std::uint8_t; using int16 = std::int16_t; using uint16 = std::uint16_t; diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h index 0f78e0f728585ab27a8116a4707ac9614a6ea060..dd6932ffe7b7a6f1101f146ce6472b0df4cbda3b 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h @@ -1696,15 +1696,15 @@ inline void DepthwiseConv(const uint8* input_data, const Dims<4>& input_dims, #ifdef __aarch64__ // Call kernel optimized for depthwise convolutions using 3x3 filters if // parameters are supported. - if (Fast3by3FilterKernelSupported(input_dims, filter_dims, stride_width, - stride_height, pad_width, pad_height, - depth_multiplier, output_dims)) { - 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); + if (Fast3x3FilterKernelSupported(input_dims, filter_dims, stride_width, + stride_height, pad_width, pad_height, + depth_multiplier, output_dims)) { + DepthwiseConv3x3Filter(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 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 index a349892076fcc4989e2f4cad188b383d2b31d470..55e0d5c3aa9ebb8b46403550e190b00a54cb53e5 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h @@ -1,4 +1,4 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -40,412 +40,4380 @@ inline void preload_l1_keep(const uint8* ptr) { // NEON intrinsics vector data types. // See: https://bugs.llvm.org/show_bug.cgi?id=34945 -struct Int32x16 { - int32x4_t v0, v1, v2, v3; +struct Int32x8 { + int32x4_t low, high; }; -struct Int16x16 { - int16x8_t low, high; +struct Filter3x3x8 { + int16x8_t f0, f1, f2, f3, f4, f5, f6, f7, f8; }; -struct Int16x16x3 { - Int16x16 v0, v1, v2; +// Loads 3x3 filter of depth 8 and adds filter offsets. +inline Filter3x3x8 Load3x3Filter(const uint8* filter_ptr, int32 filter_offset, + int output_depth) { + Filter3x3x8 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; + 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 + 1 * output_depth); + temp_u8_2 = vld1_u8(filter_ptr + 2 * output_depth); + temp_u8_3 = vld1_u8(filter_ptr + 3 * output_depth); + temp_u8_4 = vld1_u8(filter_ptr + 4 * output_depth); + temp_u8_5 = vld1_u8(filter_ptr + 5 * output_depth); + temp_u8_6 = vld1_u8(filter_ptr + 6 * output_depth); + temp_u8_7 = vld1_u8(filter_ptr + 7 * output_depth); + temp_u8_8 = vld1_u8(filter_ptr + 8 * output_depth); + + filter.f0 = vreinterpretq_s16_u16(vmovl_u8(temp_u8_0)); + filter.f1 = vreinterpretq_s16_u16(vmovl_u8(temp_u8_1)); + filter.f2 = vreinterpretq_s16_u16(vmovl_u8(temp_u8_2)); + filter.f3 = vreinterpretq_s16_u16(vmovl_u8(temp_u8_3)); + filter.f4 = vreinterpretq_s16_u16(vmovl_u8(temp_u8_4)); + filter.f5 = vreinterpretq_s16_u16(vmovl_u8(temp_u8_5)); + filter.f6 = vreinterpretq_s16_u16(vmovl_u8(temp_u8_6)); + filter.f7 = vreinterpretq_s16_u16(vmovl_u8(temp_u8_7)); + filter.f8 = vreinterpretq_s16_u16(vmovl_u8(temp_u8_8)); + + filter.f0 = vaddq_s16(filter.f0, filter_offset_vec); + filter.f1 = vaddq_s16(filter.f1, filter_offset_vec); + filter.f2 = vaddq_s16(filter.f2, filter_offset_vec); + filter.f3 = vaddq_s16(filter.f3, filter_offset_vec); + filter.f4 = vaddq_s16(filter.f4, filter_offset_vec); + filter.f5 = vaddq_s16(filter.f5, filter_offset_vec); + filter.f6 = vaddq_s16(filter.f6, filter_offset_vec); + filter.f7 = vaddq_s16(filter.f7, filter_offset_vec); + filter.f8 = vaddq_s16(filter.f8, filter_offset_vec); + + return filter; +} + +// Applies activation, offset and downquantize on a set of accumulator +// registers that correspond to a 2x2 output of depth 8. +// Stores results to output. +inline void DownquantizeAndStore2x2Output( + Int32x8 acc_0, Int32x8 acc_1, Int32x8 acc_2, Int32x8 acc_3, + 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) { + using gemmlowp::RoundingDivideByPOT; + const int32x4_t output_offset_vec = vdupq_n_s32(output_offset); + const int32x4_t output_activation_min_vec = + vdupq_n_s32(output_activation_min); + const int32x4_t output_activation_max_vec = + vdupq_n_s32(output_activation_max); + + // Fixed-point multiplication. + acc_0.low = vqrdmulhq_n_s32(acc_0.low, output_multiplier); + acc_0.high = vqrdmulhq_n_s32(acc_0.high, output_multiplier); + acc_1.low = vqrdmulhq_n_s32(acc_1.low, output_multiplier); + acc_1.high = vqrdmulhq_n_s32(acc_1.high, output_multiplier); + acc_2.low = vqrdmulhq_n_s32(acc_2.low, output_multiplier); + acc_2.high = vqrdmulhq_n_s32(acc_2.high, output_multiplier); + acc_3.low = vqrdmulhq_n_s32(acc_3.low, output_multiplier); + acc_3.high = vqrdmulhq_n_s32(acc_3.high, output_multiplier); + + acc_0.low = RoundingDivideByPOT(acc_0.low, output_shift); + acc_0.high = RoundingDivideByPOT(acc_0.high, output_shift); + acc_1.low = RoundingDivideByPOT(acc_1.low, output_shift); + acc_1.high = RoundingDivideByPOT(acc_1.high, output_shift); + acc_2.low = RoundingDivideByPOT(acc_2.low, output_shift); + acc_2.high = RoundingDivideByPOT(acc_2.high, output_shift); + acc_3.low = RoundingDivideByPOT(acc_3.low, output_shift); + acc_3.high = RoundingDivideByPOT(acc_3.high, output_shift); + + // Add the output offset. + acc_0.low = vaddq_s32(acc_0.low, output_offset_vec); + acc_0.high = vaddq_s32(acc_0.high, output_offset_vec); + acc_1.low = vaddq_s32(acc_1.low, output_offset_vec); + acc_1.high = vaddq_s32(acc_1.high, output_offset_vec); + acc_2.low = vaddq_s32(acc_2.low, output_offset_vec); + acc_2.high = vaddq_s32(acc_2.high, output_offset_vec); + acc_3.low = vaddq_s32(acc_3.low, output_offset_vec); + acc_3.high = vaddq_s32(acc_3.high, output_offset_vec); + + // Apply the activation function. + acc_0.low = vmaxq_s32(acc_0.low, output_activation_min_vec); + acc_0.high = vmaxq_s32(acc_0.high, output_activation_min_vec); + acc_1.low = vmaxq_s32(acc_1.low, output_activation_min_vec); + acc_1.high = vmaxq_s32(acc_1.high, output_activation_min_vec); + acc_2.low = vmaxq_s32(acc_2.low, output_activation_min_vec); + acc_2.high = vmaxq_s32(acc_2.high, output_activation_min_vec); + acc_3.low = vmaxq_s32(acc_3.low, output_activation_min_vec); + acc_3.high = vmaxq_s32(acc_3.high, output_activation_min_vec); + + acc_0.low = vminq_s32(acc_0.low, output_activation_max_vec); + acc_0.high = vminq_s32(acc_0.high, output_activation_max_vec); + acc_1.low = vminq_s32(acc_1.low, output_activation_max_vec); + acc_1.high = vminq_s32(acc_1.high, output_activation_max_vec); + acc_2.low = vminq_s32(acc_2.low, output_activation_max_vec); + acc_2.high = vminq_s32(acc_2.high, output_activation_max_vec); + acc_3.low = vminq_s32(acc_3.low, output_activation_max_vec); + acc_3.high = vminq_s32(acc_3.high, output_activation_max_vec); + + // Saturating cast to uint8 and store to destination. + int16x4_t acc_0_low_s16 = vqmovn_s32(acc_0.low); + int16x4_t acc_0_high_s16 = vqmovn_s32(acc_0.high); + int16x4_t acc_1_low_s16 = vqmovn_s32(acc_1.low); + int16x4_t acc_1_high_s16 = vqmovn_s32(acc_1.high); + int16x4_t acc_2_low_s16 = vqmovn_s32(acc_2.low); + int16x4_t acc_2_high_s16 = vqmovn_s32(acc_2.high); + int16x4_t acc_3_low_s16 = vqmovn_s32(acc_3.low); + int16x4_t acc_3_high_s16 = vqmovn_s32(acc_3.high); + + int16x8_t res_0_s16 = vcombine_s16(acc_0_low_s16, acc_0_high_s16); + int16x8_t res_1_s16 = vcombine_s16(acc_1_low_s16, acc_1_high_s16); + int16x8_t res_2_s16 = vcombine_s16(acc_2_low_s16, acc_2_high_s16); + int16x8_t res_3_s16 = vcombine_s16(acc_3_low_s16, acc_3_high_s16); + + uint8x8_t res_0_u8 = vqmovun_s16(res_0_s16); + uint8x8_t res_1_u8 = vqmovun_s16(res_1_s16); + uint8x8_t res_2_u8 = vqmovun_s16(res_2_s16); + uint8x8_t res_3_u8 = vqmovun_s16(res_3_s16); + + vst1_u8(output_ptr, res_0_u8); + vst1_u8(output_ptr + output_depth, res_1_u8); + vst1_u8(output_ptr + output_depth * output_width, res_2_u8); + vst1_u8(output_ptr + output_depth * output_width + output_depth, res_3_u8); +} + +inline void DownquantizeAndStore(Int32x8 acc, int32 output_offset, + int32 output_multiplier, int output_shift, + int32 output_activation_min, + int32 output_activation_max, + uint8* output_ptr) { + using gemmlowp::RoundingDivideByPOT; + const int32x4_t output_offset_vec = vdupq_n_s32(output_offset); + const int32x4_t output_activation_min_vec = + vdupq_n_s32(output_activation_min); + const int32x4_t output_activation_max_vec = + vdupq_n_s32(output_activation_max); + + acc.low = vqrdmulhq_n_s32(acc.low, output_multiplier); + acc.high = vqrdmulhq_n_s32(acc.high, output_multiplier); + + acc.low = RoundingDivideByPOT(acc.low, output_shift); + acc.high = RoundingDivideByPOT(acc.high, output_shift); + + acc.low = vaddq_s32(acc.low, output_offset_vec); + acc.high = vaddq_s32(acc.high, output_offset_vec); + + acc.low = vmaxq_s32(acc.low, output_activation_min_vec); + acc.high = vmaxq_s32(acc.high, output_activation_min_vec); + + acc.low = vminq_s32(acc.low, output_activation_max_vec); + acc.high = vminq_s32(acc.high, output_activation_max_vec); + + int16x4_t acc_low_s16 = vqmovn_s32(acc.low); + int16x4_t acc_high_s16 = vqmovn_s32(acc.high); + + int16x8_t res_s16 = vcombine_s16(acc_low_s16, acc_high_s16); + uint8x8_t res_u8 = vqmovun_s16(res_s16); + vst1_u8(output_ptr, res_u8); +} + +inline void DownquantizeAndStore2Output( + Int32x8 acc_0, Int32x8 acc_1, int32 output_offset, int32 output_multiplier, + int output_shift, int32 output_activation_min, int32 output_activation_max, + uint8* output_ptr, int output_ptr_offset) { + { + using gemmlowp::RoundingDivideByPOT; + const int32x4_t output_offset_vec = vdupq_n_s32(output_offset); + const int32x4_t output_activation_min_vec = + vdupq_n_s32(output_activation_min); + const int32x4_t output_activation_max_vec = + vdupq_n_s32(output_activation_max); + + // Fixed-point multiplication. + acc_0.low = vqrdmulhq_n_s32(acc_0.low, output_multiplier); + acc_0.high = vqrdmulhq_n_s32(acc_0.high, output_multiplier); + acc_1.low = vqrdmulhq_n_s32(acc_1.low, output_multiplier); + acc_1.high = vqrdmulhq_n_s32(acc_1.high, output_multiplier); + + acc_0.low = RoundingDivideByPOT(acc_0.low, output_shift); + acc_0.high = RoundingDivideByPOT(acc_0.high, output_shift); + acc_1.low = RoundingDivideByPOT(acc_1.low, output_shift); + acc_1.high = RoundingDivideByPOT(acc_1.high, output_shift); + + // Add the output offset. + acc_0.low = vaddq_s32(acc_0.low, output_offset_vec); + acc_0.high = vaddq_s32(acc_0.high, output_offset_vec); + acc_1.low = vaddq_s32(acc_1.low, output_offset_vec); + acc_1.high = vaddq_s32(acc_1.high, output_offset_vec); + + // Apply the activation function. + acc_0.low = vmaxq_s32(acc_0.low, output_activation_min_vec); + acc_0.high = vmaxq_s32(acc_0.high, output_activation_min_vec); + acc_1.low = vmaxq_s32(acc_1.low, output_activation_min_vec); + acc_1.high = vmaxq_s32(acc_1.high, output_activation_min_vec); + + acc_0.low = vminq_s32(acc_0.low, output_activation_max_vec); + acc_0.high = vminq_s32(acc_0.high, output_activation_max_vec); + acc_1.low = vminq_s32(acc_1.low, output_activation_max_vec); + acc_1.high = vminq_s32(acc_1.high, output_activation_max_vec); + } + + // Saturating cast to uint8 and store to destination. + int16x8_t res_0_s16; + { + int16x4_t acc_0_low_s16 = vqmovn_s32(acc_0.low); + int16x4_t acc_0_high_s16 = vqmovn_s32(acc_0.high); + res_0_s16 = vcombine_s16(acc_0_low_s16, acc_0_high_s16); + } + + int16x8_t res_1_s16; + { + int16x4_t acc_1_low_s16 = vqmovn_s32(acc_1.low); + int16x4_t acc_1_high_s16 = vqmovn_s32(acc_1.high); + res_1_s16 = vcombine_s16(acc_1_low_s16, acc_1_high_s16); + } + + uint8x8_t res_0_u8 = vqmovun_s16(res_0_s16); + uint8x8_t res_1_u8 = vqmovun_s16(res_1_s16); + vst1_u8(output_ptr, res_0_u8); + vst1_u8(output_ptr + output_ptr_offset, res_1_u8); +} + +// Performs multiply accumulate on 3 inputs of depth 8. +inline Int32x8 MultiplyAccumulateRow(Int32x8 accum, int16x8_t f0, int16x8_t f1, + int16x8_t f2, int16x8_t i0, int16x8_t i1, + int16x8_t i2) { + accum.low = vmlal_s16(accum.low, vget_low_s16(f0), vget_low_s16(i0)); + accum.high = vmlal_s16(accum.high, vget_high_s16(f0), vget_high_s16(i0)); + accum.low = vmlal_s16(accum.low, vget_low_s16(f1), vget_low_s16(i1)); + accum.high = vmlal_s16(accum.high, vget_high_s16(f1), vget_high_s16(i1)); + accum.low = vmlal_s16(accum.low, vget_low_s16(f2), vget_low_s16(i2)); + accum.high = vmlal_s16(accum.high, vget_high_s16(f2), vget_high_s16(i2)); + return accum; +} + +// Performs multiply accumulate on 3 inputs of depth 8. +inline Int32x8 MultiplyAccumulate3x3Filter(const Filter3x3x8& f, int16x8_t i0, + int16x8_t i1, int16x8_t i2, + int16x8_t i3, int16x8_t i4, + int16x8_t i5, int16x8_t i6, + int16x8_t i7, int16x8_t i8, + Int32x8 accum) { + accum.low = vmlal_s16(accum.low, vget_low_s16(f.f0), vget_low_s16(i0)); + accum.high = vmlal_s16(accum.high, vget_high_s16(f.f0), vget_high_s16(i0)); + accum.low = vmlal_s16(accum.low, vget_low_s16(f.f1), vget_low_s16(i1)); + accum.high = vmlal_s16(accum.high, vget_high_s16(f.f1), vget_high_s16(i1)); + accum.low = vmlal_s16(accum.low, vget_low_s16(f.f2), vget_low_s16(i2)); + accum.high = vmlal_s16(accum.high, vget_high_s16(f.f2), vget_high_s16(i2)); + accum.low = vmlal_s16(accum.low, vget_low_s16(f.f3), vget_low_s16(i3)); + accum.high = vmlal_s16(accum.high, vget_high_s16(f.f3), vget_high_s16(i3)); + accum.low = vmlal_s16(accum.low, vget_low_s16(f.f4), vget_low_s16(i4)); + accum.high = vmlal_s16(accum.high, vget_high_s16(f.f4), vget_high_s16(i4)); + accum.low = vmlal_s16(accum.low, vget_low_s16(f.f5), vget_low_s16(i5)); + accum.high = vmlal_s16(accum.high, vget_high_s16(f.f5), vget_high_s16(i5)); + accum.low = vmlal_s16(accum.low, vget_low_s16(f.f6), vget_low_s16(i6)); + accum.high = vmlal_s16(accum.high, vget_high_s16(f.f6), vget_high_s16(i6)); + accum.low = vmlal_s16(accum.low, vget_low_s16(f.f7), vget_low_s16(i7)); + accum.high = vmlal_s16(accum.high, vget_high_s16(f.f7), vget_high_s16(i7)); + accum.low = vmlal_s16(accum.low, vget_low_s16(f.f8), vget_low_s16(i8)); + accum.high = vmlal_s16(accum.high, vget_high_s16(f.f8), vget_high_s16(i8)); + return accum; +} + +inline void DotProductAndStore(const Filter3x3x8& filter, int16x8_t i0, + int16x8_t i1, int16x8_t i2, int16x8_t i3, + int16x8_t i4, int16x8_t i5, int16x8_t i6, + int16x8_t i7, int16x8_t i8, + const int32* bias_ptr, int32 output_offset, + int32 output_multiplier, int output_shift, + int32 output_activation_min, + int32 output_activation_max, uint8* output_ptr) { + Int32x8 acc; + acc.low = vld1q_s32(bias_ptr); + acc.high = vld1q_s32(bias_ptr + 4); + + acc = MultiplyAccumulate3x3Filter(filter, i0, i1, i2, i3, i4, i5, i6, i7, i8, + acc); + + DownquantizeAndStore(acc, output_offset, output_multiplier, output_shift, + output_activation_min, output_activation_max, + output_ptr); +} + +// Performs multiply-accumulate on a 3x4 input for 2 horizontal outputs. +inline void DotProductAndStore2xStride1( + const Filter3x3x8& filter, int16x8_t i0, int16x8_t i1, int16x8_t i2, + int16x8_t i3, int16x8_t i4, int16x8_t i5, int16x8_t i6, int16x8_t i7, + int16x8_t i8, int16x8_t i9, int16x8_t i10, int16x8_t i11, + 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_ptr_offset) { + Int32x8 acc_0, acc_1; + acc_0.low = vld1q_s32(bias_ptr); + acc_1.low = vld1q_s32(bias_ptr); + acc_0.high = vld1q_s32(bias_ptr + 4); + acc_1.high = vld1q_s32(bias_ptr + 4); + + acc_0 = MultiplyAccumulate3x3Filter(filter, i0, i1, i2, i4, i5, i6, i8, i9, + i10, acc_0); + acc_1 = MultiplyAccumulate3x3Filter(filter, i1, i2, i3, i5, i6, i7, i9, i10, + i11, acc_1); + DownquantizeAndStore2Output(acc_0, acc_1, output_offset, output_multiplier, + output_shift, output_activation_min, + output_activation_max, output_ptr, + output_ptr_offset); +} + +// Performs multiply-accumulate on a 4x3 input for 2 vertical outputs. +inline void DotProductAndStore2yStride1( + const Filter3x3x8& filter, int16x8_t i0, int16x8_t i1, int16x8_t i2, + int16x8_t i3, int16x8_t i4, int16x8_t i5, int16x8_t i6, int16x8_t i7, + int16x8_t i8, int16x8_t i9, int16x8_t i10, int16x8_t i11, + 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_ptr_offset) { + Int32x8 acc_0, acc_1; + acc_0.low = vld1q_s32(bias_ptr); + acc_1.low = vld1q_s32(bias_ptr); + acc_0.high = vld1q_s32(bias_ptr + 4); + acc_1.high = vld1q_s32(bias_ptr + 4); + + acc_0 = MultiplyAccumulate3x3Filter(filter, i0, i1, i2, i3, i4, i5, i6, i7, + i8, acc_0); + acc_1 = MultiplyAccumulate3x3Filter(filter, i3, i4, i5, i6, i7, i8, i9, i10, + i11, acc_1); + DownquantizeAndStore2Output(acc_0, acc_1, output_offset, output_multiplier, + output_shift, output_activation_min, + output_activation_max, output_ptr, + output_ptr_offset); +} + +// A kernel that is optimized on the number of output cells in the x and y +// direction, and the stride. Assumes 3x3 filters of 8 depth. +template +struct ConvKernel3x3FilterDepth8 {}; + +template <> +struct ConvKernel3x3FilterDepth8<8, 8, 1, 1> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + Filter3x3x8 filter = Load3x3Filter(filter_ptr, filter_offset, output_depth); + + const int16x8_t input_offset_vec = vdupq_n_s16(input_offset); + const int output_row_size = output_depth * output_width; + + // To process 8x8 outputs using a 3x3 filter, we require 10x10 inputs. + // Load inputs for the first 2 filters on the top left, then slide to + // the right, down, left, down, right, etc. in a snake-like path. This + // minimizes the total number of loads. + // + // INPUT OUTPUT + // |\----------------\ |\------------\ + // | \ \ | \ \ + // | \----------------\ | \------------\ + // | | 0 ... 9 | | | 0 ... 7 | + // | | 10 ... 19 | ---> | | 8 ... 15 | + // | | 20 ... 29 | \ | .. ... .. | + // \ | .. ... .. | \| 56 ... 63 | + // \| 90 ... 109 | |------------| + // |----------------| + // + // The first set of loads corresponds to: + // + // INPUT OUTPUT + // |\----------------- |\----------- + // | \ | \ + // | \----------------- | \---------- + // | | 0 1 2 3 ... | | 0 1 ... + // | | 10 11 12 13 ... ---> | | .. ... + // | | 20 21 22 23 ... | .. ... + // | | .. ... ... + // + // The next set of loads correspond to a sliding window to the right. + // It loads inputs 4, 5, 14, 15, 23, 24 and keeps 2, 3, 12, 13, and 22: + // + // INPUT OUTPUT + // |\------------------- |\------------- + // | \ | \ + // | \------------------- | \------------ + // | | .. 2 3 4 5 ... | | .. 2 3 ... + // | | .. 12 13 14 15 ... ---> | | .. ... + // | | .. 21 22 23 24 ... | .. ... + // | | .. ... ... + // + // And so on... + + int16x8_t input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11; + + // Load inputs for 1x2 outputs starting from the top left. Referring to the + // indexes in the diagram above, this corresponds to outputs (0) and (1). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth); + + // Slide to the right for outputs x = [2, 3], y = 0. Referring to the + // indexes in the diagram above, this corresponds to outputs (2) and (3). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 4 * input_depth; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_2, input_3, input_0, input_1, input_6, input_7, input_4, + input_5, input_10, input_11, input_8, input_9, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 2 * output_depth, output_depth); + + // Slide to the right again for outputs x = [4, 5], y = 0. Referring to the + // indexes in the diagram above, this corresponds to outputs (4) and (5). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 6 * input_depth; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 4 * output_depth, output_depth); + + // Slide to the right one last time for outputs x = [6, 7], y = 0. + // Referring to the indexes in the diagram above, this corresponds to + // outputs (6) and (7). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 8 * input_depth; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_2, input_3, input_0, input_1, input_6, input_7, input_4, + input_5, input_10, input_11, input_8, input_9, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 6 * output_depth, output_depth); + + // Slide to down for outputs x = [6, 7], y = 1. Referring to the indexes in + // the diagram above, this corresponds to outputs (14) and (15). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr + 6 * input_depth + 3 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_6, input_7, input_4, input_5, input_10, input_11, input_8, + input_9, input_2, input_3, input_0, input_1, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 6 * output_depth + output_row_size, + output_depth); + + // Slide left for outputs x = [4, 5], y = 1. Referring to the indexes in + // the diagram above, this corresponds to outputs (12) and (13). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 4 * input_depth + input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_4, input_5, input_6, input_7, input_8, input_9, input_10, + input_11, input_0, input_1, input_2, input_3, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 4 * output_depth + output_row_size, + output_depth); + + // Slide left again for outputs x = [2, 3], y = 1. Referring to the indexes + // in the diagram above, this corresponds to outputs (10) and (11). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 2 * input_depth + input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_6, input_7, input_4, input_5, input_10, input_11, input_8, + input_9, input_2, input_3, input_0, input_1, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 2 * output_depth + output_row_size, + output_depth); + + // Slide left one more time for outputs x = [0, 1], y = 1. Referring to the + // indexes in the diagram above, this corresponds to outputs (8) and (9). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_4, input_5, input_6, input_7, input_8, input_9, input_10, + input_11, input_0, input_1, input_2, input_3, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + output_row_size, output_depth); + + // Slide down for outputs x = [0, 1], y = 2. Referring to the + // indexes in the diagram above, this corresponds to outputs (16) and (17). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr + 4 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_8, input_9, input_10, input_11, input_0, input_1, input_2, + input_3, input_4, input_5, input_6, input_7, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 2 * output_row_size, output_depth); + + // Slide right for outputs x = [2, 3], y = 2. Referring to the + // indexes in the diagram above, this corresponds to outputs (18) and (19). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 4 * input_depth + 2 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_10, input_11, input_8, input_9, input_2, input_3, input_0, + input_1, input_6, input_7, input_4, input_5, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 2 * output_depth + 2 * output_row_size, output_depth); + + // Slide right for outputs x = [4, 5], y = 2. Referring to the + // indexes in the diagram above, this corresponds to outputs (20) and (21). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 6 * input_depth + 2 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_8, input_9, input_10, input_11, input_0, input_1, input_2, + input_3, input_4, input_5, input_6, input_7, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 4 * output_depth + 2 * output_row_size, output_depth); + + // Slide right one more time for outputs x = [6, 7], y = 2. Referring to the + // indexes in the diagram above, this corresponds to outputs (22) and (23). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 8 * input_depth + 2 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_10, input_11, input_8, input_9, input_2, input_3, input_0, + input_1, input_6, input_7, input_4, input_5, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 6 * output_depth + 2 * output_row_size, output_depth); + + // Slide down for outputs x = [6, 7], y = 3. Referring to the indexes in + // the diagram above, this corresponds to outputs (30) and (31). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr + 6 * input_depth + 5 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_2, input_3, input_0, input_1, input_6, input_7, input_4, + input_5, input_10, input_11, input_8, input_9, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 6 * output_depth + 3 * output_row_size, output_depth); + + // Slide left for outputs x = [4, 5], y = 3. Referring to the indexes in + // the diagram above, this corresponds to outputs (28) and (29). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 4 * input_depth + 3 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 4 * output_depth + 3 * output_row_size, output_depth); + + // Slide left for outputs x = [2, 3], y = 3. Referring to the indexes in + // the diagram above, this corresponds to outputs (26) and (27). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 2 * input_depth + 3 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_2, input_3, input_0, input_1, input_6, input_7, input_4, + input_5, input_10, input_11, input_8, input_9, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 2 * output_depth + 3 * output_row_size, output_depth); + + // Slide left one more time for outputs x = [0, 1], y = 3. Referring to the + // indexes in the diagram above, this corresponds to outputs (24) and (25). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 3 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 3 * output_row_size, output_depth); + + // Slide down for outputs x = [0, 1], y = 4. Referring to the indexes in + // the diagram above, this corresponds to outputs (32) and (33). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr + 6 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_4, input_5, input_6, input_7, input_8, input_9, input_10, + input_11, input_0, input_1, input_2, input_3, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 4 * output_row_size, output_depth); + + // Slide right for outputs x = [2, 3], y = 4. Referring to the indexes in + // the diagram above, this corresponds to outputs (34) and (35). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 4 * input_depth + 4 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_6, input_7, input_4, input_5, input_10, input_11, input_8, + input_9, input_2, input_3, input_0, input_1, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 2 * output_depth + 4 * output_row_size, output_depth); + + // Slide right for outputs x = [4, 5], y = 4. Referring to the indexes in + // the diagram above, this corresponds to outputs (36) and (37). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 6 * input_depth + 4 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_4, input_5, input_6, input_7, input_8, input_9, input_10, + input_11, input_0, input_1, input_2, input_3, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 4 * output_depth + 4 * output_row_size, output_depth); + + // Slide right one more time for outputs x = [6, 7], y = 4. Referring to the + // indexes in the diagram above, this corresponds to outputs (38) and (39). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 8 * input_depth + 4 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_6, input_7, input_4, input_5, input_10, input_11, input_8, + input_9, input_2, input_3, input_0, input_1, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 6 * output_depth + 4 * output_row_size, output_depth); + + // Slide down for outputs x = [6, 7], y = 5. Referring to the indexes in + // the diagram above, this corresponds to outputs (46) and (47). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr + 6 * input_depth + 7 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_10, input_11, input_8, input_9, input_2, input_3, input_0, + input_1, input_6, input_7, input_4, input_5, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 6 * output_depth + 5 * output_row_size, output_depth); + + // Slide left for outputs x = [4, 5], y = 5. Referring to the indexes in + // the diagram above, this corresponds to outputs (44) and (45). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 4 * input_depth + 5 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_8, input_9, input_10, input_11, input_0, input_1, input_2, + input_3, input_4, input_5, input_6, input_7, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 4 * output_depth + 5 * output_row_size, output_depth); + + // Slide left for outputs x = [2, 3], y = 5. Referring to the indexes in + // the diagram above, this corresponds to outputs (42) and (43). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 2 * input_depth + 5 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_10, input_11, input_8, input_9, input_2, input_3, input_0, + input_1, input_6, input_7, input_4, input_5, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 2 * output_depth + 5 * output_row_size, output_depth); + + // Slide left one more time for outputs x = [0, 1], y = 5. Referring to the + // indexes in the diagram above, this corresponds to outputs (40) and (41). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 5 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_8, input_9, input_10, input_11, input_0, input_1, input_2, + input_3, input_4, input_5, input_6, input_7, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 5 * output_row_size, output_depth); + + // Slide down for outputs x = [0, 1], y = 6. Referring to the indexes in + // the diagram above, this corresponds to outputs (48) and (49). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr + 8 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 6 * output_row_size, output_depth); + + // Slide right for outputs x = [2, 3], y = 6. Referring to the indexes in + // the diagram above, this corresponds to outputs (50) and (51). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 4 * input_depth + 6 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_2, input_3, input_0, input_1, input_6, input_7, input_4, + input_5, input_10, input_11, input_8, input_9, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 2 * output_depth + 6 * output_row_size, output_depth); + + // Slide right for outputs x = [4, 5], y = 6. Referring to the indexes in + // the diagram above, this corresponds to outputs (52) and (53). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 6 * input_depth + 6 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 4 * output_depth + 6 * output_row_size, output_depth); + + // Slide right one more time for outputs x = [6, 7], y = 6. Referring to the + // indexes in the diagram above, this corresponds to outputs (54) and (55). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 8 * input_depth + 6 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_2, input_3, input_0, input_1, input_6, input_7, input_4, + input_5, input_10, input_11, input_8, input_9, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 6 * output_depth + 6 * output_row_size, output_depth); + + // Slide down for outputs x = [6, 7], y = 7. Referring to the indexes in the + // diagram above, this corresponds to outputs (62) and (63). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr + 6 * input_depth + 9 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_6, input_7, input_4, input_5, input_10, input_11, input_8, + input_9, input_2, input_3, input_0, input_1, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 6 * output_depth + 7 * output_row_size, output_depth); + + // Slide left for outputs x = [4, 5], y = 7. Referring to the indexes in the + // diagram above, this corresponds to outputs (60) and (61). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 4 * input_depth + 7 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_4, input_5, input_6, input_7, input_8, input_9, input_10, + input_11, input_0, input_1, input_2, input_3, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 4 * output_depth + 7 * output_row_size, output_depth); + + // Slide left for outputs x = [2, 3], y = 7. Referring to the indexes in the + // diagram above, this corresponds to outputs (58) and (59). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 2 * input_depth + 7 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_6, input_7, input_4, input_5, input_10, input_11, input_8, + input_9, input_2, input_3, input_0, input_1, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 2 * output_depth + 7 * output_row_size, output_depth); + + // Slide left one more time for outputs x = [0, 1], y = 7. Referring to the + // indexes in the diagram above, this corresponds to outputs (56) and (57). + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 7 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_4, input_5, input_6, input_7, input_8, input_9, input_10, + input_11, input_0, input_1, input_2, input_3, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 7 * output_row_size, output_depth); + } +}; + +template <> +struct ConvKernel3x3FilterDepth8<4, 4, 1, 1> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + Filter3x3x8 filter = Load3x3Filter(filter_ptr, filter_offset, output_depth); + + const int16x8_t input_offset_vec = vdupq_n_s16(input_offset); + const int output_row_size = output_depth * output_width; + + // To process 4x4 outputs using a 3x3 filter, we require 6x6 inputs. + // Load inputs for the first 2 filters on the top left, then slide to + // the right, down, left, down, right, etc. in a snake-like path. This + // minimizes the total number of loads. + int16x8_t input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11; + + // Load inputs for 1x2 outputs starting from the top left. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth); + + // Now load 1x2 inputs on the top right. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 4 * input_depth; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_2, input_3, input_0, input_1, input_6, input_7, input_4, + input_5, input_10, input_11, input_8, input_9, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 2 * output_depth, output_depth); + + // Now load next inputs when sliding window down. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr + 2 * input_depth + 3 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_6, input_7, input_4, input_5, input_10, input_11, input_8, + input_9, input_2, input_3, input_0, input_1, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 2 * output_depth + output_row_size, + output_depth); + + // Now load next inputs when sliding window left. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_4, input_5, input_6, input_7, input_8, input_9, input_10, + input_11, input_0, input_1, input_2, input_3, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + output_row_size, output_depth); + + // Now load next inputs when sliding window down. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr + 4 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_8, input_9, input_10, input_11, input_0, input_1, input_2, + input_3, input_4, input_5, input_6, input_7, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 2 * output_row_size, output_depth); + + // Now load next inputs when sliding window right. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 4 * input_depth + 2 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_10, input_11, input_8, input_9, input_2, input_3, input_0, + input_1, input_6, input_7, input_4, input_5, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 2 * output_depth + 2 * output_row_size, output_depth); + + // Now load next inputs when sliding window down. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr + 2 * input_depth + 5 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_2, input_3, input_0, input_1, input_6, input_7, input_4, + input_5, input_10, input_11, input_8, input_9, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, + output_ptr + 2 * output_depth + 3 * output_row_size, output_depth); + + // Now load next inputs when sliding window left. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 3 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 3 * output_row_size, output_depth); + } +}; + +template <> +struct ConvKernel3x3FilterDepth8<4, 2, 1, 1> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + Filter3x3x8 filter = Load3x3Filter(filter_ptr, filter_offset, output_depth); + + const int16x8_t input_offset_vec = vdupq_n_s16(input_offset); + const int output_row_size = output_depth * output_width; + + int16x8_t input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11; + + // Load inputs for 1x2 outputs starting from the top. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth); + + output_ptr += output_row_size; + + // Now load next inputs one row down. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr + 3 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_4, input_5, input_6, input_7, input_8, input_9, input_10, + input_11, input_0, input_1, input_2, input_3, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth); + + output_ptr += output_row_size; + + // Now load next row. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr + 4 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_8, input_9, input_10, input_11, input_0, input_1, input_2, + input_3, input_4, input_5, input_6, input_7, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth); + + output_ptr += output_row_size; + + // Now load last row. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr + 5 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth); + } +}; + +template <> +struct ConvKernel3x3FilterDepth8<4, 1, 1, 1> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + Filter3x3x8 filter = Load3x3Filter(filter_ptr, filter_offset, output_depth); + + const int16x8_t input_offset_vec = vdupq_n_s16(input_offset); + const int output_row_size = output_depth * output_width; + + int16x8_t input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11; + + // Load inputs for 2x1 outputs starting from the top. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_3 = vld1_u8(ptr); + temp_4 = vld1_u8(ptr + input_depth); + temp_5 = vld1_u8(ptr + 2 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_3 = vld1_u8(ptr); + temp_4 = vld1_u8(ptr + input_depth); + temp_5 = vld1_u8(ptr + 2 * input_depth); + + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + } + + DotProductAndStore2yStride1( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_row_size); + + // Load inputs for bottom 2 rows. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 4 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_3 = vld1_u8(ptr); + temp_4 = vld1_u8(ptr + input_depth); + temp_5 = vld1_u8(ptr + 2 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + } + + DotProductAndStore2yStride1( + filter, input_6, input_7, input_8, input_9, input_10, input_11, input_0, + input_1, input_2, input_3, input_4, input_5, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 2 * output_row_size, + output_row_size); + } +}; + +template <> +struct ConvKernel3x3FilterDepth8<2, 2, 1, 1> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + Filter3x3x8 filter = Load3x3Filter(filter_ptr, filter_offset, output_depth); + + Int32x8 acc_0, acc_1, acc_2, acc_3; + + acc_0.low = vld1q_s32(bias_ptr); + acc_1.low = vld1q_s32(bias_ptr); + acc_2.low = vld1q_s32(bias_ptr); + acc_3.low = vld1q_s32(bias_ptr); + + bias_ptr += 4; + acc_0.high = vld1q_s32(bias_ptr); + acc_1.high = vld1q_s32(bias_ptr); + acc_2.high = vld1q_s32(bias_ptr); + acc_3.high = vld1q_s32(bias_ptr); + + const int16x8_t input_offset_vec = vdupq_n_s16(input_offset); + + // Add scope for input registers to help the compiler know that it is + // not needed. + { + // To process 2x2 outputs using a 3x3 filter, we require 4x4 inputs. + // Load inputs for the top two filters first. + int16x8_t input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11; + + const uint8* ptr = input_ptr; + + // Load top 3 rows. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + } + + // Multiply-accum for top-left output. + acc_0 = MultiplyAccumulate3x3Filter(filter, input_0, input_1, input_2, + input_4, input_5, input_6, input_8, + input_9, input_10, acc_0); + + // Multiply-accum for top-right output. + acc_1 = MultiplyAccumulate3x3Filter(filter, input_1, input_2, input_3, + input_5, input_6, input_7, input_9, + input_10, input_11, acc_1); + + // Now load the bottom row. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + } + + // Multiply-accum for bottom-left output. + acc_2 = MultiplyAccumulate3x3Filter(filter, input_4, input_5, input_6, + input_8, input_9, input_10, input_0, + input_1, input_2, acc_2); + + // Multiply-accum for bottom-right output. + acc_3 = MultiplyAccumulate3x3Filter(filter, input_5, input_6, input_7, + input_9, input_10, input_11, input_1, + input_2, input_3, acc_3); + } + + DownquantizeAndStore2x2Output(acc_0, acc_1, acc_2, acc_3, output_offset, + output_multiplier, output_shift, + output_activation_min, output_activation_max, + output_ptr, output_depth, output_width); + } +}; + +template <> +struct ConvKernel3x3FilterDepth8<2, 4, 1, 1> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + Filter3x3x8 filter = Load3x3Filter(filter_ptr, filter_offset, output_depth); + + const int16x8_t input_offset_vec = vdupq_n_s16(input_offset); + const int output_row_size = output_depth * output_width; + + int16x8_t input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11; + + // Load inputs for 1x2 outputs starting from the top left. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth); + + // Now load 1x2 inputs on the top right. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + 4 * input_depth; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_2, input_3, input_0, input_1, input_6, input_7, input_4, + input_5, input_10, input_11, input_8, input_9, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 2 * output_depth, output_depth); + + // Now load next inputs when sliding window down. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr + 2 * input_depth + 3 * input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_6, input_7, input_4, input_5, input_10, input_11, input_8, + input_9, input_2, input_3, input_0, input_1, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 2 * output_depth + output_row_size, + output_depth); + + // Now load next inputs when sliding window left. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_4, input_5, input_6, input_7, input_8, input_9, input_10, + input_11, input_0, input_1, input_2, input_3, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + output_row_size, output_depth); + } +}; + +template <> +struct ConvKernel3x3FilterDepth8<1, 4, 1, 1> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + Filter3x3x8 filter = Load3x3Filter(filter_ptr, filter_offset, output_depth); + + const int16x8_t input_offset_vec = vdupq_n_s16(input_offset); + + int16x8_t input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11; + + // Load inputs for 1x2 outputs starting from the left. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3; + + const uint8* ptr = input_ptr; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth); + + // Now load 1x2 inputs on the right. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr + input_depth * 4; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_2 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + } + + DotProductAndStore2xStride1( + filter, input_2, input_3, input_0, input_1, input_6, input_7, input_4, + input_5, input_10, input_11, input_8, input_9, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr + 2 * output_depth, output_depth); + } +}; + +template <> +struct ConvKernel3x3FilterDepth8<2, 1, 1, 1> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + Filter3x3x8 filter = Load3x3Filter(filter_ptr, filter_offset, output_depth); + + // To process 2x1 outputs using a 3x3 filter, we require 4x3 inputs. + // Load all inputs at the beginning. + int16x8_t input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11; + + // Load inputs for 1x2 outputs starting from the top left. + { + const int16x8_t input_offset_vec = vdupq_n_s16(input_offset); + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5; + + const uint8* ptr = input_ptr; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_3 = vld1_u8(ptr); + temp_4 = vld1_u8(ptr + input_depth); + temp_5 = vld1_u8(ptr + 2 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_3 = vld1_u8(ptr); + temp_4 = vld1_u8(ptr + input_depth); + temp_5 = vld1_u8(ptr + 2 * input_depth); + + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_10 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_11 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + input_10 = vaddq_s16(input_10, input_offset_vec); + input_11 = vaddq_s16(input_11, input_offset_vec); + } + + DotProductAndStore2yStride1( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9, input_10, input_11, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth * output_width); + } +}; + +template <> +struct ConvKernel3x3FilterDepth8<4, 2, 2, 2> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + const int output_row_size = output_depth * output_width; + + Filter3x3x8 filter = Load3x3Filter(filter_ptr, filter_offset, output_depth); + + Int32x8 acc_0, acc_1; + acc_0.low = vld1q_s32(bias_ptr); + acc_1.low = vld1q_s32(bias_ptr); + acc_0.high = vld1q_s32(bias_ptr + 4); + acc_1.high = vld1q_s32(bias_ptr + 4); + + const int16x8_t input_offset_vec = vdupq_n_s16(input_offset); + + int16x8_t input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9; + + const uint8* ptr = input_ptr; + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4; + + // Load first 2 rows. + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + temp_4 = vld1_u8(ptr + 4 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + temp_4 = vld1_u8(ptr + 4 * input_depth); + + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + + acc_0 = MultiplyAccumulateRow(acc_0, filter.f0, filter.f1, filter.f2, + input_0, input_1, input_2); + + acc_1 = MultiplyAccumulateRow(acc_1, filter.f0, filter.f1, filter.f2, + input_2, input_3, input_4); + + acc_0 = MultiplyAccumulateRow(acc_0, filter.f3, filter.f4, filter.f5, + input_5, input_6, input_7); + + acc_1 = MultiplyAccumulateRow(acc_1, filter.f3, filter.f4, filter.f5, + input_7, input_8, input_9); + + // Load next 2 rows. + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + temp_4 = vld1_u8(ptr + 4 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + temp_4 = vld1_u8(ptr + 4 * input_depth); + + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + + acc_0 = MultiplyAccumulateRow(acc_0, filter.f6, filter.f7, filter.f8, + input_0, input_1, input_2); + + acc_1 = MultiplyAccumulateRow(acc_1, filter.f6, filter.f7, filter.f8, + input_2, input_3, input_4); + + DownquantizeAndStore2Output( + acc_0, acc_1, output_offset, output_multiplier, output_shift, + output_activation_min, output_activation_max, output_ptr, output_depth); + + output_ptr += output_row_size; + + // Moving onto the next row of outputs. + acc_0.low = vld1q_s32(bias_ptr); + acc_1.low = vld1q_s32(bias_ptr); + acc_0.high = vld1q_s32(bias_ptr + 4); + acc_1.high = vld1q_s32(bias_ptr + 4); + + acc_0 = MultiplyAccumulateRow(acc_0, filter.f0, filter.f1, filter.f2, + input_0, input_1, input_2); + + acc_1 = MultiplyAccumulateRow(acc_1, filter.f0, filter.f1, filter.f2, + input_2, input_3, input_4); + + acc_0 = MultiplyAccumulateRow(acc_0, filter.f3, filter.f4, filter.f5, + input_5, input_6, input_7); + + acc_1 = MultiplyAccumulateRow(acc_1, filter.f3, filter.f4, filter.f5, + input_7, input_8, input_9); + + // Load next 2 rows. + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + temp_4 = vld1_u8(ptr + 4 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + temp_4 = vld1_u8(ptr + 4 * input_depth); + + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + + acc_0 = MultiplyAccumulateRow(acc_0, filter.f6, filter.f7, filter.f8, + input_0, input_1, input_2); + + acc_1 = MultiplyAccumulateRow(acc_1, filter.f6, filter.f7, filter.f8, + input_2, input_3, input_4); + + DownquantizeAndStore2Output( + acc_0, acc_1, output_offset, output_multiplier, output_shift, + output_activation_min, output_activation_max, output_ptr, output_depth); + + output_ptr += output_row_size; + + // Moving onto the next row of outputs. + acc_0.low = vld1q_s32(bias_ptr); + acc_1.low = vld1q_s32(bias_ptr); + acc_0.high = vld1q_s32(bias_ptr + 4); + acc_1.high = vld1q_s32(bias_ptr + 4); + + acc_0 = MultiplyAccumulateRow(acc_0, filter.f0, filter.f1, filter.f2, + input_0, input_1, input_2); + + acc_1 = MultiplyAccumulateRow(acc_1, filter.f0, filter.f1, filter.f2, + input_2, input_3, input_4); + + acc_0 = MultiplyAccumulateRow(acc_0, filter.f3, filter.f4, filter.f5, + input_5, input_6, input_7); + + acc_1 = MultiplyAccumulateRow(acc_1, filter.f3, filter.f4, filter.f5, + input_7, input_8, input_9); + + // Load next 2 rows. + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + temp_4 = vld1_u8(ptr + 4 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + temp_4 = vld1_u8(ptr + 4 * input_depth); + + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + + acc_0 = MultiplyAccumulateRow(acc_0, filter.f6, filter.f7, filter.f8, + input_0, input_1, input_2); + + acc_1 = MultiplyAccumulateRow(acc_1, filter.f6, filter.f7, filter.f8, + input_2, input_3, input_4); + + DownquantizeAndStore2Output( + acc_0, acc_1, output_offset, output_multiplier, output_shift, + output_activation_min, output_activation_max, output_ptr, output_depth); + + output_ptr += output_row_size; + + // Moving onto the next row of outputs. + acc_0.low = vld1q_s32(bias_ptr); + acc_1.low = vld1q_s32(bias_ptr); + acc_0.high = vld1q_s32(bias_ptr + 4); + acc_1.high = vld1q_s32(bias_ptr + 4); + + acc_0 = MultiplyAccumulateRow(acc_0, filter.f0, filter.f1, filter.f2, + input_0, input_1, input_2); + + acc_1 = MultiplyAccumulateRow(acc_1, filter.f0, filter.f1, filter.f2, + input_2, input_3, input_4); + + acc_0 = MultiplyAccumulateRow(acc_0, filter.f3, filter.f4, filter.f5, + input_5, input_6, input_7); + + acc_1 = MultiplyAccumulateRow(acc_1, filter.f3, filter.f4, filter.f5, + input_7, input_8, input_9); + + // Load last row. + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + temp_4 = vld1_u8(ptr + 4 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + + acc_0 = MultiplyAccumulateRow(acc_0, filter.f6, filter.f7, filter.f8, + input_0, input_1, input_2); + + acc_1 = MultiplyAccumulateRow(acc_1, filter.f6, filter.f7, filter.f8, + input_2, input_3, input_4); + + DownquantizeAndStore2Output( + acc_0, acc_1, output_offset, output_multiplier, output_shift, + output_activation_min, output_activation_max, output_ptr, output_depth); + } +}; + +template <> +struct ConvKernel3x3FilterDepth8<4, 4, 2, 2> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + // Reuse 4x2 kernel twice. + ConvKernel3x3FilterDepth8<4, 2, 2, 2>::Run( + input_ptr, input_depth, input_offset, input_row_size, filter_ptr, + filter_offset, bias_ptr, output_offset, output_multiplier, output_shift, + output_activation_min, output_activation_max, output_ptr, output_depth, + output_width); + + ConvKernel3x3FilterDepth8<4, 2, 2, 2>::Run( + input_ptr + 4 * input_depth, input_depth, input_offset, input_row_size, + filter_ptr, filter_offset, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, + output_ptr + 2 * output_depth, output_depth, output_width); + } +}; + +template <> +struct ConvKernel3x3FilterDepth8<4, 1, 2, 2> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + const int output_row_size = output_depth * output_width; + + Filter3x3x8 filter = Load3x3Filter(filter_ptr, filter_offset, output_depth); + + const int16x8_t input_offset_vec = vdupq_n_s16(input_offset); + int16x8_t input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8; + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5, temp_6, temp_7, + temp_8; + + const uint8* ptr = input_ptr; + + // Load all inputs for top output. + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_3 = vld1_u8(ptr); + temp_4 = vld1_u8(ptr + input_depth); + temp_5 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_6 = vld1_u8(ptr); + temp_7 = vld1_u8(ptr + input_depth); + temp_8 = vld1_u8(ptr + 2 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_6)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_7)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_8)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + + DotProductAndStore( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, output_ptr); + + // Second output. + output_ptr += output_row_size; + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_3 = vld1_u8(ptr); + temp_4 = vld1_u8(ptr + input_depth); + temp_5 = vld1_u8(ptr + 2 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + + DotProductAndStore( + filter, input_6, input_7, input_8, input_0, input_1, input_2, input_3, + input_4, input_5, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, output_ptr); + + // Third output. + output_ptr += output_row_size; + + ptr += input_row_size; + temp_6 = vld1_u8(ptr); + temp_7 = vld1_u8(ptr + input_depth); + temp_8 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_6)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_7)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_8)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + + DotProductAndStore( + filter, input_3, input_4, input_5, input_6, input_7, input_8, input_0, + input_1, input_2, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, output_ptr); + + // Fourth output. + output_ptr += output_row_size; + + ptr += input_row_size; + temp_3 = vld1_u8(ptr); + temp_4 = vld1_u8(ptr + input_depth); + temp_5 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_6 = vld1_u8(ptr); + temp_7 = vld1_u8(ptr + input_depth); + temp_8 = vld1_u8(ptr + 2 * input_depth); + + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_6)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_7)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_8)); + + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + + DotProductAndStore( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, output_ptr); + } +}; + +template <> +struct ConvKernel3x3FilterDepth8<2, 2, 2, 2> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + Filter3x3x8 filter = Load3x3Filter(filter_ptr, filter_offset, output_depth); + + Int32x8 acc_0, acc_1, acc_2, acc_3; + acc_0.low = vld1q_s32(bias_ptr); + acc_1.low = vld1q_s32(bias_ptr); + acc_2.low = vld1q_s32(bias_ptr); + acc_3.low = vld1q_s32(bias_ptr); + + bias_ptr += 4; + acc_0.high = vld1q_s32(bias_ptr); + acc_1.high = vld1q_s32(bias_ptr); + acc_2.high = vld1q_s32(bias_ptr); + acc_3.high = vld1q_s32(bias_ptr); + + const int16x8_t input_offset_vec = vdupq_n_s16(input_offset); + + // Add scope for input registers to help the compiler know that it is + // not needed. + { + // To process 2x2 outputs using a 3x3 filter at stride 2, we require + // 5x5 inputs. We load the first 5x2 inputs at a time. + int16x8_t input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, input_9; + + const uint8* ptr = input_ptr; + + // Load inputs. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4; + + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + temp_4 = vld1_u8(ptr + 4 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + temp_4 = vld1_u8(ptr + 4 * input_depth); + + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + } + + acc_0 = MultiplyAccumulateRow(acc_0, filter.f0, filter.f1, filter.f2, + input_0, input_1, input_2); + + acc_1 = MultiplyAccumulateRow(acc_1, filter.f0, filter.f1, filter.f2, + input_2, input_3, input_4); + + acc_0 = MultiplyAccumulateRow(acc_0, filter.f3, filter.f4, filter.f5, + input_5, input_6, input_7); + + acc_1 = MultiplyAccumulateRow(acc_1, filter.f3, filter.f4, filter.f5, + input_7, input_8, input_9); + + // Load next inputs. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4; + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + temp_4 = vld1_u8(ptr + 4 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + temp_4 = vld1_u8(ptr + 4 * input_depth); + + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_9 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_9 = vaddq_s16(input_9, input_offset_vec); + } + + acc_0 = MultiplyAccumulateRow(acc_0, filter.f6, filter.f7, filter.f8, + input_0, input_1, input_2); + + acc_1 = MultiplyAccumulateRow(acc_1, filter.f6, filter.f7, filter.f8, + input_2, input_3, input_4); + + // Moving onto the two bottom outputs. + acc_2 = MultiplyAccumulateRow(acc_2, filter.f0, filter.f1, filter.f2, + input_0, input_1, input_2); + + acc_3 = MultiplyAccumulateRow(acc_3, filter.f0, filter.f1, filter.f2, + input_2, input_3, input_4); + + acc_2 = MultiplyAccumulateRow(acc_2, filter.f3, filter.f4, filter.f5, + input_5, input_6, input_7); + + acc_3 = MultiplyAccumulateRow(acc_3, filter.f3, filter.f4, filter.f5, + input_7, input_8, input_9); + + // Load last input row. + { + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4; + + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + temp_3 = vld1_u8(ptr + 3 * input_depth); + temp_4 = vld1_u8(ptr + 4 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + } + + acc_2 = MultiplyAccumulateRow(acc_2, filter.f6, filter.f7, filter.f8, + input_0, input_1, input_2); + + acc_3 = MultiplyAccumulateRow(acc_3, filter.f6, filter.f7, filter.f8, + input_2, input_3, input_4); + } + + DownquantizeAndStore2x2Output(acc_0, acc_1, acc_2, acc_3, output_offset, + output_multiplier, output_shift, + output_activation_min, output_activation_max, + output_ptr, output_depth, output_width); + } }; -struct Filter3x3x16 { - Int16x16x3 r0, r1, r2; +template <> +struct ConvKernel3x3FilterDepth8<2, 4, 2, 2> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + // Reuse 2x2 kernel twice. + ConvKernel3x3FilterDepth8<2, 2, 2, 2>::Run( + input_ptr, input_depth, input_offset, input_row_size, filter_ptr, + filter_offset, bias_ptr, output_offset, output_multiplier, output_shift, + output_activation_min, output_activation_max, output_ptr, output_depth, + output_width); + + ConvKernel3x3FilterDepth8<2, 2, 2, 2>::Run( + input_ptr + 4 * input_depth, input_depth, input_offset, input_row_size, + filter_ptr, filter_offset, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, + output_ptr + 2 * output_depth, output_depth, output_width); + } }; -// 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; +template <> +struct ConvKernel3x3FilterDepth8<2, 1, 2, 2> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + const int output_row_size = output_depth * output_width; - 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); + Filter3x3x8 filter = Load3x3Filter(filter_ptr, filter_offset, output_depth); - 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); + const int16x8_t input_offset_vec = vdupq_n_s16(input_offset); + int16x8_t input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8; + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5, temp_6, temp_7, + temp_8; - return filter; -} + const uint8* ptr = input_ptr; -// 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; -} + // Load all inputs for top output. + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_3 = vld1_u8(ptr); + temp_4 = vld1_u8(ptr + input_depth); + temp_5 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_6 = vld1_u8(ptr); + temp_7 = vld1_u8(ptr + input_depth); + temp_8 = vld1_u8(ptr + 2 * input_depth); -// 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; -} + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_6)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_7)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_8)); -// 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); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); - 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); + DotProductAndStore( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, output_ptr); - // 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); + // Second output. + output_ptr += output_row_size; - // 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); + ptr += input_row_size; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_3 = vld1_u8(ptr); + temp_4 = vld1_u8(ptr + input_depth); + temp_5 = vld1_u8(ptr + 2 * input_depth); - 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); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); - // 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)); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + + DotProductAndStore( + filter, input_6, input_7, input_8, input_0, input_1, input_2, input_3, + input_4, input_5, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, output_ptr); + } +}; + +template <> +struct ConvKernel3x3FilterDepth8<1, 2, 2, 2> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + Filter3x3x8 filter = Load3x3Filter(filter_ptr, filter_offset, output_depth); + + const int16x8_t input_offset_vec = vdupq_n_s16(input_offset); + int16x8_t input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8; + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5, temp_6, temp_7, + temp_8; + + const uint8* ptr = input_ptr; + + // Load all inputs for top output. + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_3 = vld1_u8(ptr); + temp_4 = vld1_u8(ptr + input_depth); + temp_5 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_6 = vld1_u8(ptr); + temp_7 = vld1_u8(ptr + input_depth); + temp_8 = vld1_u8(ptr + 2 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_6)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_7)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_8)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + + DotProductAndStore( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, output_ptr); + + // Second output. + output_ptr += output_depth; + + ptr = input_ptr + 3 * input_depth; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + ptr += input_row_size; + temp_3 = vld1_u8(ptr); + temp_4 = vld1_u8(ptr + input_depth); + ptr += input_row_size; + temp_6 = vld1_u8(ptr); + temp_7 = vld1_u8(ptr + input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_6)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_7)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + + DotProductAndStore( + filter, input_2, input_0, input_1, input_5, input_3, input_4, input_8, + input_6, input_7, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, output_ptr); + } +}; + +template <> +struct ConvKernel3x3FilterDepth8<1, 4, 2, 2> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + Filter3x3x8 filter = Load3x3Filter(filter_ptr, filter_offset, output_depth); + + const int16x8_t input_offset_vec = vdupq_n_s16(input_offset); + int16x8_t input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8; + uint8x8_t temp_0, temp_1, temp_2, temp_3, temp_4, temp_5, temp_6, temp_7, + temp_8; + + const uint8* ptr = input_ptr; + + // Load all inputs for top output. + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + temp_2 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_3 = vld1_u8(ptr); + temp_4 = vld1_u8(ptr + input_depth); + temp_5 = vld1_u8(ptr + 2 * input_depth); + ptr += input_row_size; + temp_6 = vld1_u8(ptr); + temp_7 = vld1_u8(ptr + input_depth); + temp_8 = vld1_u8(ptr + 2 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_6)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_7)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_8)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + + DotProductAndStore( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, output_ptr); + + // Second output. + output_ptr += output_depth; + + ptr = input_ptr + 3 * input_depth; + temp_0 = vld1_u8(ptr); + temp_1 = vld1_u8(ptr + input_depth); + ptr += input_row_size; + temp_3 = vld1_u8(ptr); + temp_4 = vld1_u8(ptr + input_depth); + ptr += input_row_size; + temp_6 = vld1_u8(ptr); + temp_7 = vld1_u8(ptr + input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_6)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_7)); + + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + + DotProductAndStore( + filter, input_2, input_0, input_1, input_5, input_3, input_4, input_8, + input_6, input_7, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, output_ptr); + + // Third output. + output_ptr += output_depth; + + ptr = input_ptr + 5 * input_depth; + temp_2 = vld1_u8(ptr); + temp_0 = vld1_u8(ptr + input_depth); + ptr += input_row_size; + temp_5 = vld1_u8(ptr); + temp_3 = vld1_u8(ptr + input_depth); + ptr += input_row_size; + temp_8 = vld1_u8(ptr); + temp_6 = vld1_u8(ptr + input_depth); + + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_8)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_6)); + + input_2 = vaddq_s16(input_2, input_offset_vec); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + + DotProductAndStore( + filter, input_1, input_2, input_0, input_4, input_5, input_3, input_7, + input_8, input_6, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, output_ptr); + + // Fourth output. + output_ptr += output_depth; + + ptr = input_ptr + 7 * input_depth; + temp_1 = vld1_u8(ptr); + temp_2 = vld1_u8(ptr + input_depth); + ptr += input_row_size; + temp_4 = vld1_u8(ptr); + temp_5 = vld1_u8(ptr + input_depth); + ptr += input_row_size; + temp_7 = vld1_u8(ptr); + temp_8 = vld1_u8(ptr + input_depth); + + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_7)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_8)); + + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + + DotProductAndStore( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, output_ptr); + } +}; + +template +struct ConvKernel3x3FilterDepth8<1, 1, kFixedStrideWidth, kFixedStrideHeight> { + static inline void Run(const uint8* input_ptr, int input_depth, + int32 input_offset, int input_row_size, + const uint8* filter_ptr, int32 filter_offset, + 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) { + Filter3x3x8 filter = Load3x3Filter(filter_ptr, filter_offset, output_depth); + + int16x8_t input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8; + + uint8x8_t temp_0 = vld1_u8(input_ptr); + uint8x8_t temp_1 = vld1_u8(input_ptr + input_depth); + uint8x8_t temp_2 = vld1_u8(input_ptr + 2 * input_depth); + + input_ptr += input_row_size; + uint8x8_t temp_3 = vld1_u8(input_ptr); + uint8x8_t temp_4 = vld1_u8(input_ptr + input_depth); + uint8x8_t temp_5 = vld1_u8(input_ptr + 2 * input_depth); + + input_ptr += input_row_size; + uint8x8_t temp_6 = vld1_u8(input_ptr); + uint8x8_t temp_7 = vld1_u8(input_ptr + input_depth); + uint8x8_t temp_8 = vld1_u8(input_ptr + 2 * input_depth); + + input_0 = vreinterpretq_s16_u16(vmovl_u8(temp_0)); + input_1 = vreinterpretq_s16_u16(vmovl_u8(temp_1)); + input_2 = vreinterpretq_s16_u16(vmovl_u8(temp_2)); + input_3 = vreinterpretq_s16_u16(vmovl_u8(temp_3)); + input_4 = vreinterpretq_s16_u16(vmovl_u8(temp_4)); + input_5 = vreinterpretq_s16_u16(vmovl_u8(temp_5)); + input_6 = vreinterpretq_s16_u16(vmovl_u8(temp_6)); + input_7 = vreinterpretq_s16_u16(vmovl_u8(temp_7)); + input_8 = vreinterpretq_s16_u16(vmovl_u8(temp_8)); + + const int16x8_t input_offset_vec = vdupq_n_s16(input_offset); + input_0 = vaddq_s16(input_0, input_offset_vec); + input_1 = vaddq_s16(input_1, input_offset_vec); + input_2 = vaddq_s16(input_2, input_offset_vec); + input_3 = vaddq_s16(input_3, input_offset_vec); + input_4 = vaddq_s16(input_4, input_offset_vec); + input_5 = vaddq_s16(input_5, input_offset_vec); + input_6 = vaddq_s16(input_6, input_offset_vec); + input_7 = vaddq_s16(input_7, input_offset_vec); + input_8 = vaddq_s16(input_8, input_offset_vec); + + DotProductAndStore( + filter, input_0, input_1, input_2, input_3, input_4, input_5, input_6, + input_7, input_8, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, output_ptr); + } +}; + +inline void ShuffleInput(const uint8* input_ptr, int input_depth, + int input_width, int input_height, int output_depth, + int output_width, int output_height, + uint8* output_ptr) { + const int input_row_size = input_depth * input_width; + + for (int y = 0; y < output_height; y++) { + const uint8* ptr = input_ptr; + for (int x = 0; x < output_width; x++) { + memcpy(output_ptr, ptr, output_depth); + output_ptr += output_depth; + ptr += input_depth; + } + input_ptr += input_row_size; + } } -// 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 ConvRow3x3FilterDepth8 {}; + +template +struct ConvRow3x3FilterDepth8<1, kFixedStrideWidth, kFixedStrideHeight> { + static inline void Run(const uint8* input_data, int start_x, int start_y, + int input_depth, int input_width, int input_height, + int input_row_size, int32 input_offset, + const uint8* filter_data, int32 filter_offset, + const int32* bias_data, int32 output_offset, + int32 output_multiplier, int output_shift, + int32 output_activation_min, + int32 output_activation_max, uint8* output_data, + int output_depth, int output_width, + uint8* shuffle_workspace) { + int out_x = start_x; + + // 1x4 at a time. + for (; out_x <= output_width - 4; out_x += 4) { + const int32* bias_ptr = bias_data; + const uint8* filter_ptr = filter_data; + + const uint8* input_ptr = input_data; + uint8* output_ptr = output_data; + + for (int depth = 0; depth <= output_depth - 8; depth += 8) { + ConvKernel3x3FilterDepth8<1, 4, kFixedStrideWidth, kFixedStrideHeight>:: + Run(input_ptr, input_depth, input_offset, input_row_size, + filter_ptr, filter_offset, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth, output_width); + + input_ptr += 8; + output_ptr += 8; + filter_ptr += 8; + bias_ptr += 8; + } + + input_data += 4 * kFixedStrideWidth * input_depth; + output_data += 4 * output_depth; + } + + // 1x1 at a time. + for (; out_x < output_width; out_x++) { + const int32* bias_ptr = bias_data; + const uint8* filter_ptr = filter_data; + + const uint8* input_ptr = input_data; + uint8* output_ptr = output_data; + + for (int depth = 0; depth <= output_depth - 8; depth += 8) { + ConvKernel3x3FilterDepth8<1, 1, kFixedStrideWidth, kFixedStrideHeight>:: + Run(input_ptr, input_depth, input_offset, input_row_size, + filter_ptr, filter_offset, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth, output_width); + + input_ptr += 8; + output_ptr += 8; + filter_ptr += 8; + bias_ptr += 8; + } + + input_data += kFixedStrideWidth * input_depth; + output_data += output_depth; + } + } +}; + +template +struct ConvRow3x3FilterDepth8<2, kFixedStrideWidth, kFixedStrideHeight> { + static inline void Run(const uint8* input_data, int start_x, int start_y, + int input_depth, int input_width, int input_height, + int input_row_size, int32 input_offset, + const uint8* filter_data, int32 filter_offset, + const int32* bias_data, int32 output_offset, + int32 output_multiplier, int output_shift, + int32 output_activation_min, + int32 output_activation_max, uint8* output_data, + int output_depth, int output_width, + uint8* shuffle_workspace) { + int out_x = start_x; + + // 2x4 at a time. + for (; out_x <= output_width - 4; out_x += 4) { + const int32* bias_ptr = bias_data; + const uint8* filter_ptr = filter_data; + + const uint8* input_ptr = input_data; + uint8* output_ptr = output_data; + + for (int depth = 0; depth <= output_depth - 8; depth += 8) { + ConvKernel3x3FilterDepth8<2, 4, kFixedStrideWidth, kFixedStrideHeight>:: + Run(input_ptr, input_depth, input_offset, input_row_size, + filter_ptr, filter_offset, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth, output_width); + + input_ptr += 8; + output_ptr += 8; + filter_ptr += 8; + bias_ptr += 8; + } + + input_data += 4 * kFixedStrideWidth * input_depth; + output_data += 4 * output_depth; + } + + // 2x2 at a time. + for (; out_x <= output_width - 2; out_x += 2) { + const int32* bias_ptr = bias_data; + const uint8* filter_ptr = filter_data; + + const uint8* input_ptr = input_data; + uint8* output_ptr = output_data; + + for (int depth = 0; depth <= output_depth - 8; depth += 8) { + ConvKernel3x3FilterDepth8<2, 2, kFixedStrideWidth, kFixedStrideHeight>:: + Run(input_ptr, input_depth, input_offset, input_row_size, + filter_ptr, filter_offset, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth, output_width); + + input_ptr += 8; + output_ptr += 8; + filter_ptr += 8; + bias_ptr += 8; + } + + input_data += 2 * kFixedStrideWidth * input_depth; + output_data += 2 * output_depth; + } + + // 2x1 at a time. + for (; out_x < output_width; out_x++) { + const int32* bias_ptr = bias_data; + const uint8* filter_ptr = filter_data; + + const uint8* input_ptr = input_data; + uint8* output_ptr = output_data; + + for (int depth = 0; depth <= output_depth - 8; depth += 8) { + ConvKernel3x3FilterDepth8<2, 1, kFixedStrideWidth, kFixedStrideHeight>:: + Run(input_ptr, input_depth, input_offset, input_row_size, + filter_ptr, filter_offset, bias_ptr, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth, output_width); + + input_ptr += 8; + output_ptr += 8; + filter_ptr += 8; + bias_ptr += 8; + } + + input_data += kFixedStrideWidth * input_depth; + output_data += output_depth; + } + } +}; + +template <> +struct ConvRow3x3FilterDepth8<4, 1, 1> { + static inline void Run(const uint8* input_data, int start_x, int start_y, + int input_depth, int input_width, int input_height, + int input_row_size, int32 input_offset, + const uint8* filter_data, int32 filter_offset, + const int32* bias_data, int32 output_offset, + int32 output_multiplier, int output_shift, + int32 output_activation_min, + int32 output_activation_max, uint8* output_data, + int output_depth, int output_width, + uint8* shuffle_workspace) { + int out_x = start_x; + + // 4x4 at a time. + for (; out_x <= output_width - 4; out_x += 4) { + const int32* bias_ptr = bias_data; + const uint8* filter_ptr = filter_data; + + const uint8* input_ptr = input_data; + uint8* output_ptr = output_data; + + for (int depth = 0; depth <= output_depth - 8; depth += 8) { + ConvKernel3x3FilterDepth8<4, 4, 1, 1>::Run( + input_ptr, input_depth, input_offset, input_row_size, filter_ptr, + filter_offset, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, + output_ptr, output_depth, output_width); + + input_ptr += 8; + output_ptr += 8; + filter_ptr += 8; + bias_ptr += 8; + } + + input_data += 4 * input_depth; + output_data += 4 * output_depth; + } + + // Handle the rest of the right side. + // 4x2 at a time. + for (; out_x <= output_width - 2; out_x += 2) { + const int32* bias_ptr = bias_data; + const uint8* filter_ptr = filter_data; + + const uint8* input_ptr = input_data; + uint8* output_ptr = output_data; + + for (int depth = 0; depth <= output_depth - 8; depth += 8) { + ConvKernel3x3FilterDepth8<4, 2, 1, 1>::Run( + input_ptr, input_depth, input_offset, input_row_size, filter_ptr, + filter_offset, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, + output_ptr, output_depth, output_width); + + input_ptr += 8; + output_ptr += 8; + filter_ptr += 8; + bias_ptr += 8; + } + + input_data += 2 * input_depth; + output_data += 2 * output_depth; + } + + // 4x1 at a time. + for (; out_x < output_width; out_x++) { + const int32* bias_ptr = bias_data; + const uint8* filter_ptr = filter_data; + + const uint8* input_ptr = input_data; + uint8* output_ptr = output_data; + + for (int depth = 0; depth <= output_depth - 8; depth += 8) { + ConvKernel3x3FilterDepth8<4, 1, 1, 1>::Run( + input_ptr, input_depth, input_offset, input_row_size, filter_ptr, + filter_offset, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, + output_ptr, output_depth, output_width); + + input_ptr += 8; + output_ptr += 8; + filter_ptr += 8; + bias_ptr += 8; + } + + input_data += input_depth; + output_data += output_depth; + } + } +}; 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); +struct ConvRow3x3FilterDepth8<4, 2, 2> { + // The buffer size of the shuffled input. + static inline constexpr int ShuffleWorkspaceSize() { return 64 * 9 * 9; } + + static inline void Run(const uint8* input_data, int start_x, int start_y, + int input_depth, int input_width, int input_height, + int input_row_size, int32 input_offset, + const uint8* filter_data, int32 filter_offset, + const int32* bias_data, int32 output_offset, + int32 output_multiplier, int output_shift, + int32 output_activation_min, + int32 output_activation_max, uint8* output_data, + int output_depth, int output_width, + uint8* shuffle_workspace) { + // Branch and cache misses increase substantially with stride 2 kernels. + // Adding prefetching reduces latency by as much as 2x. + const int i0 = 0; + const int i1 = input_depth; + const int i2 = 2 * input_depth; + const int i3 = 3 * input_depth; + const int i4 = 4 * input_depth; + const int i5 = 5 * input_depth; + const int i6 = 6 * input_depth; + const int i7 = 7 * input_depth; + const int i8 = 8 * input_depth; + +#define DEPTHWISECONV_PRELOAD_ROW(input_ptr, i) \ + preload_l1_keep(input_ptr + i * input_row_size + i0); \ + preload_l1_keep(input_ptr + i * input_row_size + i1); \ + preload_l1_keep(input_ptr + i * input_row_size + i2); \ + preload_l1_keep(input_ptr + i * input_row_size + i3); \ + preload_l1_keep(input_ptr + i * input_row_size + i4); \ + preload_l1_keep(input_ptr + i * input_row_size + i5); \ + preload_l1_keep(input_ptr + i * input_row_size + i6); \ + preload_l1_keep(input_ptr + i * input_row_size + i7); \ + preload_l1_keep(input_ptr + i * input_row_size + i8); + + int out_x = start_x; + // 4x4 at a time. + for (; out_x <= output_width - 4; out_x += 4) { + const int32* bias_ptr = bias_data; + const uint8* filter_ptr = filter_data; + + const uint8* input_ptr = input_data; + uint8* output_ptr = output_data; + + int depth = 0; + for (; depth <= output_depth - 64; depth += 64) { + // Preload 9x9 input. + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 0); + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 1); + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 2); + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 3); + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 4); + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 5); + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 6); + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 7); + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 8); + + // For a large input window (64x9x9) that is small enough to fit in L1 + // cache, copy the input into a separate buffer and run the kernel on + // this new buffer. This reduces the likelihood of cache misses when + // the kernel is loading input data. If this size is ever changed, + // update the ShuffleWorkspaceSize() function to return the new size. + ShuffleInput(input_ptr, input_depth, input_width, input_height, 64, 9, + 9, shuffle_workspace); + const uint8* shuffled_ptr = &shuffle_workspace[0]; + + for (int micro_depth = 0; micro_depth <= 64 - 8; micro_depth += 8) { + ConvKernel3x3FilterDepth8<4, 4, 2, 2>::Run( + shuffled_ptr, 64, input_offset, 64 * 9, filter_ptr, filter_offset, + bias_ptr, output_offset, output_multiplier, output_shift, + output_activation_min, output_activation_max, output_ptr, + output_depth, output_width); + + shuffled_ptr += 8; + output_ptr += 8; + filter_ptr += 8; + bias_ptr += 8; + } + input_ptr += 64; + } + + // Preload 9x9 input one more time for the rest of the depth. + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 0); + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 1); + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 2); + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 3); + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 4); + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 5); + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 6); + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 7); + DEPTHWISECONV_PRELOAD_ROW(input_ptr, 8); + + for (; depth <= output_depth - 8; depth += 8) { + ConvKernel3x3FilterDepth8<4, 4, 2, 2>::Run( + input_ptr, input_depth, input_offset, input_row_size, filter_ptr, + filter_offset, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, + output_ptr, output_depth, output_width); + + input_ptr += 8; + output_ptr += 8; + filter_ptr += 8; + bias_ptr += 8; + } + + input_data += 4 * 2 * input_depth; + output_data += 4 * output_depth; + } + +#undef DEPTHWISECONV_PRELOAD_ROW + + // Handle the rest of the right side. + // 4x2 at a time. + for (; out_x <= output_width - 2; out_x += 2) { + const int32* bias_ptr = bias_data; + const uint8* filter_ptr = filter_data; + + const uint8* input_ptr = input_data; + uint8* output_ptr = output_data; + + for (int depth = 0; depth <= output_depth - 8; depth += 8) { + ConvKernel3x3FilterDepth8<4, 2, 2, 2>::Run( + input_ptr, input_depth, input_offset, input_row_size, filter_ptr, + filter_offset, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, + output_ptr, output_depth, output_width); + + input_ptr += 8; + output_ptr += 8; + filter_ptr += 8; + bias_ptr += 8; + } + + input_data += 2 * 2 * input_depth; + output_data += 2 * output_depth; + } + + // 4x1 at a time. + for (; out_x < output_width; out_x++) { + const int32* bias_ptr = bias_data; + const uint8* filter_ptr = filter_data; + + const uint8* input_ptr = input_data; + uint8* output_ptr = output_data; + + for (int depth = 0; depth <= output_depth - 8; depth += 8) { + ConvKernel3x3FilterDepth8<4, 1, 2, 2>::Run( + input_ptr, input_depth, input_offset, input_row_size, filter_ptr, + filter_offset, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, + output_ptr, output_depth, output_width); + + input_ptr += 8; + output_ptr += 8; + filter_ptr += 8; + bias_ptr += 8; + } + + input_data += 2 * input_depth; + output_data += output_depth; + } } }; 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); +struct ConvRow3x3FilterDepth8<8, 2, 2> { + static inline void Run(const uint8* input_data, int start_x, int start_y, + int input_depth, int input_width, int input_height, + int input_row_size, int32 input_offset, + const uint8* filter_data, int32 filter_offset, + const int32* bias_data, int32 output_offset, + int32 output_multiplier, int output_shift, + int32 output_activation_min, + int32 output_activation_max, uint8* output_data, + int output_depth, int output_width, + uint8* shuffle_workspace) { + // Reuse 4 row kernels twice. + ConvRow3x3FilterDepth8<4, 2, 2>::Run( + input_data, start_x, start_y, input_depth, input_width, input_height, + input_row_size, input_offset, filter_data, filter_offset, bias_data, + output_offset, output_multiplier, output_shift, output_activation_min, + output_activation_max, output_data, output_depth, output_width, + shuffle_workspace); + + ConvRow3x3FilterDepth8<4, 2, 2>::Run( + input_data + 2 * 4 * input_row_size, start_x, start_y + 4, input_depth, + input_width, input_height, input_row_size, input_offset, filter_data, + filter_offset, bias_data, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, + output_data + 4 * output_depth * output_width, output_depth, + output_width, shuffle_workspace); } }; 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); +struct ConvRow3x3FilterDepth8<8, 1, 1> { + // The buffer size of the shuffled input. + static inline constexpr int ShuffleWorkspaceSize() { return 64 * 10 * 10; } + + static inline void Run(const uint8* input_data, int start_x, int start_y, + int input_depth, int input_width, int input_height, + int input_row_size, int32 input_offset, + const uint8* filter_data, int32 filter_offset, + const int32* bias_data, int32 output_offset, + int32 output_multiplier, int output_shift, + int32 output_activation_min, + int32 output_activation_max, uint8* output_data, + int output_depth, int output_width, + uint8* shuffle_workspace) { + int out_x = start_x; + // 8x8 at a time. + for (; out_x <= output_width - 8; out_x += 8) { + const int32* bias_ptr = bias_data; + const uint8* filter_ptr = filter_data; + + const uint8* input_ptr = input_data; + uint8* output_ptr = output_data; + + int depth = 0; + for (; depth <= output_depth - 64; depth += 64) { + // For a large input window (64x10x10) that is small enough to fit in L1 + // cache, copy the input into a separate buffer and run the kernel on + // this new buffer. This reduces the likelihood of cache misses when + // the kernel is loading input data. If the size of the input window + // changes, update the function ShuffleWorkspaceSize() with the new + // size. + ShuffleInput(input_ptr, input_depth, input_width, input_height, 64, 10, + 10, shuffle_workspace); + const uint8* shuffled_ptr = shuffle_workspace; + + for (int micro_depth = 0; micro_depth <= 64 - 8; micro_depth += 8) { + ConvKernel3x3FilterDepth8<8, 8, 1, 1>::Run( + shuffled_ptr, 64, input_offset, 64 * 10, filter_ptr, + filter_offset, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, + output_ptr, output_depth, output_width); + + shuffled_ptr += 8; + output_ptr += 8; + filter_ptr += 8; + bias_ptr += 8; + } + input_ptr += 64; + } + + for (; depth <= output_depth - 8; depth += 8) { + ConvKernel3x3FilterDepth8<8, 8, 1, 1>::Run( + input_ptr, input_depth, input_offset, input_row_size, filter_ptr, + filter_offset, bias_ptr, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, + output_ptr, output_depth, output_width); + + input_ptr += 8; + output_ptr += 8; + filter_ptr += 8; + bias_ptr += 8; + } + + input_data += 8 * input_depth; + output_data += 8 * output_depth; + } + + // Handle the rest of the right side by re-using 4 row kernels twice. + ConvRow3x3FilterDepth8<4, 1, 1>::Run( + input_data, out_x, start_y, input_depth, input_width, input_height, + input_row_size, input_offset, filter_data, filter_offset, bias_data, + output_offset, output_multiplier, output_shift, output_activation_min, + output_activation_max, output_data, output_depth, output_width, + shuffle_workspace); + + ConvRow3x3FilterDepth8<4, 1, 1>::Run( + input_data + 4 * input_row_size, out_x, start_y + 4, input_depth, + input_width, input_height, input_row_size, input_offset, filter_data, + filter_offset, bias_data, output_offset, output_multiplier, + output_shift, output_activation_min, output_activation_max, + output_data + 4 * output_depth * output_width, output_depth, + output_width, shuffle_workspace); } }; -inline bool Fast3by3FilterKernelSupported(const Dims<4>& input_dims, - const Dims<4>& filter_dims, - int stride_width, int stride_height, - int pad_width, int pad_height, - int depth_multiplier, - const Dims<4>& output_dims) { +inline bool Fast3x3FilterKernelSupported(const Dims<4>& input_dims, + const Dims<4>& filter_dims, + int stride_width, int stride_height, + int pad_width, int pad_height, + int depth_multiplier, + const Dims<4>& output_dims) { const int input_height = ArraySize(input_dims, 2); const int input_width = ArraySize(input_dims, 1); const int input_depth = ArraySize(input_dims, 0); @@ -458,14 +4426,15 @@ inline bool Fast3by3FilterKernelSupported(const Dims<4>& input_dims, depth_multiplier == 1 && (stride_width == 1 || stride_width == 2) && (stride_height == 1 || stride_height == 2) && - pad_width == 0 && pad_height == 0 && (input_depth % 16) == 0; + (stride_width == stride_height) && pad_width == 0 && + pad_height == 0 && (input_depth % 8) == 0; if (!supported) { return false; } - // Handle case where padding is zero but type is not kValid. This would - // require special boundary case handling that is not supported yet. + // Handle case where padding is zero but padding type is not kValid. + // This would require special boundary case handling that is not supported. const int out_x = output_width - 1; const int out_y = output_height - 1; @@ -481,7 +4450,7 @@ inline bool Fast3by3FilterKernelSupported(const Dims<4>& input_dims, return in_x_end <= input_width && in_y_end <= input_height; } -inline void DepthwiseConv3by3FilterDepth16( +inline void DepthwiseConv3x3Filter( 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, @@ -500,241 +4469,109 @@ inline void DepthwiseConv3by3FilterDepth16( 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. + // of 1, 3x3 filter, no padding and 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 || stride_width == 2); TFLITE_DCHECK(stride_height == 1 || stride_height == 2); + TFLITE_DCHECK(stride_width == 1 || stride_width == 2); + TFLITE_DCHECK(stride_width == stride_height); - // 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 input_row_size = input_depth * (input_width + 2 * pad_width); + const int output_row_size = output_depth * output_width; + const int input_batch_size = input_row_size * (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; - } + using conv_row_func_t = decltype(&ConvRow3x3FilterDepth8<1, 1, 1>::Run); + conv_row_func_t conv_1_output_row = ConvRow3x3FilterDepth8<1, 1, 1>::Run; + conv_row_func_t conv_2_output_rows = ConvRow3x3FilterDepth8<2, 1, 1>::Run; + conv_row_func_t conv_4_output_rows = ConvRow3x3FilterDepth8<4, 1, 1>::Run; + conv_row_func_t conv_8_output_rows = ConvRow3x3FilterDepth8<8, 1, 1>::Run; + + if (stride_width == 2) { + conv_1_output_row = ConvRow3x3FilterDepth8<1, 2, 2>::Run; + conv_2_output_rows = ConvRow3x3FilterDepth8<2, 2, 2>::Run; + conv_4_output_rows = ConvRow3x3FilterDepth8<4, 2, 2>::Run; + conv_8_output_rows = ConvRow3x3FilterDepth8<8, 2, 2>::Run; } - using dot_product_func_t = - decltype(&ConvKernel3x3FilterDepth16<1, 2, 1>::Run); - dot_product_func_t dot_product_func = nullptr; + // Allocate maximum memory needed for shuffled input. + // TODO(mariewhite): The size of this workspace is small enough to be + // allocated on the stack. Eventually we will want to move it to the heap + // and have it allocated outside of this function, like the im2col_array used + // in gemmlowp. +#define DEPTHWISECONV_SHUFFLE_WORKSPACE_SIZE 10 * 10 * 64 + uint8 shuffle_workspace[DEPTHWISECONV_SHUFFLE_WORKSPACE_SIZE]; - if (stride_width == 1 && stride_height == 1) { - dot_product_func = ConvKernel3x3FilterDepth16<1, 2, 1>::Run; - } else { - dot_product_func = ConvKernel3x3FilterDepth16<1, 2, 2>::Run; - } + // Make sure the kernels using this buffer will not run out of bounds. + static_assert(ConvRow3x3FilterDepth8<8, 1, 1>::ShuffleWorkspaceSize() <= + DEPTHWISECONV_SHUFFLE_WORKSPACE_SIZE, + "Shuffle workspace size is too small."); + static_assert(ConvRow3x3FilterDepth8<4, 2, 2>::ShuffleWorkspaceSize() <= + DEPTHWISECONV_SHUFFLE_WORKSPACE_SIZE, + "Shuffle workspace size is too small."); - // 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; - const int i12 = 4 * input_row_width; - const int i13 = 4 * input_row_width + input_depth; - const int i14 = 4 * input_row_width + 2 * input_depth; +#undef DEPTHWISECONV_SHUFFLE_WORKSPACE_SIZE 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; - - // 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); - - if (stride_height == 2) { - preload_l1_keep(input_ptr + i12); - preload_l1_keep(input_ptr + i13); - preload_l1_keep(input_ptr + i14); - } - } 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); - - if (stride_height == 2) { - preload_l1_keep(input_ptr + i12); - } - } + const uint8* input_ptr = input_data + b * input_batch_size; + uint8* output_ptr = output_data + b * output_batch_size; - uint8* output_ptr = output_data + depth + (out_x * output_depth) + - (output_depth * output_width * out_y) + - out_batch_offset; - - for (; out_x < out_x_end; out_x += num_x_outputs) { - dot_product_func(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); - preload_l1_keep(input_ptr + i13); - preload_l1_keep(input_ptr + i14); - } - } + int out_y = 0; - // 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 kSame padding. - } - } + // Handle 8 rows at a time. + for (; out_y <= output_height - 8; out_y += 8) { + conv_8_output_rows(input_ptr, 0, out_y, input_depth, input_width, + input_height, input_row_size, input_offset, + filter_data, filter_offset, bias_data, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth, + output_width, shuffle_workspace); - // 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; - - 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); - } else { - preload_l1_keep(input_ptr + i0); - preload_l1_keep(input_ptr + i3); - preload_l1_keep(input_ptr + i6); - } + input_ptr += 8 * stride_height * input_row_size; + output_ptr += 8 * output_row_size; + } - uint8* output_ptr = output_data + depth + (out_x * output_depth) + - (output_depth * output_width * out_y) + - out_batch_offset; + // Handle 4 rows at a time. + for (; out_y <= output_height - 4; out_y += 4) { + conv_4_output_rows(input_ptr, 0, out_y, input_depth, input_width, + input_height, input_row_size, input_offset, + filter_data, filter_offset, bias_data, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth, + output_width, shuffle_workspace); - 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 += 4 * stride_height * input_row_size; + output_ptr += 4 * output_row_size; + } - input_ptr += input_ptr_x_increment; - output_ptr += output_depth; - - if (stride_width == 1) { - preload_l1_keep(input_ptr + i2); - preload_l1_keep(input_ptr + i5); - preload_l1_keep(input_ptr + i8); - } 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); - } - } - } - filter_ptr += 16; - bias_ptr += 16; + // Handle 2 rows at a time. + for (; out_y <= output_height - 2; out_y += 2) { + conv_2_output_rows(input_ptr, 0, out_y, input_depth, input_width, + input_height, input_row_size, input_offset, + filter_data, filter_offset, bias_data, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth, + output_width, shuffle_workspace); + + input_ptr += 2 * stride_height * input_row_size; + output_ptr += 2 * output_row_size; + } + + // Handle one row at a time. + for (; out_y < output_height; out_y++) { + conv_1_output_row(input_ptr, 0, out_y, input_depth, input_width, + input_height, input_row_size, input_offset, filter_data, + filter_offset, bias_data, output_offset, + output_multiplier, output_shift, output_activation_min, + output_activation_max, output_ptr, output_depth, + output_width, shuffle_workspace); + + input_ptr += stride_height * input_row_size; + output_ptr += output_row_size; } } } diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h index 5acf1eaede5aa6c61cf345b60009369bf7d4cdde..d27c6ccf3d404cb18460876fc2961d8447d5f819 100644 --- a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h +++ b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h @@ -30,6 +30,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/quantization_util.h" #include "tensorflow/contrib/lite/kernels/internal/round.h" #include "tensorflow/contrib/lite/kernels/internal/types.h" @@ -1202,6 +1203,142 @@ void FullyConnected(const uint8* input_data, const Dims<4>& input_dims, output_activation_max, output_data, output_dims, gemm_context); } +inline void ExperimentalShuffledFullyConnected( + const uint8* input_data, const Dims<4>& input_dims, + const uint8* shuffled_weights_data, const Dims<4>& weights_dims, + const int32* bias_data, const Dims<4>& bias_dims, 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( + "ExperimentalShuffledFullyConnected/8bit"); + (void)gemm_context; // only used in optimized code. + TFLITE_DCHECK_EQ(output_activation_min, -32768); + TFLITE_DCHECK_EQ(output_activation_max, 32767); + // 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(weights_dims, 1, output_dims, 0); + const int accum_depth = ArraySize(weights_dims, 0); + TFLITE_DCHECK(IsPackedWithoutStrides(input_dims)); + TFLITE_DCHECK(IsPackedWithoutStrides(weights_dims)); + // The experimental shuffling is an optimization for matrix*vector product. + // We aren't interested in supporting non-matrix*vector-product cases, i.e. + // batches>1. + TFLITE_DCHECK_EQ(batches, 1); + // Shuffled weights have had their sign bit (0x80) pre-flipped (xor'd) + // so that just reinterpreting them as int8 values is equivalent to + // subtracting 128 from them, thus implementing for free the subtraction of + // the zero_point value 128. + const int8* shuffled_weights_ptr = + reinterpret_cast(shuffled_weights_data); +#if defined USE_NEON + // We'll only need to xor signbit to the input activation values, as + // that xor-ing is pre-built into the shuffled weights values. + const uint8x16_t signbit = vdupq_n_u8(0x80); + const int right_shift = output_shift > 0 ? output_shift : 0; + const int left_shift = output_shift > 0 ? 0 : -output_shift; + for (int c = 0; c < output_depth; c += 4) { + // Accumulation loop. + int32x4_t row_accum0 = vdupq_n_s32(0); + int32x4_t row_accum1 = vdupq_n_s32(0); + int32x4_t row_accum2 = vdupq_n_s32(0); + int32x4_t row_accum3 = vdupq_n_s32(0); + for (int d = 0; d < accum_depth; d += 16) { + int8x16_t weights0 = vld1q_s8(shuffled_weights_ptr + 0); + int8x16_t weights1 = vld1q_s8(shuffled_weights_ptr + 16); + int8x16_t weights2 = vld1q_s8(shuffled_weights_ptr + 32); + int8x16_t weights3 = vld1q_s8(shuffled_weights_ptr + 48); + shuffled_weights_ptr += 64; + int8x16_t input = + vreinterpretq_s8_u8(veorq_u8(signbit, vld1q_u8(input_data + d))); + int16x8_t local_accum0 = + vmull_s8(vget_low_s8(weights0), vget_low_s8(input)); + int16x8_t local_accum1 = + vmull_s8(vget_low_s8(weights1), vget_low_s8(input)); + int16x8_t local_accum2 = + vmull_s8(vget_low_s8(weights2), vget_low_s8(input)); + int16x8_t local_accum3 = + vmull_s8(vget_low_s8(weights3), vget_low_s8(input)); + local_accum0 = + vmlal_s8(local_accum0, vget_high_s8(weights0), vget_high_s8(input)); + local_accum1 = + vmlal_s8(local_accum1, vget_high_s8(weights1), vget_high_s8(input)); + local_accum2 = + vmlal_s8(local_accum2, vget_high_s8(weights2), vget_high_s8(input)); + local_accum3 = + vmlal_s8(local_accum3, vget_high_s8(weights3), vget_high_s8(input)); + row_accum0 = vpadalq_s16(row_accum0, local_accum0); + row_accum1 = vpadalq_s16(row_accum1, local_accum1); + row_accum2 = vpadalq_s16(row_accum2, local_accum2); + row_accum3 = vpadalq_s16(row_accum3, local_accum3); + } + // 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(row_accum0), vget_high_s32(row_accum0)); + pairwise_reduced_acc_1 = + vpadd_s32(vget_low_s32(row_accum1), vget_high_s32(row_accum1)); + pairwise_reduced_acc_2 = + vpadd_s32(vget_low_s32(row_accum2), vget_high_s32(row_accum2)); + pairwise_reduced_acc_3 = + vpadd_s32(vget_low_s32(row_accum3), vget_high_s32(row_accum3)); + 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_data + c); + reduced = vaddq_s32(reduced, bias_vec); + reduced = vshlq_s32(reduced, vdupq_n_s32(left_shift)); + // Multiply by the fixed-point multiplier. + reduced = vqrdmulhq_n_s32(reduced, output_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_data + c, res16); + } +#else + for (int c = 0; c < output_depth; c += 4) { + // Internal accumulation. + // Initialize accumulator with the bias-value. + int32 accum[4] = {0}; + // Accumulation loop. + for (int d = 0; d < accum_depth; d += 16) { + for (int i = 0; i < 4; i++) { + for (int j = 0; j < 16; j++) { + int8 input_val = input_data[d + j] - 128; + int8 weights_val = *shuffled_weights_ptr++; + accum[i] += weights_val * input_val; + } + } + } + for (int i = 0; i < 4; i++) { + // Add bias value + int acc = accum[i] + bias_data[c + i]; + // 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). + acc = + MultiplyByQuantizedMultiplier(acc, output_multiplier, -output_shift); + // Saturate, cast to int16, and store to output array. + acc = std::max(acc, output_activation_min); + acc = std::min(acc, output_activation_max); + output_data[c + i] = acc; + } + } +#endif +} + template inline void ExtractPatchIntoBufferColumn( const Dims<4>& input_dims, int w, int h, int b, int kheight, int kwidth, @@ -3937,7 +4074,7 @@ inline void Softmax(const uint8* input_data, const Dims<4>& input_dims, using FixedPointAccum = gemmlowp::FixedPoint; using FixedPoint0 = gemmlowp::FixedPoint; - gemmlowp::ScopedProfilingLabel label("Softmax/8bit"); +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); @@ -4135,6 +4272,7 @@ inline void Softmax(const uint8* input_data, const Dims<4>& input_dims, // optimized yet. inline void LogSoftmax(const float* input_data, const Dims<4>& input_dims, float* output_data, const Dims<4>& output_dims) { + gemmlowp::ScopedProfilingLabel label("LogSoftmax"); 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); @@ -4168,6 +4306,94 @@ inline void LogSoftmax(const float* input_data, const Dims<4>& input_dims, } } +// Currently just a copy of the reference code. +inline void LogSoftmax(const uint8* input_data, const Dims<4>& input_dims, + int32 input_multiplier, int32 input_left_shift, + int32 reverse_scaling_divisor, + int32 reverse_scaling_right_shift, int diff_min, + uint8* output_data, const Dims<4>& output_dims) { + gemmlowp::ScopedProfilingLabel label("LogSoftmax/Uint8"); + // The representation chosen for the input to the exp() function is Q5.26. + // We need to leave extra space since values that we skip might be as large as + // -32 before multiplying by input_beta_multiplier, and therefore as large as + // -16 afterwards. Note that exp(-8) is definitely not insignificant to + // accumulation, but exp(-16) definitely is. + static constexpr int kScaledDiffIntegerBits = 5; + static constexpr int kAccumulationIntegerBits = 12; + static constexpr int kOutputIntegerBits = 4; + using FixedPointScaledDiff = + gemmlowp::FixedPoint; + using FixedPointAccum = gemmlowp::FixedPoint; + using FixedPoint0 = gemmlowp::FixedPoint; + + const int outer_size = MatchingFlatSizeSkipDim(input_dims, 0, output_dims); + const int depth = MatchingArraySize(input_dims, 0, output_dims, 0); + + for (int i = 0; i < outer_size; ++i) { + uint8 max_in_row = 0; + for (int c = 0; c < depth; ++c) { + max_in_row = std::max(max_in_row, input_data[i * depth + c]); + } + + FixedPointAccum sum_of_exps = FixedPointAccum::Zero(); + for (int c = 0; c < depth; ++c) { + int32 input_diff = + static_cast(input_data[i * depth + c]) - max_in_row; + if (input_diff >= diff_min) { + const int32 input_diff_rescaled = + MultiplyByQuantizedMultiplierGreaterThanOne( + input_diff, input_multiplier, input_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)); + } + } + + // TODO(b/77858996): Implement fixed-point log(). + // Not a fully-quantized implementation: floating-point log(). + const float float_log_sum_of_exps = + std::log(static_cast(sum_of_exps.raw()) / + (1 << (31 - kAccumulationIntegerBits))); + const int32 fixed_log_sum_of_exps = static_cast(TfLiteRound( + float_log_sum_of_exps * (1 << (31 - kScaledDiffIntegerBits)))); + + // rescaled_diff_min is smallest representable in + // Q(kScaledDiffIntegerBits).(31-kScaledDiffIntegerBits) plus the + // log-sub-exps that will be subtracted in the loop. + // + // The thresholds diff_min, etc are negative. + const int rescaled_diff_min = + fixed_log_sum_of_exps + std::numeric_limits::lowest(); + const int adjusted_diff_min = + std::max(diff_min - 1, // Note use of > below instead of >= above. + MultiplyByQuantizedMultiplierSmallerThanOne( + rescaled_diff_min, reverse_scaling_divisor, + reverse_scaling_right_shift)); + + for (int c = 0; c < depth; ++c) { + int32 input_diff = + static_cast(input_data[i * depth + c]) - max_in_row; + if (input_diff > adjusted_diff_min) { + const int32 input_diff_rescaled = + MultiplyByQuantizedMultiplierGreaterThanOne( + input_diff, input_multiplier, input_left_shift); + int32 unsat_output = + gemmlowp::RoundingDivideByPOT( + (input_diff_rescaled - fixed_log_sum_of_exps), + 31 - kScaledDiffIntegerBits - kOutputIntegerBits) + + 255; + + output_data[i * depth + c] = static_cast( + std::max(std::min(unsat_output, static_cast(255)), 0)); + } else { + // Set output to smallest value. + output_data[i * depth + c] = 0; + } + } + } +} + inline void Logistic(const float* input_data, const Dims<4>& input_dims, float* output_data, const Dims<4>& output_dims) { gemmlowp::ScopedProfilingLabel label("Logistic"); @@ -4181,7 +4407,7 @@ inline void Logistic(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) { - gemmlowp::ScopedProfilingLabel label("Logistic"); + gemmlowp::ScopedProfilingLabel label("Logistic/Uint8"); /* batches */ MatchingArraySize(input_dims, 3, output_dims, 3); /* height */ MatchingArraySize(input_dims, 2, output_dims, 2); /* width */ MatchingArraySize(input_dims, 1, output_dims, 1); @@ -4661,66 +4887,23 @@ inline void Dequantize(const uint8* input_data, const Dims<4>& input_dims, } inline void FakeQuant(const float* input_data, const Dims<4>& input_dims, - float rmin, float rmax, float* output_data, + float rmin, float rmax, int num_bits, float* output_data, const Dims<4>& output_dims) { gemmlowp::ScopedProfilingLabel label("FakeQuant"); // 0 should always be a representable value. Let's assume that the initial // min,max range contains 0. - TFLITE_DCHECK_LE(rmin, 0.); - TFLITE_DCHECK_GE(rmax, 0.); - - // Determine quantization parameters: zero_point, scale. - using Integer = uint8; - const Integer qmin = std::numeric_limits::min(); - const Integer qmax = std::numeric_limits::max(); - const float qmin_float = qmin; - const float qmax_float = qmax; - int32 zero_point = 0; - float scale = 0.f; - // If rmin==rmax, both must be zero per the above assertion, - // so we are done. - if (rmin != rmax) { - // First determine the scale. - scale = (rmax - rmin) / (qmax_float - qmin_float); - - // 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 float zero_point_from_min = qmin_float - rmin / scale; - const float zero_point_from_max = qmax_float - rmax / scale; - const float zero_point_from_min_error = - std::abs(qmin_float) + std::abs(rmin / scale); - const float zero_point_from_max_error = - std::abs(qmax_float) + std::abs(rmax / scale); - - const float zero_point_float = - 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). - if (zero_point_float < qmin_float) { - zero_point = qmin; - } else if (zero_point_float > qmax_float) { - zero_point = qmax; - } else { - zero_point = static_cast(TfLiteRound(zero_point_float)); - } - // The zero point should always be in the range of quantized value, - // [qmin, qmax]. - TFLITE_DCHECK_GE(zero_point, qmin); - TFLITE_DCHECK_LE(zero_point, qmax); - } + TFLITE_DCHECK_LE(rmin, 0.0f); + TFLITE_DCHECK_GE(rmax, 0.0f); + TFLITE_DCHECK_LT(rmin, rmax); + + // Code matches tensorflow's FakeQuantWithMinMaxArgsFunctor. + int quant_min = 0; + int quant_max = (1 << num_bits) - 1; + float nudged_min, nudged_max, nudged_scale; + NudgeQuantizationRange(rmin, rmax, quant_min, quant_max, &nudged_min, + &nudged_max, &nudged_scale); + const float inv_nudged_scale = 1.0f / nudged_scale; const int batches = MatchingArraySize(input_dims, 3, output_dims, 3); const int height = MatchingArraySize(input_dims, 2, output_dims, 2); @@ -4731,11 +4914,12 @@ inline void FakeQuant(const float* input_data, const Dims<4>& input_dims, for (int x = 0; x < width; ++x) { for (int c = 0; c < depth; ++c) { const float src_val = input_data[Offset(input_dims, c, x, y, b)]; - const float unclamped_quantized_val = - TfLiteRound(zero_point + src_val / scale); - const float quantized_val = std::min( - qmax_float, std::max(qmin_float, unclamped_quantized_val)); - const float dst_val = scale * (quantized_val - zero_point); + const float clamped = + std::min(nudged_max, std::max(nudged_min, src_val)); + const float clamped_shifted = clamped - nudged_min; + const float dst_val = + TfLiteRound(clamped_shifted * inv_nudged_scale) * nudged_scale + + nudged_min; output_data[Offset(output_dims, c, x, y, b)] = dst_val; } } @@ -5164,6 +5348,7 @@ inline void SpaceToBatchND(const T* input_data, const Dims<4>& input_dims, const int32* paddings_data, const Dims<4>& paddings_dims, T* output_data, const Dims<4>& output_dims) { + // Unoptimized - Straight copy from reference ops. gemmlowp::ScopedProfilingLabel label("SpaceToBatchND"); const int output_batch_size = ArraySize(output_dims, 3); @@ -5205,29 +5390,76 @@ inline void SpaceToBatchND(const T* input_data, const Dims<4>& input_dims, } } +// Helper methods for BatchToSpaceND. +// `spatial_index_dim` specifies post-crop offset index in this spatial +// dimension, i.e. spatial offset introduced by flattening batch to spatial +// dimension minus the crop size at beginning. `block_shape_dim` is the block +// size in current dimension. `input_dim` and `output_dim` are input and output +// size of BatchToSpaceND operation in current dimension. +// Output start index is inclusive and end index is exclusive. +inline void GetIndexRange(int spatial_index_dim, int block_shape_dim, + int input_dim, int output_dim, int* start_index, + int* end_index) { + // (*start_index) * block_shape_dim is effectively rounded up to the next + // multiple of block_shape_dim by the integer division. + *start_index = + std::max(0, (-spatial_index_dim + block_shape_dim - 1) / block_shape_dim); + // Similarly, (*end_index) * block_shape_dim is rounded up too (note that + // end_index is exclusive). + *end_index = std::min( + input_dim, + (output_dim - spatial_index_dim + block_shape_dim - 1) / block_shape_dim); +} + template inline void BatchToSpaceND(const T* input_data, const Dims<4>& input_dims, const int32* block_shape_data, - const Dims<4>& block_shape_dims, T* output_data, - const Dims<4>& output_dims) { + const Dims<4>& block_shape_dims, + const int32* crops_data, const Dims<4>& crops_dims, + T* output_data, const Dims<4>& output_dims) { gemmlowp::ScopedProfilingLabel label("BatchToSpaceND"); const int output_batch_size = ArraySize(output_dims, 3); + const int output_height = ArraySize(output_dims, 2); + const int output_width = ArraySize(output_dims, 1); const int input_batch_size = ArraySize(input_dims, 3); const int input_height = ArraySize(input_dims, 2); const int input_width = ArraySize(input_dims, 1); const int depth = ArraySize(input_dims, 0); const int block_shape_width = block_shape_data[1]; const int block_shape_height = block_shape_data[0]; + const int crops_top = crops_data[0]; + const int crops_left = crops_data[2]; for (int in_batch = 0; in_batch < input_batch_size; ++in_batch) { - for (int in_h = 0; in_h < input_height; ++in_h) { - for (int in_w = 0; in_w < input_width; ++in_w) { - int out_batch = in_batch % output_batch_size; - int out_w = in_w * block_shape_width + - (in_batch / output_batch_size) % block_shape_width; - int out_h = in_h * block_shape_height + - (in_batch / output_batch_size) / block_shape_width; + const int out_batch = in_batch % output_batch_size; + const int spatial_offset = in_batch / output_batch_size; + + int in_h_start = 0; + int in_h_end = 0; + // GetIndexRange ensures start and end indices are in [0, output_height). + GetIndexRange(spatial_offset / block_shape_width - crops_top, + block_shape_height, input_height, output_height, &in_h_start, + &in_h_end); + + for (int in_h = in_h_start; in_h < in_h_end; ++in_h) { + const int out_h = in_h * block_shape_height + + spatial_offset / block_shape_width - crops_top; + TFLITE_DCHECK_GE(out_h, 0); + TFLITE_DCHECK_LT(out_h, output_height); + + int in_w_start = 0; + int in_w_end = 0; + // GetIndexRange ensures start and end indices are in [0, output_width). + GetIndexRange(spatial_offset % block_shape_width - crops_left, + block_shape_width, input_width, output_width, &in_w_start, + &in_w_end); + + for (int in_w = in_w_start; in_w < in_w_end; ++in_w) { + const int out_w = in_w * block_shape_width + + spatial_offset % block_shape_width - crops_left; + TFLITE_DCHECK_GE(out_w, 0); + TFLITE_DCHECK_LT(out_w, output_width); T* out = output_data + Offset(output_dims, 0, out_w, out_h, out_batch); const T* in = input_data + Offset(input_dims, 0, in_w, in_h, in_batch); memcpy(out, in, depth * sizeof(T)); @@ -5240,7 +5472,7 @@ template inline void Pad(const T* input_data, const Dims<4>& input_dims, const std::vector& left_paddings, const std::vector& right_paddings, T* output_data, - const Dims<4>& output_dims) { + const Dims<4>& output_dims, const int32_t pad_value) { gemmlowp::ScopedProfilingLabel label("Pad"); const int output_batch = ArraySize(output_dims, 3); const int output_height = ArraySize(output_dims, 2); @@ -5260,27 +5492,27 @@ inline void Pad(const T* input_data, const Dims<4>& input_dims, const int input_depth = ArraySize(input_dims, 0); if (left_b_padding != 0) { - memset(output_data, 0, + memset(output_data, pad_value, left_b_padding * output_height * output_width * output_depth * sizeof(T)); } for (int out_b = left_b_padding; out_b < output_batch - right_b_padding; ++out_b) { if (left_h_padding != 0) { - memset(output_data + Offset(output_dims, 0, 0, 0, out_b), 0, + memset(output_data + Offset(output_dims, 0, 0, 0, out_b), pad_value, left_h_padding * output_width * output_depth * sizeof(T)); } for (int out_h = left_h_padding; out_h < output_height - right_h_padding; ++out_h) { if (left_w_padding != 0) { - memset(output_data + Offset(output_dims, 0, 0, out_h, out_b), 0, + memset(output_data + Offset(output_dims, 0, 0, out_h, out_b), pad_value, left_w_padding * output_depth * sizeof(T)); } for (int out_w = left_w_padding; out_w < output_width - right_w_padding; ++out_w) { if (left_d_padding != 0) { - memset(output_data + Offset(output_dims, 0, out_w, out_h, out_b), 0, - left_d_padding * sizeof(T)); + memset(output_data + Offset(output_dims, 0, out_w, out_h, out_b), + pad_value, left_d_padding * sizeof(T)); } T* out = output_data + @@ -5294,20 +5526,21 @@ inline void Pad(const T* input_data, const Dims<4>& input_dims, memset( output_data + Offset(output_dims, output_depth - right_d_padding, out_w, out_h, out_b), - 0, right_d_padding * sizeof(T)); + pad_value, right_d_padding * sizeof(T)); } } if (right_w_padding != 0) { memset( output_data + Offset(output_dims, 0, output_width - right_w_padding, out_h, out_b), - 0, right_w_padding * output_depth * sizeof(T)); + pad_value, right_w_padding * output_depth * sizeof(T)); } } if (right_h_padding != 0) { memset(output_data + Offset(output_dims, 0, 0, output_height - right_h_padding, out_b), - 0, right_h_padding * output_width * output_depth * sizeof(T)); + pad_value, + right_h_padding * output_width * output_depth * sizeof(T)); } } if (right_b_padding != 0) { @@ -5319,43 +5552,136 @@ inline void Pad(const T* input_data, const Dims<4>& input_dims, } } +template +inline void Pad(const T* input_data, const Dims<4>& input_dims, + const std::vector& left_paddings, + const std::vector& right_paddings, T* output_data, + const Dims<4>& output_dims) { + Pad(input_data, input_dims, left_paddings, right_paddings, output_data, + output_dims, 0); +} + +// UNOPTIMIZED COPY of StridedSlice from reference_ops.h (see comments there). + +// Use until std::clamp() is available from C++17. +inline int Clamp(const int v, const int lo, const int hi) { + TFLITE_DCHECK(!(hi < lo)); + if (hi < v) return hi; + if (v < lo) return lo; + return v; +} + +inline int StartForAxis(int begin_mask, const std::vector& start_indices, + const std::vector& strides, + const Dims<4>& input_shape, int axis) { + // Begin with the specified index + int start = start_indices[axis]; + + // begin_mask override + if (begin_mask & 1 << axis) { + if (strides[axis] > 0) { + // Forward iteration - use the first element. These values will get + // clamped below (Note: We could have set them to 0 and axis_size-1, but + // use lowest() and max() to maintain symmetry with StopForAxis()) + start = std::numeric_limits::lowest(); + } else { + // Backward iteration - use the last element. + start = std::numeric_limits::max(); + } + } + + // Handle negative indices + int axis_size = input_shape.sizes[axis]; + if (start < 0) { + start += axis_size; + } + + // Clamping + start = Clamp(start, 0, axis_size - 1); + + return start; +} + +inline int StopForAxis(int end_mask, const std::vector& stop_indices, + const std::vector& strides, + const Dims<4>& input_shape, int axis) { + // Begin with the specified index + int stop = stop_indices[axis]; + + // end_mask override + if (end_mask & (1 << axis)) { + if (strides[axis] > 0) { + // Forward iteration - use the last element. These values will get + // clamped below + stop = std::numeric_limits::max(); + } else { + // Backward iteration - use the first element. + stop = std::numeric_limits::lowest(); + } + } + + // Handle negative indices + int axis_size = input_shape.sizes[axis]; + if (stop < 0) { + stop += axis_size; + } + + // Clamping + // Because the end index points one past the last element, we need slightly + // different clamping ranges depending on the direction. + if (strides[axis] > 0) { + // Forward iteration + stop = Clamp(stop, 0, axis_size); + } else { + // Backward iteration + stop = Clamp(stop, -1, axis_size - 1); + } + + return stop; +} + +inline bool LoopCondition(int index, int stop, int stride) { + // True when we have reached the end of an axis and should loop. + return stride > 0 ? index >= stop : index <= stop; +} + template inline void StridedSlice(const T* input_data, const Dims<4>& input_dims, int begin_mask, int end_mask, - const std::vector& starts, - const std::vector& stops, + const std::vector& start_indices, + const std::vector& stop_indices, const std::vector& strides, T* output_data, const Dims<4>& output_dims) { - gemmlowp::ScopedProfilingLabel label("StridedSlice"); - const int start_b = (begin_mask & 8) ? 0 : starts[3]; - const int stop_b = (end_mask & 8) ? input_dims.sizes[3] : stops[3]; - const int start_h = (begin_mask & 4) ? 0 : starts[2]; - const int stop_h = (end_mask & 4) ? input_dims.sizes[2] : stops[2]; - const int start_w = (begin_mask & 2) ? 0 : starts[1]; - const int stop_w = (end_mask & 2) ? input_dims.sizes[1] : stops[1]; - const int start_d = (begin_mask & 1) ? 0 : starts[0]; - const int stop_d = (end_mask & 1) ? input_dims.sizes[0] : stops[0]; + TFLITE_DCHECK_EQ(start_indices.size(), 4); + TFLITE_DCHECK_EQ(stop_indices.size(), 4); + TFLITE_DCHECK_EQ(strides.size(), 4); + const int start_b = + StartForAxis(begin_mask, start_indices, strides, input_dims, 3); + const int stop_b = + StopForAxis(end_mask, stop_indices, strides, input_dims, 3); + const int start_h = + StartForAxis(begin_mask, start_indices, strides, input_dims, 2); + const int stop_h = + StopForAxis(end_mask, stop_indices, strides, input_dims, 2); + const int start_w = + StartForAxis(begin_mask, start_indices, strides, input_dims, 1); + const int stop_w = + StopForAxis(end_mask, stop_indices, strides, input_dims, 1); + const int start_d = + StartForAxis(begin_mask, start_indices, strides, input_dims, 0); + const int stop_d = + StopForAxis(end_mask, stop_indices, strides, input_dims, 0); T* out_ptr = output_data; - if (strides[0] == 0) { - for (int in_b = start_b; in_b < stop_b; in_b += strides[3]) { - for (int in_h = start_h; in_h < stop_h; in_h += strides[2]) { - for (int in_w = start_w; in_w < stop_w; in_w += strides[1]) { - const int len = stop_d - start_d; - memcpy(out_ptr, - input_data + Offset(input_dims, start_d, in_w, in_h, in_b), - len * sizeof(T)); - out_ptr += len; - } - } - } - } else { - for (int in_b = start_b; in_b < stop_b; in_b += strides[3]) { - for (int in_h = start_h; in_h < stop_h; in_h += strides[2]) { - for (int in_w = start_w; in_w < stop_w; in_w += strides[1]) { - for (int in_d = start_d; in_d < stop_d; in_d += strides[0]) { - *out_ptr++ = input_data[Offset(input_dims, in_d, in_w, in_h, in_b)]; - } + for (int in_b = start_b; !LoopCondition(in_b, stop_b, strides[3]); + in_b += strides[3]) { + for (int in_h = start_h; !LoopCondition(in_h, stop_h, strides[2]); + in_h += strides[2]) { + for (int in_w = start_w; !LoopCondition(in_w, stop_w, strides[1]); + in_w += strides[1]) { + for (int in_d = start_d; !LoopCondition(in_d, stop_d, strides[0]); + in_d += strides[0]) { + *out_ptr++ = input_data[Offset(input_dims, in_d, in_w, in_h, in_b)]; } } } diff --git a/tensorflow/contrib/lite/kernels/internal/quantization_util.cc b/tensorflow/contrib/lite/kernels/internal/quantization_util.cc index 18be6777a5caeb45a4ffabd8b7f1793de7b053f8..b0951aac8cbb98a181d9dcaef88770fadfc74f62 100644 --- a/tensorflow/contrib/lite/kernels/internal/quantization_util.cc +++ b/tensorflow/contrib/lite/kernels/internal/quantization_util.cc @@ -78,6 +78,22 @@ void PreprocessSoftmaxScaling(double beta, double input_scale, quantized_multiplier, left_shift); } +void PreprocessLogSoftmaxScaling(double beta, double input_scale, + int input_integer_bits, + int32_t* quantized_multiplier, int* left_shift, + int32_t* reverse_scaling_divisor, + int* reverse_scaling_right_shift) { + PreprocessSoftmaxScaling(beta, input_scale, input_integer_bits, + quantized_multiplier, left_shift); + + // Also calculate what amounts to the inverse scaling factor for the input. + const double real_reverse_scaling_divisor = + (1 << (31 - *left_shift)) / static_cast(*quantized_multiplier); + tflite::QuantizeMultiplierSmallerThanOne(real_reverse_scaling_divisor, + reverse_scaling_divisor, + reverse_scaling_right_shift); +} + int CalculateInputRadius(int input_integer_bits, int input_left_shift) { const double max_input_rescaled = 1.0 * ((1 << input_integer_bits) - 1) * (1ll << (31 - input_integer_bits)) / @@ -88,4 +104,25 @@ int CalculateInputRadius(int input_integer_bits, int input_left_shift) { return static_cast(std::floor(max_input_rescaled)); } +void NudgeQuantizationRange(const float min, const float max, + const int quant_min, const int quant_max, + float* nudged_min, float* nudged_max, + float* scale) { + // This code originates from tensorflow/core/kernels/fake_quant_ops_functor.h. + const float quant_min_float = static_cast(quant_min); + const float quant_max_float = static_cast(quant_max); + *scale = (max - min) / (quant_max_float - quant_min_float); + const float zero_point_from_min = quant_min_float - min / *scale; + uint16 nudged_zero_point; + if (zero_point_from_min < quant_min_float) { + nudged_zero_point = static_cast(quant_min); + } else if (zero_point_from_min > quant_max_float) { + nudged_zero_point = static_cast(quant_max); + } else { + nudged_zero_point = static_cast(TfLiteRound(zero_point_from_min)); + } + *nudged_min = (quant_min_float - nudged_zero_point) * (*scale); + *nudged_max = (quant_max_float - nudged_zero_point) * (*scale); +} + } // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/quantization_util.h b/tensorflow/contrib/lite/kernels/internal/quantization_util.h index 9a04b76e56b2527b06f5b0ec1e75e991fd1cbdea..4a217515f142b2451ebd61e423871b95cdc09748 100644 --- a/tensorflow/contrib/lite/kernels/internal/quantization_util.h +++ b/tensorflow/contrib/lite/kernels/internal/quantization_util.h @@ -196,7 +196,12 @@ void QuantizeMultiplier(double double_multiplier, int32_t* quantized_multiplier, void PreprocessSoftmaxScaling(double beta, double input_scale, int input_integer_bits, int32_t* quantized_multiplier, int* left_shift); - +// Like PreprocessSoftmaxScaling, but inverse scaling factors also calculated. +void PreprocessLogSoftmaxScaling(double beta, double input_scale, + int input_integer_bits, + int32_t* quantized_multiplier, int* left_shift, + int32_t* reverse_scaling_divisor, + int* reverse_scaling_right_shift); // Calculate the largest input that will result in a within-bounds intermediate // result within MultiplyByQuantizedMultiplierGreaterThanOne. In other words, // it must not overflow before we reduce the value by multiplication by the @@ -204,6 +209,14 @@ void PreprocessSoftmaxScaling(double beta, double input_scale, // Softmax. int CalculateInputRadius(int input_integer_bits, int input_left_shift); +// Nudges a min/max quantization range to ensure zero is zero. +// Gymnastics with nudged zero point is to ensure that real zero maps to +// an integer, which is required for e.g. zero-padding in convolutional layers. +// Outputs nudged_min, nudged_max, nudged_scale. +void NudgeQuantizationRange(const float min, const float max, + const int quant_min, const int quant_max, + float* nudged_min, float* nudged_max, float* scale); + } // namespace tflite #endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_QUANTIZATION_UTIL_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h index 31e190e2480f6af389a65b66a5c65374e1e42ee1..9ad125b8eba7f6457d582cd3ee512b87234a0e95 100644 --- a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h +++ b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h @@ -27,6 +27,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/quantization_util.h" #include "tensorflow/contrib/lite/kernels/internal/round.h" #include "tensorflow/contrib/lite/kernels/internal/types.h" @@ -601,6 +602,67 @@ inline void FullyConnected(const uint8* input_data, const Dims<4>& input_dims, } } +inline void ExperimentalShuffledFullyConnected( + const uint8* input_data, const Dims<4>& input_dims, + const uint8* shuffled_weights_data, const Dims<4>& weights_dims, + const int32* bias_data, const Dims<4>& bias_dims, 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); + // 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(weights_dims, 1, output_dims, 0); + const int accum_depth = ArraySize(weights_dims, 0); + TFLITE_DCHECK(IsPackedWithoutStrides(input_dims)); + TFLITE_DCHECK(IsPackedWithoutStrides(weights_dims)); + // The experimental shuffling is an optimization for matrix*vector product. + // We aren't interested in supporting non-matrix*vector-product cases, i.e. + // batches>1. + TFLITE_DCHECK_EQ(batches, 1); + // Shuffled weights have had their sign bit (0x80) pre-flipped (xor'd) + // so that just reinterpreting them as int8 values is equivalent to + // subtracting 128 from them, thus implementing for free the subtraction of + // the zero_point value 128. + const int8* shuffled_weights_ptr = + reinterpret_cast(shuffled_weights_data); + for (int c = 0; c < output_depth; c += 4) { + // Internal accumulation. + // Initialize accumulator with the bias-value. + int32 accum[4] = {0}; + // Accumulation loop. + for (int d = 0; d < accum_depth; d += 16) { + for (int i = 0; i < 4; i++) { + for (int j = 0; j < 16; j++) { + int8 input_val = input_data[d + j] - 128; + int8 weights_val = *shuffled_weights_ptr++; + accum[i] += weights_val * input_val; + } + } + } + for (int i = 0; i < 4; i++) { + // Add bias value + int acc = accum[i] + bias_data[c + i]; + // 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). + acc = + MultiplyByQuantizedMultiplier(acc, output_multiplier, -output_shift); + // Saturate, cast to int16, and store to output array. + acc = std::max(acc, output_activation_min); + acc = std::min(acc, output_activation_max); + output_data[c + i] = acc; + } + } +} + // legacy, for compatibility with old checked-in code template void FullyConnected(const uint8* input_data, const Dims<4>& input_dims, @@ -635,27 +697,14 @@ void NonGlobalBatchNormalization( const Dims<4>& offset_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, mean_dims, 2, multiplier_dims, 2, - offset_dims, 2, output_dims, 2); - const int width = - MatchingArraySize(input_dims, 1, mean_dims, 1, multiplier_dims, 1, - offset_dims, 1, output_dims, 1); - const int depth = - MatchingArraySize(input_dims, 0, mean_dims, 0, multiplier_dims, 0, - offset_dims, 0, output_dims, 0); + const int inner_size = MatchingFlatSizeSkipDim( + input_dims, 3, mean_dims, multiplier_dims, offset_dims, output_dims); for (int b = 0; b < batches; ++b) { - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - for (int c = 0; c < depth; ++c) { - output_data[Offset(output_dims, c, x, y, b)] = ActivationFunction( - (input_data[Offset(input_dims, c, x, y, b)] - - mean_data[Offset(mean_dims, c, x, y, 0)]) * - multiplier_data[Offset(multiplier_dims, c, x, y, 0)] + - offset_data[Offset(offset_dims, c, x, y, 0)]); - } - } + for (int i = 0; i < inner_size; ++i) { + output_data[b * inner_size + i] = ActivationFunction( + (input_data[b * inner_size + i] - mean_data[i]) * multiplier_data[i] + + offset_data[i]); } } } @@ -669,87 +718,52 @@ void GlobalBatchNormalization(const float* input_data, const float* offset_data, const Dims<4>& offset_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 outer_size = MatchingFlatSizeSkipDim(input_dims, 0, output_dims); const int depth = MatchingArraySize(input_dims, 0, mean_dims, 0, multiplier_dims, 0, offset_dims, 0, output_dims, 0); - for (int b = 0; b < batches; ++b) { - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - for (int c = 0; c < depth; ++c) { - output_data[Offset(output_dims, c, x, y, b)] = ActivationFunction( - (input_data[Offset(input_dims, c, x, y, b)] - - mean_data[Offset(mean_dims, c, 0, 0, 0)]) * - multiplier_data[Offset(multiplier_dims, c, 0, 0, 0)] + - offset_data[Offset(offset_dims, c, 0, 0, 0)]); - } - } + for (int i = 0; i < outer_size; ++i) { + for (int c = 0; c < depth; ++c) { + output_data[depth * i + c] = ActivationFunction( + (input_data[depth * i + c] - mean_data[c]) * multiplier_data[c] + + offset_data[c]); } } } inline void Relu(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) { - for (int c = 0; c < depth; ++c) { - float val = input_data[Offset(input_dims, c, x, y, b)]; - const float lower = 0; - float clamped = val < lower ? lower : val; - output_data[Offset(output_dims, c, x, y, b)] = clamped; - } - } - } + const int flat_size = MatchingFlatSize(input_dims, output_dims); + for (int i = 0; i < flat_size; ++i) { + const float val = input_data[i]; + const float lower = 0; + const float clamped = val < lower ? lower : val; + output_data[i] = clamped; } } inline void Relu1(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) { - for (int c = 0; c < depth; ++c) { - float val = input_data[Offset(input_dims, c, x, y, b)]; - const float upper = 1; - const float lower = -1; - float clamped = val > upper ? upper : val < lower ? lower : val; - output_data[Offset(output_dims, c, x, y, b)] = clamped; - } - } - } + const int flat_size = MatchingFlatSize(input_dims, output_dims); + for (int i = 0; i < flat_size; ++i) { + const float val = input_data[i]; + const float upper = 1; + const float lower = -1; + const float clamped = val > upper ? upper : val < lower ? lower : val; + output_data[i] = clamped; } } inline void Relu6(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) { - for (int c = 0; c < depth; ++c) { - float val = input_data[Offset(input_dims, c, x, y, b)]; - const float upper = 6; - const float lower = 0; - float clamped = val > upper ? upper : val < lower ? lower : val; - output_data[Offset(output_dims, c, x, y, b)] = clamped; - } - } - } + const int flat_size = MatchingFlatSize(input_dims, output_dims); + for (int i = 0; i < flat_size; ++i) { + const float val = input_data[i]; + const float upper = 6; + const float lower = 0; + const float clamped = val > upper ? upper : val < lower ? lower : val; + output_data[i] = clamped; } } @@ -757,24 +771,17 @@ template void L2Normalization(const float* input_data, const Dims<4>& input_dims, float* output_data, const Dims<4>& output_dims) { static_assert(Ac == FusedActivationFunctionType::kNone, ""); - 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 outer_size = MatchingFlatSizeSkipDim(input_dims, 0, output_dims); 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) { - float squared_l2_norm = 0; - for (int c = 0; c < depth; ++c) { - float val = input_data[Offset(input_dims, c, x, y, b)]; - squared_l2_norm += val * val; - } - float l2_norm = std::sqrt(squared_l2_norm); - 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)] / l2_norm; - } - } + for (int i = 0; i < outer_size; ++i) { + float squared_l2_norm = 0; + for (int c = 0; c < depth; ++c) { + const float val = input_data[depth * i + c]; + squared_l2_norm += val * val; + } + const float l2_norm = std::sqrt(squared_l2_norm); + for (int c = 0; c < depth; ++c) { + output_data[depth * i + c] = input_data[depth * i + c] / l2_norm; } } } @@ -859,26 +866,11 @@ inline void Add(const float* input1_data, const Dims<4>& input1_dims, const float* input2_data, const Dims<4>& input2_dims, float output_activation_min, float output_activation_max, float* output_data, const Dims<4>& output_dims) { - const int batches = - MatchingArraySize(input1_dims, 3, input2_dims, 3, output_dims, 3); - const int height = - MatchingArraySize(input1_dims, 2, input2_dims, 2, output_dims, 2); - const int width = - MatchingArraySize(input1_dims, 1, input2_dims, 1, output_dims, 1); - const int depth = - MatchingArraySize(input1_dims, 0, input2_dims, 0, output_dims, 0); - for (int b = 0; b < batches; ++b) { - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - for (int c = 0; c < depth; ++c) { - output_data[Offset(output_dims, c, x, y, b)] = - ActivationFunctionWithMinMax( - input1_data[Offset(input1_dims, c, x, y, b)] + - input2_data[Offset(input2_dims, c, x, y, b)], - output_activation_min, output_activation_max); - } - } - } + const int flat_size = MatchingFlatSize(input1_dims, input2_dims, output_dims); + for (int i = 0; i < flat_size; ++i) { + output_data[i] = ActivationFunctionWithMinMax( + input1_data[i] + input2_data[i], output_activation_min, + output_activation_max); } } @@ -1141,26 +1133,11 @@ inline void Mul(const float* input1_data, const Dims<4>& input1_dims, const float* input2_data, const Dims<4>& input2_dims, float output_activation_min, float output_activation_max, float* output_data, const Dims<4>& output_dims) { - const int batches = - MatchingArraySize(input1_dims, 3, input2_dims, 3, output_dims, 3); - const int height = - MatchingArraySize(input1_dims, 2, input2_dims, 2, output_dims, 2); - const int width = - MatchingArraySize(input1_dims, 1, input2_dims, 1, output_dims, 1); - const int depth = - MatchingArraySize(input1_dims, 0, input2_dims, 0, output_dims, 0); - for (int b = 0; b < batches; ++b) { - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - for (int c = 0; c < depth; ++c) { - output_data[Offset(output_dims, c, x, y, b)] = - ActivationFunctionWithMinMax( - input1_data[Offset(input1_dims, c, x, y, b)] * - input2_data[Offset(input2_dims, c, x, y, b)], - output_activation_min, output_activation_max); - } - } - } + const int flat_size = MatchingFlatSize(input1_dims, input2_dims, output_dims); + for (int i = 0; i < flat_size; ++i) { + output_data[i] = ActivationFunctionWithMinMax( + input1_data[i] * input2_data[i], output_activation_min, + output_activation_max); } } @@ -1339,6 +1316,33 @@ inline void BroadcastMul(const uint8* input1_data, const Dims<4>& input1_dims, output_data, output_dims); } +inline void Div(const float* input1_data, const Dims<4>& input1_dims, + const float* input2_data, const Dims<4>& input2_dims, + float output_activation_min, float output_activation_max, + float* output_data, const Dims<4>& output_dims) { + const int batches = + MatchingArraySize(input1_dims, 3, input2_dims, 3, output_dims, 3); + const int height = + MatchingArraySize(input1_dims, 2, input2_dims, 2, output_dims, 2); + const int width = + MatchingArraySize(input1_dims, 1, input2_dims, 1, output_dims, 1); + const int depth = + MatchingArraySize(input1_dims, 0, input2_dims, 0, output_dims, 0); + for (int b = 0; b < batches; ++b) { + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + for (int c = 0; c < depth; ++c) { + output_data[Offset(output_dims, c, x, y, b)] = + ActivationFunctionWithMinMax( + input1_data[Offset(input1_dims, c, x, y, b)] / + input2_data[Offset(input2_dims, c, x, y, b)], + output_activation_min, output_activation_max); + } + } + } + } +} + // 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 @@ -1380,57 +1384,15 @@ void BroadcastDiv(const T* input1_data, const Dims<4>& input1_dims, } } -inline void Div(const float* input1_data, const Dims<4>& input1_dims, - const float* input2_data, const Dims<4>& input2_dims, - float output_activation_min, float output_activation_max, - float* output_data, const Dims<4>& output_dims) { - const int batches = - MatchingArraySize(input1_dims, 3, input2_dims, 3, output_dims, 3); - const int height = - MatchingArraySize(input1_dims, 2, input2_dims, 2, output_dims, 2); - const int width = - MatchingArraySize(input1_dims, 1, input2_dims, 1, output_dims, 1); - const int depth = - MatchingArraySize(input1_dims, 0, input2_dims, 0, output_dims, 0); - for (int b = 0; b < batches; ++b) { - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - for (int c = 0; c < depth; ++c) { - output_data[Offset(output_dims, c, x, y, b)] = - ActivationFunctionWithMinMax( - input1_data[Offset(input1_dims, c, x, y, b)] / - input2_data[Offset(input2_dims, c, x, y, b)], - output_activation_min, output_activation_max); - } - } - } - } -} - inline void Sub(const float* input1_data, const Dims<4>& input1_dims, const float* input2_data, const Dims<4>& input2_dims, float output_activation_min, float output_activation_max, float* output_data, const Dims<4>& output_dims) { - const int batches = - MatchingArraySize(input1_dims, 3, input2_dims, 3, output_dims, 3); - const int height = - MatchingArraySize(input1_dims, 2, input2_dims, 2, output_dims, 2); - const int width = - MatchingArraySize(input1_dims, 1, input2_dims, 1, output_dims, 1); - const int depth = - MatchingArraySize(input1_dims, 0, input2_dims, 0, output_dims, 0); - for (int b = 0; b < batches; ++b) { - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - for (int c = 0; c < depth; ++c) { - output_data[Offset(output_dims, c, x, y, b)] = - ActivationFunctionWithMinMax( - input1_data[Offset(input1_dims, c, x, y, b)] - - input2_data[Offset(input2_dims, c, x, y, b)], - output_activation_min, output_activation_max); - } - } - } + const int flat_size = MatchingFlatSize(input1_dims, input2_dims, output_dims); + for (int i = 0; i < flat_size; ++i) { + output_data[i] = ActivationFunctionWithMinMax( + input1_data[i] - input2_data[i], output_activation_min, + output_activation_max); } } @@ -1812,15 +1774,9 @@ void LstmCell(const uint8* input_data_uint8, const Dims<4>& input_dims, (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); + const int outer_size = + MatchingFlatSizeSkipDim(input_dims, 0, prev_activ_dims, prev_state_dims, + output_state_dims, output_activ_dims); TFLITE_CHECK_EQ(ArraySize(weights_dims, 2), 1); TFLITE_CHECK_EQ(ArraySize(weights_dims, 3), 1); const int input_depth = ArraySize(input_dims, 0); @@ -1836,9 +1792,7 @@ void LstmCell(const uint8* input_data_uint8, const Dims<4>& input_dims, 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_batches = FlatSizeSkipDim(activ_temp_dims, 0); const int fc_output_depth = MatchingArraySize(weights_dims, 1, activ_temp_dims, 0); const int fc_accum_depth = ArraySize(weights_dims, 0); @@ -1883,7 +1837,6 @@ void LstmCell(const uint8* input_data_uint8, const Dims<4>& input_dims, // 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 @@ -2418,28 +2371,20 @@ inline void LocalResponseNormalization(const float* input_data, float bias, float alpha, float beta, 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 outer_size = MatchingFlatSizeSkipDim(input_dims, 0, output_dims); 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 int begin_input_c = std::max(0, c - range); - const int end_input_c = std::min(depth, c + range); - float accum = 0.f; - for (int input_c = begin_input_c; input_c < end_input_c; ++input_c) { - const float input_val = - input_data[Offset(input_dims, input_c, x, y, b)]; - accum += input_val * input_val; - } - const float multiplier = std::pow(bias + alpha * accum, -beta); - output_data[Offset(output_dims, c, x, y, b)] = - input_data[Offset(input_dims, c, x, y, b)] * multiplier; - } + for (int i = 0; i < outer_size; ++i) { + for (int c = 0; c < depth; ++c) { + const int begin_input_c = std::max(0, c - range); + const int end_input_c = std::min(depth, c + range); + float accum = 0.f; + for (int input_c = begin_input_c; input_c < end_input_c; ++input_c) { + const float input_val = input_data[i * depth + input_c]; + accum += input_val * input_val; } + const float multiplier = std::pow(bias + alpha * accum, -beta); + output_data[i * depth + c] = input_data[i * depth + c] * multiplier; } } } @@ -2447,37 +2392,28 @@ inline void LocalResponseNormalization(const float* input_data, inline void Softmax(const float* input_data, const Dims<4>& input_dims, float beta, 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 outer_size = MatchingFlatSizeSkipDim(input_dims, 0, output_dims); 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: - // exp(x[i])/sum(exp(x[i])) == 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)]); - } + for (int i = 0; i < outer_size; ++i) { + // Find max element value which we'll use to ensure numerical stability + // taking advantage of the following equality: + // exp(x[i])/sum(exp(x[i])) == 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[i * depth + c]); + } - // 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) * - beta); - } + // Compute sum. + float sum = 0.f; + for (int c = 0; c < depth; ++c) { + sum += std::exp((input_data[i * depth + c] - max) * beta); + } - // Compute result. - for (int c = 0; c < depth; ++c) { - output_data[Offset(output_dims, c, x, y, b)] = - std::exp((input_data[Offset(input_dims, c, x, y, b)] - max) * - beta) / - sum; - } - } + // Compute result. + for (int c = 0; c < depth; ++c) { + output_data[i * depth + c] = + std::exp((input_data[i * depth + c] - max) * beta) / sum; } } } @@ -2498,73 +2434,63 @@ inline void Softmax(const uint8* input_data, const Dims<4>& input_dims, using FixedPointAccum = gemmlowp::FixedPoint; using FixedPoint0 = gemmlowp::FixedPoint; - 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 outer_size = MatchingFlatSizeSkipDim(input_dims, 0, output_dims); const int depth = MatchingArraySize(input_dims, 0, output_dims, 0); - for (int b = 0; b < batches; ++b) { - for (int x = 0; x < width; ++x) { - for (int y = 0; y < height; ++y) { - uint8 max_in_row = 0; - for (int c = 0; c < depth; ++c) { - max_in_row = - std::max(max_in_row, input_data[Offset(input_dims, c, x, y, b)]); - } + for (int i = 0; i < outer_size; ++i) { + uint8 max_in_row = 0; + for (int c = 0; c < depth; ++c) { + max_in_row = std::max(max_in_row, input_data[i * depth + c]); + } + + FixedPointAccum sum_of_exps = FixedPointAccum::Zero(); + for (int c = 0; c < depth; ++c) { + int32 input_diff = + static_cast(input_data[i * depth + 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)); + } + } + + int32 fixed_sum_of_exps = sum_of_exps.raw(); + int headroom_plus_one = + CountLeadingZeros(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)); + + for (int c = 0; c < depth; ++c) { + int32 input_diff = + static_cast(input_data[i * depth + 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[i * depth + c] = static_cast( + std::max(std::min(unsat_output, static_cast(255)), 0)); - FixedPointAccum sum_of_exps = FixedPointAccum::Zero(); - for (int c = 0; c < depth; ++c) { - int32 input_diff = - static_cast(input_data[Offset(input_dims, c, x, y, b)]) - - max_in_row; - if (input_diff >= diff_min) { - const int32 input_diff_rescaled = - MultiplyByQuantizedMultiplierGreaterThanOne( - input_diff, input_beta_multiplier, input_beta_left_shift); - const FixedPointScaledDiff scaled_diff_f8 = - FixedPointScaledDiff::FromRaw(input_diff_rescaled); - sum_of_exps = - sum_of_exps + gemmlowp::Rescale( - exp_on_negative_values(scaled_diff_f8)); - } - } - - int32 fixed_sum_of_exps = sum_of_exps.raw(); - int headroom_plus_one = - CountLeadingZeros(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)); - - 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)] = static_cast( - std::max(std::min(unsat_output, static_cast(255)), 0)); - - } else { - output_data[Offset(output_dims, c, x, y, b)] = 0; - } - } + } else { + output_data[i * depth + c] = 0; } } } @@ -2572,109 +2498,170 @@ 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 outer_size = MatchingFlatSizeSkipDim(input_dims, 0, output_dims); 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)]); - } + for (int i = 0; i < outer_size; ++i) { + // 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[i * depth + c]); + } - // 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 sum. + float sum = 0.f; + for (int c = 0; c < depth; ++c) { + sum += std::exp(input_data[i * depth + c] - 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; - } - } + // Compute result. + const float log_sum = std::log(sum); + for (int c = 0; c < depth; ++c) { + output_data[i * depth + c] = input_data[i * depth + c] - 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); - const int height = MatchingArraySize(input_dims, 2, output_dims, 2); - const int width = MatchingArraySize(input_dims, 1, output_dims, 1); +inline void LogSoftmax(const uint8* input_data, const Dims<4>& input_dims, + int32 input_multiplier, int32 input_left_shift, + int32 reverse_scaling_divisor, + int32 reverse_scaling_right_shift, int diff_min, + uint8* output_data, const Dims<4>& output_dims) { + // The representation chosen for the input to the exp() function is Q5.26. + // We need to leave extra space since values that we skip might be as large as + // -32 before multiplying by input_beta_multiplier, and therefore as large as + // -16 afterwards. Note that exp(-8) is definitely not insignificant to + // accumulation, but exp(-16) definitely is. + static constexpr int kScaledDiffIntegerBits = 5; + static constexpr int kAccumulationIntegerBits = 12; + static constexpr int kOutputIntegerBits = 4; + using FixedPointScaledDiff = + gemmlowp::FixedPoint; + using FixedPointAccum = gemmlowp::FixedPoint; + using FixedPoint0 = gemmlowp::FixedPoint; + + const int outer_size = MatchingFlatSizeSkipDim(input_dims, 0, output_dims); 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) { - float val = input_data[Offset(input_dims, c, x, y, b)]; - float result = 1.f / (1.f + std::exp(-val)); - output_data[Offset(output_dims, c, x, y, b)] = result; - } + + for (int i = 0; i < outer_size; ++i) { + uint8 max_in_row = 0; + for (int c = 0; c < depth; ++c) { + max_in_row = std::max(max_in_row, input_data[i * depth + c]); + } + + FixedPointAccum sum_of_exps = FixedPointAccum::Zero(); + for (int c = 0; c < depth; ++c) { + int32 input_diff = + static_cast(input_data[i * depth + c]) - max_in_row; + if (input_diff >= diff_min) { + const int32 input_diff_rescaled = + MultiplyByQuantizedMultiplierGreaterThanOne( + input_diff, input_multiplier, input_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)); + } + } + + // TODO(b/77858996): Implement fixed-point log(). + // Not a fully-quantized implementation: floating-point log(). + const float float_log_sum_of_exps = + std::log(static_cast(sum_of_exps.raw()) / + (1 << (31 - kAccumulationIntegerBits))); + const int32 fixed_log_sum_of_exps = static_cast(TfLiteRound( + float_log_sum_of_exps * (1 << (31 - kScaledDiffIntegerBits)))); + + // rescaled_diff_min is smallest representable in + // Q(kScaledDiffIntegerBits).(31-kScaledDiffIntegerBits) plus the + // log-sub-exps that will be subtracted in the loop. + // + // The thresholds diff_min, etc are negative. + const int rescaled_diff_min = + fixed_log_sum_of_exps + std::numeric_limits::lowest(); + const int adjusted_diff_min = + std::max(diff_min - 1, // Note use of > below instead of >= above. + MultiplyByQuantizedMultiplierSmallerThanOne( + rescaled_diff_min, reverse_scaling_divisor, + reverse_scaling_right_shift)); + + for (int c = 0; c < depth; ++c) { + int32 input_diff = + static_cast(input_data[i * depth + c]) - max_in_row; + if (input_diff > adjusted_diff_min) { + const int32 input_diff_rescaled = + MultiplyByQuantizedMultiplierGreaterThanOne( + input_diff, input_multiplier, input_left_shift); + int32 unsat_output = + gemmlowp::RoundingDivideByPOT( + (input_diff_rescaled - fixed_log_sum_of_exps), + 31 - kScaledDiffIntegerBits - kOutputIntegerBits) + + 255; + + output_data[i * depth + c] = static_cast( + std::max(std::min(unsat_output, static_cast(255)), 0)); + } else { + // Set output to smallest value. + output_data[i * depth + c] = 0; } } } } +inline void Logistic(const float* input_data, const Dims<4>& input_dims, + float* output_data, const Dims<4>& output_dims) { + const int flat_size = MatchingFlatSize(output_dims, input_dims); + + for (int i = 0; i < flat_size; i++) { + float val = input_data[i]; + float result = 1.f / (1.f + std::exp(-val)); + output_data[i] = result; + } +} + inline void Logistic(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 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::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); - } - output_data[Offset(output_dims, c, x, y, b)] = output_val; - } + const int flat_size = MatchingFlatSize(output_dims, input_dims); + + for (int i = 0; i < flat_size; i++) { + const uint8 input_val_u8 = input_data[i]; + 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::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); } + output_data[i] = output_val; } } 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); + const int flat_size = MatchingFlatSize(output_dims, input_dims); for (int i = 0; i < flat_size; i++) { // F0 uses 0 integer bits, range [-1, 1]. @@ -2692,20 +2679,12 @@ inline void Logistic(const int16* input_data, const Dims<4>& input_dims, 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); - 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) { - float val = input_data[Offset(input_dims, c, x, y, b)]; - float result = std::tanh(val); - output_data[Offset(output_dims, c, x, y, b)] = result; - } - } - } + const int flat_size = MatchingFlatSize(output_dims, input_dims); + + for (int i = 0; i < flat_size; i++) { + float val = input_data[i]; + float result = std::tanh(val); + output_data[i] = result; } } @@ -2714,47 +2693,38 @@ inline void Tanh(const uint8* input_data, const Dims<4>& input_dims, 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; - } + const int flat_size = MatchingFlatSize(output_dims, input_dims); + + for (int i = 0; i < flat_size; i++) { + const uint8 input_val_u8 = input_data[i]; + 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[i] = output_val; } } @@ -2766,8 +2736,7 @@ inline void Tanh(const int16* input_data, const Dims<4>& input_dims, 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); + const int flat_size = MatchingFlatSize(output_dims, input_dims); // F0 uses 0 integer bits, range [-1, 1]. // This is the return type of math functions such as tanh, logistic, @@ -2795,138 +2764,62 @@ inline void Tanh(const int16* input_data, const Dims<4>& input_dims, 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) { - 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) { - int32 val = input_data[Offset(input_dims, c, x, y, b)]; - float result = static_cast(scale * (val - zero_point)); - output_data[Offset(output_dims, c, x, y, b)] = result; - } - } - } + const int flat_size = MatchingFlatSize(output_dims, input_dims); + + for (int i = 0; i < flat_size; i++) { + int32 val = input_data[i]; + float result = static_cast(scale * (val - zero_point)); + output_data[i] = result; } } inline void FakeQuant(const float* input_data, const Dims<4>& input_dims, - float rmin, float rmax, float* output_data, + float rmin, float rmax, int num_bits, float* output_data, const Dims<4>& output_dims) { // 0 should always be a representable value. Let's assume that the initial // min,max range contains 0. - TFLITE_DCHECK_LE(rmin, 0.); - TFLITE_DCHECK_GE(rmax, 0.); - - // Determine quantization parameters: zero_point, scale. - using Integer = uint8; - const Integer qmin = std::numeric_limits::min(); - const Integer qmax = std::numeric_limits::max(); - const float qmin_float = qmin; - const float qmax_float = qmax; - int32 zero_point = 0; - float scale = 0.f; - // If rmin==rmax, both must be zero per the above assertion, - // so we are done. - if (rmin != rmax) { - // First determine the scale. - scale = (rmax - rmin) / (qmax_float - qmin_float); - - // 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 float zero_point_from_min = qmin_float - rmin / scale; - const float zero_point_from_max = qmax_float - rmax / scale; - const float zero_point_from_min_error = - std::abs(qmin_float) + std::abs(rmin / scale); - const float zero_point_from_max_error = - std::abs(qmax_float) + std::abs(rmax / scale); - - const float zero_point_float = - 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). - if (zero_point_float < qmin_float) { - zero_point = qmin; - } else if (zero_point_float > qmax_float) { - zero_point = qmax; - } else { - zero_point = static_cast(TfLiteRound(zero_point_float)); - } - // The zero point should always be in the range of quantized value, - // [qmin, qmax]. - TFLITE_DCHECK_GE(zero_point, qmin); - TFLITE_DCHECK_LE(zero_point, qmax); - } - - 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 float src_val = input_data[Offset(input_dims, c, x, y, b)]; - const float unclamped_quantized_val = - TfLiteRound(zero_point + src_val / scale); - const float quantized_val = std::min( - qmax_float, std::max(qmin_float, unclamped_quantized_val)); - const float dst_val = scale * (quantized_val - zero_point); - output_data[Offset(output_dims, c, x, y, b)] = dst_val; - } - } - } + TFLITE_DCHECK_LE(rmin, 0.0f); + TFLITE_DCHECK_GE(rmax, 0.0f); + TFLITE_DCHECK_LT(rmin, rmax); + + // Code matches tensorflow's FakeQuantWithMinMaxArgsFunctor. + int quant_min = 0; + int quant_max = (1 << num_bits) - 1; + float nudged_min, nudged_max, nudged_scale; + NudgeQuantizationRange(rmin, rmax, quant_min, quant_max, &nudged_min, + &nudged_max, &nudged_scale); + const float inv_nudged_scale = 1.0f / nudged_scale; + + const int flat_size = MatchingFlatSize(output_dims, input_dims); + for (int i = 0; i < flat_size; i++) { + const float src_val = input_data[i]; + const float clamped = std::min(nudged_max, std::max(nudged_min, src_val)); + const float clamped_shifted = clamped - nudged_min; + const float dst_val = + TfLiteRound(clamped_shifted * inv_nudged_scale) * nudged_scale + + nudged_min; + output_data[i] = dst_val; } } template inline void Cast(const SrcT* input_data, const Dims<4>& input_dims, DstT* 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) { - for (int c = 0; c < depth; ++c) { - int offset = Offset(input_dims, c, x, y, b); - output_data[offset] = static_cast(input_data[offset]); - } - } - } + const int flat_size = MatchingFlatSize(output_dims, input_dims); + + for (int i = 0; i < flat_size; i++) { + int offset = i; + output_data[offset] = static_cast(input_data[offset]); } } inline void Floor(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) { - for (int c = 0; c < depth; ++c) { - int offset = Offset(input_dims, c, x, y, b); - output_data[offset] = std::floor(input_data[offset]); - } - } - } + const int flat_size = MatchingFlatSize(output_dims, input_dims); + + for (int i = 0; i < flat_size; i++) { + int offset = i; + output_data[offset] = std::floor(input_data[offset]); } } @@ -3056,24 +2949,37 @@ inline void SpaceToBatchND(const T* input_data, const Dims<4>& input_dims, template inline void BatchToSpaceND(const T* input_data, const Dims<4>& input_dims, const int32* block_shape_data, - const Dims<4>& block_shape_dims, T* output_data, - const Dims<4>& output_dims) { + const Dims<4>& block_shape_dims, + const int32* crops_data, const Dims<4>& crops_dims, + T* output_data, const Dims<4>& output_dims) { const int output_batch_size = ArraySize(output_dims, 3); + const int output_height = ArraySize(output_dims, 2); + const int output_width = ArraySize(output_dims, 1); const int input_batch_size = ArraySize(input_dims, 3); const int input_height = ArraySize(input_dims, 2); const int input_width = ArraySize(input_dims, 1); const int depth = ArraySize(input_dims, 0); const int block_shape_width = block_shape_data[1]; const int block_shape_height = block_shape_data[0]; + const int crops_top = crops_data[0]; + const int crops_left = crops_data[2]; for (int in_batch = 0; in_batch < input_batch_size; ++in_batch) { + const int out_batch = in_batch % output_batch_size; + const int spatial_offset = in_batch / output_batch_size; for (int in_h = 0; in_h < input_height; ++in_h) { + const int out_h = in_h * block_shape_height + + spatial_offset / block_shape_width - crops_top; + if (out_h < 0 || out_h >= output_height) { + continue; + } for (int in_w = 0; in_w < input_width; ++in_w) { - int out_batch = in_batch % output_batch_size; - int out_w = in_w * block_shape_width + - (in_batch / output_batch_size) % block_shape_width; - int out_h = in_h * block_shape_height + - (in_batch / output_batch_size) / block_shape_width; + const int out_w = in_w * block_shape_width + + spatial_offset % block_shape_width - crops_left; + + if (out_w < 0 || out_w >= output_width) { + continue; + } T* out = output_data + Offset(output_dims, 0, out_w, out_h, out_batch); const T* in = input_data + Offset(input_dims, 0, in_w, in_h, in_batch); memcpy(out, in, depth * sizeof(T)); @@ -3086,7 +2992,7 @@ template inline void Pad(const T* input_data, const Dims<4>& input_dims, const std::vector& left_paddings, const std::vector& right_paddings, T* output_data, - const Dims<4>& output_dims) { + const Dims<4>& output_dims, const int32_t pad_value) { const int output_batch = ArraySize(output_dims, 3); const int output_height = ArraySize(output_dims, 2); const int output_width = ArraySize(output_dims, 1); @@ -3116,7 +3022,7 @@ inline void Pad(const T* input_data, const Dims<4>& input_dims, out_w >= output_width - right_w_padding || out_d < left_d_padding || out_d >= output_depth - right_d_padding) { - *out_ptr++ = 0; + *out_ptr++ = static_cast(pad_value); } else { *out_ptr++ = *in_ptr++; } @@ -3126,59 +3032,148 @@ inline void Pad(const T* input_data, const Dims<4>& input_dims, } } -inline bool LoopCondition(int index, int stop, int stride) { - return stride > 0 ? index < stop : index > stop; +template +inline void Pad(const T* input_data, const Dims<4>& input_dims, + const std::vector& left_paddings, + const std::vector& right_paddings, T* output_data, + const Dims<4>& output_dims) { + Pad(input_data, input_dims, left_paddings, right_paddings, output_data, + output_dims, 0); } -inline int StartIndex(int start, int stride, int dim, bool masked) { - return masked ? (stride > 0 ? 0 : dim - 1) : start; +// STRIDED SLICE +// The functions below for StridedSlice are mirrored in a number of places: +// +// propagate_fixed_sizes.cc +// propagate_shapes.cc +// resolve_constant_strided_slice.cc +// optimized_ops.h +// +// It is designed for an arbitrary number of dimensions, even though dimensions +// here are fixed at 4. This is because we expect to eventually support +// arbitrary dimensionality. Also note that the axis orders are reversed for +// runtime ops, and so the indices and masks must be as well too. +// +// Be warned this code involves some rather subtle logic of python slicing. The +// best "ground truth" is to compare results to actual python execution. + +// Use until std::clamp() is available from C++17. +inline int Clamp(const int v, const int lo, const int hi) { + TFLITE_DCHECK(!(hi < lo)); + if (hi < v) return hi; + if (v < lo) return lo; + return v; +} + +inline int StartForAxis(int begin_mask, const std::vector& start_indices, + const std::vector& strides, + const Dims<4>& input_shape, int axis) { + // Begin with the specified index + int start = start_indices[axis]; + + // begin_mask override + if (begin_mask & 1 << axis) { + if (strides[axis] > 0) { + // Forward iteration - use the first element. These values will get + // clamped below (Note: We could have set them to 0 and axis_size-1, but + // use lowest() and max() to maintain symmetry with StopForAxis()) + start = std::numeric_limits::lowest(); + } else { + // Backward iteration - use the last element. + start = std::numeric_limits::max(); + } + } + + // Handle negative indices + int axis_size = input_shape.sizes[axis]; + if (start < 0) { + start += axis_size; + } + + // Clamping + start = Clamp(start, 0, axis_size - 1); + + return start; +} + +inline int StopForAxis(int end_mask, const std::vector& stop_indices, + const std::vector& strides, + const Dims<4>& input_shape, int axis) { + // Begin with the specified index + int stop = stop_indices[axis]; + + // end_mask override + if (end_mask & (1 << axis)) { + if (strides[axis] > 0) { + // Forward iteration - use the last element. These values will get + // clamped below + stop = std::numeric_limits::max(); + } else { + // Backward iteration - use the first element. + stop = std::numeric_limits::lowest(); + } + } + + // Handle negative indices + int axis_size = input_shape.sizes[axis]; + if (stop < 0) { + stop += axis_size; + } + + // Clamping + // Because the end index points one past the last element, we need slightly + // different clamping ranges depending on the direction. + if (strides[axis] > 0) { + // Forward iteration + stop = Clamp(stop, 0, axis_size); + } else { + // Backward iteration + stop = Clamp(stop, -1, axis_size - 1); + } + + return stop; } -inline int StopIndex(int start, int stop, int stride, int dim, bool masked, - bool shrink_axis_masked) { - return shrink_axis_masked ? stride > 0 ? start + 1 : start - 1 - : masked ? (stride > 0 ? dim : -1) : stop; +inline bool LoopCondition(int index, int stop, int stride) { + // True when we have reached the end of an axis and should loop. + return stride > 0 ? index >= stop : index <= stop; } template inline void StridedSlice(const T* input_data, const Dims<4>& input_dims, - int begin_mask, int end_mask, int shrink_axis_mask, - const std::vector& starts, - const std::vector& stops, + int begin_mask, int end_mask, + const std::vector& start_indices, + const std::vector& stop_indices, const std::vector& strides, T* output_data, const Dims<4>& output_dims) { - TFLITE_DCHECK_EQ(starts.size(), 4); - TFLITE_DCHECK_EQ(stops.size(), 4); + TFLITE_DCHECK_EQ(start_indices.size(), 4); + TFLITE_DCHECK_EQ(stop_indices.size(), 4); TFLITE_DCHECK_EQ(strides.size(), 4); const int start_b = - StartIndex(starts[3], strides[3], input_dims.sizes[3], begin_mask & 8); + StartForAxis(begin_mask, start_indices, strides, input_dims, 3); const int stop_b = - StopIndex(start_b, stops[3], strides[3], input_dims.sizes[3], - end_mask & 8, shrink_axis_mask & 8); + StopForAxis(end_mask, stop_indices, strides, input_dims, 3); const int start_h = - StartIndex(starts[2], strides[2], input_dims.sizes[2], begin_mask & 4); + StartForAxis(begin_mask, start_indices, strides, input_dims, 2); const int stop_h = - StopIndex(start_h, stops[2], strides[2], input_dims.sizes[2], - end_mask & 4, shrink_axis_mask & 4); + StopForAxis(end_mask, stop_indices, strides, input_dims, 2); const int start_w = - StartIndex(starts[1], strides[1], input_dims.sizes[1], begin_mask & 2); + StartForAxis(begin_mask, start_indices, strides, input_dims, 1); const int stop_w = - StopIndex(start_w, stops[1], strides[1], input_dims.sizes[1], - end_mask & 2, shrink_axis_mask & 2); + StopForAxis(end_mask, stop_indices, strides, input_dims, 1); const int start_d = - StartIndex(starts[0], strides[0], input_dims.sizes[0], begin_mask & 1); + StartForAxis(begin_mask, start_indices, strides, input_dims, 0); const int stop_d = - StopIndex(start_d, stops[0], strides[0], input_dims.sizes[0], - end_mask & 1, shrink_axis_mask & 1); + StopForAxis(end_mask, stop_indices, strides, input_dims, 0); T* out_ptr = output_data; - for (int in_b = start_b; LoopCondition(in_b, stop_b, strides[3]); + for (int in_b = start_b; !LoopCondition(in_b, stop_b, strides[3]); in_b += strides[3]) { - for (int in_h = start_h; LoopCondition(in_h, stop_h, strides[2]); + for (int in_h = start_h; !LoopCondition(in_h, stop_h, strides[2]); in_h += strides[2]) { - for (int in_w = start_w; LoopCondition(in_w, stop_w, strides[1]); + for (int in_w = start_w; !LoopCondition(in_w, stop_w, strides[1]); in_w += strides[1]) { - for (int in_d = start_d; LoopCondition(in_d, stop_d, strides[0]); + for (int in_d = start_d; !LoopCondition(in_d, stop_d, strides[0]); in_d += strides[0]) { *out_ptr++ = input_data[Offset(input_dims, in_d, in_w, in_h, in_b)]; } @@ -3187,18 +3182,6 @@ inline void StridedSlice(const T* input_data, const Dims<4>& input_dims, } } -template -inline void StridedSlice(const T* input_data, const Dims<4>& input_dims, - int begin_mask, int end_mask, - const std::vector& starts, - const std::vector& stops, - const std::vector& strides, T* output_data, - const Dims<4>& output_dims) { - StridedSlice(input_data, input_dims, begin_mask, end_mask, - /*shrink_axis_mask=*/0, starts, stops, strides, output_data, - output_dims); -} - template inline void Slice(const T* input_data, const Dims<4>& input_dims, const std::vector& begin, const std::vector& size, @@ -3375,23 +3358,11 @@ template void TensorFlowMinimum(const T* input1_data, const Dims<4>& input1_dims, const T* input2_data, T* output_data, const Dims<4>& output_dims) { - int batches = MatchingArraySize(input1_dims, 3, output_dims, 3); - int input_height = MatchingArraySize(input1_dims, 2, output_dims, 2); - int input_width = MatchingArraySize(input1_dims, 1, output_dims, 1); - int depth = MatchingArraySize(input1_dims, 0, output_dims, 0); + const int flat_size = MatchingFlatSize(output_dims, input1_dims); auto min_value = input2_data[0]; - - for (int b = 0; b < batches; b++) { - for (int y = 0; y < input_height; y++) { - for (int x = 0; x < input_width; x++) { - for (int c = 0; c < depth; c++) { - int offset = Offset(input1_dims, c, x, y, b); - output_data[offset] = - input1_data[offset] > min_value ? min_value : input1_data[offset]; - } - } - } + for (int i = 0; i < flat_size; i++) { + output_data[i] = input1_data[i] > min_value ? min_value : input1_data[i]; } } @@ -3399,30 +3370,19 @@ template void TensorFlowMaximum(const T* input1_data, const Dims<4>& input1_dims, const T* input2_data, T* output_data, const Dims<4>& output_dims) { - int batches = MatchingArraySize(input1_dims, 3, output_dims, 3); - int input_height = MatchingArraySize(input1_dims, 2, output_dims, 2); - int input_width = MatchingArraySize(input1_dims, 1, output_dims, 1); - int depth = MatchingArraySize(input1_dims, 0, output_dims, 0); + const int flat_size = MatchingFlatSize(output_dims, input1_dims); auto max_value = input2_data[0]; - - for (int b = 0; b < batches; b++) { - for (int y = 0; y < input_height; y++) { - for (int x = 0; x < input_width; x++) { - for (int c = 0; c < depth; c++) { - int offset = Offset(input1_dims, c, x, y, b); - output_data[offset] = - input1_data[offset] < max_value ? max_value : input1_data[offset]; - } - } - } + for (int i = 0; i < flat_size; i++) { + output_data[i] = input1_data[i] < max_value ? max_value : input1_data[i]; } } -template -void TensorFlowMaximum(const T* input1_data, const Dims<4>& input1_dims, - const T* input2_data, const Dims<4>& input2_dims, - T* output_data, const Dims<4>& output_dims) { +template +void TensorFlowMaximumMinimum(const T* input1_data, const Dims<4>& input1_dims, + const T* input2_data, const Dims<4>& input2_dims, + T* output_data, const Dims<4>& output_dims, + Op op) { NdArrayDesc<4> desc1; NdArrayDesc<4> desc2; NdArrayDescsForElementwiseBroadcast(input1_dims, input2_dims, &desc1, &desc2); @@ -3436,7 +3396,7 @@ void TensorFlowMaximum(const T* input1_data, const Dims<4>& input1_dims, auto in2_idx = SubscriptToIndex(desc2, c, x, y, b); auto in1_val = input1_data[in1_idx]; auto in2_val = input2_data[in2_idx]; - output_data[out_idx] = in1_val > in2_val ? in1_val : in2_val; + output_data[out_idx] = op(in1_val, in2_val); } } } @@ -3455,25 +3415,20 @@ void ArgMax(const T3* axis, const T1* input_data, const Dims<4>& input_dims, // input dimensions here. We enforce the constraint that the last dimension // must always be 1. TFLITE_DCHECK_EQ(ArraySize(output_dims, 0), 1); - 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 outer_size = MatchingFlatSizeSkipDim(input_dims, 0, output_dims); const int depth = ArraySize(input_dims, 0); - for (int b = 0; b < batches; ++b) { - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - auto max_value = input_data[Offset(input_dims, 0, x, y, b)]; - int max_index = 0; - for (int d = 1; d < depth; ++d) { - const auto& curr_value = input_data[Offset(input_dims, d, x, y, b)]; - if (curr_value > max_value) { - max_value = curr_value; - max_index = d; - } - } - output_data[Offset(output_dims, 0, x, y, b)] = max_index; + + for (int i = 0; i < outer_size; ++i) { + auto max_value = input_data[i * depth]; + int max_index = 0; + for (int d = 1; d < depth; ++d) { + const auto& curr_value = input_data[i * depth + d]; + if (curr_value > max_value) { + max_value = curr_value; + max_index = d; } } + output_data[i] = max_index; } } @@ -3523,11 +3478,11 @@ inline void TransposeConv(const float* input_data, const Dims<4>& input_dims, // 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. + // 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; } @@ -3567,6 +3522,51 @@ inline void TransposeConv(const float* input_data, const Dims<4>& input_dims, } } +template +inline void Less(int64_t num_elements, const T* input1, const T* input2, + bool* output) { + for (int64_t i = 0; i < num_elements; ++i) { + output[i] = input1[i] < input2[i]; + } +} + +template +inline void Less(const T* input1_data, const Dims<4>& input1_dims, + const T* input2_data, const Dims<4>& input2_dims, + bool* output_data, const Dims<4>& output_dims) { + const int64_t batches = + MatchingArraySize(input1_dims, 3, input2_dims, 3, output_dims, 3); + const int64_t height = + MatchingArraySize(input1_dims, 2, input2_dims, 2, output_dims, 2); + const int64_t width = + MatchingArraySize(input1_dims, 1, input2_dims, 1, output_dims, 1); + const int64_t depth = + MatchingArraySize(input1_dims, 0, input2_dims, 0, output_dims, 0); + Less(batches * height * width * depth, input1_data, input2_data, output_data); +} + +template +inline void BroadcastLess(T1* input1_data, const Dims<4>& input1_dims, + T2* input2_data, const Dims<4>& input2_dims, + bool* output_data, const Dims<4>& output_dims) { + gemmlowp::ScopedProfilingLabel label("BroadcastLess"); + NdArrayDesc<4> desc1; + NdArrayDesc<4> desc2; + NdArrayDescsForElementwiseBroadcast(input1_dims, input2_dims, &desc1, &desc2); + + 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)] = + input1_data[SubscriptToIndex(desc1, c, x, y, b)] < + input2_data[SubscriptToIndex(desc2, c, x, y, b)]; + } + } + } + } +} + } // namespace reference_ops } // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/tensor.h b/tensorflow/contrib/lite/kernels/internal/tensor.h index 4bce2ffaaf326cf083a76c76adb093f3ac2e8850..62cea143e6afc0631493012be26808a89eb03138 100644 --- a/tensorflow/contrib/lite/kernels/internal/tensor.h +++ b/tensorflow/contrib/lite/kernels/internal/tensor.h @@ -44,6 +44,11 @@ inline int64_t* GetTensorData(TfLiteTensor* tensor) { return tensor != nullptr ? tensor->data.i64 : nullptr; } +template <> +inline bool* GetTensorData(TfLiteTensor* tensor) { + return tensor != nullptr ? tensor->data.b : nullptr; +} + inline int RemapDim(int max_dimensions, int d) { return max_dimensions - d - 1; } diff --git a/tensorflow/contrib/lite/kernels/internal/types.h b/tensorflow/contrib/lite/kernels/internal/types.h index 293538fcbb6406d6065d8efd25adb3b163638c92..3290c364c18224edb733c177ad72bf86b6892434 100644 --- a/tensorflow/contrib/lite/kernels/internal/types.h +++ b/tensorflow/contrib/lite/kernels/internal/types.h @@ -130,14 +130,125 @@ int MatchingArraySize(const ArrayType1& array1, int index1, return MatchingArraySize(array1, index1, args...); } -inline int RequiredBufferSizeForDims(const Dims<4>& dims) { +template +inline int FlatSize(const Dims& dims) { int max_offset = 0; - for (int i = 0; i < 4; i++) { + for (int i = 0; i < N; i++) { max_offset += (dims.sizes[i] - 1) * dims.strides[i]; } return max_offset + 1; } +// Deprecated. Prefer FlatSize. +inline int RequiredBufferSizeForDims(const Dims<4>& dims) { + return FlatSize(dims); +} + +// Flat size calculation, checking that dimensions match with one or more other +// arrays. +template +inline int MatchingFlatSize(const Dims& dims, const Dims& check_dims_0) { + for (int i = 0; i < N; i++) { + TFLITE_DCHECK_EQ(ArraySize(dims, i), ArraySize(check_dims_0, i)); + } + return FlatSize(dims); +} + +template +inline int MatchingFlatSize(const Dims& dims, const Dims& check_dims_0, + const Dims& check_dims_1) { + for (int i = 0; i < N; i++) { + TFLITE_DCHECK_EQ(ArraySize(dims, i), ArraySize(check_dims_0, i)); + } + return MatchingFlatSize(dims, check_dims_1); +} + +template +inline int MatchingFlatSize(const Dims& dims, const Dims& check_dims_0, + const Dims& check_dims_1, + const Dims& check_dims_2) { + for (int i = 0; i < N; i++) { + TFLITE_DCHECK_EQ(ArraySize(dims, i), ArraySize(check_dims_0, i)); + } + return FlatSize(dims, check_dims_1, check_dims_2); +} + +template +inline int MatchingFlatSize(const Dims& dims, const Dims& check_dims_0, + const Dims& check_dims_1, + const Dims& check_dims_2, + const Dims& check_dims_3) { + for (int i = 0; i < N; i++) { + TFLITE_DCHECK_EQ(ArraySize(dims, i), ArraySize(check_dims_0, i)); + } + return FlatSize(dims, check_dims_1, check_dims_2, check_dims_3); +} + +// Data is required to be contiguous, and so many operators can use either the +// full array flat size or the flat size with one dimension skipped (commonly +// the depth). +template +inline int FlatSizeSkipDim(const Dims& dims, int skip_dim) { + TFLITE_DCHECK(skip_dim >= 0 && skip_dim < N); + int flat_size = 1; + for (int i = 0; i < N; i++) { + flat_size *= (i == skip_dim) ? 1 : dims.sizes[i]; + } + return flat_size; +} + +// A combination of MatchingFlatSize() and FlatSizeSkipDim(). +template +inline int MatchingFlatSizeSkipDim(const Dims& dims, int skip_dim, + const Dims& check_dims_0) { + for (int i = 0; i < N; i++) { + if (i != skip_dim) { + TFLITE_DCHECK_EQ(ArraySize(dims, i), ArraySize(check_dims_0, i)); + } + } + return FlatSizeSkipDim(dims, skip_dim); +} + +template +inline int MatchingFlatSizeSkipDim(const Dims& dims, int skip_dim, + const Dims& check_dims_0, + const Dims& check_dims_1) { + for (int i = 0; i < N; i++) { + if (i != skip_dim) { + TFLITE_DCHECK_EQ(ArraySize(dims, i), ArraySize(check_dims_0, i)); + } + } + return MatchingFlatSizeSkipDim(dims, skip_dim, check_dims_1); +} + +template +inline int MatchingFlatSizeSkipDim(const Dims& dims, int skip_dim, + const Dims& check_dims_0, + const Dims& check_dims_1, + const Dims& check_dims_2) { + for (int i = 0; i < N; i++) { + if (i != skip_dim) { + TFLITE_DCHECK_EQ(ArraySize(dims, i), ArraySize(check_dims_0, i)); + } + } + return MatchingFlatSizeSkipDim(dims, skip_dim, check_dims_1, check_dims_2); +} + +template +inline int MatchingFlatSizeSkipDim(const Dims& dims, int skip_dim, + const Dims& check_dims_0, + const Dims& check_dims_1, + const Dims& check_dims_2, + const Dims& check_dims_3) { + for (int i = 0; i < N; i++) { + if (i != skip_dim) { + TFLITE_DCHECK_EQ(ArraySize(dims, i), ArraySize(check_dims_0, i)); + } + } + return MatchingFlatSizeSkipDim(dims, skip_dim, check_dims_1, check_dims_2, + check_dims_3); +} + template bool IsPackedWithoutStrides(const Dims& dims) { int expected_stride = 1; diff --git a/tensorflow/contrib/lite/kernels/maximum.cc b/tensorflow/contrib/lite/kernels/maximum_minimum.cc similarity index 59% rename from tensorflow/contrib/lite/kernels/maximum.cc rename to tensorflow/contrib/lite/kernels/maximum_minimum.cc index 13c40603ced6338086809b908539156e2c0985e7..5a28d663c9e756040746f0a98b356afba76cceab 100644 --- a/tensorflow/contrib/lite/kernels/maximum.cc +++ b/tensorflow/contrib/lite/kernels/maximum_minimum.cc @@ -24,9 +24,9 @@ limitations under the License. namespace tflite { namespace ops { namespace builtin { -namespace maximum { +namespace maximum_minimum { -// This file has a reference implemenation of TFMaximum. +// This file has a reference implemenation of TFMaximum/TFMinimum. enum KernelType { kReference, }; @@ -35,8 +35,8 @@ constexpr int kInputTensor1 = 0; constexpr int kInputTensor2 = 1; constexpr int kOutputTensor = 0; -struct MaximumContext { - MaximumContext(TfLiteContext* context, TfLiteNode* node) { +struct OpContext { + OpContext(TfLiteContext* context, TfLiteNode* node) { input1 = GetInput(context, node, kInputTensor1); input2 = GetInput(context, node, kInputTensor2); output = GetOutput(context, node, kOutputTensor); @@ -50,7 +50,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); - MaximumContext op_context(context, node); + OpContext op_context(context, node); TF_LITE_ENSURE_EQ(context, op_context.input1->type, op_context.input2->type); op_context.output->type = op_context.input1->type; @@ -69,23 +69,49 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { return context->ResizeTensor(context, op_context.output, output_size); } -template -TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - MaximumContext op_context(context, node); +struct MaximumOp { + template + static data_type op(data_type el1, data_type el2) { + return el1 > el2 ? el1 : el2; + } +}; + +struct MinimumOp { + template + static data_type op(data_type el1, data_type el2) { + return el1 < el2 ? el1 : el2; + } +}; + +template +void TFLiteOperation(TfLiteContext* context, TfLiteNode* node, + const OpContext& op_context) { + reference_ops::TensorFlowMaximumMinimum( + GetTensorData(op_context.input1), + GetTensorDims(op_context.input1), + GetTensorData(op_context.input2), + GetTensorDims(op_context.input2), + GetTensorData(op_context.output), + GetTensorDims(op_context.output), op_type::template op); +} -#define TF_LITE_MAXIMUM(kernel_type, data_type) \ - kernel_type::TensorFlowMaximum( \ - GetTensorData(op_context.input1), \ - GetTensorDims(op_context.input1), \ - GetTensorData(op_context.input2), \ - GetTensorDims(op_context.input2), \ - GetTensorData(op_context.output), \ - GetTensorDims(op_context.output)) +template +TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { + OpContext op_context(context, node); if (kernel_type == kReference) { switch (op_context.output->type) { case kTfLiteFloat32: - TF_LITE_MAXIMUM(reference_ops, float); + TFLiteOperation(context, node, op_context); + break; + case kTfLiteUInt8: + TFLiteOperation(context, node, op_context); + break; + case kTfLiteInt32: + TFLiteOperation(context, node, op_context); + break; + case kTfLiteInt64: + TFLiteOperation(context, node, op_context); break; default: context->ReportError(context, @@ -99,19 +125,28 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { op_context.output->type); return kTfLiteError; } -#undef TF_LITE_MAXIMUM return kTfLiteOk; } -} // namespace maximum +} // namespace maximum_minimum TfLiteRegistration* Register_MAXIMUM_REF() { - static TfLiteRegistration r = {nullptr, nullptr, maximum::Prepare, - maximum::Eval}; + static TfLiteRegistration r = { + nullptr, nullptr, maximum_minimum::Prepare, + maximum_minimum::Eval}; return &r; } +TfLiteRegistration* Register_MINIMUM_REF() { + static TfLiteRegistration r = { + nullptr, nullptr, maximum_minimum::Prepare, + maximum_minimum::Eval}; + return &r; +} TfLiteRegistration* Register_MAXIMUM() { return Register_MAXIMUM_REF(); } +TfLiteRegistration* Register_MINIMUM() { return Register_MINIMUM_REF(); } } // namespace builtin } // namespace ops diff --git a/tensorflow/contrib/lite/kernels/maximum_minimum_test.cc b/tensorflow/contrib/lite/kernels/maximum_minimum_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..0752aa1804722accb1f88910fe013ffd632a4503 --- /dev/null +++ b/tensorflow/contrib/lite/kernels/maximum_minimum_test.cc @@ -0,0 +1,143 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT 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 MaxMinOpModel : public SingleOpModel { + public: + MaxMinOpModel(tflite::BuiltinOperator op, const TensorData& input1, + const TensorData& input2, const TensorType& output) { + input1_ = AddInput(input1); + input2_ = AddInput(input2); + output_ = AddOutput(output); + SetBuiltinOp(op, BuiltinOptions_MaximumMinimumOptions, + CreateMaximumMinimumOptions(builder_).Union()); + BuildInterpreter({GetShape(input1_), GetShape(input2_)}); + } + + template + void SetInput1(std::initializer_list data) { + PopulateTensor(input1_, data); + } + + template + void SetInput2(std::initializer_list data) { + PopulateTensor(input2_, data); + } + + template + std::vector GetOutput() { + return ExtractVector(output_); + } + std::vector GetOutputShape() { return GetTensorShape(output_); } + + protected: + int input1_; + int input2_; + int output_; +}; + +template +void TestModel(tflite::BuiltinOperator op, const TensorData& input1, + const TensorData& input2, const TensorData& output, + std::initializer_list input1_values, + std::initializer_list input2_values, + std::initializer_list output_values) { + MaxMinOpModel m(op, input1, input2, output.type); + m.SetInput1(input1_values); + m.SetInput2(input2_values); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray(output.shape)); + EXPECT_THAT(m.GetOutput(), ElementsAreArray(output_values)); +} + +template <> +void TestModel(tflite::BuiltinOperator op, const TensorData& input1, + const TensorData& input2, const TensorData& output, + std::initializer_list input1_values, + std::initializer_list input2_values, + std::initializer_list output_values) { + MaxMinOpModel m(op, input1, input2, output.type); + m.SetInput1(input1_values); + m.SetInput2(input2_values); + m.Invoke(); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray(output.shape)); + EXPECT_THAT(m.GetOutput(), + ElementsAreArray(ArrayFloatNear(output_values))); +} + +TEST(MaximumOpTest, FloatTest) { + std::initializer_list data1 = {1.0, 0.0, -1.0, 11.0, -2.0, -1.44}; + std::initializer_list data2 = {-1.0, 0.0, 1.0, 12.0, -3.0, -1.43}; + TestModel(BuiltinOperator_MAXIMUM, {TensorType_FLOAT32, {3, 1, 2}}, + {TensorType_FLOAT32, {3, 1, 2}}, + {TensorType_FLOAT32, {3, 1, 2}}, data1, data2, + {1.0, 0.0, 1.0, 12.0, -2.0, -1.43}); + TestModel(BuiltinOperator_MINIMUM, {TensorType_FLOAT32, {3, 1, 2}}, + {TensorType_FLOAT32, {3, 1, 2}}, + {TensorType_FLOAT32, {3, 1, 2}}, data1, data2, + {-1.0, 0.0, -1.0, 11.0, -3.0, -1.44}); +} + +TEST(MaxMinOpTest, Uint8Test) { + std::initializer_list data1 = {1, 0, 2, 11, 2, 23}; + std::initializer_list data2 = {0, 0, 1, 12, 255, 1}; + TestModel(BuiltinOperator_MAXIMUM, {TensorType_UINT8, {3, 1, 2}}, + {TensorType_UINT8, {3, 1, 2}}, + {TensorType_UINT8, {3, 1, 2}}, data1, data2, + {1, 0, 2, 12, 255, 23}); + TestModel(BuiltinOperator_MINIMUM, {TensorType_UINT8, {3, 1, 2}}, + {TensorType_UINT8, {3, 1, 2}}, + {TensorType_UINT8, {3, 1, 2}}, data1, data2, + {0, 0, 1, 11, 2, 1}); +} + +TEST(MaximumOpTest, FloatWithBroadcastTest) { + std::initializer_list data1 = {1.0, 0.0, -1.0, -2.0, -1.44, 11.0}; + std::initializer_list data2 = {0.5, 2.0}; + TestModel(BuiltinOperator_MAXIMUM, {TensorType_FLOAT32, {3, 1, 2}}, + {TensorType_FLOAT32, {2}}, {TensorType_FLOAT32, {3, 1, 2}}, + data1, data2, {1.0, 2.0, 0.5, 2.0, 0.5, 11.0}); + TestModel(BuiltinOperator_MINIMUM, {TensorType_FLOAT32, {3, 1, 2}}, + {TensorType_FLOAT32, {2}}, {TensorType_FLOAT32, {3, 1, 2}}, + data1, data2, {0.5, 0.0, -1.0, -2.0, -1.44, 2.0}); +} + +TEST(MaximumOpTest, Int32WithBroadcastTest) { + std::initializer_list data1 = {1, 0, -1, -2, 3, 11}; + std::initializer_list data2 = {2}; + TestModel(BuiltinOperator_MAXIMUM, {TensorType_INT32, {3, 1, 2}}, + {TensorType_INT32, {1}}, {TensorType_INT32, {3, 1, 2}}, + data1, data2, {2, 2, 2, 2, 3, 11}); + TestModel(BuiltinOperator_MINIMUM, {TensorType_INT32, {3, 1, 2}}, + {TensorType_INT32, {1}}, {TensorType_INT32, {3, 1, 2}}, + data1, data2, {1, 0, -1, -2, 2, 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/maximum_test.cc b/tensorflow/contrib/lite/kernels/maximum_test.cc deleted file mode 100644 index df2bf29c205e0a3ff6ea5df2bba8ca721a09e626..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/maximum_test.cc +++ /dev/null @@ -1,95 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT 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 MaximumOpModel : public SingleOpModel { - public: - MaximumOpModel(const TensorData& input1, const TensorData& input2, - const TensorType& output) { - input1_ = AddInput(input1); - input2_ = AddInput(input2); - output_ = AddOutput(output); - SetBuiltinOp(BuiltinOperator_MAXIMUM, BuiltinOptions_MaximumOptions, - CreateMaximumOptions(builder_).Union()); - BuildInterpreter({GetShape(input1_), GetShape(input2_)}); - } - - template - void SetInput1(std::initializer_list data) { - PopulateTensor(input1_, data); - } - - template - void SetInput2(std::initializer_list data) { - PopulateTensor(input2_, data); - } - - template - std::vector GetOutput() { - return ExtractVector(output_); - } - std::vector GetOutputShape() { return GetTensorShape(output_); } - - protected: - int input1_; - int input2_; - int output_; -}; - -TEST(MaximumOpTest, FloatTest) { - std::initializer_list data1 = {1.0, 0.0, -1.0, 11.0, -2.0, -1.44}; - std::initializer_list data2 = {-1.0, 0.0, 1.0, 12.0, -3.0, -1.43}; - MaximumOpModel m({TensorType_FLOAT32, {3, 1, 2}}, - {TensorType_FLOAT32, {3, 1, 2}}, TensorType_FLOAT32); - m.SetInput1(data1); - m.SetInput2(data2); - m.Invoke(); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3, 1, 2})); - EXPECT_THAT( - m.GetOutput(), - ElementsAreArray(ArrayFloatNear({1.0, 0.0, 1.0, 12.0, -2.0, -1.43}))); -} - -TEST(MaximumOpTest, FloatWithBroadcastTest) { - std::initializer_list data1 = {1.0, 0.0, -1.0, -2.0, -1.44, 11.0}; - std::initializer_list data2 = {0.5, 2.0}; - MaximumOpModel m({TensorType_FLOAT32, {3, 1, 2}}, {TensorType_FLOAT32, {2}}, - TensorType_FLOAT32); - m.SetInput1(data1); - m.SetInput2(data2); - m.Invoke(); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3, 1, 2})); - EXPECT_THAT( - m.GetOutput(), - ElementsAreArray(ArrayFloatNear({1.0, 2.0, 0.5, 2.0, 0.5, 11.0}))); -} - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/pad.cc b/tensorflow/contrib/lite/kernels/pad.cc index c29da3862e84d6756bf5ef34b2ca06307b0a065d..4f9449a225c66a0fb2a9285e6aff3a1f7147f5dd 100644 --- a/tensorflow/contrib/lite/kernels/pad.cc +++ b/tensorflow/contrib/lite/kernels/pad.cc @@ -119,39 +119,46 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { after_padding.push_back(paddings_data[idx * 2 + 1]); } -#define TF_LITE_PAD(type, scalar) \ +#define TF_LITE_PAD(type, scalar, pad_value) \ type::Pad(GetTensorData(op_context.input), \ GetTensorDims(op_context.input), before_padding, after_padding, \ GetTensorData(op_context.output), \ - GetTensorDims(op_context.output)) + GetTensorDims(op_context.output), pad_value) switch (op_context.input->type) { case kTfLiteFloat32: if (kernel_type == kReference) { - TF_LITE_PAD(reference_ops, float); + TF_LITE_PAD(reference_ops, float, 0); } else if (kernel_type == kGenericOptimized) { - TF_LITE_PAD(optimized_ops, float); + TF_LITE_PAD(optimized_ops, float, 0); } break; case kTfLiteUInt8: + // Quantized Pad requires that 0 is represented in the quantized range. + TF_LITE_ENSURE(context, op_context.output->params.zero_point >= + std::numeric_limits::min()); + TF_LITE_ENSURE(context, op_context.output->params.zero_point <= + std::numeric_limits::max()); if (kernel_type == kReference) { - TF_LITE_PAD(reference_ops, uint8_t); + TF_LITE_PAD(reference_ops, uint8_t, + op_context.output->params.zero_point); } else if (kernel_type == kGenericOptimized) { - TF_LITE_PAD(optimized_ops, uint8_t); + TF_LITE_PAD(optimized_ops, uint8_t, + op_context.output->params.zero_point); } break; case kTfLiteInt32: if (kernel_type == kReference) { - TF_LITE_PAD(reference_ops, int32_t); + TF_LITE_PAD(reference_ops, int32_t, 0); } else if (kernel_type == kGenericOptimized) { - TF_LITE_PAD(optimized_ops, int32_t); + TF_LITE_PAD(optimized_ops, int32_t, 0); } break; case kTfLiteInt64: if (kernel_type == kReference) { - TF_LITE_PAD(reference_ops, int64_t); + TF_LITE_PAD(reference_ops, int64_t, 0); } else if (kernel_type == kGenericOptimized) { - TF_LITE_PAD(optimized_ops, int64_t); + TF_LITE_PAD(optimized_ops, int64_t, 0); } break; default: diff --git a/tensorflow/contrib/lite/kernels/pad_test.cc b/tensorflow/contrib/lite/kernels/pad_test.cc index 28834ad0719291b2e868bca2d86a6685e6eb9962..c06237e5720874e66c5953edab2d3749cc88af28 100644 --- a/tensorflow/contrib/lite/kernels/pad_test.cc +++ b/tensorflow/contrib/lite/kernels/pad_test.cc @@ -22,6 +22,7 @@ namespace tflite { namespace { using ::testing::ElementsAreArray; +using ::testing::Matcher; class PadOpModel : public SingleOpModel { public: @@ -29,6 +30,10 @@ class PadOpModel : public SingleOpModel { PopulateTensor(input_, data); } + void SetQuantizedInput(std::initializer_list data) { + QuantizeAndPopulate(input_, data); + } + void SetPaddings(std::initializer_list paddings) { PopulateTensor(paddings_, paddings); } @@ -36,6 +41,11 @@ class PadOpModel : public SingleOpModel { std::vector GetOutput() { return ExtractVector(output_); } std::vector GetOutputShape() { return GetTensorShape(output_); } + std::vector GetDequantizedOutput() { + return Dequantize(ExtractVector(output_), + GetScale(output_), GetZeroPoint(output_)); + } + protected: int input_; int output_; @@ -50,16 +60,17 @@ class PadOpModel : public SingleOpModel { // m.Invoke(); class PadOpConstModel : public PadOpModel { public: - PadOpConstModel(std::initializer_list input_shape, + PadOpConstModel(const TensorData& input, std::initializer_list paddings_shape, - std::initializer_list paddings) { - input_ = AddInput(TensorType_FLOAT32); + std::initializer_list paddings, + const TensorData& output) { + input_ = AddInput(input); paddings_ = AddConstInput(TensorType_INT32, paddings, paddings_shape); - output_ = AddOutput(TensorType_FLOAT32); + output_ = AddOutput(output); SetBuiltinOp(BuiltinOperator_PAD, BuiltinOptions_PadOptions, CreatePadOptions(builder_).Union()); - BuildInterpreter({input_shape}); + BuildInterpreter({input.shape}); } }; @@ -72,40 +83,45 @@ class PadOpConstModel : public PadOpModel { // m.Invoke(); class PadOpDynamicModel : public PadOpModel { public: - PadOpDynamicModel(std::initializer_list input_shape, - std::initializer_list paddings_shape) { - input_ = AddInput(TensorType_FLOAT32); + PadOpDynamicModel(const TensorData& input, + std::initializer_list paddings_shape, + const TensorData& output) { + input_ = AddInput(input); paddings_ = AddInput(TensorType_INT32); - output_ = AddOutput(TensorType_FLOAT32); + output_ = AddOutput(output); SetBuiltinOp(BuiltinOperator_PAD, BuiltinOptions_PadOptions, CreatePadOptions(builder_).Union()); - BuildInterpreter({input_shape, paddings_shape}); + BuildInterpreter({input.shape, paddings_shape}); } }; TEST(PadOpTest, TooManyDimensions) { EXPECT_DEATH( - PadOpConstModel({1, 2, 3, 4, 5, 6, 7, 8, 9}, {9, 2}, - {1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9}), + PadOpConstModel({TensorType_FLOAT32, {1, 2, 3, 4, 5, 6, 7, 8, 9}}, {9, 2}, + {1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9}, + {TensorType_FLOAT32}), "dims != 4"); } TEST(PadOpTest, UnequalDimensions) { - EXPECT_DEATH(PadOpConstModel({1, 1, 2, 1}, {3, 2}, {1, 1, 2, 2, 3, 3}), + EXPECT_DEATH(PadOpConstModel({TensorType_FLOAT32, {1, 1, 2, 1}}, {3, 2}, + {1, 1, 2, 2, 3, 3}, {TensorType_FLOAT32}), "3 != 4"); } TEST(PadOpTest, InvalidPadValue) { EXPECT_DEATH( - PadOpConstModel({1, 1, 2, 1}, {4, 2}, {0, 0, 1, -1, 2, -1, 0, 0}), + PadOpConstModel({TensorType_FLOAT32, {1, 1, 2, 1}}, {4, 2}, + {0, 0, 1, -1, 2, -1, 0, 0}, {TensorType_FLOAT32}), "Pad value has to be greater than equal to 0."); } TEST(PadOpTest, SimpleConstTest) { // Padding is represented as four 2-D lists representing above padding and // below padding (i.e. {{0, 0}, {1, 1}, {1, 1}, {0, 0}}). - PadOpConstModel m({1, 2, 2, 1}, {4, 2}, {0, 0, 1, 1, 1, 1, 0, 0}); + PadOpConstModel m({TensorType_FLOAT32, {1, 2, 2, 1}}, {4, 2}, + {0, 0, 1, 1, 1, 1, 0, 0}, {TensorType_FLOAT32}); m.SetInput({1, 2, 3, 4}); m.Invoke(); EXPECT_THAT(m.GetOutput(), ElementsAreArray({0, 0, 0, 0, 0, 1, 2, 0, 0, 3, 4, @@ -114,7 +130,8 @@ TEST(PadOpTest, SimpleConstTest) { } TEST(PadOpTest, SimpleDynamicTest) { - PadOpDynamicModel m({1, 2, 2, 1}, {4, 2}); + PadOpDynamicModel m({TensorType_FLOAT32, {1, 2, 2, 1}}, {4, 2}, + {TensorType_FLOAT32}); m.SetInput({1, 2, 3, 4}); m.SetPaddings({0, 0, 1, 1, 1, 1, 0, 0}); m.Invoke(); @@ -124,7 +141,8 @@ TEST(PadOpTest, SimpleDynamicTest) { } TEST(PadOpTest, AdvancedConstTest) { - PadOpConstModel m({1, 2, 3, 1}, {4, 2}, {0, 0, 0, 2, 1, 3, 0, 0}); + PadOpConstModel m({TensorType_FLOAT32, {1, 2, 3, 1}}, {4, 2}, + {0, 0, 0, 2, 1, 3, 0, 0}, {TensorType_FLOAT32}); m.SetInput({1, 2, 3, 4, 5, 6}); m.Invoke(); EXPECT_THAT(m.GetOutput(), @@ -134,7 +152,8 @@ TEST(PadOpTest, AdvancedConstTest) { } TEST(PadOpTest, AdvancedDynamicTest) { - PadOpDynamicModel m({1, 2, 3, 1}, {4, 2}); + PadOpDynamicModel m({TensorType_FLOAT32, {1, 2, 3, 1}}, {4, 2}, + {TensorType_FLOAT32}); m.SetInput({1, 2, 3, 4, 5, 6}); m.SetPaddings({0, 0, 0, 2, 1, 3, 0, 0}); m.Invoke(); @@ -144,6 +163,80 @@ TEST(PadOpTest, AdvancedDynamicTest) { EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 7, 1})); } +class QuantizedPadOpTest : public ::testing::Test { + protected: + std::vector> DequantizedArrayNear( + const std::vector& values, const float min, const float max) { + const float quantization_tolerance = (max - min) / 255.0; + return ArrayFloatNear(values, quantization_tolerance); + } +}; + +TEST_F(QuantizedPadOpTest, ZeroNotInQuantizationRange) { + // The test_util and actual quantization code currently ensure that the range + // must include zero, but if that ever changes, this test will catch it. + EXPECT_DEATH(PadOpConstModel m({TensorType_UINT8, {1, 2, 2, 1}, 1.0, 2.0}, + {4, 2}, {0, 0, 1, 1, 1, 1, 0, 0}, + {TensorType_UINT8, {}, 1.0, 2.0}), + ".*Check failed: f_min <= 0.*"); +} + +TEST_F(QuantizedPadOpTest, SimpleConstTest) { + // Padding is represented as four 2-D lists representing above padding and + // below padding (i.e. {{0, 0}, {1, 1}, {1, 1}, {0, 0}}). + PadOpConstModel m({TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0}, {4, 2}, + {0, 0, 1, 1, 1, 1, 0, 0}, + {TensorType_UINT8, {}, -1.0, 1.0}); + m.SetQuantizedInput({-0.8, 0.2, 0.9, 0.7}); + m.Invoke(); + EXPECT_THAT(m.GetDequantizedOutput(), + ElementsAreArray(DequantizedArrayNear( + {0, 0, 0, 0, 0, -0.8, 0.2, 0, 0, 0.9, 0.7, 0, 0, 0, 0, 0}, + -1.0, 1.0))); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1})); +} + +TEST_F(QuantizedPadOpTest, SimpleDynamicTest) { + PadOpDynamicModel m({TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0}, {4, 2}, + {TensorType_UINT8, {}, -1.0, 1.0}); + m.SetQuantizedInput({-0.8, 0.2, 0.9, 0.7}); + m.SetPaddings({0, 0, 1, 1, 1, 1, 0, 0}); + m.Invoke(); + EXPECT_THAT(m.GetDequantizedOutput(), + ElementsAreArray(DequantizedArrayNear( + {0, 0, 0, 0, 0, -0.8, 0.2, 0, 0, 0.9, 0.7, 0, 0, 0, 0, 0}, + -1.0, 1.0))); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1})); +} + +TEST_F(QuantizedPadOpTest, AdvancedConstTest) { + PadOpConstModel m({TensorType_UINT8, {1, 2, 3, 1}, -1.0, 1.0}, {4, 2}, + {0, 0, 0, 2, 1, 3, 0, 0}, + {TensorType_UINT8, {}, -1.0, 1.0}); + m.SetQuantizedInput({-0.8, 0.2, 0.9, 0.7, 0.1, -0.3}); + m.Invoke(); + EXPECT_THAT(m.GetDequantizedOutput(), + ElementsAreArray(DequantizedArrayNear( + {0, -0.8, 0.2, 0.9, 0, 0, 0, 0, 0.7, 0.1, -0.3, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, + -1.0, 1.0))); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 7, 1})); +} + +TEST_F(QuantizedPadOpTest, AdvancedDynamicTest) { + PadOpDynamicModel m({TensorType_UINT8, {1, 2, 3, 1}, -1.0, 1.0}, {4, 2}, + {TensorType_UINT8, {}, -1.0, 1.0}); + m.SetQuantizedInput({-0.8, 0.2, 0.9, 0.7, 0.1, -0.3}); + m.SetPaddings({0, 0, 0, 2, 1, 3, 0, 0}); + m.Invoke(); + EXPECT_THAT(m.GetDequantizedOutput(), + ElementsAreArray(DequantizedArrayNear( + {0, -0.8, 0.2, 0.9, 0, 0, 0, 0, 0.7, 0.1, -0.3, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, + -1.0, 1.0))); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 7, 1})); +} + } // namespace } // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/padding.h b/tensorflow/contrib/lite/kernels/padding.h index 40b8476b3779c66e31a04856bce8aebd378f1e5f..e81b970e0fb149e8c5d95ed12622917fdc336f7a 100644 --- a/tensorflow/contrib/lite/kernels/padding.h +++ b/tensorflow/contrib/lite/kernels/padding.h @@ -17,9 +17,10 @@ limitations under the License. namespace tflite { -inline int ComputePadding(int stride, int in_size, int filter_size, - int out_size) { - int padding = ((out_size - 1) * stride + filter_size - in_size) / 2; +inline int ComputePadding(int stride, int dilation_rate, int in_size, + int filter_size, int out_size) { + int effective_filter_size = (filter_size - 1) * dilation_rate + 1; + int padding = ((out_size - 1) * stride + effective_filter_size - in_size) / 2; return padding > 0 ? padding : 0; } diff --git a/tensorflow/contrib/lite/kernels/pooling.cc b/tensorflow/contrib/lite/kernels/pooling.cc index b79880110897a1438a589d97363fd861c61667e7..0bf27c34c1337b4ae4b8b73ee2dafcc931c7ce3c 100644 --- a/tensorflow/contrib/lite/kernels/pooling.cc +++ b/tensorflow/contrib/lite/kernels/pooling.cc @@ -94,9 +94,9 @@ TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteNode* node) { int outHeight = computeOutSize(height, params->filter_height, params->stride_height); - data->padding.height = ComputePadding(params->stride_height, height, + data->padding.height = ComputePadding(params->stride_height, 1, height, params->filter_height, outHeight); - data->padding.width = ComputePadding(params->stride_width, width, + data->padding.width = ComputePadding(params->stride_width, 1, width, params->filter_width, outWidth); if (input->type == kTfLiteUInt8) { diff --git a/tensorflow/contrib/lite/kernels/register.cc b/tensorflow/contrib/lite/kernels/register.cc index 0f98154b904b1f776016e6bbee3263027f815244..b07e7b6ff32e9ea513e60619078e09b7e6d1db72 100644 --- a/tensorflow/contrib/lite/kernels/register.cc +++ b/tensorflow/contrib/lite/kernels/register.cc @@ -77,6 +77,9 @@ TfLiteRegistration* Register_CAST(); TfLiteRegistration* Register_DEQUANTIZE(); TfLiteRegistration* Register_PRELU(); TfLiteRegistration* Register_MAXIMUM(); +TfLiteRegistration* Register_MINIMUM(); +TfLiteRegistration* Register_ARG_MAX(); +TfLiteRegistration* Register_LESS(); BuiltinOpResolver::BuiltinOpResolver() { AddBuiltin(BuiltinOperator_RELU, Register_RELU()); @@ -135,6 +138,9 @@ BuiltinOpResolver::BuiltinOpResolver() { AddBuiltin(BuiltinOperator_DEQUANTIZE, Register_DEQUANTIZE()); AddBuiltin(BuiltinOperator_PRELU, Register_PRELU()); AddBuiltin(BuiltinOperator_MAXIMUM, Register_MAXIMUM()); + AddBuiltin(BuiltinOperator_MINIMUM, Register_MINIMUM()); + AddBuiltin(BuiltinOperator_ARG_MAX, Register_ARG_MAX()); + AddBuiltin(BuiltinOperator_LESS, Register_LESS()); // TODO(andrewharp, ahentz): Move these somewhere more appropriate so that // custom ops aren't always included by default. diff --git a/tensorflow/contrib/lite/kernels/strided_slice.cc b/tensorflow/contrib/lite/kernels/strided_slice.cc index e6d5c300dcd47821b0572e3239b36f14bd6ea3d0..40ac436b7dcabe7a12166e5381f0381941a204d3 100644 --- a/tensorflow/contrib/lite/kernels/strided_slice.cc +++ b/tensorflow/contrib/lite/kernels/strided_slice.cc @@ -87,6 +87,8 @@ inline int32_t ClampedIndex(int32_t index, int dim, bool pos_stride) { std::min(std::max(index, -dim), dim - 1), dim)); } +// TODO(b/77971377) this logic should be removed, as it's a duplication of +// StartForAxis() & StopForAxis() in kernels/internal/reference/reference_ops.h 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; @@ -188,8 +190,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { 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)); + starts.emplace_back(GetTensorData(op_context.begin)[idx]); + stops.emplace_back(GetTensorData(op_context.end)[idx]); strides.emplace_back(GetTensorData(op_context.strides)[idx]); } @@ -202,15 +204,13 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { int begin_mask = ReverseMaskBits(op_context.params->begin_mask, op_context.dims); 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), begin_mask, end_mask, shrink_axis_mask, \ - starts, stops, strides, GetTensorData(op_context.output), \ - GetTensorDims(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, starts, stops, strides, \ + GetTensorData(op_context.output), \ + GetTensorDims(op_context.output)) switch (op_context.input->type) { case kTfLiteFloat32: diff --git a/tensorflow/contrib/lite/kernels/strided_slice_test.cc b/tensorflow/contrib/lite/kernels/strided_slice_test.cc index 22d7b097cbd4e1349516eae9fd378aa186e06de7..cc39179bc705aa1083e74b06f8f7f3fb45e9f616 100644 --- a/tensorflow/contrib/lite/kernels/strided_slice_test.cc +++ b/tensorflow/contrib/lite/kernels/strided_slice_test.cc @@ -377,29 +377,18 @@ TEST(StridedSliceOpTest, In1D_ShrinkAxisMask1) { StridedSliceOpModel<> m({4}, {1}, {1}, {1}, 0, 0, 0, 0, 1); m.SetInput({1, 2, 3, 4}); m.SetBegin({1}); - m.SetEnd({3}); + m.SetEnd({2}); m.SetStrides({1}); m.Invoke(); EXPECT_TRUE(m.GetOutputShape().empty()); EXPECT_THAT(m.GetOutput(), ElementsAreArray({2})); } -TEST(StridedSliceOpTest, In1D_EmptyOutputShrinkAxisMask1) { - StridedSliceOpModel<> m({4}, {1}, {1}, {1}, 0, 0, 0, 0, 1); - m.SetInput({1, 2, 3, 4}); - m.SetBegin({2}); - m.SetEnd({1}); - m.SetStrides({1}); - m.Invoke(); - EXPECT_TRUE(m.GetOutputShape().empty()); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({3})); -} - TEST(StridedSliceOpTest, In1D_BeginMaskShrinkAxisMask1) { StridedSliceOpModel<> m({4}, {1}, {1}, {1}, 1, 0, 0, 0, 1); m.SetInput({1, 2, 3, 4}); m.SetBegin({1}); - m.SetEnd({3}); + m.SetEnd({1}); m.SetStrides({1}); m.Invoke(); EXPECT_TRUE(m.GetOutputShape().empty()); @@ -421,7 +410,7 @@ TEST(StridedSliceOpTest, In2D_ShrinkAxisMask1) { StridedSliceOpModel<> m({2, 3}, {2}, {2}, {2}, 0, 0, 0, 0, 1); m.SetInput({1, 2, 3, 4, 5, 6}); m.SetBegin({0, 0}); - m.SetEnd({2, 3}); + m.SetEnd({1, 3}); m.SetStrides({1, 1}); m.Invoke(); EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3})); @@ -432,7 +421,7 @@ TEST(StridedSliceOpTest, In2D_ShrinkAxisMask2) { StridedSliceOpModel<> m({2, 3}, {2}, {2}, {2}, 0, 0, 0, 0, 2); m.SetInput({1, 2, 3, 4, 5, 6}); m.SetBegin({0, 0}); - m.SetEnd({2, 3}); + m.SetEnd({2, 1}); m.SetStrides({1, 1}); m.Invoke(); EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2})); @@ -443,7 +432,7 @@ TEST(StridedSliceOpTest, In2D_ShrinkAxisMask3) { StridedSliceOpModel<> m({2, 3}, {2}, {2}, {2}, 0, 0, 0, 0, 3); m.SetInput({1, 2, 3, 4, 5, 6}); m.SetBegin({0, 0}); - m.SetEnd({2, 3}); + m.SetEnd({1, 1}); m.SetStrides({1, 1}); m.Invoke(); EXPECT_TRUE(m.GetOutputShape().empty()); @@ -454,7 +443,7 @@ TEST(StridedSliceOpTest, In3D_IdentityShrinkAxis1) { StridedSliceOpModel<> m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 1); m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); m.SetBegin({0, 0, 0}); - m.SetEnd({2, 3, 2}); + m.SetEnd({1, 3, 2}); m.SetStrides({1, 1, 1}); m.Invoke(); EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3, 2})); @@ -465,7 +454,7 @@ TEST(StridedSliceOpTest, In3D_IdentityShrinkAxis2) { StridedSliceOpModel<> m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 2); m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); m.SetBegin({0, 0, 0}); - m.SetEnd({2, 3, 2}); + m.SetEnd({2, 1, 2}); m.SetStrides({1, 1, 1}); m.Invoke(); EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2})); @@ -476,7 +465,7 @@ TEST(StridedSliceOpTest, In3D_IdentityShrinkAxis3) { StridedSliceOpModel<> m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 3); m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); m.SetBegin({0, 0, 0}); - m.SetEnd({2, 3, 2}); + m.SetEnd({1, 1, 2}); m.SetStrides({1, 1, 1}); m.Invoke(); EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2})); @@ -487,7 +476,7 @@ TEST(StridedSliceOpTest, In3D_IdentityShrinkAxis4) { StridedSliceOpModel<> m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 4); m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); m.SetBegin({0, 0, 0}); - m.SetEnd({2, 3, 2}); + m.SetEnd({2, 3, 1}); m.SetStrides({1, 1, 1}); m.Invoke(); EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 3})); @@ -498,7 +487,7 @@ TEST(StridedSliceOpTest, In3D_IdentityShrinkAxis5) { StridedSliceOpModel<> m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 5); m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); m.SetBegin({0, 0, 0}); - m.SetEnd({2, 3, 2}); + m.SetEnd({1, 3, 1}); m.SetStrides({1, 1, 1}); m.Invoke(); EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3})); @@ -509,7 +498,7 @@ TEST(StridedSliceOpTest, In3D_IdentityShrinkAxis6) { StridedSliceOpModel<> m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 6); m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); m.SetBegin({0, 0, 0}); - m.SetEnd({2, 3, 2}); + m.SetEnd({2, 1, 1}); m.SetStrides({1, 1, 1}); m.Invoke(); EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2})); @@ -520,7 +509,7 @@ TEST(StridedSliceOpTest, In3D_IdentityShrinkAxis7) { StridedSliceOpModel<> m({2, 3, 2}, {3}, {3}, {3}, 0, 0, 0, 0, 7); m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); m.SetBegin({0, 0, 0}); - m.SetEnd({2, 3, 2}); + m.SetEnd({1, 1, 1}); m.SetStrides({1, 1, 1}); m.Invoke(); EXPECT_TRUE(m.GetOutputShape().empty()); @@ -553,7 +542,7 @@ TEST(StridedSliceOpTest, In3D_IdentityShrinkAxis1Uint8) { 0, 0, 1); m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); m.SetBegin({0, 0, 0}); - m.SetEnd({2, 3, 2}); + m.SetEnd({1, 3, 2}); m.SetStrides({1, 1, 1}); m.Invoke(); EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3, 2})); diff --git a/tensorflow/contrib/lite/kernels/sub.cc b/tensorflow/contrib/lite/kernels/sub.cc index 66b06aeaec52dd3d2d98acfec8218ffdd0ae6bf3..7c60a4fdbffdc96b8967f52f8dbab3e18ecbcc0a 100644 --- a/tensorflow/contrib/lite/kernels/sub.cc +++ b/tensorflow/contrib/lite/kernels/sub.cc @@ -174,7 +174,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { EvalQuantized(context, node, params, data, input1, input2, output); } else { - context->ReportError(context, "Inputs and outputs not all float types."); + context->ReportError(context, + "Inputs and outputs not all float|uint8 types."); return kTfLiteError; } diff --git a/tensorflow/contrib/lite/model.cc b/tensorflow/contrib/lite/model.cc index 3448de68e8e678e3b5316e82617a5b15a5e7b8b4..2dd6d67e078619df41524e8242a0475320c02013 100644 --- a/tensorflow/contrib/lite/model.cc +++ b/tensorflow/contrib/lite/model.cc @@ -57,6 +57,9 @@ TfLiteStatus ConvertTensorType(TensorType tensor_type, TfLiteType* type, case TensorType_STRING: *type = kTfLiteString; break; + case TensorType_BOOL: + *type = kTfLiteBool; + break; default: error_reporter->Report("Unimplemented data type %s (%d) in tensor\n", EnumNameTensorType(tensor_type), tensor_type); @@ -261,13 +264,11 @@ T* MallocPOD() { // Parse the appropriate data out of the op. // // This handles builtin data explicitly as there are flatbuffer schemas. -// -// Returns memory that must be feed. -// -// TODO(nupurgarg): Pass in void ** and return TfLiteStatus to ensure program -// crashes if error reporter is called. -void* ParseOpData(const Operator* op, BuiltinOperator op_type, - ErrorReporter* error_reporter) { +// If it returns kTfLiteOk, it passes the data out with `builtin_data`, which +// need to be released by calling `free`.` +// If it returns kTfLiteError, `builtin_data` will be `nullptr`. +TfLiteStatus ParseOpData(const Operator* op, BuiltinOperator op_type, + ErrorReporter* error_reporter, void** builtin_data) { auto parse_padding = [](Padding padding) { switch (padding) { case Padding_SAME: @@ -316,7 +317,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, } }; - void* builtin_data = nullptr; + *builtin_data = nullptr; switch (op_type) { case BuiltinOperator_CALL: // TODO(aselle): Implement call in BuiltinOptions, but nullptrs are @@ -332,8 +333,10 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->stride_height = conv_params->stride_h(); params->activation = parse_activation(conv_params->fused_activation_function()); + params->dilation_width_factor = conv_params->dilation_w_factor(); + params->dilation_height_factor = conv_params->dilation_h_factor(); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_TANH: @@ -358,10 +361,11 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, ConvertTensorType(schema_params->out_data_type(), ¶ms->out_data_type, error_reporter); if (in_status != kTfLiteOk || out_status != kTfLiteOk) { - break; + free(params); + return kTfLiteError; } } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_LSH_PROJECTION: { @@ -370,7 +374,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, if (auto* lshParams = op->builtin_options_as_LSHProjectionOptions()) { params->type = parseLSHProjectionType(lshParams->type()); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_AVERAGE_POOL_2D: @@ -386,7 +390,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->activation = parse_activation(pool_params->fused_activation_function()); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_DEPTHWISE_CONV_2D: { @@ -400,7 +404,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->activation = parse_activation(conv_params->fused_activation_function()); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_SVDF: { @@ -410,7 +414,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->activation = parse_activation(svdf_params->fused_activation_function()); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN: @@ -422,7 +426,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, parse_activation(sequence_rnn_params->fused_activation_function()); params->time_major = sequence_rnn_params->time_major(); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_RNN: { @@ -431,7 +435,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->activation = parse_activation(rnn_params->fused_activation_function()); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_EMBEDDING_LOOKUP: @@ -444,7 +448,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, op->builtin_options_as_EmbeddingLookupSparseOptions()) { params->combiner = parseCombinerType(embedding_params->combiner()); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_FULLY_CONNECTED: { @@ -455,7 +459,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->activation = parse_activation( fully_connected_params->fused_activation_function()); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_HASHTABLE_LOOKUP: @@ -466,7 +470,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, if (auto* softmax_params = op->builtin_options_as_SoftmaxOptions()) { params->beta = softmax_params->beta(); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_CONCATENATION: { @@ -478,7 +482,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, parse_activation(concatenation_params->fused_activation_function()); params->axis = concatenation_params->axis(); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_MUL: { @@ -487,7 +491,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->activation = parse_activation(schema_params->fused_activation_function()); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_ADD: { @@ -496,7 +500,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->activation = parse_activation(schema_params->fused_activation_function()); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_DIV: { @@ -505,7 +509,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->activation = parse_activation(schema_params->fused_activation_function()); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_SUB: { @@ -514,7 +518,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->activation = parse_activation(schema_params->fused_activation_function()); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_L2_NORMALIZATION: { @@ -523,7 +527,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->activation = parse_activation(schema_params->fused_activation_function()); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_LOCAL_RESPONSE_NORMALIZATION: { @@ -535,7 +539,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->alpha = schema_params->alpha(); params->beta = schema_params->beta(); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM: @@ -548,7 +552,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->cell_clip = lstm_params->cell_clip(); params->proj_clip = lstm_params->proj_clip(); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_RESIZE_BILINEAR: { @@ -557,7 +561,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, op->builtin_options_as_ResizeBilinearOptions()) { params->align_corners = schema_params->align_corners(); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_PAD: { @@ -571,7 +575,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->shape, error_reporter); params->num_dimensions = new_shape->Length(); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_SKIP_GRAM: { @@ -581,7 +585,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->max_skip_size = skip_gram_params->max_skip_size(); params->include_all_ngrams = skip_gram_params->include_all_ngrams(); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_SPACE_TO_DEPTH: { @@ -589,7 +593,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, if (auto* schema_params = op->builtin_options_as_SpaceToDepthOptions()) { params->block_size = schema_params->block_size(); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_GATHER: { @@ -599,7 +603,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->axis = gather_params->axis(); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_SPACE_TO_BATCH_ND: { @@ -616,7 +620,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, if (auto* schema_params = op->builtin_options_as_MeanOptions()) { params->keep_dims = schema_params->keep_dims(); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_SPLIT: { @@ -624,7 +628,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, if (auto* schema_params = op->builtin_options_as_SplitOptions()) { params->num_splits = schema_params->num_splits(); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_SQUEEZE: { @@ -635,7 +639,7 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->squeeze_dims, error_reporter); params->num_squeeze_dims = squeeze_dims->Length(); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); break; } case BuiltinOperator_STRIDED_SLICE: { @@ -647,19 +651,32 @@ void* ParseOpData(const Operator* op, BuiltinOperator op_type, params->new_axis_mask = schema_params->new_axis_mask(); params->shrink_axis_mask = schema_params->shrink_axis_mask(); } - builtin_data = reinterpret_cast(params); + *builtin_data = reinterpret_cast(params); + break; + } + case BuiltinOperator_MAXIMUM: + case BuiltinOperator_MINIMUM: { + break; + } + case BuiltinOperator_ARG_MAX: { + auto* params = MallocPOD(); + if (auto* schema_params = op->builtin_options_as_ArgMaxOptions()) { + ConvertTensorType(schema_params->output_type(), ¶ms->output_type, + error_reporter); + } + *builtin_data = reinterpret_cast(params); break; } - case BuiltinOperator_MAXIMUM: { + case BuiltinOperator_LESS: { break; } case BuiltinOperator_DELEGATE: { // TODO(ycling): Revisit when supporting saving delegated models. error_reporter->Report("DELEGATE op shouldn't exist in model."); - break; + return kTfLiteError; } } - return builtin_data; + return kTfLiteOk; } } // namespace @@ -699,10 +716,13 @@ TfLiteStatus InterpreterBuilder::ParseNodes( reinterpret_cast(op->custom_options()->data()), op->custom_options()->size(), nullptr, reg); } else { + void* builtin_data = nullptr; + TF_LITE_ENSURE_STATUS( + ParseOpData(op, op_type, error_reporter_, &builtin_data)); interpreter->AddNodeWithParameters( FlatBufferIntArrayToVector(op->inputs()), - FlatBufferIntArrayToVector(op->outputs()), nullptr, 0, - ParseOpData(op, op_type, error_reporter_), reg); + FlatBufferIntArrayToVector(op->outputs()), nullptr, 0, builtin_data, + reg); } } diff --git a/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h b/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h index 85aca3687402a89b557d76ab5ace80dea8f8b23d..ace4827d8ce2150cde69a0b0c2d5ca39203193bd 100644 --- a/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h +++ b/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h @@ -34,10 +34,13 @@ limitations under the License. inline void* loadLibrary(const char* name) { // TODO: change RTLD_LOCAL? Assumes there can be multiple instances of nn // api RT - void* handle = dlopen(name, RTLD_LAZY | RTLD_LOCAL); + void* handle = nullptr; +#ifdef __ANDROID__ + handle = dlopen(name, RTLD_LAZY | RTLD_LOCAL); if (handle == nullptr) { NNAPI_LOG("nnapi error: unable to open library %s", name); } +#endif return handle; } diff --git a/tensorflow/contrib/lite/nnapi_delegate.cc b/tensorflow/contrib/lite/nnapi_delegate.cc index bc13444dc70f27e3360774e843985b6294be6996..eab82ea8ef23542cb3fc490d913313d4c757e466 100644 --- a/tensorflow/contrib/lite/nnapi_delegate.cc +++ b/tensorflow/contrib/lite/nnapi_delegate.cc @@ -351,6 +351,9 @@ void AddOpsAndParams(tflite::Interpreter* interpreter, case tflite::BuiltinOperator_CAST: case tflite::BuiltinOperator_PRELU: case tflite::BuiltinOperator_MAXIMUM: + case tflite::BuiltinOperator_MINIMUM: + case tflite::BuiltinOperator_ARG_MAX: + case tflite::BuiltinOperator_LESS: 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/optional_debug_tools.cc b/tensorflow/contrib/lite/optional_debug_tools.cc index 1f762e6688d0cc2a91417b9d82201446e3060a6f..e1366639c78a4e90740aaf42a9ba5770ec65cb78 100644 --- a/tensorflow/contrib/lite/optional_debug_tools.cc +++ b/tensorflow/contrib/lite/optional_debug_tools.cc @@ -48,6 +48,8 @@ const char* TensorTypeName(TfLiteType type) { return "kTfLiteInt64"; case kTfLiteString: return "kTfLiteString"; + case kTfLiteBool: + return "kTfLiteBool"; } return "(invalid)"; } diff --git a/tensorflow/contrib/lite/profiling/BUILD b/tensorflow/contrib/lite/profiling/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..15999e5d4188db1e191936ae6d84faf8cce5ca6e --- /dev/null +++ b/tensorflow/contrib/lite/profiling/BUILD @@ -0,0 +1,44 @@ +package(default_visibility = ["//visibility:public"]) + +licenses(["notice"]) # Apache 2.0 + +common_copts = [ + "-Wall", +] + +cc_library( + name = "profiler", + hdrs = ["profiler.h"], + copts = common_copts, + deps = [":profile_buffer"], +) + +cc_test( + name = "profiler_test", + srcs = ["profiler_test.cc"], + copts = ["-DTFLITE_PROFILING_ENABLED"], + defines = ["TFLITE_PROFILING_ENABLED"], + deps = [ + ":profiler", + "//tensorflow/contrib/lite/testing:util", + "@com_google_googletest//:gtest", + ], +) + +cc_library( + name = "profile_buffer", + hdrs = ["profile_buffer.h"], + copts = common_copts, +) + +cc_test( + name = "profile_buffer_test", + srcs = ["profile_buffer_test.cc"], + copts = ["-DTFLITE_PROFILING_ENABLED"], + defines = ["TFLITE_PROFILING_ENABLED"], + deps = [ + ":profile_buffer", + "//tensorflow/contrib/lite/testing:util", + "@com_google_googletest//:gtest", + ], +) diff --git a/tensorflow/contrib/lite/profiling/profile_buffer.h b/tensorflow/contrib/lite/profiling/profile_buffer.h new file mode 100644 index 0000000000000000000000000000000000000000..3bfe02571ba59f05ff316d327d2a964f6d5b4e1e --- /dev/null +++ b/tensorflow/contrib/lite/profiling/profile_buffer.h @@ -0,0 +1,150 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT 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_PROFILING_PROFILE_BUFFER_H_ +#define TENSORFLOW_CONTRIB_LITE_PROFILING_PROFILE_BUFFER_H_ + +#include +#include + +namespace tflite { +namespace profiling { + +// A profiling event. +struct ProfileEvent { + // Describes the type of event. + // The event_metadata field may contain additional data for interpreting + // the event. + enum class EventType { + // Default event type, the metadata field has no special significance. + DEFAULT = 0, + // The event is an operator invocation and the event_metadata field is the + // index of operator node. + OPERATOR_INVOKE_EVENT = 1 + }; + + // Label of the event. This usually describes the event. + const char* tag; + // Timestamp in microseconds when the event began. + int64_t begin_timestamp_ms; + // Timestamp in microseconds when the event ended. + int64_t end_timestamp_ms; + // The field containing the type of event. This must be one of the event types + // in EventType. + EventType event_type; + // Extra data describing the details of the event. + uint32_t event_metadata; +}; +} // namespace profiling +} // namespace tflite + +#ifdef TFLITE_PROFILING_ENABLED + +#include +#include + +namespace tflite { +namespace profiling { +constexpr uint32_t kInvalidEventHandle = static_cast(~0) - 1; + +// A ring buffer of profile events. +// This class is not thread safe. +class ProfileBuffer { + public: + ProfileBuffer(uint32_t max_num_entries, bool enabled) + : enabled_(enabled), current_index_(0), event_buffer_(max_num_entries) {} + + // Adds an event to the buffer with begin timestamp set to the current + // timestamp. Returns a handle to event that can be used to call EndEvent. If + // buffer is disabled this has no affect. + // The tag of the event should remain valid till the buffer is valid. + uint32_t BeginEvent(const char* tag, ProfileEvent::EventType event_type, + uint32_t event_metadata) { + if (!enabled_) { + return kInvalidEventHandle; + } + int64_t timestamp = NowMicros(); + int index = current_index_ % event_buffer_.size(); + event_buffer_[index].tag = tag; + event_buffer_[index].event_type = event_type; + event_buffer_[index].event_metadata = event_metadata; + event_buffer_[index].begin_timestamp_ms = timestamp; + event_buffer_[index].end_timestamp_ms = 0; + current_index_++; + return index; + } + + // Sets the enabled state of buffer to |enabled| + void SetEnabled(bool enabled) { enabled_ = enabled; } + + // Sets the end timestamp for event for the handle to current time. + // If the buffer is disabled or previous event has been overwritten this + // operation has not effect. + void EndEvent(uint32_t event_handle) { + if (!enabled_ || event_handle == kInvalidEventHandle || + event_handle > current_index_) { + return; + } + const uint32_t max_size = event_buffer_.size(); + if (current_index_ > (max_size + event_handle)) { + // Ignore, buffer has already overflowed. + return; + } + + int event_index = event_handle % max_size; + event_buffer_[event_index].end_timestamp_ms = NowMicros(); + } + + // Returns the size of the buffer. + size_t Size() const { + return (current_index_ >= event_buffer_.size()) ? event_buffer_.size() + : current_index_; + } + + // Resets the buffer. + void Reset() { + enabled_ = false; + current_index_ = 0; + } + + // Returns the profile event at the given index. If the index is invalid a + // nullptr is returned. The return event may get overwritten if more events + // are added to buffer. + const struct ProfileEvent* const At(int index) const { + size_t size = Size(); + if (index >= size) { + return nullptr; + } + const uint32_t max_size = event_buffer_.size(); + uint32_t start = + (current_index_ > max_size) ? current_index_ % max_size : max_size; + index = (index + start) % max_size; + return &event_buffer_[index]; + } + + private: + static int64_t NowMicros() { + // TODO(shashishekhar): Refactor this to a separate file. + struct timeval tv; + gettimeofday(&tv, nullptr); + return static_cast(tv.tv_sec) * 1000000 + tv.tv_usec; + } + bool enabled_; + uint32_t current_index_; + std::vector event_buffer_; +}; +} // namespace profiling +} // namespace tflite +#endif // TFLITE_PROFILING_ENABLED +#endif // TENSORFLOW_CONTRIB_LITE_PROFILING_PROFILE_BUFFER_H_ diff --git a/tensorflow/contrib/lite/profiling/profile_buffer_test.cc b/tensorflow/contrib/lite/profiling/profile_buffer_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..0c5f0cd31495d67beb867beb835cfc897d7d6265 --- /dev/null +++ b/tensorflow/contrib/lite/profiling/profile_buffer_test.cc @@ -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 +#include "tensorflow/contrib/lite/profiling/profile_buffer.h" +#include "tensorflow/contrib/lite/testing/util.h" + +namespace tflite { +namespace profiling { + +namespace { + +std::vector GetProfileEvents(const ProfileBuffer& buffer) { + std::vector events; + for (auto i = 0; i < buffer.Size(); i++) { + events.push_back(buffer.At(i)); + } + return events; +} + +TEST(ProfileBufferTest, Empty) { + ProfileBuffer buffer(/*max_size*/ 0, /*enabled*/ true); + EXPECT_EQ(0, buffer.Size()); +} + +TEST(ProfileBufferTest, AddEvent) { + ProfileBuffer buffer(/*max_size*/ 10, /*enabled*/ true); + EXPECT_EQ(0, buffer.Size()); + auto event_handle = buffer.BeginEvent( + "hello", ProfileEvent::EventType::DEFAULT, /* event_metadata */ 42); + + EXPECT_GE(event_handle, 0); + EXPECT_EQ(1, buffer.Size()); + + auto event = GetProfileEvents(buffer)[0]; + EXPECT_EQ(event->tag, "hello"); + EXPECT_GT(event->begin_timestamp_ms, 0); + EXPECT_EQ(event->event_type, ProfileEvent::EventType::DEFAULT); + EXPECT_EQ(event->event_metadata, 42); + + buffer.EndEvent(event_handle); + EXPECT_EQ(1, buffer.Size()); + EXPECT_GE(event->end_timestamp_ms, event->begin_timestamp_ms); +} + +TEST(ProfileBufferTest, OverFlow) { + const int max_size = 4; + ProfileBuffer buffer{max_size, true}; + std::vector eventNames = {"first", "second", "third", "fourth"}; + for (int i = 0; i < 2 * max_size; i++) { + buffer.BeginEvent(eventNames[i % 4].c_str(), + ProfileEvent::EventType::DEFAULT, i); + size_t expected_size = std::min(i + 1, max_size); + EXPECT_EQ(expected_size, buffer.Size()); + } + EXPECT_EQ(max_size, buffer.Size()); + for (int j = 0; j < buffer.Size(); ++j) { + auto event = buffer.At(j); + EXPECT_EQ(eventNames[j % 4], event->tag); + EXPECT_EQ(ProfileEvent::EventType::DEFAULT, event->event_type); + EXPECT_EQ(4 + j, event->event_metadata); + } +} + +TEST(ProfileBufferTest, Enable) { + ProfileBuffer buffer(/*max_size*/ 10, /*enabled*/ false); + EXPECT_EQ(0, buffer.Size()); + auto event_handle = buffer.BeginEvent( + "hello", ProfileEvent::EventType::DEFAULT, /* event_metadata */ 42); + EXPECT_EQ(kInvalidEventHandle, event_handle); + EXPECT_EQ(0, buffer.Size()); + buffer.SetEnabled(true); + event_handle = buffer.BeginEvent("hello", ProfileEvent::EventType::DEFAULT, + /* event_metadata */ 42); + EXPECT_GE(event_handle, 0); + EXPECT_EQ(1, buffer.Size()); +} + +} // namespace +} // namespace profiling +} // namespace tflite + +int main(int argc, char** argv) { + ::tflite::LogToStderr(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/profiling/profiler.h b/tensorflow/contrib/lite/profiling/profiler.h new file mode 100644 index 0000000000000000000000000000000000000000..dfa98a6708edc874708537dc324d8c340664dc63 --- /dev/null +++ b/tensorflow/contrib/lite/profiling/profiler.h @@ -0,0 +1,174 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT 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_PROFILING_PROFILER_H_ +#define TENSORFLOW_CONTRIB_LITE_PROFILING_PROFILER_H_ + +#include + +#include "tensorflow/contrib/lite/profiling/profile_buffer.h" + +#ifdef TFLITE_PROFILING_ENABLED + +namespace tflite { +namespace profiling { +class ScopedProfile; +class ScopedOperatorProfile; + +// Controls whether profiling is enabled or disabled and collects profiles. +// TFLite is used on platforms that don't have posix threads, so the profiler is +// kept as simple as possible. It is designed to be used only on a single +// thread. +// +// Profiles are collected using Scoped*Profile objects that begin and end a +// profile event. +// An example usage is shown in the example below: +// +// Say Worker class has a DoWork method and we are interested in profiling +// the overall execution time for DoWork and time spent in Task1 and Task2 +// functions. +// +// class Worker { +// public: +// void DoWork() { +// ScopedProfile(&controller, "DoWork"); +// Task1(); +// Task2(); +// ..... +// } +// +// void Task1() { +// ScopedProfile(&controller, "Task1"); +// .... +// } +// +// void Task2() { +// ScopedProfile(&controller, "Task2"); +// } +// +// Profiler profiler; +// } +// +// We instrument the functions that need to be profiled. +// +// Profile can be collected by enable profiling and then getting profile +// events. +// +// void ProfileWorker() { +// Worker worker; +// worker.profiler.EnableProfiling(); +// worker.DoWork(); +// worker.profiler.DisableProfiling(); +// // Profiling is complete, extract profiles. +// auto profile_events = worker.profiler.GetProfiles(); +// } +// +// +class Profiler { + public: + Profiler() : buffer_(1024, false) {} + + void StartProfiling() { buffer_.SetEnabled(true); } + void StopProfiling() { buffer_.SetEnabled(false); } + void Reset() { buffer_.Reset(); } + std::vector GetProfileEvents() { + std::vector profile_events; + profile_events.reserve(buffer_.Size()); + for (int i = 0; i < buffer_.Size(); i++) { + profile_events.push_back(buffer_.At(i)); + } + return profile_events; + } + + private: + friend class ScopedProfile; + friend class ScopedOperatorProfile; + ProfileBuffer* GetProfileBuffer() { return &buffer_; } + ProfileBuffer buffer_; +}; + +class ScopedProfile { + public: + // Adds a profile event to profile that begins with the construction + // of object and ends when the object goes out of scope. + // The lifetime of tag should be at least the lifetime of profiler. + ScopedProfile(Profiler* profiler, const char* tag) { + if (profiler) { + buffer_ = profiler->GetProfileBuffer(); + event_handle_ = + buffer_->BeginEvent(tag, ProfileEvent::EventType::DEFAULT, 0); + } + } + ~ScopedProfile() { + if (buffer_) { + buffer_->EndEvent(event_handle_); + } + } + + private: + ProfileBuffer* buffer_; + int32_t event_handle_; +}; + +class ScopedOperatorProfile { + public: + // Adds a profile event to profile that begins with the construction + // of object and ends when the object goes out of scope. + // The lifetime of tag should be at least the lifetime of profiler. + ScopedOperatorProfile(Profiler* profiler, const char* tag, int node_index) { + if (profiler) { + buffer_ = profiler->GetProfileBuffer(); + event_handle_ = buffer_->BeginEvent( + tag, ProfileEvent::EventType::OPERATOR_INVOKE_EVENT, node_index); + } + } + + ~ScopedOperatorProfile() { + if (buffer_) { + buffer_->EndEvent(event_handle_); + } + } + + private: + ProfileBuffer* buffer_; + int32_t event_handle_; +}; + +} // namespace profiling +} // namespace tflite + +#define SCOPED_OPERATOR_PROFILE(profiler, node_index) \ + tflite::profiling::ScopedOperatorProfile _profile((profiler), "OpInvoke", \ + (node_index)) +#else + +namespace tflite { +namespace profiling { +// A noop version of profiler when profiling is disabled. +class Profiler { + public: + Profiler() {} + void StartProfiling() {} + void StopProfiling() {} + void Reset() {} + std::vector GetProfileEvents() { return {}; } +}; +} // namespace profiling +} // namespace tflite + +#define SCOPED_OPERATOR_PROFILE(profiler, node_index) + +#endif // TFLITE_PROFILING_ENABLED + +#endif // TENSORFLOW_CONTRIB_LITE_PROFILING_PROFILER_H_ diff --git a/tensorflow/contrib/lite/profiling/profiler_test.cc b/tensorflow/contrib/lite/profiling/profiler_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..994523a8fb74b511a18f2ec53147ce81323bec1b --- /dev/null +++ b/tensorflow/contrib/lite/profiling/profiler_test.cc @@ -0,0 +1,105 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT 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 // NOLINT(build/c++11) +#include +#include // NOLINT(build/c++11) + +#include +#include +#include "tensorflow/contrib/lite/profiling/profiler.h" +#include "tensorflow/contrib/lite/testing/util.h" + +namespace tflite { +namespace profiling { +namespace { + +void AssertDurationOfEventAroundMs(const ProfileEvent* event, + double expected_ms, double eps_ms) { + double duration_ms = + (event->end_timestamp_ms - event->begin_timestamp_ms) / 1e3; + EXPECT_NEAR(expected_ms, duration_ms, eps_ms); +} + +void SleepForQuarterSecond(Profiler* profiler) { + ScopedProfile profile(profiler, "SleepForQuarter"); + std::this_thread::sleep_for(std::chrono::milliseconds(250)); +} + +void ChildFunction(Profiler* profiler) { + ScopedProfile profile(profiler, "Child"); + SleepForQuarterSecond(profiler); +} + +void ParentFunction(Profiler* profiler) { + ScopedProfile profile(profiler, "Parent"); + for (int i = 0; i < 2; i++) { + ChildFunction(profiler); + } +} + +TEST(ProfilerTest, NoProfilesAreCollectedWhenDisabled) { + Profiler profiler; + ParentFunction(&profiler); + auto profile_events = profiler.GetProfileEvents(); + EXPECT_EQ(0, profile_events.size()); +} + +TEST(ProfilingTest, ProfilesAreCollected) { + Profiler profiler; + profiler.StartProfiling(); + ParentFunction(&profiler); + profiler.StopProfiling(); + auto profile_events = profiler.GetProfileEvents(); + // ParentFunction calls the ChildFunction 2 times. + // Each ChildFunction calls SleepForQuarterSecond once. + // We expect 1 entry for ParentFunction, 2 for ChildFunction and 2 for + // SleepForQuarterSecond: Total: 1+ 2 + 2 = 5 + // Profiles should look like: + // Parent ~ 500 ms (due to 2 Child calls) + // - Child ~ 250 ms (due to SleepForQuarter calls) + // - SleepForQuarter ~ 250ms + // - Child ~ 250 ms (due to SleepForQuarter calls) + // - SleepForQuarter ~ 250ms + // + ASSERT_EQ(5, profile_events.size()); + EXPECT_EQ("Parent", profile_events[0]->tag); + EXPECT_EQ("Child", profile_events[1]->tag); + EXPECT_EQ("SleepForQuarter", profile_events[2]->tag); + EXPECT_EQ("Child", profile_events[3]->tag); + EXPECT_EQ("SleepForQuarter", profile_events[4]->tag); + + AssertDurationOfEventAroundMs(profile_events[0], /*expected_ms*/ 500, + /*eps_ms*/ 2); + AssertDurationOfEventAroundMs(profile_events[1], /*expected_ms*/ 250, + /*eps_ms*/ 2); + AssertDurationOfEventAroundMs(profile_events[2], /*expected_ms*/ 250, + /*eps_ms*/ 2); + AssertDurationOfEventAroundMs(profile_events[3], /*expected_ms*/ 250, + /*eps_ms*/ 2); + AssertDurationOfEventAroundMs(profile_events[4], /*expected_ms*/ 250, + /*eps_ms*/ 2); +} + +} // namespace +} // namespace profiling +} // namespace tflite + +int main(int argc, char** argv) { + ::tflite::LogToStderr(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/contrib/lite/python/BUILD b/tensorflow/contrib/lite/python/BUILD index e735062a7f2749c1e1c43e9c5f4971b3c7383387..926896d609d83aac3b875d33dfe3c4dc7ae89ccd 100644 --- a/tensorflow/contrib/lite/python/BUILD +++ b/tensorflow/contrib/lite/python/BUILD @@ -97,6 +97,19 @@ py_binary( ], ) +py_binary( + name = "create_custom_op", + srcs = ["create_custom_op.py"], + srcs_version = "PY2AND3", + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/contrib/framework:framework_py", + "//tensorflow/core:protos_all_py", + "//tensorflow/python:platform", + "@absl_py//absl/flags", + ], +) + py_test( name = "convert_saved_model_test", srcs = ["convert_saved_model_test.py"], @@ -106,8 +119,13 @@ py_test( deps = [ ":convert_saved_model", "//tensorflow/python:client_testlib", + "//tensorflow/python:layers", + "//tensorflow/python:nn", "//tensorflow/python:platform_test", "//tensorflow/python:session", + "//tensorflow/python/estimator:estimator_py", + "//tensorflow/python/keras", + "//tensorflow/python/ops/losses", "//tensorflow/python/saved_model", ], ) diff --git a/tensorflow/contrib/lite/python/convert_saved_model_test.py b/tensorflow/contrib/lite/python/convert_saved_model_test.py index d87fbeb91cc3d2779c0ae01aff488f88bd340c1c..734e42d619bdb79de0306a94e304ce46065d14d4 100644 --- a/tensorflow/contrib/lite/python/convert_saved_model_test.py +++ b/tensorflow/contrib/lite/python/convert_saved_model_test.py @@ -25,21 +25,21 @@ from __future__ import print_function import os from tensorflow.contrib.lite.python import convert_saved_model -from tensorflow.python import estimator from tensorflow.python import keras -from tensorflow.python import layers -from tensorflow.python import losses -from tensorflow.python import nn -from tensorflow.python import saved_model -from tensorflow.python import train from tensorflow.python.client import session +from tensorflow.python.estimator import estimator_lib as estimator from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import test_util +from tensorflow.python.layers import layers from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops +from tensorflow.python.ops import nn from tensorflow.python.ops import random_ops +from tensorflow.python.ops.losses import losses from tensorflow.python.platform import test +from tensorflow.python.saved_model import saved_model +from tensorflow.python.training import training as train class ConvertSavedModelTestBasicGraph(test_util.TensorFlowTestCase): diff --git a/tensorflow/contrib/lite/python/create_custom_op.py b/tensorflow/contrib/lite/python/create_custom_op.py new file mode 100644 index 0000000000000000000000000000000000000000..830f95358c455047db2cbad15cfed8c221e95dca --- /dev/null +++ b/tensorflow/contrib/lite/python/create_custom_op.py @@ -0,0 +1,111 @@ +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +r"""Replaces a subgraph of a TensorFlow GraphDef with a single node. + +In conjunction with TOCO's --allow_custom_op this script allows selected +portions of a TensorFlow GraphDef to be executed by custom code. + +Example: + +bazel run tensorflow/contrib/lite/python:create_custom_op -- \ + --input_graph=/tmp/input.pb \ + --output_graph=/tmp/output.pb \ + --inputs=concat,concat_1 \ + --outputs=detection_classes \ + --op_definition='op:"PostProcessing" attr{key:"num" value:{i:10}}' + +The above will identify a subgraph starting at nodes 'concat' and 'concat_1', +and ending at 'detection_classes'. All nodes in between will be removed and +replaced by a new op called 'PostProcessing'. + +""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function +import uuid as _uuid +from absl import app +from absl import flags +from google.protobuf import text_format +from tensorflow.contrib.framework.python.framework.graph_util import fuse_op +from tensorflow.core.framework import graph_pb2 +from tensorflow.core.framework import node_def_pb2 +from tensorflow.core.framework import types_pb2 +from tensorflow.python.platform import gfile + +FLAGS = flags.FLAGS + +flags.DEFINE_string("input_graph", "", "Binary graphdef to load.") +flags.DEFINE_string("output_graph", "", "Resulting binary graphdef.") + +flags.DEFINE_string("inputs", "", + "Comma-separated list of inputs to the subgraph.") +flags.DEFINE_string("outputs", "", + "Comma-separated list of outputs of the subgraph.") +flags.DEFINE_string("op_definition", "", + "A text NodeDef defining the contents of the custom op.") + + +def _read_graph_def(filename): + if not gfile.Exists(filename): + raise ValueError("Input graph file '" + filename + "' does not exist!") + + graph_def = graph_pb2.GraphDef() + with gfile.FastGFile(filename, "rb") as f: + graph_def.ParseFromString(f.read()) + return graph_def + + +def _write_graph_def(graph_def, filename): + if not filename: + raise ValueError("Output graph file not specified") + + with gfile.Open(filename, "wb") as f: + f.write(graph_def.SerializeToString()) + + +def _collapse_subgraph(graph_def, inputs, outputs, op_definition): + """Substitute a custom op for the subgraph delimited by inputs and outputs.""" + name = _uuid.uuid1().hex + # We need a default type, but it can be changed using 'op_definition'. + default_type = types_pb2.DT_FLOAT + new_graph = fuse_op( + graph_def=graph_def, + input_nodes=inputs, + output_nodes=outputs, + output_dtypes=[default_type for _ in outputs], + output_quantized=False, + op_name=name, + op_type="CustomTfLiteOp") + node_def = node_def_pb2.NodeDef() + text_format.Parse(op_definition, node_def) + for node in new_graph.node: + if node.name == name: + node.MergeFrom(node_def) + return new_graph + + +def main(argv): + del argv # unused + graph = _read_graph_def(filename=flags.FLAGS.input_graph) + graph = _collapse_subgraph( + graph_def=graph, + inputs=flags.FLAGS.inputs.split(","), + outputs=flags.FLAGS.outputs.split(","), + op_definition=flags.FLAGS.op_definition) + _write_graph_def(graph_def=graph, filename=flags.FLAGS.output_graph) + + +if __name__ == "__main__": + app.run(main) diff --git a/tensorflow/contrib/lite/python/interpreter.py b/tensorflow/contrib/lite/python/interpreter.py index b8638007f7e49737726d9939a00e8cb1d6a41281..cb9c0d3121895595ffce91e254bea3f527714809 100644 --- a/tensorflow/contrib/lite/python/interpreter.py +++ b/tensorflow/contrib/lite/python/interpreter.py @@ -121,8 +121,8 @@ class Interpreter(object): 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') + if not self._interpreter.ResizeInputTensor(input_index, tensor_size): + raise ValueError('Failed to resize input') def get_output_details(self): """Gets model output details. diff --git a/tensorflow/contrib/lite/python/interpreter_test.py b/tensorflow/contrib/lite/python/interpreter_test.py index cd2386f5263f24e1e034015ec6880e71f0608c7c..f802edf020db8a9d4e7bb890aadaae7e34e983a8 100644 --- a/tensorflow/contrib/lite/python/interpreter_test.py +++ b/tensorflow/contrib/lite/python/interpreter_test.py @@ -81,6 +81,9 @@ class InterpreterTest(test_util.TensorFlowTestCase): test_input = np.array([[1, 2, 3, 4]], dtype=np.uint8) expected_output = np.array([[4, 3, 2, 1]], dtype=np.uint8) + interpreter.resize_tensor_input(input_details[0]['index'], + np.array(test_input.shape, dtype=np.int32)) + interpreter.allocate_tensors() interpreter.set_tensor(input_details[0]['index'], test_input) interpreter.invoke() diff --git a/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.cc b/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.cc index 35ad226b78c906f0819afd5b029a1a0d438d69af..04fc098129854e168d68de3b308eabbcaa968ea8 100644 --- a/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.cc +++ b/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.cc @@ -72,6 +72,8 @@ int TfLiteTypeToPyArrayType(TfLiteType tf_lite_type) { return NPY_INT64; case kTfLiteString: return NPY_OBJECT; + case kTfLiteBool: + return NPY_BOOL; case kTfLiteNoType: return -1; } @@ -90,6 +92,8 @@ TfLiteType TfLiteTypeFromPyArray(PyArrayObject* array) { return kTfLiteUInt8; case NPY_INT64: return kTfLiteInt64; + case NPY_BOOL: + return kTfLiteBool; case NPY_OBJECT: case NPY_STRING: case NPY_UNICODE: @@ -186,7 +190,7 @@ bool InterpreterWrapper::ResizeInputTensor(int i, PyObject* value) { std::vector dims(PyArray_SHAPE(array)[0]); memcpy(dims.data(), PyArray_BYTES(array), dims.size() * sizeof(int)); - return interpreter_->ResizeInputTensor(i, dims); + return (interpreter_->ResizeInputTensor(i, dims) == kTfLiteOk); } std::string InterpreterWrapper::TensorName(int i) const { diff --git a/tensorflow/contrib/lite/python/lite.py b/tensorflow/contrib/lite/python/lite.py index ed6dd036f9fd9f39b74e902498d815793943924b..cf50f9d4d65cb7a36af8f82e2d29babbc9884d23 100644 --- a/tensorflow/contrib/lite/python/lite.py +++ b/tensorflow/contrib/lite/python/lite.py @@ -145,7 +145,8 @@ def toco_convert(input_data, input_format=TENSORFLOW_GRAPHDEF, output_format=TFLITE, quantized_input_stats=None, - drop_control_dependency=True): + drop_control_dependency=True, + allow_custom_ops=None): """Convert a model using TOCO from `input_format` to `output_format`. Typically this is to convert from TensorFlow GraphDef to TFLite, in which @@ -178,9 +179,12 @@ def toco_convert(input_data, toco = _toco_flags_pb2.TocoFlags() toco.input_format = input_format toco.output_format = output_format + toco.inference_type = inference_type toco.drop_control_dependency = drop_control_dependency + if allow_custom_ops is not None: + toco.allow_custom_ops = allow_custom_ops + model = _model_flags_pb2.ModelFlags() - toco.inference_type = inference_type for idx, input_tensor in enumerate(input_tensors): if input_tensor.dtype == _dtypes.float32: tflite_input_type = FLOAT diff --git a/tensorflow/contrib/lite/schema/BUILD b/tensorflow/contrib/lite/schema/BUILD index 246ec85fe47e496e157a91ab4ff84f6b1eeab4a4..9717a4a1a496b888348514584888e62c4e3703b4 100644 --- a/tensorflow/contrib/lite/schema/BUILD +++ b/tensorflow/contrib/lite/schema/BUILD @@ -63,6 +63,9 @@ cc_test( "schema.fbs", "schema_v3.fbs", ], + tags = [ + "tflite_not_portable_android", + ], deps = [ "//tensorflow/core:lib_platform", "@com_google_googletest//:gtest", diff --git a/tensorflow/contrib/lite/schema/schema.fbs b/tensorflow/contrib/lite/schema/schema.fbs index c63bfb28cc66494c3a42250632485c67861c5e9b..2b62c257d8410f9af1b250c9d108eba6737a9efe 100644 --- a/tensorflow/contrib/lite/schema/schema.fbs +++ b/tensorflow/contrib/lite/schema/schema.fbs @@ -33,6 +33,7 @@ enum TensorType : byte { UINT8 = 3, INT64 = 4, STRING = 5, + BOOL = 6, } // Parameters for converting a quantized tensor back to float. Given a @@ -132,6 +133,9 @@ enum BuiltinOperator : byte { CAST = 53, PRELU = 54, MAXIMUM = 55, + ARG_MAX = 56, + MINIMUM = 57, + LESS = 58, } // Options for the builtin operators. @@ -174,7 +178,9 @@ union BuiltinOptions { LogSoftmaxOptions, CastOptions, DequantizeOptions, - MaximumOptions, + MaximumMinimumOptions, + ArgMaxOptions, + LessOptions, } enum Padding : byte { SAME, VALID } @@ -193,6 +199,8 @@ table Conv2DOptions { stride_w:int; stride_h:int; fused_activation_function:ActivationFunctionType; + dilation_w_factor:int = 1; + dilation_h_factor:int = 1; } table Pool2DOptions { @@ -388,7 +396,14 @@ table CastOptions { table DequantizeOptions { } -table MaximumOptions { +table MaximumMinimumOptions { +} + +table ArgMaxOptions { + output_type : TensorType; +} + +table LessOptions { } // An OperatorCode can be an enum value (BuiltinOperator) if the operator is a diff --git a/tensorflow/contrib/lite/schema/schema_generated.h b/tensorflow/contrib/lite/schema/schema_generated.h index 0735be5c8f1b1c8a87c3d47839ce54595d58af7d..0b9961d606d6095b0dc693df631e4efffcb0e35e 100755 --- a/tensorflow/contrib/lite/schema/schema_generated.h +++ b/tensorflow/contrib/lite/schema/schema_generated.h @@ -145,8 +145,14 @@ struct CastOptionsT; struct DequantizeOptions; struct DequantizeOptionsT; -struct MaximumOptions; -struct MaximumOptionsT; +struct MaximumMinimumOptions; +struct MaximumMinimumOptionsT; + +struct ArgMaxOptions; +struct ArgMaxOptionsT; + +struct LessOptions; +struct LessOptionsT; struct OperatorCode; struct OperatorCodeT; @@ -170,18 +176,20 @@ enum TensorType { TensorType_UINT8 = 3, TensorType_INT64 = 4, TensorType_STRING = 5, + TensorType_BOOL = 6, TensorType_MIN = TensorType_FLOAT32, - TensorType_MAX = TensorType_STRING + TensorType_MAX = TensorType_BOOL }; -inline TensorType (&EnumValuesTensorType())[6] { +inline TensorType (&EnumValuesTensorType())[7] { static TensorType values[] = { TensorType_FLOAT32, TensorType_FLOAT16, TensorType_INT32, TensorType_UINT8, TensorType_INT64, - TensorType_STRING + TensorType_STRING, + TensorType_BOOL }; return values; } @@ -194,6 +202,7 @@ inline const char **EnumNamesTensorType() { "UINT8", "INT64", "STRING", + "BOOL", nullptr }; return names; @@ -259,11 +268,14 @@ enum BuiltinOperator { BuiltinOperator_CAST = 53, BuiltinOperator_PRELU = 54, BuiltinOperator_MAXIMUM = 55, + BuiltinOperator_ARG_MAX = 56, + BuiltinOperator_MINIMUM = 57, + BuiltinOperator_LESS = 58, BuiltinOperator_MIN = BuiltinOperator_ADD, - BuiltinOperator_MAX = BuiltinOperator_MAXIMUM + BuiltinOperator_MAX = BuiltinOperator_LESS }; -inline BuiltinOperator (&EnumValuesBuiltinOperator())[54] { +inline BuiltinOperator (&EnumValuesBuiltinOperator())[57] { static BuiltinOperator values[] = { BuiltinOperator_ADD, BuiltinOperator_AVERAGE_POOL_2D, @@ -318,7 +330,10 @@ inline BuiltinOperator (&EnumValuesBuiltinOperator())[54] { BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM, BuiltinOperator_CAST, BuiltinOperator_PRELU, - BuiltinOperator_MAXIMUM + BuiltinOperator_MAXIMUM, + BuiltinOperator_ARG_MAX, + BuiltinOperator_MINIMUM, + BuiltinOperator_LESS }; return values; } @@ -381,6 +396,9 @@ inline const char **EnumNamesBuiltinOperator() { "CAST", "PRELU", "MAXIMUM", + "ARG_MAX", + "MINIMUM", + "LESS", nullptr }; return names; @@ -431,12 +449,14 @@ enum BuiltinOptions { BuiltinOptions_LogSoftmaxOptions = 36, BuiltinOptions_CastOptions = 37, BuiltinOptions_DequantizeOptions = 38, - BuiltinOptions_MaximumOptions = 39, + BuiltinOptions_MaximumMinimumOptions = 39, + BuiltinOptions_ArgMaxOptions = 40, + BuiltinOptions_LessOptions = 41, BuiltinOptions_MIN = BuiltinOptions_NONE, - BuiltinOptions_MAX = BuiltinOptions_MaximumOptions + BuiltinOptions_MAX = BuiltinOptions_LessOptions }; -inline BuiltinOptions (&EnumValuesBuiltinOptions())[40] { +inline BuiltinOptions (&EnumValuesBuiltinOptions())[42] { static BuiltinOptions values[] = { BuiltinOptions_NONE, BuiltinOptions_Conv2DOptions, @@ -477,7 +497,9 @@ inline BuiltinOptions (&EnumValuesBuiltinOptions())[40] { BuiltinOptions_LogSoftmaxOptions, BuiltinOptions_CastOptions, BuiltinOptions_DequantizeOptions, - BuiltinOptions_MaximumOptions + BuiltinOptions_MaximumMinimumOptions, + BuiltinOptions_ArgMaxOptions, + BuiltinOptions_LessOptions }; return values; } @@ -523,7 +545,9 @@ inline const char **EnumNamesBuiltinOptions() { "LogSoftmaxOptions", "CastOptions", "DequantizeOptions", - "MaximumOptions", + "MaximumMinimumOptions", + "ArgMaxOptions", + "LessOptions", nullptr }; return names; @@ -690,8 +714,16 @@ template<> struct BuiltinOptionsTraits { static const BuiltinOptions enum_value = BuiltinOptions_DequantizeOptions; }; -template<> struct BuiltinOptionsTraits { - static const BuiltinOptions enum_value = BuiltinOptions_MaximumOptions; +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_MaximumMinimumOptions; +}; + +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_ArgMaxOptions; +}; + +template<> struct BuiltinOptionsTraits { + static const BuiltinOptions enum_value = BuiltinOptions_LessOptions; }; struct BuiltinOptionsUnion { @@ -1029,13 +1061,29 @@ struct BuiltinOptionsUnion { return type == BuiltinOptions_DequantizeOptions ? reinterpret_cast(value) : nullptr; } - MaximumOptionsT *AsMaximumOptions() { - return type == BuiltinOptions_MaximumOptions ? - reinterpret_cast(value) : nullptr; + MaximumMinimumOptionsT *AsMaximumMinimumOptions() { + return type == BuiltinOptions_MaximumMinimumOptions ? + reinterpret_cast(value) : nullptr; + } + const MaximumMinimumOptionsT *AsMaximumMinimumOptions() const { + return type == BuiltinOptions_MaximumMinimumOptions ? + reinterpret_cast(value) : nullptr; + } + ArgMaxOptionsT *AsArgMaxOptions() { + return type == BuiltinOptions_ArgMaxOptions ? + reinterpret_cast(value) : nullptr; } - const MaximumOptionsT *AsMaximumOptions() const { - return type == BuiltinOptions_MaximumOptions ? - reinterpret_cast(value) : nullptr; + const ArgMaxOptionsT *AsArgMaxOptions() const { + return type == BuiltinOptions_ArgMaxOptions ? + reinterpret_cast(value) : nullptr; + } + LessOptionsT *AsLessOptions() { + return type == BuiltinOptions_LessOptions ? + reinterpret_cast(value) : nullptr; + } + const LessOptionsT *AsLessOptions() const { + return type == BuiltinOptions_LessOptions ? + reinterpret_cast(value) : nullptr; } }; @@ -1430,11 +1478,15 @@ struct Conv2DOptionsT : public flatbuffers::NativeTable { int32_t stride_w; int32_t stride_h; ActivationFunctionType fused_activation_function; + int32_t dilation_w_factor; + int32_t dilation_h_factor; Conv2DOptionsT() : padding(Padding_SAME), stride_w(0), stride_h(0), - fused_activation_function(ActivationFunctionType_NONE) { + fused_activation_function(ActivationFunctionType_NONE), + dilation_w_factor(0), + dilation_h_factor(0) { } }; @@ -1444,7 +1496,9 @@ struct Conv2DOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { VT_PADDING = 4, VT_STRIDE_W = 6, VT_STRIDE_H = 8, - VT_FUSED_ACTIVATION_FUNCTION = 10 + VT_FUSED_ACTIVATION_FUNCTION = 10, + VT_DILATION_W_FACTOR = 12, + VT_DILATION_H_FACTOR = 14 }; Padding padding() const { return static_cast(GetField(VT_PADDING, 0)); @@ -1458,12 +1512,20 @@ struct Conv2DOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { ActivationFunctionType fused_activation_function() const { return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); } + int32_t dilation_w_factor() const { + return GetField(VT_DILATION_W_FACTOR, 0); + } + int32_t dilation_h_factor() const { + return GetField(VT_DILATION_H_FACTOR, 0); + } bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && VerifyField(verifier, VT_PADDING) && VerifyField(verifier, VT_STRIDE_W) && VerifyField(verifier, VT_STRIDE_H) && VerifyField(verifier, VT_FUSED_ACTIVATION_FUNCTION) && + VerifyField(verifier, VT_DILATION_W_FACTOR) && + VerifyField(verifier, VT_DILATION_H_FACTOR) && verifier.EndTable(); } Conv2DOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; @@ -1486,6 +1548,12 @@ struct Conv2DOptionsBuilder { void add_fused_activation_function(ActivationFunctionType fused_activation_function) { fbb_.AddElement(Conv2DOptions::VT_FUSED_ACTIVATION_FUNCTION, static_cast(fused_activation_function), 0); } + void add_dilation_w_factor(int32_t dilation_w_factor) { + fbb_.AddElement(Conv2DOptions::VT_DILATION_W_FACTOR, dilation_w_factor, 0); + } + void add_dilation_h_factor(int32_t dilation_h_factor) { + fbb_.AddElement(Conv2DOptions::VT_DILATION_H_FACTOR, dilation_h_factor, 0); + } explicit Conv2DOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) : fbb_(_fbb) { start_ = fbb_.StartTable(); @@ -1503,8 +1571,12 @@ inline flatbuffers::Offset CreateConv2DOptions( Padding padding = Padding_SAME, int32_t stride_w = 0, int32_t stride_h = 0, - ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE) { + ActivationFunctionType fused_activation_function = ActivationFunctionType_NONE, + int32_t dilation_w_factor = 0, + int32_t dilation_h_factor = 0) { Conv2DOptionsBuilder builder_(_fbb); + builder_.add_dilation_h_factor(dilation_h_factor); + builder_.add_dilation_w_factor(dilation_w_factor); builder_.add_stride_h(stride_h); builder_.add_stride_w(stride_w); builder_.add_fused_activation_function(fused_activation_function); @@ -3806,45 +3878,139 @@ inline flatbuffers::Offset CreateDequantizeOptions( flatbuffers::Offset CreateDequantizeOptions(flatbuffers::FlatBufferBuilder &_fbb, const DequantizeOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); -struct MaximumOptionsT : public flatbuffers::NativeTable { - typedef MaximumOptions TableType; - MaximumOptionsT() { +struct MaximumMinimumOptionsT : public flatbuffers::NativeTable { + typedef MaximumMinimumOptions TableType; + MaximumMinimumOptionsT() { } }; -struct MaximumOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { - typedef MaximumOptionsT NativeTableType; +struct MaximumMinimumOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef MaximumMinimumOptionsT NativeTableType; bool Verify(flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && verifier.EndTable(); } - MaximumOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; - void UnPackTo(MaximumOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; - static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const MaximumOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + MaximumMinimumOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(MaximumMinimumOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const MaximumMinimumOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); }; -struct MaximumOptionsBuilder { +struct MaximumMinimumOptionsBuilder { flatbuffers::FlatBufferBuilder &fbb_; flatbuffers::uoffset_t start_; - explicit MaximumOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + explicit MaximumMinimumOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) : fbb_(_fbb) { start_ = fbb_.StartTable(); } - MaximumOptionsBuilder &operator=(const MaximumOptionsBuilder &); - flatbuffers::Offset Finish() { + MaximumMinimumOptionsBuilder &operator=(const MaximumMinimumOptionsBuilder &); + flatbuffers::Offset Finish() { const auto end = fbb_.EndTable(start_); - auto o = flatbuffers::Offset(end); + auto o = flatbuffers::Offset(end); return o; } }; -inline flatbuffers::Offset CreateMaximumOptions( +inline flatbuffers::Offset CreateMaximumMinimumOptions( flatbuffers::FlatBufferBuilder &_fbb) { - MaximumOptionsBuilder builder_(_fbb); + MaximumMinimumOptionsBuilder builder_(_fbb); + return builder_.Finish(); +} + +flatbuffers::Offset CreateMaximumMinimumOptions(flatbuffers::FlatBufferBuilder &_fbb, const MaximumMinimumOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct ArgMaxOptionsT : public flatbuffers::NativeTable { + typedef ArgMaxOptions TableType; + TensorType output_type; + ArgMaxOptionsT() + : output_type(TensorType_FLOAT32) { + } +}; + +struct ArgMaxOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef ArgMaxOptionsT NativeTableType; + enum { + VT_OUTPUT_TYPE = 4 + }; + TensorType output_type() const { + return static_cast(GetField(VT_OUTPUT_TYPE, 0)); + } + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + VerifyField(verifier, VT_OUTPUT_TYPE) && + verifier.EndTable(); + } + ArgMaxOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(ArgMaxOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const ArgMaxOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct ArgMaxOptionsBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + void add_output_type(TensorType output_type) { + fbb_.AddElement(ArgMaxOptions::VT_OUTPUT_TYPE, static_cast(output_type), 0); + } + explicit ArgMaxOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + ArgMaxOptionsBuilder &operator=(const ArgMaxOptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +inline flatbuffers::Offset CreateArgMaxOptions( + flatbuffers::FlatBufferBuilder &_fbb, + TensorType output_type = TensorType_FLOAT32) { + ArgMaxOptionsBuilder builder_(_fbb); + builder_.add_output_type(output_type); return builder_.Finish(); } -flatbuffers::Offset CreateMaximumOptions(flatbuffers::FlatBufferBuilder &_fbb, const MaximumOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +flatbuffers::Offset CreateArgMaxOptions(flatbuffers::FlatBufferBuilder &_fbb, const ArgMaxOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); + +struct LessOptionsT : public flatbuffers::NativeTable { + typedef LessOptions TableType; + LessOptionsT() { + } +}; + +struct LessOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { + typedef LessOptionsT NativeTableType; + bool Verify(flatbuffers::Verifier &verifier) const { + return VerifyTableStart(verifier) && + verifier.EndTable(); + } + LessOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const; + void UnPackTo(LessOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const; + static flatbuffers::Offset Pack(flatbuffers::FlatBufferBuilder &_fbb, const LessOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); +}; + +struct LessOptionsBuilder { + flatbuffers::FlatBufferBuilder &fbb_; + flatbuffers::uoffset_t start_; + explicit LessOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb) + : fbb_(_fbb) { + start_ = fbb_.StartTable(); + } + LessOptionsBuilder &operator=(const LessOptionsBuilder &); + flatbuffers::Offset Finish() { + const auto end = fbb_.EndTable(start_); + auto o = flatbuffers::Offset(end); + return o; + } +}; + +inline flatbuffers::Offset CreateLessOptions( + flatbuffers::FlatBufferBuilder &_fbb) { + LessOptionsBuilder builder_(_fbb); + return builder_.Finish(); +} + +flatbuffers::Offset CreateLessOptions(flatbuffers::FlatBufferBuilder &_fbb, const LessOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr); struct OperatorCodeT : public flatbuffers::NativeTable { typedef OperatorCode TableType; @@ -4077,8 +4243,14 @@ struct Operator FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table { const DequantizeOptions *builtin_options_as_DequantizeOptions() const { return builtin_options_type() == BuiltinOptions_DequantizeOptions ? static_cast(builtin_options()) : nullptr; } - const MaximumOptions *builtin_options_as_MaximumOptions() const { - return builtin_options_type() == BuiltinOptions_MaximumOptions ? static_cast(builtin_options()) : nullptr; + const MaximumMinimumOptions *builtin_options_as_MaximumMinimumOptions() const { + return builtin_options_type() == BuiltinOptions_MaximumMinimumOptions ? static_cast(builtin_options()) : nullptr; + } + const ArgMaxOptions *builtin_options_as_ArgMaxOptions() const { + return builtin_options_type() == BuiltinOptions_ArgMaxOptions ? static_cast(builtin_options()) : nullptr; + } + const LessOptions *builtin_options_as_LessOptions() const { + return builtin_options_type() == BuiltinOptions_LessOptions ? static_cast(builtin_options()) : nullptr; } const flatbuffers::Vector *custom_options() const { return GetPointer *>(VT_CUSTOM_OPTIONS); @@ -4258,8 +4430,16 @@ template<> inline const DequantizeOptions *Operator::builtin_options_as inline const MaximumOptions *Operator::builtin_options_as() const { - return builtin_options_as_MaximumOptions(); +template<> inline const MaximumMinimumOptions *Operator::builtin_options_as() const { + return builtin_options_as_MaximumMinimumOptions(); +} + +template<> inline const ArgMaxOptions *Operator::builtin_options_as() const { + return builtin_options_as_ArgMaxOptions(); +} + +template<> inline const LessOptions *Operator::builtin_options_as() const { + return builtin_options_as_LessOptions(); } struct OperatorBuilder { @@ -4729,6 +4909,8 @@ inline void Conv2DOptions::UnPackTo(Conv2DOptionsT *_o, const flatbuffers::resol { 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 = dilation_w_factor(); _o->dilation_w_factor = _e; }; + { auto _e = dilation_h_factor(); _o->dilation_h_factor = _e; }; } inline flatbuffers::Offset Conv2DOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const Conv2DOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { @@ -4743,12 +4925,16 @@ inline flatbuffers::Offset CreateConv2DOptions(flatbuffers::FlatB auto _stride_w = _o->stride_w; auto _stride_h = _o->stride_h; auto _fused_activation_function = _o->fused_activation_function; + auto _dilation_w_factor = _o->dilation_w_factor; + auto _dilation_h_factor = _o->dilation_h_factor; return tflite::CreateConv2DOptions( _fbb, _padding, _stride_w, _stride_h, - _fused_activation_function); + _fused_activation_function, + _dilation_w_factor, + _dilation_h_factor); } inline Pool2DOptionsT *Pool2DOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { @@ -5796,26 +5982,75 @@ inline flatbuffers::Offset CreateDequantizeOptions(flatbuffer _fbb); } -inline MaximumOptionsT *MaximumOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { - auto _o = new MaximumOptionsT(); +inline MaximumMinimumOptionsT *MaximumMinimumOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new MaximumMinimumOptionsT(); UnPackTo(_o, _resolver); return _o; } -inline void MaximumOptions::UnPackTo(MaximumOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { +inline void MaximumMinimumOptions::UnPackTo(MaximumMinimumOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { (void)_o; (void)_resolver; } -inline flatbuffers::Offset MaximumOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const MaximumOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { - return CreateMaximumOptions(_fbb, _o, _rehasher); +inline flatbuffers::Offset MaximumMinimumOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const MaximumMinimumOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateMaximumMinimumOptions(_fbb, _o, _rehasher); } -inline flatbuffers::Offset CreateMaximumOptions(flatbuffers::FlatBufferBuilder &_fbb, const MaximumOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { +inline flatbuffers::Offset CreateMaximumMinimumOptions(flatbuffers::FlatBufferBuilder &_fbb, const MaximumMinimumOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { (void)_rehasher; (void)_o; - struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const MaximumOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; - return tflite::CreateMaximumOptions( + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const MaximumMinimumOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + return tflite::CreateMaximumMinimumOptions( + _fbb); +} + +inline ArgMaxOptionsT *ArgMaxOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new ArgMaxOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void ArgMaxOptions::UnPackTo(ArgMaxOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; + { auto _e = output_type(); _o->output_type = _e; }; +} + +inline flatbuffers::Offset ArgMaxOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const ArgMaxOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateArgMaxOptions(_fbb, _o, _rehasher); +} + +inline flatbuffers::Offset CreateArgMaxOptions(flatbuffers::FlatBufferBuilder &_fbb, const ArgMaxOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const ArgMaxOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + auto _output_type = _o->output_type; + return tflite::CreateArgMaxOptions( + _fbb, + _output_type); +} + +inline LessOptionsT *LessOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const { + auto _o = new LessOptionsT(); + UnPackTo(_o, _resolver); + return _o; +} + +inline void LessOptions::UnPackTo(LessOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const { + (void)_o; + (void)_resolver; +} + +inline flatbuffers::Offset LessOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const LessOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) { + return CreateLessOptions(_fbb, _o, _rehasher); +} + +inline flatbuffers::Offset CreateLessOptions(flatbuffers::FlatBufferBuilder &_fbb, const LessOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) { + (void)_rehasher; + (void)_o; + struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const LessOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va; + return tflite::CreateLessOptions( _fbb); } @@ -6151,8 +6386,16 @@ inline bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, const void *ob auto ptr = reinterpret_cast(obj); return verifier.VerifyTable(ptr); } - case BuiltinOptions_MaximumOptions: { - auto ptr = reinterpret_cast(obj); + case BuiltinOptions_MaximumMinimumOptions: { + auto ptr = reinterpret_cast(obj); + return verifier.VerifyTable(ptr); + } + case BuiltinOptions_ArgMaxOptions: { + auto ptr = reinterpret_cast(obj); + return verifier.VerifyTable(ptr); + } + case BuiltinOptions_LessOptions: { + auto ptr = reinterpret_cast(obj); return verifier.VerifyTable(ptr); } default: return false; @@ -6325,8 +6568,16 @@ inline void *BuiltinOptionsUnion::UnPack(const void *obj, BuiltinOptions type, c auto ptr = reinterpret_cast(obj); return ptr->UnPack(resolver); } - case BuiltinOptions_MaximumOptions: { - auto ptr = reinterpret_cast(obj); + case BuiltinOptions_MaximumMinimumOptions: { + auto ptr = reinterpret_cast(obj); + return ptr->UnPack(resolver); + } + case BuiltinOptions_ArgMaxOptions: { + auto ptr = reinterpret_cast(obj); + return ptr->UnPack(resolver); + } + case BuiltinOptions_LessOptions: { + auto ptr = reinterpret_cast(obj); return ptr->UnPack(resolver); } default: return nullptr; @@ -6487,9 +6738,17 @@ inline flatbuffers::Offset BuiltinOptionsUnion::Pack(flatbuffers::FlatBuff auto ptr = reinterpret_cast(value); return CreateDequantizeOptions(_fbb, ptr, _rehasher).Union(); } - case BuiltinOptions_MaximumOptions: { - auto ptr = reinterpret_cast(value); - return CreateMaximumOptions(_fbb, ptr, _rehasher).Union(); + case BuiltinOptions_MaximumMinimumOptions: { + auto ptr = reinterpret_cast(value); + return CreateMaximumMinimumOptions(_fbb, ptr, _rehasher).Union(); + } + case BuiltinOptions_ArgMaxOptions: { + auto ptr = reinterpret_cast(value); + return CreateArgMaxOptions(_fbb, ptr, _rehasher).Union(); + } + case BuiltinOptions_LessOptions: { + auto ptr = reinterpret_cast(value); + return CreateLessOptions(_fbb, ptr, _rehasher).Union(); } default: return 0; } @@ -6649,8 +6908,16 @@ inline BuiltinOptionsUnion::BuiltinOptionsUnion(const BuiltinOptionsUnion &u) FL value = new DequantizeOptionsT(*reinterpret_cast(u.value)); break; } - case BuiltinOptions_MaximumOptions: { - value = new MaximumOptionsT(*reinterpret_cast(u.value)); + case BuiltinOptions_MaximumMinimumOptions: { + value = new MaximumMinimumOptionsT(*reinterpret_cast(u.value)); + break; + } + case BuiltinOptions_ArgMaxOptions: { + value = new ArgMaxOptionsT(*reinterpret_cast(u.value)); + break; + } + case BuiltinOptions_LessOptions: { + value = new LessOptionsT(*reinterpret_cast(u.value)); break; } default: @@ -6850,8 +7117,18 @@ inline void BuiltinOptionsUnion::Reset() { delete ptr; break; } - case BuiltinOptions_MaximumOptions: { - auto ptr = reinterpret_cast(value); + case BuiltinOptions_MaximumMinimumOptions: { + auto ptr = reinterpret_cast(value); + delete ptr; + break; + } + case BuiltinOptions_ArgMaxOptions: { + auto ptr = reinterpret_cast(value); + delete ptr; + break; + } + case BuiltinOptions_LessOptions: { + auto ptr = reinterpret_cast(value); delete ptr; break; } diff --git a/tensorflow/contrib/lite/testing/BUILD b/tensorflow/contrib/lite/testing/BUILD index 62f20638bac943e9f674087e46c18233e8b09d63..bd888a415b035917e64932103740890876eb493e 100644 --- a/tensorflow/contrib/lite/testing/BUILD +++ b/tensorflow/contrib/lite/testing/BUILD @@ -18,6 +18,7 @@ gen_zipped_test_files( name = "optest", files = [ "add.zip", + "arg_max.zip", "avg_pool.zip", "batch_to_space_nd.zip", "concat.zip", @@ -33,14 +34,15 @@ gen_zipped_test_files( "global_batch_norm.zip", "l2_pool.zip", "l2norm.zip", + "less.zip", "local_response_norm.zip", "log_softmax.zip", "max_pool.zip", "maximum.zip", "mean.zip", + "minimum.zip", "mul.zip", "pad.zip", - "prelu.zip", "relu.zip", "relu1.zip", "relu6.zip", @@ -160,6 +162,9 @@ cc_test( size = "small", srcs = ["tflite_driver_test.cc"], data = ["//tensorflow/contrib/lite:testdata/multi_add.bin"], + tags = [ + "tflite_not_portable_android", + ], deps = [ ":tflite_driver", "@com_google_googletest//:gtest_main", @@ -194,7 +199,6 @@ cc_library( cc_library( name = "util", - testonly = 1, hdrs = ["util.h"], ) @@ -249,7 +253,6 @@ cc_test( cc_library( name = "generate_testspec", - testonly = 1, srcs = ["generate_testspec.cc"], hdrs = ["generate_testspec.h"], deps = [ @@ -275,7 +278,6 @@ cc_test( cc_library( name = "tflite_diff_util", - testonly = 1, srcs = ["tflite_diff_util.cc"], hdrs = ["tflite_diff_util.h"], deps = [ @@ -293,7 +295,6 @@ cc_library( cc_library( name = "tflite_diff_flags", - testonly = 1, hdrs = ["tflite_diff_flags.h"], deps = [ ":split", @@ -336,6 +337,15 @@ tf_cc_test( ], ) +cc_binary( + name = "tflite_diff", + srcs = ["tflite_diff_example_test.cc"], + deps = [ + ":tflite_diff_flags", + ":tflite_diff_util", + ], +) + tf_cc_test( name = "generated_examples_zip_test", size = "large", diff --git a/tensorflow/contrib/lite/testing/generate_examples.py b/tensorflow/contrib/lite/testing/generate_examples.py index f919517e9333331fd3f1fb3908c8679c2b9c21b6..f72a4e0d8cbc89c9de5ee0df61f78d7d32bde73e 100644 --- a/tensorflow/contrib/lite/testing/generate_examples.py +++ b/tensorflow/contrib/lite/testing/generate_examples.py @@ -93,9 +93,6 @@ KNOWN_BUGS = { r"softmax.*input_shape=\[1,3,4,3\]": "67749831", # SpaceToDepth only supports float32. r"space_to_depth.*(float16|int32|uint8|int64)": "68018134", - # BatchToSpaceND doesn't support cropping. This catches test cases with - # const tensors as crops. - r"batch_to_space_nd.*crops=\[\[1,1\],\[1,1\]\]": "70594634", # BatchToSpaceND only supports 4D tensors. r"batch_to_space_nd.*input_shape=\[8,2,2,2,1,1\]": "70594733", # Div will use floordiv. @@ -104,6 +101,10 @@ KNOWN_BUGS = { r"strided_slice.*begin=\[0\].*end=\[1\].*": "73170889", # No support for SplitV r"split.*num_or_size_splits=\[2,2\]": "73377559", + # Needs support for dimensions other than the last one in argmax. + r"arg_max.*axis=0.*": "77546240", + r"arg_max.*axis=1.*": "77546240", + r"arg_max.*axis=2.*": "77546240", } @@ -626,54 +627,6 @@ def make_relu6_tests(zip_path): make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) -def make_prelu_tests(zip_path): - """Make a set of tests to do PReLU.""" - - test_parameters = [{ - # The canonical case for image processing is having a 4D `input` (NHWC) - # and `shared_axes`=[1, 2], so the alpha parameter is per channel. - "input_shape": [[1, 10, 10, 3], [3, 3, 3, 3]], - "shared_axes": [[1, 2], [1]], - }] - - def build_graph(parameters): - """Build the graph for the test case.""" - - input_tensor = tf.placeholder( - dtype=tf.float32, name="input", shape=parameters["input_shape"]) - prelu = tf.keras.layers.PReLU(shared_axes=parameters["shared_axes"]) - out = prelu(input_tensor) - return [input_tensor], [out] - - def build_inputs(parameters, sess, inputs, outputs): - """Build the inputs for the test case.""" - - input_shape = parameters["input_shape"] - input_values = create_tensor_data( - np.float32, input_shape, min_value=-10, max_value=10) - shared_axes = parameters["shared_axes"] - - alpha_shape = [] - for dim in range(1, len(input_shape)): - alpha_shape.append(1 if dim in shared_axes else input_shape[dim]) - - alpha_values = create_tensor_data(np.float32, alpha_shape) - - with tf.variable_scope("", reuse=True): - alpha = tf.get_variable("p_re_lu/alpha") - sess.run(alpha.assign(alpha_values)) - - return [input_values], sess.run( - outputs, feed_dict=dict(zip(inputs, [input_values]))) - - make_zip_of_tests( - zip_path, - test_parameters, - build_graph, - build_inputs, - use_frozen_graph=True) - - # This function tests various TensorFLow functions that generates Const op, # including `tf.ones`, `tf.zeros` and random functions. def make_constant_tests(zip_path): @@ -906,6 +859,41 @@ def make_maximum_tests(zip_path): make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) +def make_minimum_tests(zip_path): + """Make a set of tests to do minimum.""" + + test_parameters = [{ + "input_dtype": [tf.float32], + "input_shape_1": [[3], [1, 100], [4, 2, 3], [5, 224, 224, 3]], + "input_shape_2": [[3], [1, 100], [4, 2, 3], [5, 224, 224, 3]], + }] + + def build_graph(parameters): + """Build the minimum op testing graph.""" + input_tensor_1 = tf.placeholder( + dtype=parameters["input_dtype"], + name="input_1", + shape=parameters["input_shape_1"]) + input_tensor_2 = tf.placeholder( + dtype=parameters["input_dtype"], + name="input_2", + shape=parameters["input_shape_2"]) + + out = tf.minimum(input_tensor_1, input_tensor_2) + return [input_tensor_1, input_tensor_2], [out] + + def build_inputs(parameters, sess, inputs, outputs): + values = [ + create_tensor_data(parameters["input_dtype"], + parameters["input_shape_1"]), + create_tensor_data(parameters["input_dtype"], + parameters["input_shape_2"]) + ] + 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) @@ -1051,6 +1039,7 @@ def make_conv_tests(zip_path): "input_shape": [[1, 3, 4, 3]], "filter_shape": [[1, 1, 3, 2]], "strides": [[1, 1, 1, 1], [1, 2, 3, 1]], + "dilations": [[1, 1, 1, 1], [1, 3, 2, 1], [1, 2, 2, 1]], "padding": ["SAME", "VALID"], "data_format": ["NHWC"], # TODO(aselle): NCHW would be good "constant_filter": [True, False], @@ -1059,6 +1048,7 @@ def make_conv_tests(zip_path): "input_shape": [[2, 14, 14, 2]], "filter_shape": [[6, 6, 2, 2]], "strides": [[1, 1, 1, 1], [1, 2, 3, 1]], + "dilations": [[1, 1, 1, 1], [1, 2, 2, 1]], "padding": ["SAME", "VALID"], "data_format": ["NHWC"], # TODO(aselle): NCHW would be good "constant_filter": [True, False], @@ -1084,6 +1074,7 @@ def make_conv_tests(zip_path): input_tensor, filter_input, strides=parameters["strides"], + dilations=parameters["dilations"], padding=parameters["padding"], data_format=parameters["data_format"]) return input_tensors, [out] @@ -1604,7 +1595,7 @@ def make_batch_to_space_nd_tests(zip_path): test_parameters = [ { "dtype": [tf.float32, tf.int64, tf.int32], - "input_shape": [[12, 2, 2, 1]], + "input_shape": [[12, 3, 3, 1]], "block_shape": [[1, 4], [2, 2], [3, 4]], "crops": [[[0, 0], [0, 0]], [[1, 1], [1, 1]]], "constant_block_shape": [True, False], @@ -1767,19 +1758,7 @@ def make_strided_slice_tests(zip_path): "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], - }, + # TODO(b/73170889) Restore test paramaters removed in cl/191608113. # 2-D { "dtype": [tf.float32, tf.int32, tf.int64], @@ -1954,7 +1933,7 @@ def make_l2_pool(input_tensor, ksize, strides, padding, data_format): def make_topk_tests(zip_path): - """Make a set of tests to do gather.""" + """Make a set of tests to do topk.""" test_parameters = [{ "input_dtype": [tf.float32, tf.int32], @@ -1962,7 +1941,7 @@ def make_topk_tests(zip_path): }] def build_graph(parameters): - """Build the gather op testing graph.""" + """Build the topk op testing graph.""" input_value = tf.placeholder( dtype=parameters["input_dtype"], name="input", @@ -1979,6 +1958,69 @@ def make_topk_tests(zip_path): make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) + +def make_arg_max_tests(zip_path): + """Make a set of tests to do arg_max.""" + + test_parameters = [{ + "input_dtype": [tf.float32, tf.int32], + "input_shape": [[1, 1, 1, 3], [2, 3, 4, 5], [2, 3, 3], [5, 5], [10]], + "axis": [0, 1, 2, 3], + "output_type": [tf.int32, tf.int64], + }] + + def build_graph(parameters): + """Build the topk op testing graph.""" + input_value = tf.placeholder( + dtype=parameters["input_dtype"], + name="input", + shape=parameters["input_shape"]) + axis = tf.constant(parameters["axis"], name="axis") + out = tf.arg_max(input_value, axis, output_type=parameters["output_type"]) + return [input_value], [out] + + 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) + + +def make_less_tests(zip_path): + """Make a set of tests to do less.""" + + test_parameters = [{ + "input_dtype": [tf.float32, tf.int32, tf.int64], + "input_shape_pair": [([1, 1, 1, 3], [1, 1, 1, 3]), + ([2, 3, 4, 5], [2, 3, 4, 5]), ([2, 3, 3], [2, 3]), + ([5, 5], [1]), ([10], [2, 4, 10])], + }] + + def build_graph(parameters): + """Build the less op testing graph.""" + input_value1 = tf.placeholder( + dtype=parameters["input_dtype"], + name="input1", + shape=parameters["input_shape_pair"][0]) + input_value2 = tf.placeholder( + dtype=parameters["input_dtype"], + name="input2", + shape=parameters["input_shape_pair"][1]) + out = tf.less(input_value1, input_value2) + return [input_value1, input_value2], [out] + + def build_inputs(parameters, sess, inputs, outputs): + input_value1 = create_tensor_data(parameters["input_dtype"], + parameters["input_shape_pair"][0]) + input_value2 = create_tensor_data(parameters["input_dtype"], + parameters["input_shape_pair"][1]) + return [input_value1, input_value2], sess.run( + outputs, feed_dict=dict(zip(inputs, [input_value1, input_value2]))) + + make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs) + # Toco binary path provided by the generate rule. bin_path = None diff --git a/tensorflow/contrib/lite/testing/generate_testspec.cc b/tensorflow/contrib/lite/testing/generate_testspec.cc index eb3deafb6986e877f0a553a8b6f712102af4caca..6580845af42b3cdded19b578b41c682089aaf9ef 100644 --- a/tensorflow/contrib/lite/testing/generate_testspec.cc +++ b/tensorflow/contrib/lite/testing/generate_testspec.cc @@ -22,7 +22,22 @@ limitations under the License. namespace tflite { namespace testing { -void GenerateTestSpecFromTensorflowModel( +template +void GenerateCsv(const std::vector& shape, float min, float max, + string* out) { + auto random_float = [](int min, int max) { + static unsigned int seed; + return min + (max - min) * static_cast(rand_r(&seed)) / RAND_MAX; + }; + + std::function random_t = [&](int) { + return static_cast(random_float(min, max)); + }; + std::vector data = GenerateRandomTensor(shape, random_t); + *out = Join(data.data(), data.size(), ","); +} + +bool 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, @@ -31,12 +46,6 @@ void GenerateTestSpecFromTensorflowModel( 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()); @@ -46,15 +55,25 @@ void GenerateTestSpecFromTensorflowModel( 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(), ","); + case tensorflow::DT_FLOAT: + GenerateCsv(shape, -0.5, 0.5, &input_values[i]); + break; + case tensorflow::DT_UINT8: + GenerateCsv(shape, 0, 255, &input_values[i]); + break; + case tensorflow::DT_INT32: + GenerateCsv(shape, -100, 100, &input_values[i]); + break; + case tensorflow::DT_INT64: + GenerateCsv(shape, -100, 100, &input_values[i]); + break; + case tensorflow::DT_BOOL: + GenerateCsv(shape, 0.01, 1.99, &input_values[i]); break; - } default: - - fprintf(stderr, "Unsupported type %d when generating testspec\n", type); - return; + fprintf(stderr, "Unsupported type %d (%s) when generating testspec.\n", + type, input_layer_type[i].c_str()); + return false; } } @@ -82,6 +101,8 @@ void GenerateTestSpecFromTensorflowModel( stream << " output: \"" << runner.ReadOutput(i) << "\"\n"; } stream << "}\n"; + + return true; } } // namespace testing diff --git a/tensorflow/contrib/lite/testing/generate_testspec.h b/tensorflow/contrib/lite/testing/generate_testspec.h index 3529ee709b66625fff6e2a35b78e47f3778f0fe7..6e31a853c3f7f82a89126ff83af784ffd418741a 100644 --- a/tensorflow/contrib/lite/testing/generate_testspec.h +++ b/tensorflow/contrib/lite/testing/generate_testspec.h @@ -34,7 +34,7 @@ namespace testing { // 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( +bool 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, diff --git a/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc b/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc index 6697b86e798756bf3273e36dc105eee17d146aa6..9da8bd7a28891fc4a534710fd0fd8344cde1b197 100644 --- a/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc +++ b/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc @@ -94,6 +94,11 @@ std::map kBrokenTests = { // No support for axis!=0 in GatherV2. {R"(^\/gather.*axis=1)", "76910444"}, + + // No support for arbitrary dimensions in ArgMax. + {R"(^\/arg_max.*axis=0)", "77546240"}, + {R"(^\/arg_max.*axis=1)", "77546240"}, + {R"(^\/arg_max.*axis=2)", "77546240"}, }; // Allows test data to be unzipped into a temporary directory and makes @@ -236,6 +241,7 @@ TEST_P(OpsTest, RunStuff) { ::testing::ValuesIn(UnarchiveZipAndFindTestNames(#zip_base ".zip"))); INSTANTIATE_TESTS(add) +INSTANTIATE_TESTS(arg_max) INSTANTIATE_TESTS(avg_pool) INSTANTIATE_TESTS(batch_to_space_nd) INSTANTIATE_TESTS(concat) @@ -256,11 +262,12 @@ INSTANTIATE_TESTS(log_softmax) INSTANTIATE_TESTS(maximum) INSTANTIATE_TESTS(max_pool) INSTANTIATE_TESTS(mean) +INSTANTIATE_TESTS(minimum) INSTANTIATE_TESTS(mul) INSTANTIATE_TESTS(pad) INSTANTIATE_TESTS(relu) INSTANTIATE_TESTS(relu1) -INSTANTIATE_TESTS(prelu) +// INSTANTIATE_TESTS(prelu) INSTANTIATE_TESTS(relu6) INSTANTIATE_TESTS(reshape) INSTANTIATE_TESTS(resize_bilinear) @@ -273,6 +280,7 @@ INSTANTIATE_TESTS(squeeze) INSTANTIATE_TESTS(strided_slice) INSTANTIATE_TESTS(sub) INSTANTIATE_TESTS(transpose) +INSTANTIATE_TESTS(less) } // namespace testing } // namespace tflite diff --git a/tensorflow/contrib/lite/testing/split.h b/tensorflow/contrib/lite/testing/split.h index 428cfda4f216f0ee6409a32c43a4cf91ecc11922..896f2949efa6aeda76940bae18a11dccf3c2f01b 100644 --- a/tensorflow/contrib/lite/testing/split.h +++ b/tensorflow/contrib/lite/testing/split.h @@ -80,6 +80,16 @@ inline std::vector Split(const string& s, const string& delimiter) { return fields; } +template <> +inline std::vector Split(const string& s, const string& delimiter) { + std::vector fields; + for (const auto& p : SplitToPos(s, delimiter)) { + fields.push_back( + static_cast(strtol(s.data() + p.first, nullptr, 10))); + } + return fields; +} + } // namespace testing } // namespace tflite diff --git a/tensorflow/contrib/lite/testing/split_test.cc b/tensorflow/contrib/lite/testing/split_test.cc index 3d1e25d9c7dab50984928adfe0d7392675578662..76b918cbcd83ef43c52057b84bcc2a8f4ff6b8f7 100644 --- a/tensorflow/contrib/lite/testing/split_test.cc +++ b/tensorflow/contrib/lite/testing/split_test.cc @@ -52,6 +52,11 @@ TEST(SplitTest, SplitUint8) { EXPECT_THAT(Split("1,-1,258", ","), ElementsAre(1, 255, 2)); } +TEST(SplitTest, SplitBool) { + EXPECT_THAT(Split("1, 0, 0, 1", ","), + ElementsAre(true, false, false, true)); +} + } // namespace } // namespace testing } // namespace tflite diff --git a/tensorflow/contrib/lite/testing/tf_driver.cc b/tensorflow/contrib/lite/testing/tf_driver.cc index 2c253bb1983e5ddc5bc12858c929585d1bcee710..7b295875aab12bf48da2341ce05dd53442464cf0 100644 --- a/tensorflow/contrib/lite/testing/tf_driver.cc +++ b/tensorflow/contrib/lite/testing/tf_driver.cc @@ -87,10 +87,9 @@ TfDriver::TfDriver(const std::vector& input_layer, 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"); + Invalidate("Failed to find the model " + bin_file_path); return; } @@ -121,6 +120,10 @@ void TfDriver::SetInput(int id, const string& csv_values) { FillTensorWithData(&tensor, csv_values); break; } + case tensorflow::DT_UINT8: { + FillTensorWithData(&tensor, csv_values); + break; + } default: fprintf(stderr, "Unsupported type %d in SetInput\n", input_types_[id]); Invalidate("Unsupported tensor data type"); @@ -162,6 +165,8 @@ string TfDriver::ReadOutput(int id) { return TensorDataToCsvString(output_tensors_[id]); case tensorflow::DT_INT32: return TensorDataToCsvString(output_tensors_[id]); + case tensorflow::DT_UINT8: + return TensorDataToCsvString(output_tensors_[id]); default: fprintf(stderr, "Unsupported type %d in ResetTensor\n", input_types_[id]); Invalidate("Unsupported tensor data type"); diff --git a/tensorflow/contrib/lite/testing/tflite_diff_example_test.cc b/tensorflow/contrib/lite/testing/tflite_diff_example_test.cc index 3817e68111dbaaf2a38ceff9fbc38f30f303cb5f..5afa0f800cdaa8bf70a11cb6e2ac64ace8138e79 100644 --- a/tensorflow/contrib/lite/testing/tflite_diff_example_test.cc +++ b/tensorflow/contrib/lite/testing/tflite_diff_example_test.cc @@ -19,10 +19,13 @@ limitations under the License. int main(int argc, char** argv) { ::tflite::testing::DiffOptions options = ::tflite::testing::ParseTfliteDiffFlags(&argc, argv); + if (options.tensorflow_model.empty()) return 1; + int failure_count = 0; for (int i = 0; i < 100; i++) { if (!tflite::testing::RunDiffTest(options)) { - return 1; + ++failure_count; } } - return 0; + fprintf(stderr, "Num errors: %d\n", failure_count); + return failure_count != 0 ? 1 : 0; } diff --git a/tensorflow/contrib/lite/testing/tflite_diff_flags.h b/tensorflow/contrib/lite/testing/tflite_diff_flags.h index 5f1129d501b7235f1202b704cf36904e07b8720e..706108ed73bb3fd9bd784cffffe322d6981433e6 100644 --- a/tensorflow/contrib/lite/testing/tflite_diff_flags.h +++ b/tensorflow/contrib/lite/testing/tflite_diff_flags.h @@ -51,9 +51,11 @@ DiffOptions ParseTfliteDiffFlags(int* argc, char** argv) { "output_1,output_2"), }; + bool no_inputs = *argc == 1; bool success = tensorflow::Flags::Parse(argc, argv, flags); - if (!success || (*argc == 2 && !strcmp(argv[1], "--helpfull"))) { + if (!success || no_inputs || (*argc == 2 && !strcmp(argv[1], "--helpfull"))) { fprintf(stderr, "%s", tensorflow::Flags::Usage(argv[0], flags).c_str()); + return {}; } return {values.tensorflow_model, diff --git a/tensorflow/contrib/lite/testing/tflite_diff_util.cc b/tensorflow/contrib/lite/testing/tflite_diff_util.cc index 9ef4e1f66c7d31c746c18d63495e760585d4af9e..f601d3752ddb5df9f2b5ac73d9bc303efaade4a5 100644 --- a/tensorflow/contrib/lite/testing/tflite_diff_util.cc +++ b/tensorflow/contrib/lite/testing/tflite_diff_util.cc @@ -27,13 +27,13 @@ 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); + if (!GenerateTestSpecFromTensorflowModel( + tflite_stream, options.tensorflow_model, options.tflite_model, + options.input_layer, options.input_layer_type, + options.input_layer_shape, options.output_layer)) + return false; 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 diff --git a/tensorflow/contrib/lite/testing/tflite_driver.cc b/tensorflow/contrib/lite/testing/tflite_driver.cc index c399f4f2b78d7420ac6ea7098ed44b2122216279..58fe5bd6e40b3d5979d64fae659eb39bfe87c265 100644 --- a/tensorflow/contrib/lite/testing/tflite_driver.cc +++ b/tensorflow/contrib/lite/testing/tflite_driver.cc @@ -42,6 +42,10 @@ template <> uint8_t Value(const TfLitePtrUnion& data, int index) { return data.uint8[index]; } +template <> +bool Value(const TfLitePtrUnion& data, int index) { + return data.b[index]; +} template void SetTensorData(const std::vector& values, TfLitePtrUnion* data) { @@ -79,6 +83,8 @@ class TfLiteDriver::Expectation { return TypedCheck(verbose, tensor); case kTfLiteUInt8: return TypedCheck(verbose, tensor); + case kTfLiteBool: + return TypedCheck(verbose, tensor); default: fprintf(stderr, "Unsupported type %d in Check\n", tensor.type); return false; @@ -143,7 +149,6 @@ void TfLiteDriver::AllocateTensors() { void TfLiteDriver::LoadModel(const string& bin_file_path) { if (!IsValid()) return; - std::cout << std::endl << "Loading model: " << bin_file_path << std::endl; model_ = FlatBufferModel::BuildFromFile(GetFullPath(bin_file_path).c_str()); if (!model_) { @@ -204,6 +209,12 @@ void TfLiteDriver::SetInput(int id, const string& csv_values) { SetTensorData(values, &tensor->data); break; } + case kTfLiteBool: { + const auto& values = testing::Split(csv_values, ","); + if (!CheckSizes(tensor->bytes, values.size())) return; + SetTensorData(values, &tensor->data); + break; + } default: fprintf(stderr, "Unsupported type %d in SetInput\n", tensor->type); Invalidate("Unsupported tensor data type"); @@ -232,6 +243,9 @@ void TfLiteDriver::SetExpectation(int id, const string& csv_values) { case kTfLiteUInt8: expected_output_[id]->SetData(csv_values); break; + case kTfLiteBool: + expected_output_[id]->SetData(csv_values); + break; default: fprintf(stderr, "Unsupported type %d in SetExpectation\n", tensor->type); Invalidate("Unsupported tensor data type"); diff --git a/tensorflow/contrib/lite/toco/BUILD b/tensorflow/contrib/lite/toco/BUILD index 8a35fb9034ca9cd1b9eb87956aed1eb96485dc9b..3f73ef620e121bf3568db55e8a9fbe799900bc90 100644 --- a/tensorflow/contrib/lite/toco/BUILD +++ b/tensorflow/contrib/lite/toco/BUILD @@ -219,6 +219,8 @@ cc_library( "graph_transformations/drop_fake_quant.cc", "graph_transformations/drop_im2col_arrays.cc", "graph_transformations/ensure_bias_vectors.cc", + "graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc", + "graph_transformations/experimental_shuffle_fc_weights.cc", "graph_transformations/fuse_activation_functions.cc", "graph_transformations/fuse_binary_into_following_affine.cc", "graph_transformations/fuse_binary_into_preceding_affine.cc", @@ -237,7 +239,11 @@ cc_library( "graph_transformations/merge_reshape_into_preceding_transpose.cc", "graph_transformations/propagate_activation_function_into_constants.cc", "graph_transformations/propagate_array_data_types.cc", + "graph_transformations/propagate_default_min_max.cc", + "graph_transformations/propagate_fake_quant_num_bits.cc", "graph_transformations/propagate_fixed_sizes.cc", + "graph_transformations/quantization_util.cc", + "graph_transformations/quantization_util.h", "graph_transformations/quantize.cc", "graph_transformations/read_fake_quant_min_max.cc", "graph_transformations/remove_final_dequantize_op.cc", @@ -246,9 +252,11 @@ cc_library( "graph_transformations/remove_trivial_binary.cc", "graph_transformations/remove_trivial_concatenation.cc", "graph_transformations/remove_trivial_concatenation_input.cc", + "graph_transformations/remove_trivial_fake_quant.cc", "graph_transformations/remove_trivial_passthrough.cc", "graph_transformations/remove_trivial_passthrough.h", "graph_transformations/remove_trivial_quantized_activation_func.cc", + "graph_transformations/remove_trivial_quantized_min_max.cc", "graph_transformations/remove_trivial_reshape.cc", "graph_transformations/remove_trivial_slice.cc", "graph_transformations/remove_unused_op.cc", @@ -263,6 +271,7 @@ cc_library( "graph_transformations/resolve_constant_gather.cc", "graph_transformations/resolve_constant_random_uniform.cc", "graph_transformations/resolve_constant_range.cc", + "graph_transformations/resolve_constant_reshape.cc", "graph_transformations/resolve_constant_shape_or_rank.cc", "graph_transformations/resolve_constant_stack.cc", "graph_transformations/resolve_constant_strided_slice.cc", @@ -298,7 +307,7 @@ cc_library( ":runtime", ":toco_port", ":tooling_util", - ":types_proto_cc", + "//tensorflow/contrib/lite/kernels/internal:quantization_util", "//tensorflow/core:lib", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", @@ -373,7 +382,6 @@ 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 7a7059e3572fbd8ca83fa149cbc0ee5d07883c45..fe30b88344c5340dc8647cd89e244987c86e47fe 100644 --- a/tensorflow/contrib/lite/toco/args.h +++ b/tensorflow/contrib/lite/toco/args.h @@ -227,6 +227,8 @@ struct ParsedTocoFlags { // TODO(aselle): command_line_flags doesn't support doubles Arg default_ranges_min = Arg(0.); Arg default_ranges_max = Arg(0.); + Arg default_int16_ranges_min = Arg(0.); + Arg default_int16_ranges_max = Arg(0.); Arg inference_type; Arg inference_input_type; Arg drop_fake_quant = Arg(false); @@ -237,6 +239,8 @@ struct ParsedTocoFlags { Arg input_types; Arg debug_disable_recurrent_cell_fusion = Arg(false); Arg drop_control_dependency = Arg(false); + Arg propagate_fake_quant_num_bits = Arg(false); + Arg allow_nudging_weights_to_use_fast_gemm_kernel = Arg(false); }; } // namespace toco diff --git a/tensorflow/contrib/lite/toco/dump_graphviz.cc b/tensorflow/contrib/lite/toco/dump_graphviz.cc index c8352741b44cd627ff9edb9c4677b994c4cb9a09..c289ddcd929c74e9a5b0eaf2e0c8fc1a6a01d320 100644 --- a/tensorflow/contrib/lite/toco/dump_graphviz.cc +++ b/tensorflow/contrib/lite/toco/dump_graphviz.cc @@ -95,10 +95,8 @@ Color GetColorForArray(const Model& model, const string& array_name) { array_name == dump_options.graphviz_last_array) { return Color(0x9E, 0x9E, 0x9E); } - for (const string& output_array : model.flags.output_arrays()) { - if (array_name == output_array) { - return Color(0x9E, 0x9E, 0x9E); - } + if (IsOutputArray(model, array_name)) { + return Color(0x9E, 0x9E, 0x9E); } // Remaining arrays are intermediate activation arrays. // Lighter tone of the same grey as for input/output arrays: @@ -119,6 +117,12 @@ void AppendArrayVal(string* string, Array const& array, int index) { return; } AppendF(string, "%d", data[index]); + } else if (array.buffer->type == ArrayDataType::kInt16) { + const auto& data = array.GetBuffer().data; + if (index >= data.size()) { + return; + } + AppendF(string, "%d", data[index]); } else if (array.buffer->type == ArrayDataType::kInt32) { const auto& data = array.GetBuffer().data; if (index >= data.size()) { diff --git a/tensorflow/contrib/lite/toco/export_tensorflow.cc b/tensorflow/contrib/lite/toco/export_tensorflow.cc index 4a85f3c5a409f17f5e824434cbc6a22f4a554dea..99ccfaea648077b7b72af30b32dd53b42b85d3a2 100644 --- a/tensorflow/contrib/lite/toco/export_tensorflow.cc +++ b/tensorflow/contrib/lite/toco/export_tensorflow.cc @@ -883,6 +883,9 @@ void ConvertFakeQuantOperator(const FakeQuantOperator& src_op, CHECK(src_op.minmax); (*fakequant_op->mutable_attr())["min"].set_f(src_op.minmax->min); (*fakequant_op->mutable_attr())["max"].set_f(src_op.minmax->max); + if (src_op.num_bits) { + (*fakequant_op->mutable_attr())["num_bits"].set_i(src_op.num_bits); + } } void ConvertMaxPoolOperator(const MaxPoolOperator& src_op, diff --git a/tensorflow/contrib/lite/toco/graph_transformations/ensure_bias_vectors.cc b/tensorflow/contrib/lite/toco/graph_transformations/ensure_bias_vectors.cc index badefeca883b1e1d67f7de5276389c5e6e7f7cd3..708ecf6e0a96811ab274fbb25f748f562cd3afad 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/ensure_bias_vectors.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/ensure_bias_vectors.cc @@ -47,7 +47,7 @@ bool EnsureBiasVectors::Run(Model* model, std::size_t op_index) { op->type == OperatorType::kDepthwiseConv || op->type == OperatorType::kFullyConnected) { if (ProcessLinearOperator(model, op)) { - AddMessageF("Added bias vector to %s", LogName(*op)); + AddMessageF("Added bias vector to %s as %s", LogName(*op), op->inputs[2]); return true; } } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc b/tensorflow/contrib/lite/toco/graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc new file mode 100644 index 0000000000000000000000000000000000000000..394fa349e2663e2806344f27a96a5132a2d4a810 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.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 +#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 { + +// === Summary === +// +// TLDR: Some of our 8-bit arithmetic operations require uint8 weight values +// to avoid the value 0, thus ranging only in [1, 255]. This enables faster +// runtime arithmetic kernels on ARM NEON. This is not relevant on most +// other hardware architectures, and will cease to be relevant on ARM NEON +// in the future. These topics are elaborated below ("Context"). +// +// Having just one isolated uint8 value equal to 0 is fine. The bad case is when +// two uint8 values are both zero and are less than 16 bytes apart. +// +// By default, toco generates a fatal error when that happens. The user may opt +// in to more lax behavior by passing +// --allow_nudging_weights_to_use_fast_gemm_kernel. +// This causes toco to nudge such bad 0 values into the value 1, thus avoiding +// the problem in exchange for compromising on accuracy. +// +// The present graph transformation implements both the default fatal-erroring +// behavior, and, when allow_nudging_weights is set, also the lax nudging +// behavior. +// +// +// === Context === +// +// Since March 2017, we have been using a trick to perform faster +// 8bit matrix multiplications, to our knowledge first implemented in gemmlowp +// here: +// https://github.com/google/gemmlowp/commit/25b2989415b99e797e1ab977837111b2e231f81f +// +// This trick is explained in Appendix B of our paper, +// https://arxiv.org/abs/1712.05877 +// +// Here is the relevant paragraph: +// +// For efficient NEON implementation of the matrix multiplication’s +// core accumulation, we use the following trick. +// In the multiply-add operation in (10), we first change the +// operands’ type from uint8 to int8 (which can be done by +// subtracting 128 from the quantized values and zero-points). +// Thus the core multiply-add becomes +// +// int32 += int8 * int8. (B.1) +// +// As mentioned in section 3, with a minor tweak of the quantized +// training process, we can ensure that the weights, once +// quantized as int8 values, never take the value −128. Hence, +// the product in (B.1) is never −128 ∗ −128, and is therefore +// always less than 2^14 in absolute value. Hence, (B.1) +// can accumulate two products on a local int16 accumulator +// before that needs to be accumulated into the true int32 accumulator. +// This allows the use of an 8-way SIMD multiplication +// (SMULL on int8 operands), followed by an 8-way +// SIMD multiply-add (SMLAL on int8 operands), followed +// by a pairwise-add-and-accumulate into the int32 accumulators +// (SADALP). +// +// As that paragraph notes, quantized training should be suitably modified to +// ensure that quantized uint8 weights value only range in [1, 255]. So the +// problem that we are dealing with is only about the existing 8-bit quantized +// models that haven't been trained specifically to get 8-bit weights only in +// [1, 255]. +// +// This spreadsheet shows the speed benefit of this trick across many existing +// ARM-architecture CPUs: +// +// https://docs.google.com/spreadsheets/d/1-0LjdMvW0XtH1bYknC0bQINoFaxjTuL9eplZZcitykI/edit?usp=sharing +// +// Compare Row 18 (fast int8 trick) to Row 20 (regular uint8 kernel). +// +// The introduction of the 'dotprod' extension to ARM NEON, specifically the +// SDOT instruction, renders this eventually moot. See the experimental +// kernels contributed by ARM here, +// +// https://github.com/google/gemmlowp/pull/116 +// +// However, as of April 2018, there don't seem to be any commercially available +// CPU supporting these instructions (yet); we are waiting for +// Cortex-A{75,55}-r1 to become available; the "-r1" is key here. Even if such +// CPUs become available soon, it will presumably take years for them to +// overtake the large volume of existing CPUs not supporting these new +// instructions, especially in current and future low-end devices. All in all, +// we can foresee these 'fast int8 kernels' to remain important to have into +// the 2020s. +// +bool EnsureUint8WeightsSafeForFastInt8Kernels::Run(Model* model, + std::size_t op_index) { + const auto& op = *model->operators[op_index]; + int weights_index = 0; + switch (op.type) { + case OperatorType::kConv: + weights_index = 1; + break; + case OperatorType::kLstmCell: + weights_index = 2; + break; + case OperatorType::kFullyConnected: { + weights_index = 1; + const auto& fc_op = static_cast(op); + CHECK(!fc_op.experimental_shuffled_weights) + << "This graph transformation expects to run before FC weights get " + "shuffled."; + break; + } + default: + // Other operator types are unaffected by this graph transformation, + // because their runtime implementations don't use the fast int8 trick. + // In particular that's the case of DepthwiseConv at the moment. + // We have to update this logic when that changes, e.g. if in the future + // some DepthwiseConv kernel wants to use the trick. + // + // The reason why that's not so likely, hence why it's fairly safe to + // stay conservative in the list of operators that we handle here, is that + // the fast int8 kernel trick is only applicable to ops that either are + // implemented as a GEMM, or use symmetric ranges for both weights and + // activations. The reason why GEMM is special (can use the trick even + // without symmetric ranges) is that it is so arithmetic-intense that + // it can use techniques reducing its implementation to the symmetric + // ranges case, with limited relative overhead (O(N^2) overhead vs + // O(N^3) GEMM cost). See https://arxiv.org/pdf/1712.05877, section + // 2.3 Efficient handling of zero-points. + // + // That's why at the moment we only handle operators that use a GEMM + // (Conv, fully-connected --- note that LSTM merely wraps a + // fully-connected operator). + return false; + } + + const string& name = op.inputs[weights_index]; + auto& array = model->GetArray(name); + if (!array.buffer) { + return false; + } + if (array.data_type != ArrayDataType::kUint8) { + return false; + } + auto& buffer_data = array.GetMutableBuffer().data; + + int count_bad = 0; + int index_of_previous_bad_value = 0; + bool changed = false; + + for (int i = 0; i < buffer_data.size(); i++) { + if (buffer_data[i] == 0) { + count_bad++; + if (count_bad > 1) { + const int distance = i - index_of_previous_bad_value; + // Semi-arbitrary threshold. The idea is that trouble only occurs + // when two bad values are very close to each other so that they + // are jointly used within registers inside some GEMM kernel. + // The details of that depend on the kernel. Our current fast ARM64 + // kernel, for instance, only has an issue when the distance between + // consecutive bad values is exactly 8. We do not want to track such + // kernel details too closely here, so we pick a threshold that's + // a bit larger than that, to give us room to change kernels in the + // future without worrying. + static constexpr int kMinDistanceBetweenBadValues = 16; + if (distance < kMinDistanceBetweenBadValues) { + if (allow_nudging_weights()) { + buffer_data[i] = 1; + changed = true; + continue; + } + LOG(FATAL) << "Bad value for " << name << " at index " << i + << ", previous bad value at index " + << index_of_previous_bad_value << ", distance=" << distance + << ", kMinDistanceBetweenBadValues=" + << kMinDistanceBetweenBadValues << ". Consider passing " + << "--allow_nudging_weights_to_use_fast_gemm_kernel " + << "if you don't care about accuracy."; + } + } + index_of_previous_bad_value = i; + } + } + + if (changed) { + AddMessageF("Tweaked weights values for %s", LogName(op)); + } + + return changed; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/experimental_shuffle_fc_weights.cc b/tensorflow/contrib/lite/toco/graph_transformations/experimental_shuffle_fc_weights.cc new file mode 100644 index 0000000000000000000000000000000000000000..f098981a5cf4b91df4c7798bd3db8563705a3bd0 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/experimental_shuffle_fc_weights.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 +#include +#include +#include + +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +bool ExperimentalShuffleFCWeights::Run(Model* model, std::size_t op_index) { + Operator* op = model->operators[op_index].get(); + if (op->type != OperatorType::kFullyConnected) { + return false; + } + FullyConnectedOperator* fc_op = static_cast(op); + // Exit if this FC op already has shuffled weights + if (fc_op->experimental_shuffled_weights) { + return false; + } + const Array& input_array = model->GetArray(fc_op->inputs[0]); + const string& weights_name = fc_op->inputs[1]; + Array& weights_array = model->GetArray(weights_name); + const Array& output_array = model->GetArray(fc_op->outputs[0]); + // Exit if this FC op isn't quantized with uint8 inputs and int16 outputs, + // the only case where we are currently interested in providing a fast path + // with shuffled weights. + if (input_array.data_type != ArrayDataType::kUint8 || + weights_array.data_type != ArrayDataType::kUint8 || + output_array.data_type != ArrayDataType::kInt16 || + !input_array.quantization_params || !weights_array.quantization_params || + !output_array.quantization_params) { + return false; + } + // Exit if the shapes aren't known + if (!input_array.has_shape() || !weights_array.has_shape()) { + return false; + } + // Exit if, based on the known shapes, this FC op is not a GEMV. + // The shuffling of FC weights is only useful to enable fast GEMV paths. + const Shape& input_shape = input_array.shape(); + for (int i = 0; i < input_shape.dimensions_count() - 1; i++) { + if (input_shape.dims(i) != 1) { + // The input activations, shaped as a matrix, have multiple columns. + // This FC op isn't a matrix*vector multiplication. + AddMessageF( + "Not applying experimental shuffling to the weights of %s because " + "it's not a matrix*vector product", + LogName(*op)); + return false; + } + } + // Exit if the weights shape isn't an integral multiple of the shuffled + // block shape, 4x16. We don't want to have to write code dealing with + // odd sizes, that would go un-exercised at the moment as the models + // for which we need this shuffling have shapes that are multiples of that + // 4x16 block size. In fact, much of the rationale for this shuffling is + // to avoid cache aliasin issue with large power-of-two depths, with our + // models motivating this shuffling having FC weights shapes like + // 4096x2048. Thus, if some model doesn't get the shuffling because of that + // size requirement, that might be just fine --- that model might just not + // suffer from that cache aliasing issue that we have with large powers of + // two. + const Shape& weights_shape = weights_array.shape(); + if (weights_shape.dimensions_count() != 2) { + return false; + } + const int rows = weights_shape.dims(0); + const int cols = weights_shape.dims(1); + if (rows % 4 || cols % 16) { + AddMessageF( + "Not applying experimental shuffling to the weights of %s because its " + "shape isn't a multiple of the shuffling block shape, 4x16", + LogName(*op)); + return false; + } + // Exit if the weights aren't already a constant array. + if (!weights_array.buffer) { + return false; + } + // Exit if the weights are used by more than one op. + if (CountOpsWithInput(*model, weights_name) != 1) { + AddMessageF( + "Not applying experimental shuffling to the weights of %s because that " + "array is consumed by other operators", + LogName(*op)); + return false; + } + // Compute the shuffled weights + auto& weights_data = + weights_array.GetMutableBuffer().data; + CHECK_EQ(rows * cols, weights_data.size()); + std::vector shuffled_data(weights_data.size()); + uint8* shuffled_data_ptr = shuffled_data.data(); + for (int r = 0; r < rows; r += 4) { + for (int c = 0; c < cols; c += 16) { + for (int i = 0; i < 4; i++) { + const uint8* src_data_ptr = weights_data.data() + (r + i) * cols + c; + for (int j = 0; j < 16; j++) { + uint8 src_val = *src_data_ptr++; + // Flip the sign bit, so that the runtime will only need to + // reinterpret these uint8 values as int8, getting for free the + // subtraction of the zero_point value 128. + uint8 dst_val = src_val ^ 0x80; + *shuffled_data_ptr++ = dst_val; + } + } + } + } + CHECK_EQ(shuffled_data_ptr, shuffled_data.data() + rows * cols); + // Switch this FC op to using the shuffled weights. + weights_data = std::move(shuffled_data); + fc_op->experimental_shuffled_weights = true; + AddMessageF("Applied experimental shuffling to the weights of %s", + LogName(*op)); + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h b/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h index 27c5044bb3e06e4a052ff0c4984226fb9d113f95..72ffd51db45d0808feab3d07436c13db2420a680 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h +++ b/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h @@ -135,6 +135,7 @@ DECLARE_GRAPH_TRANSFORMATION(IdentifyDilatedConv) DECLARE_GRAPH_TRANSFORMATION(MakeInitialDequantizeOperator) DECLARE_GRAPH_TRANSFORMATION(PropagateActivationFunctionIntoConstants) DECLARE_GRAPH_TRANSFORMATION(PropagateArrayDataTypes) +DECLARE_GRAPH_TRANSFORMATION(PropagateFakeQuantNumBits); DECLARE_GRAPH_TRANSFORMATION(PropagateFixedSizes) DECLARE_GRAPH_TRANSFORMATION(HardcodeMinMax) DECLARE_GRAPH_TRANSFORMATION(Quantize) @@ -144,8 +145,10 @@ DECLARE_GRAPH_TRANSFORMATION(RemoveTensorFlowIdentity) DECLARE_GRAPH_TRANSFORMATION(RemoveTrivialBinaryOperator) DECLARE_GRAPH_TRANSFORMATION(RemoveTrivialConcatenation) DECLARE_GRAPH_TRANSFORMATION(RemoveTrivialConcatenationInput) +DECLARE_GRAPH_TRANSFORMATION(RemoveTrivialFakeQuant) DECLARE_GRAPH_TRANSFORMATION(RemoveTrivialSlice) DECLARE_GRAPH_TRANSFORMATION(RemoveTrivialQuantizedActivationFunc) +DECLARE_GRAPH_TRANSFORMATION(RemoveTrivialQuantizedMinMax) DECLARE_GRAPH_TRANSFORMATION(RemoveUnusedOp) DECLARE_GRAPH_TRANSFORMATION(ResolveBatchNormalization) DECLARE_GRAPH_TRANSFORMATION(ResolveConstantBinaryOperator) @@ -162,8 +165,8 @@ DECLARE_GRAPH_TRANSFORMATION(ResolveTensorFlowMerge) DECLARE_GRAPH_TRANSFORMATION(ResolveSqueezeAttributes) DECLARE_GRAPH_TRANSFORMATION(ResolveTensorFlowSwitch) DECLARE_GRAPH_TRANSFORMATION(ResolveTensorFlowTile) -DECLARE_GRAPH_TRANSFORMATION(ResolveConstantFakeQuant) DECLARE_GRAPH_TRANSFORMATION(ResolveConstantConcatenation) +DECLARE_GRAPH_TRANSFORMATION(ResolveConstantReshape) DECLARE_GRAPH_TRANSFORMATION(ResolveConstantTranspose) DECLARE_GRAPH_TRANSFORMATION(DropFakeQuant) DECLARE_GRAPH_TRANSFORMATION(UnfuseActivationFunctions) @@ -185,6 +188,25 @@ DECLARE_GRAPH_TRANSFORMATION(ResolveConstantGather) DECLARE_GRAPH_TRANSFORMATION(ResolveMultiplyByZero) DECLARE_GRAPH_TRANSFORMATION(Dequantize) DECLARE_GRAPH_TRANSFORMATION(UnpartitionEmbeddingLookup) +DECLARE_GRAPH_TRANSFORMATION(ExperimentalShuffleFCWeights) + +class PropagateDefaultMinMax : public GraphTransformation { + public: + bool Run(Model* model, std::size_t op_index) override; + const char* Name() const override { return "PropagateDefaultMinMax"; } + + bool has_any_ranges_defined() const { return !type_ranges_.empty(); } + void DefineTypeRange(ArrayDataType data_type, double min, double max) { + MinMax minmax; + minmax.min = min; + minmax.max = max; + type_ranges_.emplace_back(data_type, minmax); + } + + private: + bool SetArrayMinMax(const string& array_name, Array* array); + std::vector> type_ranges_; +}; class ResolveReshapeAttributes : public GraphTransformation { public: @@ -207,6 +229,36 @@ class RemoveTrivialReshape : public GraphTransformation { bool treat_expand_dims_as_trivial_ = false; }; +class ResolveConstantFakeQuant : public GraphTransformation { + public: + bool Run(Model* model, std::size_t op_index) override; + const char* Name() const override { return "ResolveConstantFakeQuant"; } + + // True if the num_bits should adjust the final data type. + bool propagate_fake_quant_num_bits() const { + return propagate_fake_quant_num_bits_; + } + void set_propagate_fake_quant_num_bits(bool val) { + propagate_fake_quant_num_bits_ = val; + } + + private: + bool propagate_fake_quant_num_bits_ = false; +}; + +class EnsureUint8WeightsSafeForFastInt8Kernels : public GraphTransformation { + public: + bool Run(Model* model, std::size_t op_index) override; + const char* Name() const override { + return "EnsureUint8WeightsSafeForFastInt8Kernels"; + } + bool allow_nudging_weights() const { return allow_nudging_weights_; } + void set_allow_nudging_weights(bool val) { allow_nudging_weights_ = val; } + + private: + bool allow_nudging_weights_ = false; +}; + #undef DECLARE_GRAPH_TRANSFORMATION } // end namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm.cc b/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm.cc index c363b93394f0af7bcfc37c1e8be5f98aca6667ae..e9842524c829b839b97b3453a36c41efe186efbb 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm.cc @@ -306,6 +306,12 @@ bool IdentifyLstmCell::Run(Model* model, std::size_t op_index) { return false; } + if (static_cast(fully_connected) + ->experimental_shuffled_weights) { + // Not yet implemented: experimental shuffled weights in fused LSTM cell. + return false; + } + // Emplace a new LSTM cell operator auto* lstm_cell_op = new LstmCellOperator; lstm_cell_op->inputs.resize(LstmCellOperator::NUM_INPUTS); 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 183b3d3f2e0ac74099236817e1b6cf5390f2d920..45d9f73a1e6416b8f3fe3936c740da637961b7fc 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 @@ -18,6 +18,7 @@ limitations under the License. #include #include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/graph_transformations/quantization_util.h" #include "tensorflow/contrib/lite/toco/model.h" #include "tensorflow/contrib/lite/toco/model_flags.pb.h" #include "tensorflow/contrib/lite/toco/tooling_util.h" diff --git a/tensorflow/contrib/lite/toco/graph_transformations/propagate_default_min_max.cc b/tensorflow/contrib/lite/toco/graph_transformations/propagate_default_min_max.cc new file mode 100644 index 0000000000000000000000000000000000000000..50b90e7c2bfddb0382a4d44ad6c90fc7f7701273 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/propagate_default_min_max.cc @@ -0,0 +1,86 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT 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/quantization_util.h" +#include "tensorflow/contrib/lite/toco/model.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +// Propagates default min/max values to any operator input/output array that +// is missing them. +// +// When provided a set of min/max values for uint8 arrays this will rescale +// the values for other data types as required and preserving the floating point +// range within the new type. +bool PropagateDefaultMinMax::Run(Model* model, std::size_t op_index) { + const auto it = model->operators.begin() + op_index; + const auto* op = it->get(); + + bool did_change = false; + + for (const auto& input : op->inputs) { + auto& input_array = model->GetArray(input); + if (!input_array.minmax && !input_array.buffer) { + did_change |= SetArrayMinMax(input, &input_array); + } + } + + for (const auto& output : op->outputs) { + auto& output_array = model->GetArray(output); + if (!output_array.minmax && !output_array.buffer) { + did_change |= SetArrayMinMax(output, &output_array); + } + } + + return did_change; +} + +// Sets the min/max on the given array, adjusting the reference_minmax for the +// final data type of the array if it is already specified. +bool PropagateDefaultMinMax::SetArrayMinMax(const string& array_name, + Array* array) { + CHECK(!array->minmax); + + ArrayDataType quantized_data_type = + GetQuantizedDataType(*array, ArrayDataType::kUint8); + for (const auto& type_range : type_ranges_) { + if (type_range.first == quantized_data_type) { + array->GetOrCreateMinMax() = type_range.second; + break; + } + } + if (!array->minmax) { + AddMessageF( + "No defaults specified for quantized data type %s of array %s, " + "skipping", + ArrayDataTypeName(quantized_data_type), array_name); + return false; + } + + AddMessageF("Adding default minmax %g,%g to array %s when quantized as %s", + array->GetMinMax().min, array->GetMinMax().max, array_name, + ArrayDataTypeName(quantized_data_type)); + + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/propagate_fake_quant_num_bits.cc b/tensorflow/contrib/lite/toco/graph_transformations/propagate_fake_quant_num_bits.cc new file mode 100644 index 0000000000000000000000000000000000000000..0bce183c1897dfba6f2c393ffc0306c054366725 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/propagate_fake_quant_num_bits.cc @@ -0,0 +1,307 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +#include +#include +#include +#include + +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/graph_transformations/quantization_util.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 { + +void ChangeArrayDataType(GraphTransformation* transformation, Array* array, + ArrayDataType new_data_type, + const MinMax* new_minmax) { + // Ensure the array ends up in the new type (if it hasn't yet been quantized). + array->final_data_type = new_data_type; + + if (array->minmax && array->quantization_params) { + // The array is already quantized and has min/max info. + // As we are changing the data type we need to fix up the existing min/max + // to the new data type range. + + double old_quantized_min, old_quantized_max; + CHECK(GetQuantizedDataTypeNumericalRange( + array->data_type, &old_quantized_min, &old_quantized_max)) + << "Existing data type is not quantized: " + << ArrayDataTypeName(array->data_type); + double new_quantized_min, new_quantized_max; + CHECK(GetQuantizedDataTypeNumericalRange(new_data_type, &new_quantized_min, + &new_quantized_max)) + << "New data type is not quantized: " + << ArrayDataTypeName(new_data_type); + + // Compute new minmax values. + double min = (old_quantized_min - array->quantization_params->zero_point) * + array->quantization_params->scale; + double max = + (old_quantized_max + 1 - array->quantization_params->zero_point) * + array->quantization_params->scale; + max = max - 1.0 / (new_quantized_max + 1); + + auto& array_minmax = array->GetOrCreateMinMax(); + transformation->AddMessageF( + "Rescaling min/max from %g,%g (%s) to %g,%g (%s)", array_minmax.min, + array_minmax.max, ArrayDataTypeName(array->data_type), min, max, + ArrayDataTypeName(new_data_type)); + + array_minmax.min = min; + array_minmax.max = max; + GetQuantizationParamsFromMinMax( + array_minmax, array->quantization_params.get()); + + // Directly change the type as the array was already quantized. + array->data_type = new_data_type; + } else { + // Array has not yet been quantized so we can just set the final data type + // and assign the new min/max value (if provided). + CHECK(!array->quantization_params); + + if (!array->minmax && new_minmax) { + transformation->AddMessageF("Forcing new minmax to %g,%g (%s)", + new_minmax->min, new_minmax->max, + ArrayDataTypeName(new_data_type)); + auto& array_minmax = array->GetOrCreateMinMax(); + array_minmax.min = new_minmax->min; + array_minmax.max = new_minmax->max; + } + } +} + +// Returns true if the op blocks our backward recursive data type propagation. +bool DoesOpBlockBackwardPropagation(const Operator& op) { + switch (op.type) { + case OperatorType::kConcatenation: + case OperatorType::kTensorFlowConcat: + case OperatorType::kTensorFlowConcatV2: + // Concat shouldn't block propagation, but we do expect that all inputs + // have the same range. + return false; + case OperatorType::kDequantize: + // Dequantize ops are inserted between the value we care about and the + // FakeQuant so make sure we move across them. + case OperatorType::kGather: + // Gathers need their parameters changed to the appropriate data type. + case OperatorType::kTensorFlowReshape: + case OperatorType::kTranspose: + // Reshapes and transposes don't change values. + return false; + default: + return true; + } +} + +// Returns true if the input of an op blocks our backward recursive data type +// propagation. +bool DoesOpInputBlockBackwardPropagation(const Operator& op, int input_index) { + switch (op.type) { + case OperatorType::kGather: + // Ignore gather indices. + return input_index != 0; + break; + case OperatorType::kTensorFlowReshape: + case OperatorType::kTranspose: + // Ignore reshape/transpose shapes/dimensions. + return input_index != 0; + default: + return false; + } +} + +// Propagates the data type up into the input arrays if they are model inputs +// that may need their type changed. May act recursively if the inputs are +// produced by ops that we can move over (such as Dequantize). +bool RecursivelyBackwardPropagateDataType(GraphTransformation* transformation, + Model* model, Operator* op, + ArrayDataType new_data_type, + const MinMax& new_minmax) { + bool did_change = false; + for (int input_index = 0; input_index < op->inputs.size(); ++input_index) { + const auto& input = op->inputs[input_index]; + auto& input_array = model->GetArray(input); + if (input_array.final_data_type == new_data_type) { + // Final data type is already - skip. + continue; + } + + // Prevent moving into constant param args that we don't want to modify. + if (DoesOpInputBlockBackwardPropagation(*op, input_index)) { + continue; + } + + if (input_array.final_data_type != new_data_type) { + transformation->AddMessageF( + "Adjusting input final data type of array %s from %s to %s", input, + ArrayDataTypeName(input_array.final_data_type), + ArrayDataTypeName(new_data_type)); + did_change = true; + ChangeArrayDataType(transformation, &input_array, new_data_type, + &new_minmax); + + // Walk up into all ops producing the inputs to this op. + for (auto& producing_op : model->operators) { + if (!DoesOpBlockBackwardPropagation(*producing_op)) { + for (const auto& output : producing_op->outputs) { + if (input == output) { + did_change |= RecursivelyBackwardPropagateDataType( + transformation, model, producing_op.get(), new_data_type, + new_minmax); + } + } + } + } + } + } + return did_change; +} + +// Returns true if the op blocks our forward recursive data type propagation. +bool DoesOpBlockForwardPropagation(const Operator& op) { + switch (op.type) { + case OperatorType::kFakeQuant: + // Always stop at another FakeQuant, as it will likely have different + // parameters. + return true; + default: + return false; + } +} + +// Recurses down the graph setting the data type of all arrays until an operator +// that blocks propagation (like another FakeQuant) or a final_data_type is +// already specified. +bool RecursivelyForwardPropagateDataType(GraphTransformation* transformation, + Model* model, Operator* op, + ArrayDataType new_data_type) { + bool did_change = false; + for (const auto& output : op->outputs) { + auto& output_array = model->GetArray(output); + if (output_array.final_data_type == new_data_type) { + // Final data type is already - skip. + continue; + } + + if (output_array.final_data_type == ArrayDataType::kNone || + output_array.final_data_type != new_data_type) { + transformation->AddMessageF( + "Adjusting output final data type of array %s from %s to %s", output, + ArrayDataTypeName(output_array.final_data_type), + ArrayDataTypeName(new_data_type)); + did_change = true; + ChangeArrayDataType(transformation, &output_array, new_data_type, + nullptr); + + // Walk down into all ops consuming the output of this op. + for (auto& consuming_op : model->operators) { + if (!DoesOpBlockForwardPropagation(*consuming_op)) { + for (const auto& input : consuming_op->inputs) { + if (input == output) { + did_change |= RecursivelyForwardPropagateDataType( + transformation, model, consuming_op.get(), new_data_type); + } + } + } + } + } + } + return did_change; +} + +} // namespace + +// Propagates the num_bits on a FakeQuant operator into the final data types +// of inputs and outputs. For example, if FakeQuant.num_bits==16 then we know +// the output must be int16 and assume all inputs up until the preceding op are +// also 16. +// +// This can be thought of as a bidirectional flood-fill of the num_bits implied +// final_data_type that terminates at other FakeQuant ops (and a few others as +// determined by DoesOpBlockBackwardPropagation/DoesOpBlockForwardPropagation). +// Once all FakeQuant ops have been visted the arrays should all have +// appropriate final_data_types if the source graph was annotated with the +// proper FakeQuant ops. +// +// Annotating a graph requires following a few hard rules: +// - every input MUST have a FakeQuant immediately following it +// - every output MUST have a FakeQuant immediately preceding it +// - important arithmetic ops (such as FullyConnected) SHOULD have a FakeQuant +// immediately following it +// - all trained weights (RHS of FullyConnected ops, params on Gather ops, etc) +// MUST have FakeQuants between them and the consuming op +// Additional FakeQuants may be used if desired, especially in areas that may +// suffer from large precision changes - such as between a Softmax and a +// FullyConnected. Only by validating accuracy differences between float +// inference with the FakeQuant ops simulating quantization and the actually +// quantized graph can you be sure the appropriate FakeQuant ops are present. +// +// You can tell if you're missing some FakeQuants by looking for warnings from +// quantize.cc about minmax ranges being determined by the contents of constant +// arrays. This will almost never produce functional models during inference. +// +// As this op may change the data types and ranges of input and output arrays +// downstream tools must also be sure to parse the output model flags to get the +// post-Transform values that may have changed due to this transformation. +// +// This isn't a GraphTransformation in the traditional respect as it affects ops +// outside of the one under transformation. This is primarily so that we can +// utilize the graph traversal and repeated pass system underlying the +// transformation system to exhaustively find all FakeQuant ops. It also gets us +// nice logging and integration with the graphviz video dumping mode. +// In general you should not copy this style of transformation and stick to +// local-only changes as seen in the other transformations. +bool PropagateFakeQuantNumBits::Run(Model* model, std::size_t op_index) { + auto it = model->operators.begin() + op_index; + auto* op = it->get(); + if (op->type != OperatorType::kFakeQuant) { + return false; + } + auto* fakequant_op = static_cast(op); + + ArrayDataType quantized_data_type = ArrayDataType::kNone; + if (!InferQuantizedDataTypeFromFakeQuant(*fakequant_op, + &quantized_data_type)) { + AddMessageF("FakeQuant op %s num_bits=%d is out of range, ignoring", + LogName(*op), fakequant_op->num_bits); + return false; + } + const auto& final_minmax = *fakequant_op->minmax; + + AddMessageF( + "Beginning propagation of fake quant %s num_bits=%d min=%g max=%g to %s", + LogName(*op), fakequant_op->num_bits, final_minmax.min, final_minmax.max, + ArrayDataTypeName(quantized_data_type)); + + bool did_change = false; + + // Propagate the FakeQuant information backward up the graph. + // This will possibly adjust input arrays or constant types (like Gather). + did_change |= RecursivelyBackwardPropagateDataType( + this, model, op, quantized_data_type, final_minmax); + + // Propagate the FakeQuant information forward down the graph. + // This will possibly adjust output arrays. + did_change |= + RecursivelyForwardPropagateDataType(this, model, op, quantized_data_type); + + return did_change; +} + +} // namespace toco 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 a648b770f841fa755634cad2d48f031901e5a24a..ba244cf5ef56301e02e49635cc7d085ddfe931a4 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc @@ -1060,17 +1060,15 @@ void ProcessBatchToSpaceNDOperator(Model* model, BatchToSpaceNDOperator* op) { } QCHECK(crops_array.data_type == ArrayDataType::kInt32); const auto& crops_data = crops_array.GetBuffer().data; - // We don't support crops now. - QCHECK_EQ(crops_data[0], 0); - QCHECK_EQ(crops_data[1], 0); - QCHECK_EQ(crops_data[2], 0); - QCHECK_EQ(crops_data[3], 0); - + const int crops_top = crops_data[0]; + const int crops_bottom = crops_data[1]; + const int crops_left = crops_data[2]; + const int crops_right = crops_data[3]; + const int output_height = + input_height * block_height - crops_top - crops_bottom; + const int output_width = input_width * block_width - crops_left - crops_right; QCHECK_EQ(input_shape.dims(0) % (block_height * block_width), 0); - int output_height = input_height * block_height; - int output_width = input_width * block_width; - model->GetArray(op->outputs[0]) .copy_shape(Shape({input_shape.dims(0) / (block_height * block_width), output_height, output_width, input_shape.dims(3)})); @@ -1255,6 +1253,83 @@ void ProcessStackOperator(Model* model, StackOperator* op) { output_array.copy_shape(*stacked_shape); } +// These StridedSlice utility functions are essentially a COPY of those in +// reference_ops.h. See comments there. + +// Use until std::clamp() is available from C++17. +int Clamp(const int v, const int lo, const int hi) { + if (hi < v) return hi; + if (v < lo) return lo; + return v; +} + +int StartForAxis(StridedSliceOperator const& op, Shape const& input_shape, + int axis) { + // Begin with the specified index + int start = op.start_indices[axis]; + + // begin_mask override + if (op.begin_mask & 1 << axis) { + if (op.strides[axis] > 0) { + // Forward iteration - use the first element. These values will get + // clamped below (Note: We could have set them to 0 and axis_size-1, but + // use lowest() and max() to maintain symmetry with StopForAxis()) + start = std::numeric_limits::lowest(); + } else { + // Backward iteration - use the last element. + start = std::numeric_limits::max(); + } + } + + // Handle negative indices + int axis_size = input_shape.dims(axis); + if (start < 0) { + start += axis_size; + } + + // Clamping + start = Clamp(start, 0, axis_size - 1); + + return start; +} + +int StopForAxis(StridedSliceOperator const& op, Shape const& input_shape, + int axis) { + // Begin with the specified index + int stop = op.stop_indices[axis]; + + // end_mask override + if (op.end_mask & (1 << axis)) { + if (op.strides[axis] > 0) { + // Forward iteration - use the last element. These values will get + // clamped below + stop = std::numeric_limits::max(); + } else { + // Backward iteration - use the first element. + stop = std::numeric_limits::lowest(); + } + } + + // Handle negative indices + int axis_size = input_shape.dims(axis); + if (stop < 0) { + stop += axis_size; + } + + // Clamping + // Because the end index points one past the last element, we need slightly + // different clamping ranges depending on the direction. + if (op.strides[axis] > 0) { + // Forward iteration + stop = Clamp(stop, 0, axis_size); + } else { + // Backward iteration + stop = Clamp(stop, -1, axis_size - 1); + } + + return stop; +} + void ProcessStridedSliceOperator(Model* model, StridedSliceOperator* op) { CHECK_GE(op->inputs.size(), 1); CHECK_EQ(op->outputs.size(), 1); @@ -1292,43 +1367,46 @@ void ProcessStridedSliceOperator(Model* model, StridedSliceOperator* op) { return; } - int dim_count = input_array.shape().dimensions_count(); - CHECK(op->start_indices.size() == dim_count) - << ": Incorrect number of start indices supplied to StridedSlice op with " - "output \"" - << op->outputs[0] << "\". Op requires " << dim_count << " start indices"; - CHECK(op->stop_indices.size() == dim_count) - << ": Incorrect number of stop indices supplied to StridedSlice op with " - "output \"" - << op->outputs[0] << "\". Op requires " << dim_count << " stop indices"; - CHECK(op->strides.size() == dim_count) - << ": Incorrect number of strides supplied to StridedSlice op with " - " output \"" - << op->outputs[0] << "\". Op requires " << dim_count << " strides"; + int num_input_axes = input_array.shape().dimensions_count(); + CHECK_LE(op->start_indices.size(), num_input_axes) + << "StridedSlice op with output \"" << op->outputs[0] + << "\", requires no more than " << num_input_axes << " start indices"; + CHECK_LE(op->stop_indices.size(), num_input_axes) + << "StridedSlice op with output \"" << op->outputs[0] + << "\", requires no more than " << num_input_axes << " stop indices"; + CHECK_LE(op->strides.size(), num_input_axes) + << "StridedSlice op with output \"" << op->outputs[0] + << "\", requires no more than " << num_input_axes << " strides"; + for (int i = 0; i < op->strides.size(); i++) { + CHECK_NE(op->strides[i], 0) << "Strides must be non-zero. Axis " << i + << " has stride=" << op->strides[i] << "."; + } + + // The TensorFlow documentation is not explicit on how it handles fewer + // supplied indices than dimensions, but they are accepted. We emulate TF's + // behavior by fully iterating over each "forgotten" dimension. + op->PadIndices(num_input_axes); // Create output shape std::vector* dims = output_array.mutable_shape()->mutable_dims(); // Compute output shape - for (int i = 0; i < dim_count; ++i) { - const int mask = 1 << i; - int start = (op->begin_mask & mask) ? 0 : op->start_indices[i]; - if (start < 0) { - // handle negative indices - start += input_array.shape().dims(i); - } - int stop = (op->end_mask & mask) ? input_array.shape().dims(i) - : op->stop_indices[i]; - if (stop < 0) { - // handle negative indices - stop += input_array.shape().dims(i); - } - - int dim_size = ceil((stop - start) / static_cast(op->strides[i])); - dim_size = dim_size < 0 ? 0 : dim_size; - if (op->shrink_axis_mask & mask) { - CHECK_EQ(dim_size, 1) << "Output size for an axis must compute to 1 when " - "shrinking that axis"; + for (int axis = 0; axis < num_input_axes; ++axis) { + int start_index = StartForAxis(*op, input_array.shape(), axis); + int stop_index = StopForAxis(*op, input_array.shape(), axis); + int dim_size = + ceil(static_cast(stop_index - start_index) / op->strides[axis]); + + CHECK_GT(dim_size, 0) + << "Output size for an axis must be greater than 0. Axis " << axis + << " computes to size " << dim_size + << " for StridedSlice op with output \"" << op->outputs[0] << "\"."; + if (op->shrink_axis_mask & (1 << axis)) { + CHECK_EQ(dim_size, 1) + << "Output size for an axis must compute to 1 when shrinking an " + "axis. Axis " + << axis << " computes to size " << dim_size + << " for StridedSlice op with output \"" << op->outputs[0] << "\"."; } else { dims->push_back(dim_size); } @@ -1438,10 +1516,7 @@ void ProcessArgMaxOperator(Model* model, ArgMaxOperator* op) { return; } - // The current ArgMax implementation only supports 4-dimensional inputs with - // the last dimension as the axis to perform ArgMax for. const std::vector& input_dims = input_array.shape().dims(); - CHECK_EQ(input_dims.size(), 4); std::vector output_dims; output_dims.reserve(input_dims.size() - 1); diff --git a/tensorflow/contrib/lite/toco/graph_transformations/quantization_util.cc b/tensorflow/contrib/lite/toco/graph_transformations/quantization_util.cc new file mode 100644 index 0000000000000000000000000000000000000000..d74cad9a626b3a472e2740d6bdaaaf7aab5bd484 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/quantization_util.cc @@ -0,0 +1,261 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT 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/graph_transformations/quantization_util.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 InferQuantizedDataTypeFromFakeQuant( + const FakeQuantOperator& op, ArrayDataType* out_quantized_data_type) { + if (op.num_bits <= 8) { + *out_quantized_data_type = ArrayDataType::kUint8; + return true; + } else if (op.num_bits <= 16) { + *out_quantized_data_type = ArrayDataType::kInt16; + return true; + } else { + *out_quantized_data_type = ArrayDataType::kNone; + return false; + } +} + +bool GetQuantizedDataTypeNumericalRange(ArrayDataType data_type, + double* out_min_value, + double* out_max_value) { + switch (data_type) { + case ArrayDataType::kUint8: + *out_min_value = 0; + *out_max_value = 255; + return true; + case ArrayDataType::kInt16: + *out_min_value = -32768; + *out_max_value = 32767; + return true; + default: + return false; + } +} + +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); + } +} + +namespace { + +template +std::unique_ptr QuantizeBuffer( + const GenericBuffer& buffer, + const QuantizationParams& quantization_params) { + const auto inverse_scale = 1. / quantization_params.scale; + CHECK(buffer.type == ArrayDataType::kFloat); + const auto& float_buffer = + static_cast&>(buffer); + auto* quantized_buffer = new Buffer; + quantized_buffer->data.resize(float_buffer.data.size()); + for (std::size_t i = 0; i < float_buffer.data.size(); i++) { + const float src_val = float_buffer.data[i]; + double scaled_val; // Astonishingly, using 'float' degrades accuracy just + // enough to make a few tests fail! + if (quantization_params.scale == 0) { + CHECK_EQ(src_val, 0) << "The quantization scale for this array is 0, " + << "so all its values should be 0."; + scaled_val = quantization_params.zero_point; + } else { + scaled_val = quantization_params.zero_point + inverse_scale * src_val; + } + quantized_buffer->data[i] = + tflite::SafeCast>(std::round(scaled_val)); + } + return std::unique_ptr(quantized_buffer); +} + +template +void QuantizeArray(GraphTransformation* transformation, Model* model, + const string& name, + const QuantizationParams& quantization_params) { + auto& array = model->GetArray(name); + CHECK(array.data_type == ArrayDataType::kFloat); + CHECK(!array.quantization_params); + array.GetOrCreateQuantizationParams() = quantization_params; + if (array.buffer) { + array.buffer = QuantizeBuffer(*array.buffer, quantization_params); + } + array.data_type = A; + array.final_data_type = A; + transformation->AddMessageF( + "Quantized array %s to %s zero_point=%g, scale=%g", name, + ArrayDataTypeName(array.data_type), quantization_params.zero_point, + quantization_params.scale); +} + +} // namespace + +void QuantizeArray(GraphTransformation* transformation, Model* model, + const string& name, ArrayDataType quantized_data_type, + const QuantizationParams& quantization_params) { + 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); + default: + LOG(FATAL) << "Unhandled case."; + } +} + +bool IsArrayQuantizedRangeSubset(GraphTransformation* transformation, + const Array& array, double clamp_min, + double clamp_max) { + ArrayDataType quantized_data_type = + GetQuantizedDataType(array, array.data_type); + if (quantized_data_type == ArrayDataType::kNone || + quantized_data_type == ArrayDataType::kFloat) { + // The array is not (or never will be) quantized. + return false; + } + + QuantizationParams quantization_params; + if (!array.quantization_params) { + if (!array.minmax) { + transformation->AddMessageF("No quantization params and no minmax"); + return false; + } else { + // Work around cases where we are asking for this prior to the Quantize + // transformation having added the quantization_params. + GetQuantizationParams(quantized_data_type, *array.minmax, + &quantization_params); + transformation->AddMessageF( + "No quantization params - infering from data type %s with minmax " + "%g,%g as zero_point=%g, scale=%g", + ArrayDataTypeName(quantized_data_type), array.minmax->min, + array.minmax->max, quantization_params.zero_point, + quantization_params.scale); + } + } else { + quantization_params = array.GetQuantizationParams(); + } + + double quantized_min, quantized_max; + CHECK(GetQuantizedDataTypeNumericalRange(quantized_data_type, &quantized_min, + &quantized_max)) + << "Type is not quantized"; + + bool has_nontrivial_min_bound = false; + bool has_nontrivial_max_bound = false; + + double lowest_representable_output = + (quantized_min - quantization_params.zero_point) * + quantization_params.scale; + if (lowest_representable_output < clamp_min) { + has_nontrivial_min_bound = true; + transformation->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); + } + + double highest_representable_output = + (quantized_max - quantization_params.zero_point) * + quantization_params.scale; + if (highest_representable_output > clamp_max) { + has_nontrivial_max_bound = true; + transformation->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); + } + + return !has_nontrivial_min_bound && !has_nontrivial_max_bound; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/quantization_util.h b/tensorflow/contrib/lite/toco/graph_transformations/quantization_util.h new file mode 100644 index 0000000000000000000000000000000000000000..79a2ce7e50887b4608b278471da0e5e63b5673e3 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/quantization_util.h @@ -0,0 +1,73 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_QUANTIZATION_UTIL_H_ +#define TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_QUANTIZATION_UTIL_H_ + +#include "tensorflow/contrib/lite/kernels/internal/quantization_util.h" +#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/model.h" + +namespace toco { + +// Gets the target quantized data type of an array based on the fake quant op. +// For example, if the num_bits is 8 the data type will be kUint8. +bool InferQuantizedDataTypeFromFakeQuant( + const FakeQuantOperator& op, ArrayDataType* out_quantized_data_type); + +// Gets the min/max numerical range for the given quantized data type. +// For example, kUint8 will return [0,255]. +// Returns true if the ranges were set and false if the type is not quantized. +bool GetQuantizedDataTypeNumericalRange(ArrayDataType data_type, + double* out_min_value, + double* out_max_value); + +// Returns the quantized data type of an array, falling back to the provided +// default data type. +ArrayDataType GetQuantizedDataType(const Array& array, + ArrayDataType default_type); + +// Returns the quantization params for the array with the given data type and +// minmax. +void GetQuantizationParams(ArrayDataType data_type, const MinMax& minmax, + QuantizationParams* quantization_params); + +// Returns the quantization params for the data type and minmax values. +template +void GetQuantizationParamsFromMinMax(const MinMax& minmax, + QuantizationParams* quantization_params) { + using Integer = DataType; + const double rmin = minmax.min; + const double rmax = minmax.max; + *quantization_params = + ::tflite::ChooseQuantizationParams(rmin, rmax); +} + +// Quantizes an array by setting its data type and (if constant) quantizing +// all values in the array. +void QuantizeArray(GraphTransformation* transformation, Model* model, + const string& name, ArrayDataType quantized_data_type, + const QuantizationParams& quantization_params); + +// Returns true if the given array, when quantized, contains only values between +// the provided clamp min/max. +// Either clamp_min or clamp_max may be +/-infinity to indicate that the value +// is unbounded on that side. +bool IsArrayQuantizedRangeSubset(GraphTransformation* transformation, + const Array& array, double clamp_min, + double clamp_max); + +} // namespace toco + +#endif // TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_QUANTIZATION_UTIL_H_ diff --git a/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc b/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc index 5b1268f9a9bf49862e072f305836ae93cdd0344b..fa46e6bc3805d3f3a7c9223ff4b111a4ed8e8559 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc @@ -21,6 +21,7 @@ limitations under the License. #include #include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/graph_transformations/quantization_util.h" #include "tensorflow/contrib/lite/toco/model.h" #include "tensorflow/contrib/lite/toco/model_flags.pb.h" #include "tensorflow/contrib/lite/toco/tooling_util.h" @@ -44,6 +45,7 @@ bool SupportsQuantization(const Operator& op) { type == OperatorType::kTensorFlowMinimum || type == OperatorType::kTensorFlowMaximum || type == OperatorType::kLogistic || type == OperatorType::kSoftmax || + type == OperatorType::kLogSoftmax || type == OperatorType::kTensorFlowSplit || type == OperatorType::kSub || type == OperatorType::kSqueeze || type == OperatorType::kPad || type == OperatorType::kTensorFlowReshape || @@ -55,72 +57,6 @@ bool SupportsQuantization(const Operator& op) { type == OperatorType::kTranspose || type == OperatorType::kMean; } -template -std::unique_ptr QuantizeBuffer( - const GenericBuffer& buffer, - const QuantizationParams& quantization_params) { - const auto inverse_scale = 1. / quantization_params.scale; - CHECK(buffer.type == ArrayDataType::kFloat); - const auto& float_buffer = - static_cast&>(buffer); - auto* quantized_buffer = new Buffer; - quantized_buffer->data.resize(float_buffer.data.size()); - for (std::size_t i = 0; i < float_buffer.data.size(); i++) { - const float src_val = float_buffer.data[i]; - double scaled_val; // Astonishingly, using 'float' degrades accuracy just - // enough to make a few tests fail! - if (quantization_params.scale == 0) { - CHECK_EQ(src_val, 0) << "The quantization scale for this array is 0, " - << "so all its values should be 0."; - scaled_val = quantization_params.zero_point; - } else { - scaled_val = quantization_params.zero_point + inverse_scale * src_val; - } - quantized_buffer->data[i] = - tflite::SafeCast>(std::round(scaled_val)); - } - return std::unique_ptr(quantized_buffer); -} - -template -void QuantizeArray(GraphTransformation* transformation, Model* model, - const string& name, - const QuantizationParams& quantization_params) { - auto& array = model->GetArray(name); - CHECK(array.data_type == ArrayDataType::kFloat); - CHECK(!array.quantization_params); - array.GetOrCreateQuantizationParams() = quantization_params; - if (array.buffer) { - array.buffer = QuantizeBuffer(*array.buffer, quantization_params); - } - array.data_type = A; - transformation->AddMessageF("Quantized array %s", name); -} - -void QuantizeArray(GraphTransformation* transformation, Model* model, - const string& name, ArrayDataType quantized_data_type, - const QuantizationParams& quantization_params) { - 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); - default: - LOG(FATAL) << "Unhandled case."; - } -} - const MinMax& GetOrComputeMinMax(Model* model, const string& array_name) { auto& array = model->GetArray(array_name); // Normally we should have a MinMax recorded on this Array, @@ -204,70 +140,6 @@ QuantizationPoints GetQuantizationPoints(ArrayDataType data_type) { } } -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, @@ -307,6 +179,8 @@ bool ChooseQuantizationForOperatorInput( const auto& input_weights = model->GetArray(op.inputs[weights_input_index]); if (!input_activations.quantization_params || !input_weights.quantization_params) { + transformation->AddMessageF( + "Input array %s is a bias vector but has no qparams", input); return false; } const auto input_activations_scale = @@ -335,12 +209,11 @@ bool ChooseQuantizationForOperatorInput( *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 %s with zero_point=%d" - ", scale=%g", + "For input array %s with min=%g, max=%g, chose to quantize as %s (f=%s) " + "with zero_point=%d, scale=%g", input, minmax.min, minmax.max, ArrayDataTypeName(*quantized_data_type), - quantization_params->zero_point, quantization_params->scale); + ArrayDataTypeName(array.final_data_type), quantization_params->zero_point, + quantization_params->scale); return true; } @@ -394,6 +267,19 @@ bool ChooseHardcodedQuantizationForOperatorOutput( *quantization_params)); return true; } + if (op.type == OperatorType::kLogSoftmax) { + // LogSoftmax has range: [LogSoftmaxOperator::kOutputRangeMin, 0]. + *quantized_data_type = GetQuantizedDataType(array, *quantized_data_type); + const QuantizationPoints qp = GetQuantizationPoints(*quantized_data_type); + quantization_params->zero_point = qp.max_value; + quantization_params->scale = + -LogSoftmaxOperator::kOutputRangeMin / (qp.max_value + 1); + // While not strictly necessary, it is easier to interpret output data and + // quantization if the scale is similar to others (such as power of 2). + CHECK(IsExactlyRepresentable(LogSoftmaxOperator::kOutputRangeMin / 2, + *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); @@ -416,6 +302,9 @@ bool ChooseQuantizationForOperatorOutput( const auto& output = op.outputs[output_index]; auto& array = model->GetArray(output); if (array.data_type != ArrayDataType::kFloat) { + transformation->AddMessageF("Array data type already set to %s, final=%s", + ArrayDataTypeName(array.data_type), + ArrayDataTypeName(array.final_data_type)); return false; } *quantized_data_type = model->GetArray(op.inputs[0]).data_type; @@ -477,29 +366,22 @@ bool ChooseQuantizationForOperatorOutput( // Fixes array minmax info to match the quantization parameters. // This is required for when quantization parameters change for an array during // quantization (such as ChooseQuantizationForOperatorOutput). -void FixMinMaxPostQuantization(ArrayDataType quantized_data_type, +void FixMinMaxPostQuantization(GraphTransformation* transformation, + ArrayDataType quantized_data_type, const QuantizationParams& quantization_params, MinMax* minmax) { - double qmin, qmax; - switch (quantized_data_type) { - case ArrayDataType::kUint8: - qmin = 0; - qmax = 255; - break; - case ArrayDataType::kInt16: - qmin = -32768; - qmax = 32767; - break; - default: - // No update required. - return; + double quantized_min, quantized_max; + if (!GetQuantizedDataTypeNumericalRange(quantized_data_type, &quantized_min, + &quantized_max)) { + // Not quantized - no update required. + return; } // Compute new minmax values. - double min = - (qmin - quantization_params.zero_point) * quantization_params.scale; - double max = - (qmax - quantization_params.zero_point) * quantization_params.scale; + double min = (quantized_min - quantization_params.zero_point) * + quantization_params.scale; + double max = (quantized_max - quantization_params.zero_point) * + quantization_params.scale; // If we are close to the existing minmax values don't bother changing them. // This prevents propagating small floating point precision errors. @@ -507,10 +389,14 @@ void FixMinMaxPostQuantization(ArrayDataType quantized_data_type, const double width = max - min; if (std::abs(min - minmax->min) > kMinMaxThreshold * width || std::abs(max - minmax->max) > kMinMaxThreshold * width) { + transformation->AddMessageF( + "Adjusting min/max from %g,%g to %g,%g to match quantization params", + minmax->min, minmax->max, min, max); minmax->min = min; minmax->max = max; } } + } // namespace bool Quantize::Run(Model* model, std::size_t op_index) { @@ -615,10 +501,33 @@ bool Quantize::Run(Model* model, std::size_t op_index) { // 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]); + // TODO(b/78013785): never rename output arrays. + if (IsInputArray(*model, dequantize_op->inputs[0])) { + // The op input is an input array and the output is an output + // array and we can't have an array be both. Insert a copy + // op to ensure the two arrays stay separate. + AddMessageF( + "Tried to rename output array %d while removing dequant " + "op %s but array is also an input; inserting copy %s " + "-> %s", + i, LogName(*dequantize_op), model->flags.output_arrays(i), + dequantize_op->inputs[0]); + InsertCopyOperator(model, dequantize_op->inputs[0], + dequantize_op->outputs[0]); + } else { + // Op output is strictly used as an output array, so we can + // just rename the array and directly bypass the op. + AddMessageF( + "Renaming output array %d after removing dequant op %s: " + "%s -> %s", + i, LogName(*dequantize_op), model->flags.output_arrays(i), + dequantize_op->inputs[0]); + model->flags.set_output_arrays(i, dequantize_op->inputs[0]); + model->EraseArray(dequantize_op->outputs[0]); + } + break; } } - model->EraseArray(dequantize_op->outputs[0]); model->operators.erase(dequantize_it); } changed = true; @@ -661,8 +570,10 @@ bool Quantize::Run(Model* model, std::size_t op_index) { // Fix up the min/max information on the output array to match the chosen // quantization parameters. + CHECK(output_array.minmax) + << "Output array named " << output << " lacks minmax"; auto& output_minmax = output_array.GetMinMax(); - FixMinMaxPostQuantization(quantized_data_type, quantization_params, + FixMinMaxPostQuantization(this, quantized_data_type, quantization_params, &output_minmax); QuantizeArray(this, model, output, quantized_data_type, @@ -673,6 +584,7 @@ bool Quantize::Run(Model* model, std::size_t op_index) { auto& dequantized_output_array = model->GetOrCreateArray(dequantized_output); dequantized_output_array.data_type = ArrayDataType::kFloat; + dequantized_output_array.final_data_type = output_array.data_type; auto& dequantized_output_minmax = dequantized_output_array.GetOrCreateMinMax(); dequantized_output_minmax.min = output_minmax.min; @@ -689,6 +601,12 @@ bool Quantize::Run(Model* model, std::size_t op_index) { dequantize_op->outputs = {dequantized_output}; for (int i = 0; i < model->flags.output_arrays_size(); i++) { if (model->flags.output_arrays(i) == output) { + // TODO(b/78013785): never rename output arrays. + AddMessageF( + "Renaming output array %d after inserting dequant op %s: %s -> " + "%s", + i, LogName(*dequantize_op), model->flags.output_arrays(i), + dequantized_output); model->flags.set_output_arrays(i, dequantized_output); } } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_fake_quant.cc b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_fake_quant.cc new file mode 100644 index 0000000000000000000000000000000000000000..2c8d04440f251f792d2a09155dd26fc01a732109 --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_fake_quant.cc @@ -0,0 +1,86 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT 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 IsFakeQuantTrivial(GraphTransformation* transformation, const Model& model, + const FakeQuantOperator& fakequant_op) { + CHECK(fakequant_op.type == OperatorType::kFakeQuant); + + if (!fakequant_op.minmax) { + // Require ReadFakeQuantMinMax to have run. + return false; + } + + // FakeQuants are trivial if they are taking input from another identical + // FakeQuant op. + auto* producing_op = GetOpWithOutput(model, fakequant_op.inputs[0]); + if (!producing_op || producing_op->type != OperatorType::kFakeQuant) { + return false; + } + const auto& producing_fakequant_op = + *static_cast(producing_op); + if (!producing_fakequant_op.minmax) { + // Require ReadFakeQuantMinMax to have run. + return false; + } + + if (*fakequant_op.minmax == *producing_fakequant_op.minmax && + fakequant_op.num_bits == producing_fakequant_op.num_bits) { + transformation->AddMessageF( + "%s is trivial because it is preceded by an identical FakeQuant %s", + LogName(fakequant_op), LogName(producing_fakequant_op)); + return true; + } + + return false; +} + +} // namespace + +// Removes FakeQuant ops that are trivial (have no effect, are redundant, etc). +bool RemoveTrivialFakeQuant::Run(Model* model, std::size_t op_index) { + const auto op_it = model->operators.begin() + op_index; + auto* op = op_it->get(); + if (op->type != OperatorType::kFakeQuant) { + return false; + } + auto* fakequant_op = static_cast(op); + + if (!IsFakeQuantTrivial(this, *model, *fakequant_op)) { + AddMessageF("%s is not trivial", LogName(*fakequant_op)); + return false; + } + + AddMessageF("Removing trivial %s", LogName(*fakequant_op)); + + CHECK_EQ(fakequant_op->inputs.size(), 1); + return RemoveTrivialPassthroughOp(this, model, op_index); +} + +} // namespace toco 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 aa93ace03af300f9cbd3f9c6620a6a58b9329aa4..3e021b819fc82d66fb70596a62fd7cee4911d4e8 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_passthrough.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_passthrough.cc @@ -82,22 +82,13 @@ bool RemoveTrivialPassthroughOp(GraphTransformation* transformation, if (IsDiscardableArray(*model, output_name)) { transformation->AddMessageF( - "Removing %s, keeping its non-constant input array", - LogName(*passthru_op)); - for (const string& input : passthru_op->inputs) { - if (IsDiscardableArray(*model, input) && input != main_input_name && - CountOpsWithInput(*model, input) == 1) { - } - } + "Removing %s, keeping its non-constant input array %s and removing %s", + LogName(*passthru_op), main_input_name, output_name); RerouteEdges(output_name, main_input_name, model); } else if (IsDiscardableArray(*model, main_input_name)) { - transformation->AddMessageF("Removing %s, keeping its output array", - LogName(*passthru_op)); - for (const string& input : passthru_op->inputs) { - if (IsDiscardableArray(*model, input) && - (input == main_input_name || CountOpsWithInput(*model, input) == 1)) { - } - } + transformation->AddMessageF( + "Removing %s, keeping its output array %s and removing input %s", + LogName(*passthru_op), output_name, main_input_name); RerouteEdges(main_input_name, output_name, model); } else { transformation->AddMessageF( @@ -113,6 +104,16 @@ bool RemoveTrivialPassthroughOp(GraphTransformation* transformation, // Remove any array that is no longer used. for (const string& removal_candidate : removal_candidates) { bool is_referenced = false; + for (const auto& array : model->flags.input_arrays()) { + if (array.name() == removal_candidate) { + is_referenced = true; + } + } + for (const auto& array_name : model->flags.output_arrays()) { + if (array_name == removal_candidate) { + is_referenced = true; + } + } for (const auto& op : model->operators) { for (const string& input : op->inputs) { if (input == removal_candidate) { 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 9b65feaa6443cd32ac1bef961600ff225d52d4b2..752560e075a087bcc2b0a3cb19dad484fb582d42 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 @@ -18,6 +18,8 @@ limitations under the License. #include #include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/graph_transformations/quantization_util.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/toco_types.h" @@ -26,27 +28,44 @@ limitations under the License. namespace toco { -bool RemoveTrivialQuantizedActivationFunc::Run(Model* model, - std::size_t op_index) { - 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; - } - const auto& output_array = model->GetArray(op->outputs[0]); - if (!output_array.quantization_params) { - return false; - } - if (output_array.data_type != ArrayDataType::kUint8) { - return false; +namespace { + +bool IsTrivialUnfusedActivationFunc(GraphTransformation* transformation, + const Model& model, OperatorType op_type, + const string& input_array_name) { + double clamp_min; + double clamp_max; + switch (op_type) { + case OperatorType::kRelu: + clamp_min = 0.0; + clamp_max = std::numeric_limits::infinity(); + break; + case OperatorType::kRelu1: + clamp_min = -1.0; + clamp_max = 1.0; + break; + case OperatorType::kRelu6: + clamp_min = 0.0; + clamp_max = 6.0; + break; + default: + return false; } - const auto& quantization_params = output_array.GetQuantizationParams(); + const auto& input_array = model.GetArray(input_array_name); + return IsArrayQuantizedRangeSubset(transformation, input_array, clamp_min, + clamp_max); +} + +bool IsTrivialFusedActivationFunc( + GraphTransformation* transformation, const Model& model, + FusedActivationFunctionType activation_function, + const string& output_array_name) { double clamp_min; double clamp_max; - switch (op->fused_activation_function) { + switch (activation_function) { + case FusedActivationFunctionType::kNone: + return false; case FusedActivationFunctionType::kRelu: clamp_min = 0.0; clamp_max = std::numeric_limits::infinity(); @@ -61,45 +80,46 @@ bool RemoveTrivialQuantizedActivationFunc::Run(Model* model, break; default: LOG(FATAL) << "Unsupported fused activation type: " - << static_cast(op->fused_activation_function); + << static_cast(activation_function); return false; } - bool has_nontrivial_min_bound = false; - bool has_nontrivial_max_bound = false; + const auto& output_array = model.GetArray(output_array_name); + return IsArrayQuantizedRangeSubset(transformation, output_array, clamp_min, + clamp_max); +} - 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); - } - 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); - } +} // namespace - if (has_nontrivial_min_bound || has_nontrivial_max_bound) { +// Attempts to remove both fused and unfused activation functions if the +// quantization params indicate that the representable values fall inside the +// activation range. +bool RemoveTrivialQuantizedActivationFunc::Run(Model* model, + std::size_t op_index) { + const auto it = model->operators.begin() + op_index; + auto* op = it->get(); + if (op->inputs.empty()) { return false; } - op->fused_activation_function = FusedActivationFunctionType::kNone; - AddMessageF( - "Removing trivial quantized activation function on %s" - " because the output quantization parameters imply at least as tight" - " a clamp anyway.", - LogName(*op)); - return true; + if (IsTrivialUnfusedActivationFunc(this, *model, op->type, op->inputs[0])) { + AddMessageF( + "Removing trivial unfused activation function %s because the input " + "minmax imply at least as tight a clamp anyway.", + LogName(*op)); + return RemoveTrivialPassthroughOp(this, model, op_index); + } + if (IsTrivialFusedActivationFunc(this, *model, op->fused_activation_function, + op->outputs[0])) { + op->fused_activation_function = FusedActivationFunctionType::kNone; + AddMessageF( + "Removing trivial quantized activation function on %s " + "because the output quantization parameters imply at least as tight " + "a clamp anyway.", + LogName(*op)); + return true; + } + return false; } } // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_quantized_min_max.cc b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_quantized_min_max.cc new file mode 100644 index 0000000000000000000000000000000000000000..eaee1c662b7cedb2baec7be47e12e348c3e7b25c --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_quantized_min_max.cc @@ -0,0 +1,90 @@ +/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT 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/quantization_util.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/toco_types.h" +#include "tensorflow/contrib/lite/toco/tooling_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace toco { + +namespace { + +bool IsTrivialMinMax(GraphTransformation* transformation, const Model& model, + OperatorType op_type, const string& input_array_name, + const string& clamp_value_array_name) { + const auto& clamp_value_array = model.GetArray(clamp_value_array_name); + if (!IsConstantParameterArray(model, clamp_value_array_name)) { + transformation->AddMessageF("Clip value array %s is non-constant", + clamp_value_array_name); + return false; + } + const auto& clamp_value_buffer = + clamp_value_array.GetBuffer(); + CHECK_EQ(clamp_value_buffer.Length(), 1); + float clamp_value = clamp_value_buffer.data[0]; + + double clamp_min; + double clamp_max; + switch (op_type) { + case OperatorType::kTensorFlowMinimum: + clamp_min = -std::numeric_limits::infinity(); + clamp_max = clamp_value; + break; + case OperatorType::kTensorFlowMaximum: + clamp_min = clamp_value; + clamp_max = std::numeric_limits::infinity(); + break; + default: + CHECK(false); + return false; + } + + const auto& input_array = model.GetArray(input_array_name); + return IsArrayQuantizedRangeSubset(transformation, input_array, clamp_min, + clamp_max); +} + +} // namespace + +// Attempts to remove min/max functions if the quantization params indicate that +// the representable values fall inside the clip range. +bool RemoveTrivialQuantizedMinMax::Run(Model* model, std::size_t op_index) { + const auto it = model->operators.begin() + op_index; + auto* op = it->get(); + if ((op->type != OperatorType::kTensorFlowMinimum && + op->type != OperatorType::kTensorFlowMaximum) || + op->inputs.size() != 2) { + return false; + } + if (IsTrivialMinMax(this, *model, op->type, op->inputs[0], op->inputs[1])) { + AddMessageF( + "Removing trivial min/max %s because the quantization parameters imply " + "at least as tight a clamp anyway.", + LogName(*op)); + return RemoveTrivialPassthroughOp(this, model, op_index); + } + return false; +} + +} // namespace toco 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 61477d59aea2f11c6347b84d8863763a86c43558..e28d8cf01eafee64e08ac2cc4b43ea7c227456c2 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_reshape.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_reshape.cc @@ -41,8 +41,8 @@ bool IsReshapeTrivial(const Model& model, const Operator& op, ShapesAgreeUpToExtending(input_array.shape(), output_array.shape())) { transformation->AddMessageF( "%s is trivial because its input and output shapes are equal up to " - "extending " - "by 1's, and we are told to aggressively discard such Reshape ops.", + "extending by 1's, and we are told to aggressively discard such " + "Reshape ops.", LogName(op)); return true; } @@ -80,6 +80,7 @@ bool RemoveTrivialReshape::Run(Model* model, std::size_t op_index) { } if (!IsReshapeTrivial(*model, *reshape_op, this)) { + AddMessageF("%s is not trivial", LogName(*reshape_op)); return false; } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/remove_unused_op.cc b/tensorflow/contrib/lite/toco/graph_transformations/remove_unused_op.cc index aa2c293382a98b476bee783ed8e177b19d35b858..8e6aaf544aa5310b4233d93e7bc8f484f6164b8a 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/remove_unused_op.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_unused_op.cc @@ -47,7 +47,8 @@ bool RemoveUnusedOp::Run(Model* model, std::size_t op_index) { bool found_output_as_rnn_state_array = false; for (const auto& rnn_state : model->flags.rnn_states()) { if (output == rnn_state.state_array()) { - CHECK(op->type == OperatorType::kFill); + CHECK(op->type == OperatorType::kFill || + op->type == OperatorType::kTensorFlowIdentity); found_output_as_rnn_state_array = true; break; } diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_batch_normalization.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_batch_normalization.cc index fb109eb91b16e3a73005230f821c18b9ef82d2fb..2b3ee36ad10e24ab7367ca44c03a234688a63a9b 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_batch_normalization.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_batch_normalization.cc @@ -33,7 +33,7 @@ bool ResolveBatchNormalization::Run(Model* model, std::size_t op_index) { const auto* bn_op = static_cast(bn_it->get()); - const auto& mean_array = model->GetArray(bn_op->inputs[1]); + auto& mean_array = model->GetArray(bn_op->inputs[1]); const auto& multiplier_array = model->GetArray(bn_op->inputs[2]); const auto& offset_array = model->GetArray(bn_op->inputs[3]); @@ -49,6 +49,13 @@ bool ResolveBatchNormalization::Run(Model* model, std::size_t op_index) { CHECK(multiplier_array.data_type == ArrayDataType::kFloat); CHECK(offset_array.data_type == ArrayDataType::kFloat); + // This graph transformations will need to address constant buffers below, + // so we need to exit early if these buffers don't exist (i.e. if the params + // haven't yet been resolved as constants). + if (!mean_array.buffer || !multiplier_array.buffer || !offset_array.buffer) { + return false; + } + // Create the new Mul, Add operators auto* mul_op = new MulOperator; auto* add_op = new AddOperator; @@ -80,9 +87,15 @@ bool ResolveBatchNormalization::Run(Model* model, std::size_t op_index) { DCHECK_EQ(bn_it->get(), bn_op); // Create the new param arrays - const auto& mean_shape = mean_array.shape(); + auto& mean_shape = *mean_array.mutable_shape(); const auto& multiplier_shape = multiplier_array.shape(); const auto& offset_shape = offset_array.shape(); + if (mean_shape.dims().empty()) { + *mean_shape.mutable_dims() = multiplier_shape.dims(); + auto& data = mean_array.GetMutableBuffer().data; + CHECK_EQ(data.size(), 1); + data.resize(RequiredBufferSizeForShape(mean_shape), data[0]); + } CHECK(mean_shape.dims() == multiplier_shape.dims()); CHECK(mean_shape.dims() == offset_shape.dims()); const auto& param_shape = mean_shape; 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 064810b53e7c3bee4601204c9dbd976c374a6a60..d916ae0ddf017fe6a2fb2709db6e9de8c258adfc 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_concatenation.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_concatenation.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include "absl/strings/str_cat.h" +#include "absl/strings/str_join.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" @@ -105,7 +106,8 @@ void ConcatenateTensorBuffers(const std::vector& input_arrays, // already set (e.g. because of previous pass in TOCO), it doesn't change it and // returns. Otherwise it uses the input arrays min and max values to compute the // concatenated array min and max. -void SetMinMaxForConcatenedArray(const std::vector& input_arrays, +void SetMinMaxForConcatenedArray(GraphTransformation* transformation, + const std::vector& input_arrays, Array* concatenated_array) { CHECK(concatenated_array->data_type == ArrayDataType::kFloat); // If the minmax is already set, use it @@ -125,6 +127,9 @@ void SetMinMaxForConcatenedArray(const std::vector& input_arrays, MinMax& minmax = concatenated_array->GetOrCreateMinMax(); minmax.min = concat_min; minmax.max = concat_max; + + transformation->AddMessageF("Setting concatenated array min/max to %g,%g", + concat_min, concat_max); } } // namespace @@ -161,11 +166,14 @@ bool ResolveConstantConcatenation::Run(Model* model, std::size_t op_index) { input_arrays.push_back(&model->GetArray(input_name)); } + AddMessageF("Performing constant concat of %s into %s", + absl::StrJoin(concat_op->inputs, ", "), concatenated_array_name); + switch (concatenated_array.data_type) { case ArrayDataType::kFloat: ConcatenateTensorBuffers( input_arrays, concatenation_axis, &concatenated_array); - SetMinMaxForConcatenedArray(input_arrays, &concatenated_array); + SetMinMaxForConcatenedArray(this, input_arrays, &concatenated_array); break; case ArrayDataType::kUint8: ConcatenateTensorBuffers( @@ -189,13 +197,13 @@ 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 + // Check to prevent removal of shared tensors. if (CountOpsWithInput(*model, input_name) == 1) { model->EraseArray(input_name); } } - // Remove concatenate operator + // Remove concatenate operator. model->operators.erase(concat_it); return true; } 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 625d90205a801ad7c3fc1026c9cedc9b509f920d..efb7bb218421dd045e3e8e2a38b9c70989f222e1 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 @@ -18,6 +18,7 @@ limitations under the License. #include #include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" +#include "tensorflow/contrib/lite/toco/graph_transformations/quantization_util.h" #include "tensorflow/contrib/lite/toco/model.h" #include "tensorflow/contrib/lite/toco/tooling_util.h" #include "tensorflow/core/platform/logging.h" @@ -45,9 +46,29 @@ bool ResolveConstantFakeQuant::Run(Model* model, std::size_t op_index) { } const auto& input_array = model->GetArray(fakequant_op->inputs[0]); + CHECK(input_array.data_type == ArrayDataType::kFloat); + + // Determine the final data type in the same way as PropagateFakeQuantNumBits. + ArrayDataType quantized_data_type = input_array.final_data_type; + if (!InferQuantizedDataTypeFromFakeQuant(*fakequant_op, + &quantized_data_type)) { + AddMessageF("Unsupported FakeQuant num_bits=%d", fakequant_op->num_bits); + return false; + } + + AddMessageF("Resolving constant %s", LogName(*fakequant_op)); + auto& output_array = model->GetArray(fakequant_op->outputs[0]); CHECK(input_array.data_type == ArrayDataType::kFloat); output_array.data_type = ArrayDataType::kFloat; + + // We'll set the final data type to what the fake quant indicates we should + // have (and would have been set if this stayed around until + // PropagateFakeQuantNumBits). + if (propagate_fake_quant_num_bits()) { + output_array.final_data_type = quantized_data_type; + } + CHECK(!output_array.buffer); const auto& input_buffer = input_array.GetBuffer(); output_array.GetOrCreateMinMax() = *fakequant_op->minmax; @@ -66,7 +87,9 @@ bool ResolveConstantFakeQuant::Run(Model* model, std::size_t op_index) { const double dst_val = qparams.scale * (quantized_val - qparams.zero_point); output_buffer.data[i] = dst_val; } - if (CountOpsWithInput(*model, fakequant_op->inputs[0]) == 1) { + + if (IsDiscardableArray(*model, fakequant_op->inputs[0]) && + CountOpsWithInput(*model, fakequant_op->inputs[0]) == 1) { model->EraseArray(fakequant_op->inputs[0]); } model->operators.erase(fakequant_it); diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_gather.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_gather.cc index d999c2df9483e096f333c6af83e1d9fee873d4d6..debe298a5a93034bcb928d7384b5ec1fc7439e47 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_gather.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_gather.cc @@ -98,6 +98,16 @@ bool ResolveConstantGather::Run(Model* model, std::size_t op_index) { CHECK(coords_array.data_type == ArrayDataType::kInt32) << "Only int32 indices are supported"; + // Copy min/max info if present. The ranges of the selected values may be + // a subset of the original range but we want to ensure the quantization + // params stay the same. + if (input_array.minmax) { + const auto& input_minmax = input_array.GetMinMax(); + auto& output_minmax = output_array.GetOrCreateMinMax(); + output_minmax.min = input_minmax.min; + output_minmax.max = input_minmax.max; + } + CHECK(!output_array.buffer); switch (output_array.data_type) { case ArrayDataType::kFloat: diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_reshape.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_reshape.cc new file mode 100644 index 0000000000000000000000000000000000000000..7e7ad383e7789891f5396845241e70143dc8b76f --- /dev/null +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_reshape.cc @@ -0,0 +1,124 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +#include + +#include "tensorflow/contrib/lite/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 { + +// Resolves a constant reshape operation by copying the buffer. +bool ResolveConstantReshape::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::kTensorFlowReshape) { + return false; + } + const auto* op = static_cast(base_op); + + CHECK_EQ(op->inputs.size(), 2); + CHECK_EQ(op->outputs.size(), 1); + + // We require constant inputs. + if (!IsConstantParameterArray(*model, op->inputs[0]) || + !IsConstantParameterArray(*model, op->inputs[1])) { + return false; + } + + 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; + } + + const Array& input_array = model->GetArray(op->inputs[0]); + if (!ShapesAgreeUpToExtending(input_array.shape(), output_array.shape())) { + AddMessageF("Constant reshape is non-trivial (%s -> %s)", + ShapeToString(input_array.shape()), + ShapeToString(output_array.shape())); + return false; + } + + CHECK(!output_array.buffer); + switch (input_array.data_type) { + case ArrayDataType::kBool: + CopyArrayBuffer(input_array, &output_array); + break; + case ArrayDataType::kFloat: + CopyArrayBuffer(input_array, &output_array); + break; + case ArrayDataType::kInt8: + CopyArrayBuffer(input_array, &output_array); + break; + case ArrayDataType::kUint8: + CopyArrayBuffer(input_array, &output_array); + break; + case ArrayDataType::kInt16: + CopyArrayBuffer(input_array, &output_array); + break; + case ArrayDataType::kUint16: + CopyArrayBuffer(input_array, &output_array); + break; + case ArrayDataType::kInt32: + CopyArrayBuffer(input_array, &output_array); + break; + case ArrayDataType::kUint32: + CopyArrayBuffer(input_array, &output_array); + break; + case ArrayDataType::kInt64: + CopyArrayBuffer(input_array, &output_array); + break; + case ArrayDataType::kUint64: + CopyArrayBuffer(input_array, &output_array); + break; + case ArrayDataType::kString: + CopyArrayBuffer(input_array, &output_array); + break; + default: + LOG(FATAL) << "Unsupported data type: " + << ArrayDataTypeName(input_array.data_type); + return false; + } + + AddMessageF("Resolving constant reshape of %s", LogName(*op)); + + if (input_array.minmax) { + output_array.GetOrCreateMinMax() = input_array.GetMinMax(); + } + if (input_array.quantization_params) { + output_array.GetOrCreateQuantizationParams() = + input_array.GetQuantizationParams(); + } + + // Erase input arrays if no longer used. + for (const auto& input : op->inputs) { + if (IsDiscardableArray(*model, input) && + CountOpsWithInput(*model, input) == 1) { + model->EraseArray(input); + } + } + + // Erase the operator. + model->operators.erase(it); + return true; +} + +} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_strided_slice.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_strided_slice.cc index a0cfc3d59763dc1211ed4d1ac114d371a4a7ee0b..8df3c2f1955c23b301ca10f613b57a1fbe138593 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_strided_slice.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_strided_slice.cc @@ -23,40 +23,88 @@ namespace toco { namespace { +// These StridedSlice utility functions are essentially a COPY of those in +// reference_ops.h. See comments there. + +// Use until std::clamp() is available from C++17. +int Clamp(const int v, const int lo, const int hi) { + if (hi < v) return hi; + if (v < lo) return lo; + return v; +} + int StartForAxis(StridedSliceOperator const& op, Shape const& input_shape, int axis) { - int start; + // Begin with the specified index + int start = op.start_indices[axis]; + + // begin_mask override if (op.begin_mask & 1 << axis) { - // If begin mask bit is set, use the first element - start = 0; - } else { - // Otherwise, use the specified element - start = op.start_indices[axis]; - if (start < 0) { - // Handle negative indices - start += input_shape.dims(axis); + if (op.strides[axis] > 0) { + // Forward iteration - use the first element. These values will get + // clamped below (Note: We could have set them to 0 and axis_size-1, but + // use lowest() and max() to maintain symmetry with StopForAxis()) + start = std::numeric_limits::lowest(); + } else { + // Backward iteration - use the last element. + start = std::numeric_limits::max(); } } + + // Handle negative indices + int axis_size = input_shape.dims(axis); + if (start < 0) { + start += axis_size; + } + + // Clamping + start = Clamp(start, 0, axis_size - 1); + return start; } int StopForAxis(StridedSliceOperator const& op, Shape const& input_shape, int axis) { - int stop; + // Begin with the specified index + int stop = op.stop_indices[axis]; + + // end_mask override if (op.end_mask & (1 << axis)) { - // If end mask bit set, use the last element - stop = input_shape.dims(axis); - } else { - // Otherwise, use the specified element - stop = op.stop_indices[axis]; - if (stop < 0) { - // Handle negative indices - stop += input_shape.dims(axis); + if (op.strides[axis] > 0) { + // Forward iteration - use the last element. These values will get + // clamped below + stop = std::numeric_limits::max(); + } else { + // Backward iteration - use the first element. + stop = std::numeric_limits::lowest(); } } + + // Handle negative indices + int axis_size = input_shape.dims(axis); + if (stop < 0) { + stop += axis_size; + } + + // Clamping + // Because the end index points one past the last element, we need slightly + // different clamping ranges depending on the direction. + if (op.strides[axis] > 0) { + // Forward iteration + stop = Clamp(stop, 0, axis_size); + } else { + // Backward iteration + stop = Clamp(stop, -1, axis_size - 1); + } + return stop; } +bool LoopCondition(int index, int stop, int stride) { + // True when we have reached the end of an axis and should loop. + return stride > 0 ? index >= stop : index <= stop; +} + template void StridedSlice(StridedSliceOperator const& op, Array const& input_array, Array* output_array) { @@ -73,9 +121,6 @@ void StridedSlice(StridedSliceOperator const& op, Array const& input_array, int num_input_axes = op.start_indices.size(); CHECK_EQ(num_input_axes, op.stop_indices.size()); CHECK_EQ(num_input_axes, op.strides.size()); - for (int i = 0; i < op.strides.size(); i++) { - CHECK_GE(op.strides[i], 0) << "Negative strides usupported"; - } // Create a buffer for the output array std::vector>& output_data = @@ -103,13 +148,15 @@ void StridedSlice(StridedSliceOperator const& op, Array const& input_array, // Compute next source input coordinates. bool carry = true; for (int axis = 0; axis < num_input_axes; axis++) { + int stride = op.strides[axis]; // Increment this axis if we carried from the previous one if (carry) { - src_coord[axis] += op.strides[axis]; + src_coord[axis] += stride; } // Check if we've overflowed. - if (src_coord[axis] >= StopForAxis(op, input_shape, axis)) { + int stop = StopForAxis(op, input_shape, axis); + if (LoopCondition(src_coord[axis], stop, stride)) { // Reset axis and set carry src_coord[axis] = StartForAxis(op, input_shape, axis); carry = true; diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_transpose.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_transpose.cc index 4f984bfde55b3457694bb411bbfdf30723c7066e..1fd20314b14d98bd82e2b20a4e70f5d9c2c3b298 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_transpose.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_transpose.cc @@ -131,6 +131,10 @@ bool ResolveConstantTranspose::Run(Model* model, std::size_t op_index) { if (input_array.minmax) { output_array.GetOrCreateMinMax() = input_array.GetMinMax(); } + if (input_array.quantization_params) { + output_array.GetOrCreateQuantizationParams() = + input_array.GetQuantizationParams(); + } if (op->perm.empty()) { // Yield until perm has been populated by ResolveTransposeAttributes. @@ -164,6 +168,8 @@ bool ResolveConstantTranspose::Run(Model* model, std::size_t op_index) { break; } + AddMessageF("Resolving constant transpose of %s", LogName(*op)); + // Erase input arrays if no longer used. for (const auto& input : op->inputs) { if (IsDiscardableArray(*model, input) && diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_merge.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_merge.cc index 477e7f13da3d88a68547d494011cd4984936b909..38e0005890ac10410df4ddb5290be8fcc948c349 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_merge.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_tensorflow_merge.cc @@ -32,7 +32,7 @@ bool ResolveTensorFlowMerge::Run(Model* model, std::size_t op_index) { } // We need to yield until this Merge node has only 1 input, which will mean - // that that is the selected input. Other graph transformations on other nodes + // that is the selected input. Other graph transformations on other nodes // such as ResolveTensorFlowSwitch, will take care of trimming the // non-selected inputs, so that at some point there will be only 1 input left. if (merge_op->inputs.size() > 1) { diff --git a/tensorflow/contrib/lite/toco/import_tensorflow.cc b/tensorflow/contrib/lite/toco/import_tensorflow.cc index 6b62eeb6389e3343eb418b602e2dcdc171d2eb1b..155d890c9f23ba206f1f0e6db645a601308cea5b 100644 --- a/tensorflow/contrib/lite/toco/import_tensorflow.cc +++ b/tensorflow/contrib/lite/toco/import_tensorflow.cc @@ -694,6 +694,8 @@ void ConvertFakeQuantWithMinMaxArgs( minmax.min = GetFloatAttr(node, "min"); minmax.max = GetFloatAttr(node, "max"); op->outputs.push_back(node.name()); + // tf.fake_quant_with_min_max_args num_bits defaults to 8. + op->num_bits = HasAttr(node, "num_bits") ? GetIntAttr(node, "num_bits") : 8; model->operators.emplace_back(op); } @@ -711,6 +713,7 @@ void ConvertFakeQuantWithMinMaxVars( op->inputs.push_back(node.input(i)); } op->outputs.push_back(node.name()); + op->num_bits = HasAttr(node, "num_bits") ? GetIntAttr(node, "num_bits") : 8; model->operators.emplace_back(op); } diff --git a/tensorflow/contrib/lite/toco/model.h b/tensorflow/contrib/lite/toco/model.h index 56ef9fe2a88617afe53e480f68da5aa754cd12f4..482cc71d8b34705041130c7518b58f7fe183bc3c 100644 --- a/tensorflow/contrib/lite/toco/model.h +++ b/tensorflow/contrib/lite/toco/model.h @@ -24,6 +24,7 @@ limitations under the License. #include "tensorflow/contrib/lite/toco/model_flags.pb.h" #include "tensorflow/contrib/lite/toco/runtime/types.h" +#include "tensorflow/contrib/lite/toco/toco_port.h" #include "tensorflow/contrib/lite/toco/toco_types.h" #include "tensorflow/core/platform/logging.h" @@ -151,9 +152,9 @@ enum class AxesOrder { }; // The type of the scalars in an array. -// Note that that does not by itself tell whether the values in the array are -// real (are literally interpreted as real numbers) or quantized (only acquire -// a meaning as real numbers in conjunction with QuantizationParams). +// Note that the type does not by itself tell whether the values in the array +// are real (are literally interpreted as real numbers) or quantized (only +// acquire a meaning as real numbers in conjunction with QuantizationParams). // // In practice though: // float values are always real @@ -425,6 +426,7 @@ struct SpaceToDepthOperator : Operator { // input activations as a matrix, followed by a MatMul node. struct FullyConnectedOperator : Operator { FullyConnectedOperator() : Operator(OperatorType::kFullyConnected) {} + bool experimental_shuffled_weights = false; }; // Dequantization operator, converting a quantized array of integers with @@ -724,8 +726,7 @@ struct L2PoolOperator : Operator { // The expected [min, max] range of values in a given array. // Used for quantization only. // This information typically comes from special nodes found in quantized -// models, -// see FakeQuantOperator, and is used during quantization to resolve +// models, see FakeQuantOperator, and is used during quantization to resolve // actual quantization parameters (see QuantizationParams). struct MinMax { double min = 0.; @@ -753,6 +754,7 @@ inline bool operator==(const MinMax& m1, const MinMax& m2) { struct FakeQuantOperator : Operator { FakeQuantOperator() : Operator(OperatorType::kFakeQuant) {} std::unique_ptr minmax; + int num_bits = 8; }; // Element-wise division operator. @@ -844,6 +846,60 @@ struct StridedSliceOperator : Operator { int end_mask; int new_axis_mask; int shrink_axis_mask; + + StridedSliceOperator(const StridedSliceOperator& other) + : Operator(OperatorType::kStridedSlice) { + inputs = other.inputs; + outputs = other.outputs; + + start_indices = other.start_indices; + stop_indices = other.stop_indices; + strides = other.strides; + + begin_mask = other.begin_mask; + ellipsis_mask = other.ellipsis_mask; + end_mask = other.end_mask; + new_axis_mask = other.new_axis_mask; + shrink_axis_mask = other.shrink_axis_mask; + } + + void PadIndices(int dim_count) { + // Add indices and mask bits to fully include extra dimensions + CHECK_GE(dim_count, start_indices.size()); + CHECK_EQ(start_indices.size(), stop_indices.size()); + CHECK_EQ(stop_indices.size(), strides.size()); + + for (int i = start_indices.size(); i < dim_count; i++) { + start_indices.push_back(0); + stop_indices.push_back(0); + strides.push_back(1); + begin_mask |= 1 << i; + end_mask |= 1 << i; + } + } + + void ReverseIndices() { + CHECK_EQ(start_indices.size(), stop_indices.size()); + CHECK_EQ(stop_indices.size(), strides.size()); + + std::reverse(start_indices.begin(), start_indices.end()); + std::reverse(stop_indices.begin(), stop_indices.end()); + std::reverse(strides.begin(), strides.end()); + + begin_mask = toco::port::ReverseBits32(static_cast(begin_mask)) >> + (32 - start_indices.size()); + ellipsis_mask = + toco::port::ReverseBits32(static_cast(ellipsis_mask)) >> + (32 - start_indices.size()); + end_mask = toco::port::ReverseBits32(static_cast(end_mask)) >> + (32 - start_indices.size()); + new_axis_mask = + toco::port::ReverseBits32(static_cast(new_axis_mask)) >> + (32 - start_indices.size()); + shrink_axis_mask = + toco::port::ReverseBits32(static_cast(shrink_axis_mask)) >> + (32 - start_indices.size()); + } }; // Reshaping operator, reshaping its input array to a two-dimensional shape @@ -1329,6 +1385,15 @@ struct SoftmaxOperator : Operator { // TensorFlow equivalent: LogSoftmax struct LogSoftmaxOperator : Operator { LogSoftmaxOperator() : Operator(OperatorType::kLogSoftmax) {} + + // LogSoftmax can in principal have very large negative output, depending on + // the input size. However, input x_i that is less than x_max-10 is + // accumulated as exp(x_i-x_max), which is truncated to zero. + // + // Since we effectively disregard smallish inputs in the normalizing factor, + // we also drop them in the output (set to minimum output), and in doing so + // make better use of the quantization range / resolution. + static constexpr float kOutputRangeMin = -16.0; }; // Cast operator. @@ -1411,8 +1476,7 @@ struct SpaceToBatchNDOperator : Operator { }; // BatchToSpaceND operator. Rearranges data from batch into blocks of -// spatial data. Currently, only 2-d blocks are supported. Cropping is not -// supported, either, and the crops array should be all zero. +// spatial data. Currently, only 2-d blocks are supported. // // Inputs: // inputs[0]: required: the input array @@ -1522,7 +1586,7 @@ class Shape { int dims(int i) const { // Always check for out-of-bounds accesses, even in optimized builds where // standard assertions are disabled. Out-of-bounds access here is a common - // occurence. + // occurrence. CHECK_GE(i, 0); CHECK_GT(dims_.size(), i); return dims_[i]; diff --git a/tensorflow/contrib/lite/toco/tflite/operator.cc b/tensorflow/contrib/lite/toco/tflite/operator.cc index f991529569d9ab56103bf7e5f91b2d2b7f2d23fe..d2e14ac5e0d7b06451de295574f42c6139cb97a0 100644 --- a/tensorflow/contrib/lite/toco/tflite/operator.cc +++ b/tensorflow/contrib/lite/toco/tflite/operator.cc @@ -68,7 +68,9 @@ class Convolution auto activation_function = ActivationFunction::Serialize(op.fused_activation_function); return ::tflite::CreateConv2DOptions(*builder, padding, op.stride_width, - op.stride_height, activation_function); + op.stride_height, activation_function, + op.dilation_width_factor, + op.dilation_height_factor); } void ReadOptions(const TfLiteOptions& options, @@ -76,6 +78,8 @@ class Convolution op->padding.type = Padding::Deserialize(options.padding()); op->stride_width = options.stride_w(); op->stride_height = options.stride_h(); + op->dilation_width_factor = options.dilation_w_factor(); + op->dilation_height_factor = options.dilation_h_factor(); op->fused_activation_function = ActivationFunction::Deserialize(options.fused_activation_function()); } @@ -260,12 +264,15 @@ class FakeQuant : public CustomOperator { flexbuffers::Builder* fbb) const override { fbb->Float("min", op.minmax->min); fbb->Float("max", op.minmax->max); + fbb->Int("num_bits", op.num_bits); } void ReadOptions(const flexbuffers::Map& m, TocoOperator* op) const override { auto* minmax = new MinMax; minmax->min = m["min"].AsFloat(); minmax->max = m["max"].AsFloat(); op->minmax.reset(minmax); + const auto& num_bits = m["num_bits"]; + op->num_bits = num_bits.IsInt() ? num_bits.AsInt32() : 8; } }; @@ -662,6 +669,23 @@ class TopK_V2 : public BuiltinOperator { + public: + using BuiltinOperator::BuiltinOperator; + flatbuffers::Offset WriteOptions( + const TocoOperator& op, + flatbuffers::FlatBufferBuilder* builder) const override { + return ::tflite::CreateArgMaxOptions( + *builder, DataType::Serialize(op.output_data_type)); + } + + void ReadOptions(const TfLiteOptions& options, + TocoOperator* op) const override { + op->output_data_type = DataType::Deserialize(options.output_type()); + } +}; + class TensorFlowUnsupported : public BaseOperator { public: using BaseOperator::BaseOperator; @@ -834,6 +858,8 @@ std::vector> BuildOperatorList() { new Lstm(::tflite::BuiltinOperator_LSTM, OperatorType::kLstmCell)); ops.emplace_back( new Cast(::tflite::BuiltinOperator_CAST, OperatorType::kCast)); + ops.emplace_back( + new ArgMax(::tflite::BuiltinOperator_ARG_MAX, OperatorType::kArgMax)); // Custom Operators. ops.emplace_back( @@ -871,6 +897,10 @@ std::vector> BuildOperatorList() { "LOG_SOFTMAX", OperatorType::kLogSoftmax)); ops.emplace_back(new SimpleOperator( "MAXIMUM", OperatorType::kTensorFlowMaximum)); + ops.emplace_back(new SimpleOperator( + "MINIMUM", OperatorType::kTensorFlowMinimum)); + ops.emplace_back(new SimpleOperator( + "LESS", OperatorType::kTensorFlowLess)); return ops; } diff --git a/tensorflow/contrib/lite/toco/tflite/operator_test.cc b/tensorflow/contrib/lite/toco/tflite/operator_test.cc index 4783843b7fa1273201e0c31816b3e1be8e98f5d5..36ed741541eadbc9435a67bec15d389ba48350c1 100644 --- a/tensorflow/contrib/lite/toco/tflite/operator_test.cc +++ b/tensorflow/contrib/lite/toco/tflite/operator_test.cc @@ -111,6 +111,10 @@ TEST_F(OperatorTest, SimpleOperators) { OperatorType::kLogSoftmax); CheckSimpleOperator( "MAXIMUM", OperatorType::kTensorFlowMaximum); + CheckSimpleOperator( + "MINIMUM", OperatorType::kTensorFlowMinimum); + CheckSimpleOperator("LESS", + OperatorType::kTensorFlowLess); } TEST_F(OperatorTest, BuiltinAdd) { @@ -163,10 +167,12 @@ TEST_F(OperatorTest, CustomFakeQuant) { minmax->min = -10; minmax->max = 200; op.minmax.reset(minmax); + op.num_bits = 16; auto output_toco_op = SerializeAndDeserialize( GetOperator("FAKE_QUANT", OperatorType::kFakeQuant), op); EXPECT_EQ(op.minmax->min, output_toco_op->minmax->min); EXPECT_EQ(op.minmax->max, output_toco_op->minmax->max); + EXPECT_EQ(op.num_bits, output_toco_op->num_bits); } TEST_F(OperatorTest, CustomFullyConnected) { @@ -391,6 +397,13 @@ TEST_F(OperatorTest, BuiltinTopKV2) { ASSERT_NE(nullptr, output_toco_op.get()); } +TEST_F(OperatorTest, BuiltinArgMax) { + ArgMaxOperator op; + auto output_toco_op = SerializeAndDeserialize( + GetOperator("ARG_MAX", OperatorType::kArgMax), op); + EXPECT_EQ(op.output_data_type, output_toco_op->output_data_type); +} + TEST_F(OperatorTest, TensorFlowUnsupported) { TensorFlowUnsupportedOperator op; op.tensorflow_op = "MyCustomUnsupportedOp"; diff --git a/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc b/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc index cc7803dd866f0282f67d1d6f227cce0fdd8c7fd6..1611c4d0c0b148e00dbf0b21b9cd65d4c8163c2c 100644 --- a/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc +++ b/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc @@ -61,11 +61,21 @@ bool ParseTocoFlagsFromCommandLineFlags( Flag("default_ranges_min", parsed_flags.default_ranges_min.bind(), parsed_flags.default_ranges_min.default_value(), "If defined, will be used as the default value for the min bound " - "of min/max ranges used for quantization."), + "of min/max ranges used for quantization of uint8 arrays."), Flag("default_ranges_max", parsed_flags.default_ranges_max.bind(), parsed_flags.default_ranges_max.default_value(), "If defined, will be used as the default value for the max bound " - "of min/max ranges used for quantization."), + "of min/max ranges used for quantization of uint8 arrays."), + Flag("default_int16_ranges_min", + parsed_flags.default_int16_ranges_min.bind(), + parsed_flags.default_int16_ranges_min.default_value(), + "If defined, will be used as the default value for the min bound " + "of min/max ranges used for quantization of int16 arrays."), + Flag("default_int16_ranges_max", + parsed_flags.default_int16_ranges_max.bind(), + parsed_flags.default_int16_ranges_max.default_value(), + "If defined, will be used as the default value for the max bound " + "of min/max ranges used for quantization of int16 arrays."), Flag("inference_type", parsed_flags.inference_type.bind(), parsed_flags.inference_type.default_value(), "Target data type of arrays in the output file (for input_arrays, " @@ -126,6 +136,18 @@ bool ParseTocoFlagsFromCommandLineFlags( 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."), + Flag("propagate_fake_quant_num_bits", + parsed_flags.propagate_fake_quant_num_bits.bind(), + parsed_flags.propagate_fake_quant_num_bits.default_value(), + "If true, use FakeQuant* operator num_bits attributes to adjust " + "array data_types."), + Flag("allow_nudging_weights_to_use_fast_gemm_kernel", + parsed_flags.allow_nudging_weights_to_use_fast_gemm_kernel.bind(), + parsed_flags.allow_nudging_weights_to_use_fast_gemm_kernel + .default_value(), + "Some fast uint8 GEMM kernels require uint8 weights to avoid the " + "value 0. This flag allows nudging them to 1 to allow proceeding, " + "with moderate inaccuracy."), }; bool asked_for_help = *argc == 2 && (!strcmp(argv[1], "--help") || !strcmp(argv[1], "-help")); @@ -207,10 +229,16 @@ void ReadTocoFlagsFromCommandLineFlags(const ParsedTocoFlags& parsed_toco_flags, PARSE_TOCO_FLAG(IODataType, inference_input_type, FlagRequirement::kNone); READ_TOCO_FLAG(default_ranges_min, FlagRequirement::kNone); READ_TOCO_FLAG(default_ranges_max, FlagRequirement::kNone); + READ_TOCO_FLAG(default_int16_ranges_min, FlagRequirement::kNone); + READ_TOCO_FLAG(default_int16_ranges_max, FlagRequirement::kNone); READ_TOCO_FLAG(drop_fake_quant, FlagRequirement::kNone); READ_TOCO_FLAG(reorder_across_fake_quant, FlagRequirement::kNone); READ_TOCO_FLAG(allow_custom_ops, FlagRequirement::kNone); READ_TOCO_FLAG(drop_control_dependency, FlagRequirement::kNone); + READ_TOCO_FLAG(debug_disable_recurrent_cell_fusion, FlagRequirement::kNone); + READ_TOCO_FLAG(propagate_fake_quant_num_bits, FlagRequirement::kNone); + READ_TOCO_FLAG(allow_nudging_weights_to_use_fast_gemm_kernel, + FlagRequirement::kNone); // Deprecated flag handling. if (parsed_toco_flags.input_type.specified()) { diff --git a/tensorflow/contrib/lite/toco/toco_flags.proto b/tensorflow/contrib/lite/toco/toco_flags.proto index 3237147a736f97f65953ca965420fcea934820a4..a04017a6bf05fb5a1ae324051375340a0adaccc6 100644 --- a/tensorflow/contrib/lite/toco/toco_flags.proto +++ b/tensorflow/contrib/lite/toco/toco_flags.proto @@ -37,7 +37,7 @@ enum FileFormat { // of as properties of models, instead describing how models are to be // processed in the context of the present tooling job. // -// Next ID to use: 14. +// Next ID to use: 18. message TocoFlags { // Input file format optional FileFormat input_format = 1; @@ -103,8 +103,14 @@ message TocoFlags { // for experimentation purposes only and should not be used in production: // they make it easy to quantize models, but the resulting quantized model // will be inaccurate. + // + // These values only apply to arrays quantized with the kUint8 data type. optional float default_ranges_min = 5; optional float default_ranges_max = 6; + // Equivalent versions of default_ranges_min/_max for arrays quantized with + // the kInt16 data type. + optional float default_int16_ranges_min = 15; + optional float default_int16_ranges_max = 16; // Ignore and discard FakeQuant nodes. For instance, that can be used to // generate plain float code without fake-quantization from a quantized @@ -141,4 +147,18 @@ message TocoFlags { // Disables transformations that fuse subgraphs such as known LSTMs (not all // LSTMs are identified). optional bool debug_disable_recurrent_cell_fusion = 13; + + // Uses the FakeQuantWithMinMaxArgs.num_bits attribute to adjust quantized + // array data types throughout the graph. The graph must be properly annotated + // with FakeQuant* ops on at least the edges and may contain additional ops on + // the interior of the graph to widen/narrow as desired. + // + // Input and output array data types may change because of this propagation + // and users must be sure to query the final data_type values. + optional bool propagate_fake_quant_num_bits = 14; + + // Some fast uint8 GEMM kernels require uint8 weights to avoid the value 0. + // This flag allows nudging them to 1 to allow proceeding, with moderate + // inaccuracy. + optional bool allow_nudging_weights_to_use_fast_gemm_kernel = 17; } diff --git a/tensorflow/contrib/lite/toco/toco_port.h b/tensorflow/contrib/lite/toco/toco_port.h index 4be3b5a0bf00ed204a1218545d9e66f7685a50d7..2d5c231bef350884cd2b0bb62cfdbddebfed7f58 100644 --- a/tensorflow/contrib/lite/toco/toco_port.h +++ b/tensorflow/contrib/lite/toco/toco_port.h @@ -75,6 +75,14 @@ Status Exists(const string& filename, const Options& options); void CopyToBuffer(const ::Cord& src, char* dest); #endif // PLATFORM_GOOGLE void CopyToBuffer(const string& src, char* dest); + +inline uint32 ReverseBits32(uint32 n) { + n = ((n >> 1) & 0x55555555) | ((n & 0x55555555) << 1); + n = ((n >> 2) & 0x33333333) | ((n & 0x33333333) << 2); + n = ((n >> 4) & 0x0F0F0F0F) | ((n & 0x0F0F0F0F) << 4); + return (((n & 0xFF) << 24) | ((n & 0xFF00) << 8) | ((n & 0xFF0000) >> 8) | + ((n & 0xFF000000) >> 24)); +} } // namespace port inline bool ParseFromStringOverload(const std::string& in, diff --git a/tensorflow/contrib/lite/toco/toco_tooling.cc b/tensorflow/contrib/lite/toco/toco_tooling.cc index 96c5ebd64f3343d454a9c445a6edcd318b08c953..7252ec2ea4d8864334fc989c91fc8c26efe2efea 100644 --- a/tensorflow/contrib/lite/toco/toco_tooling.cc +++ b/tensorflow/contrib/lite/toco/toco_tooling.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "absl/memory/memory.h" #include "absl/strings/str_join.h" #include "tensorflow/contrib/lite/toco/allocate_transient_arrays.h" #include "tensorflow/contrib/lite/toco/dump_graphviz.h" @@ -66,6 +67,7 @@ void MakeGeneralGraphTransformationsSet( transformations->Add(new RemoveTensorFlowIdentity); transformations->Add(new RemoveTrivialConcatenation); transformations->Add(new RemoveTrivialConcatenationInput); + transformations->Add(new RemoveTrivialFakeQuant); transformations->Add(new RemoveTrivialSlice); transformations->Add(new RemoveUnusedOp); transformations->Add(new EnsureBiasVectors); @@ -83,6 +85,7 @@ void MakeGeneralGraphTransformationsSet( transformations->Add(new ResolveConstantGather); transformations->Add(new ResolveConstantRandomUniform); transformations->Add(new ResolveConstantRange); + transformations->Add(new ResolveConstantReshape); transformations->Add(new ResolveConstantStack); transformations->Add(new ResolveConstantStridedSlice); transformations->Add(new ResolveConstantTranspose); @@ -108,7 +111,6 @@ void MakeGeneralGraphTransformationsSet( transformations->Add(new ResolveMeanAttributes); transformations->Add(new ResolveConstantShapeOrRank); transformations->Add(new MakeInitialDequantizeOperator); - transformations->Add(new ResolveConstantFakeQuant); transformations->Add(new UnpartitionEmbeddingLookup); } @@ -232,6 +234,12 @@ void Transform(const TocoFlags& toco_flags, Model* model) { MakeGeneralGraphTransformationsSet(&transformations); auto* remove_trivial_reshape = new RemoveTrivialReshape; transformations.Add(remove_trivial_reshape); + auto* resolve_constant_fake_quant = new ResolveConstantFakeQuant; + if (quantize_output) { + resolve_constant_fake_quant->set_propagate_fake_quant_num_bits( + toco_flags.propagate_fake_quant_num_bits()); + } + transformations.Add(resolve_constant_fake_quant); if (SupportsFusedActivationFunction(output_format)) { transformations.Add(new FuseActivationFunctions); } else { @@ -264,25 +272,63 @@ void Transform(const TocoFlags& toco_flags, Model* model) { transformations); if (quantize_output) { + if (toco_flags.propagate_fake_quant_num_bits()) { + RunGraphTransformations(model, + "fake quant propagation graph transformations", + {new PropagateFakeQuantNumBits}); + } RunGraphTransformations(model, "pre-quantization graph transformations", - {new HardcodeMinMax, new DropFakeQuant}); + { + new HardcodeMinMax, + new DropFakeQuant, + }); } + // Fix any issues with IO edges. This must happen after any transform that + // may modify the structure of the edges. + FixEdgeArrays(model); + if (quantize_output) { + // If the user specified default min/max ranges we need to set all arrays + // that didn't either have a min/max specified or get one set via + // HardcodeMinMax or PropagateFakeQuantNumBits. This may require running + // HardcodeMinMax to move changes through the graph as we make changes. + auto propagate_default_min_max = + absl::make_unique(); if (toco_flags.has_default_ranges_min() && toco_flags.has_default_ranges_max()) { - UseDefaultMinMaxRangeValues(model, toco_flags.default_ranges_min(), - toco_flags.default_ranges_max()); - // The new MinMax info may need to be propagated a bit. + propagate_default_min_max->DefineTypeRange( + ArrayDataType::kUint8, toco_flags.default_ranges_min(), + toco_flags.default_ranges_max()); + } + if (toco_flags.has_default_int16_ranges_min() && + toco_flags.has_default_int16_ranges_max()) { + propagate_default_min_max->DefineTypeRange( + ArrayDataType::kInt16, toco_flags.default_int16_ranges_min(), + toco_flags.default_int16_ranges_max()); + } + if (propagate_default_min_max->has_any_ranges_defined()) { RunGraphTransformations( model, "default min-max range propagation graph transformations", - {new HardcodeMinMax}); + { + propagate_default_min_max.release(), + new HardcodeMinMax, + }); } + CheckIsReadyForQuantization(*model); - RunGraphTransformations( - model, "quantization graph transformations", - {new Quantize, new RemoveTrivialQuantizedActivationFunc, - new RemoveFinalDequantizeOp}); + auto* ensure_safe_for_int8_kernels = + new EnsureUint8WeightsSafeForFastInt8Kernels; + ensure_safe_for_int8_kernels->set_allow_nudging_weights( + toco_flags.allow_nudging_weights_to_use_fast_gemm_kernel()); + RunGraphTransformations(model, "quantization graph transformations", + { + new RemoveTrivialQuantizedActivationFunc, + new RemoveTrivialQuantizedMinMax, + new Quantize, + new RemoveFinalDequantizeOp, + ensure_safe_for_int8_kernels, + }); } else { GraphTransformationsSet dequantization_transformations{new Dequantize}; // Dequantize creates FakeQuant nodes. We may want to discard @@ -299,10 +345,6 @@ void Transform(const TocoFlags& toco_flags, Model* model) { EncodeConstantArraysMinMaxByWrappingThemInFakeQuantNodes(model); } - // Fix any issues with IO edges. This must happen after any transform that - // may modify the structure of the edges. - FixEdgeArrays(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 bd2d5f7df039b61d50c3f7ab9fa7d5d9efd75afb..cf2cbeedc77c382b84f80f11a371f300ea4282ae 100644 --- a/tensorflow/contrib/lite/toco/tooling_util.cc +++ b/tensorflow/contrib/lite/toco/tooling_util.cc @@ -93,9 +93,18 @@ string ArrayDataTypeName(ArrayDataType data_type) { } } -bool IsInputArray(const Model& model, const string& name) { +bool IsInputArray(const Model& model, const string& array_name) { for (const auto& input_array : model.flags.input_arrays()) { - if (input_array.name() == name) { + if (array_name == input_array.name()) { + return true; + } + } + return false; +} + +bool IsOutputArray(const Model& model, const string& array_name) { + for (const auto& output_array : model.flags.output_arrays()) { + if (array_name == output_array) { return true; } } @@ -106,10 +115,8 @@ bool IsArrayConsumed(const Model& model, const string& name) { if (GetOpWithInput(model, name)) { return true; } - for (const string& model_output : model.flags.output_arrays()) { - if (model_output == name) { - return true; - } + if (IsOutputArray(model, name)) { + return true; } for (const auto& rnn_state : model.flags.rnn_states()) { if (rnn_state.back_edge_source_array() == name) { @@ -379,6 +386,7 @@ string HelpfulOperatorTypeName(const Operator& op) { bool OperatorSupportsFusedActivation(OperatorType type) { switch (type) { case OperatorType::kConcatenation: + case OperatorType::kFakeQuant: case OperatorType::kGather: case OperatorType::kSlice: case OperatorType::kSqueeze: @@ -1064,16 +1072,38 @@ void FixEdgeArrays(Model* model) { } } +namespace { +void CopyArrayAttribs(const Array& source_array, Array* target_array) { + target_array->data_type = source_array.data_type; + target_array->final_data_type = source_array.final_data_type; + target_array->copy_shape(source_array.shape()); + + if (source_array.minmax) { + target_array->GetOrCreateMinMax() = source_array.GetMinMax(); + } else { + target_array->minmax.reset(); + } + + if (source_array.quantization_params) { + target_array->GetOrCreateQuantizationParams() = + source_array.GetQuantizationParams(); + } else { + target_array->quantization_params.reset(); + } +} +} // namespace + void InsertCopyOperator(Model* model, const string& source_array_name, const string& target_array_name) { + // Reshape to the same size. This should be a no-op. + const Array& source_array = model->GetArray(source_array_name); + std::vector shape = source_array.shape().dims(); + // Drop constant data from the target array as the copy will be done at // runtime. Array& target_array = model->GetOrCreateArray(target_array_name); target_array.buffer.reset(); - - // Reshape to the same size. This should be a no-op. - const Array& source_array = model->GetArray(source_array_name); - std::vector shape = source_array.shape().dims(); + CopyArrayAttribs(source_array, &target_array); // Insert copy operator. auto* copy_op = new TensorFlowReshapeOperator; @@ -1084,22 +1114,30 @@ void InsertCopyOperator(Model* model, const string& source_array_name, model->operators.emplace_back(copy_op); } -namespace { -template -void CopyArrayBuffer(const Array& source_array, Array* target_array) { - if (source_array.buffer) { - const auto& source_buffer = source_array.GetBuffer(); - auto& target_buffer = target_array->GetMutableBuffer(); - target_buffer.data = source_buffer.data; - } -} -} // namespace - void CloneArray(Model* model, const string& source_array_name, const string& target_array_name) { CHECK(!model->HasArray(target_array_name)); const Array& source_array = model->GetArray(source_array_name); Array& target_array = model->GetOrCreateArray(target_array_name); + CopyArrayAttribs(source_array, &target_array); + + if (source_array.minmax) { + const auto& smm = source_array.GetMinMax(); + auto& tmm = target_array.GetOrCreateMinMax(); + tmm.min = smm.min; + tmm.max = smm.max; + } + + if (source_array.quantization_params) { + const auto& sqp = source_array.GetQuantizationParams(); + auto& tqp = target_array.GetOrCreateQuantizationParams(); + tqp.zero_point = sqp.zero_point; + tqp.scale = sqp.scale; + } + + target_array.data_type = source_array.data_type; + target_array.final_data_type = source_array.final_data_type; + target_array.copy_shape(source_array.shape()); switch (source_array.data_type) { case ArrayDataType::kBool: @@ -1140,25 +1178,6 @@ void CloneArray(Model* model, const string& source_array_name, << ArrayDataTypeName(source_array.data_type); return; } - - if (source_array.minmax) { - const auto& smm = source_array.GetMinMax(); - auto& tmm = target_array.GetOrCreateMinMax(); - tmm.min = smm.min; - tmm.max = smm.max; - } - - if (source_array.quantization_params) { - const auto& sqp = source_array.GetQuantizationParams(); - auto& tqp = target_array.GetOrCreateQuantizationParams(); - tqp.zero_point = sqp.zero_point; - tqp.scale = sqp.scale; - } - - target_array.data_type = source_array.data_type; - target_array.final_data_type = source_array.final_data_type; - - target_array.copy_shape(source_array.shape()); } void MakeArrayDims(int num_dims, int batch, int height, int width, int depth, @@ -1455,28 +1474,6 @@ void CheckIsReadyForQuantization(const Model& model) { } } -void UseDefaultMinMaxRangeValues(Model* model, double default_ranges_min, - double default_ranges_max) { - for (const auto& op : model->operators) { - for (const auto& input : op->inputs) { - auto& input_array = model->GetArray(input); - if (!input_array.minmax && !input_array.buffer) { - auto& minmax = input_array.GetOrCreateMinMax(); - minmax.min = default_ranges_min; - minmax.max = default_ranges_max; - } - } - for (const auto& output : op->outputs) { - auto& output_array = model->GetArray(output); - if (!output_array.minmax && !output_array.buffer) { - auto& minmax = output_array.GetOrCreateMinMax(); - minmax.min = default_ranges_min; - minmax.max = default_ranges_max; - } - } - } -} - int ElementSize(ArrayDataType data_type) { switch (data_type) { case ArrayDataType::kBool: @@ -1525,14 +1522,9 @@ bool IsAllocatableTransientArray(const Model& model, const string& array_name) { if (model.IsOptionalArray(array_name)) return false; // The model's input and output arrays are externally allocated. // They are not transient arrays. - if (IsInputArray(model, array_name)) { + if (IsInputArray(model, array_name) || IsOutputArray(model, array_name)) { return false; } - for (const string& output_array : model.flags.output_arrays()) { - if (array_name == output_array) { - return false; - } - } const auto& array = &model.GetArray(array_name); // An array with a constant buffer isn't a transient array. if (!!array->buffer) { @@ -1910,15 +1902,8 @@ int AxesCount(AxesOrder axes_order) { } bool IsDiscardableArray(const Model& model, const string& array_name) { - for (const auto& input_array : model.flags.input_arrays()) { - if (array_name == input_array.name()) { - return false; - } - } - for (const string& output_array : model.flags.output_arrays()) { - if (array_name == output_array) { - return false; - } + if (IsInputArray(model, array_name) || IsOutputArray(model, array_name)) { + return false; } for (const auto& rnn_state : model.flags.rnn_states()) { if (!rnn_state.discardable()) { @@ -1972,8 +1957,8 @@ void CheckFinalDataTypesSatisfied(const Model& model) { CHECK(array.final_data_type == array.data_type) << "Array \"" << array_entry.first << "\" has mis-matching actual and final data types (" - << static_cast(array.data_type) << "," - << static_cast(array.final_data_type) << ")."; + << ArrayDataTypeName(array.data_type) << "," + << ArrayDataTypeName(array.final_data_type) << ")."; } } } diff --git a/tensorflow/contrib/lite/toco/tooling_util.h b/tensorflow/contrib/lite/toco/tooling_util.h index dfd81173c3d3bf31a0ce688ce5434cd37fb959c7..5cc15fa57b3ea4ecb74c3f81af14e4fea95f6232 100644 --- a/tensorflow/contrib/lite/toco/tooling_util.h +++ b/tensorflow/contrib/lite/toco/tooling_util.h @@ -28,7 +28,6 @@ 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" @@ -57,7 +56,11 @@ string LogName(const Operator& op); string ArrayDataTypeName(ArrayDataType data_type); -bool IsInputArray(const Model& model, const string& name); +// Returns true if the given array is specified as a model input array. +bool IsInputArray(const Model& model, const string& array_name); +// Returns true if the given array is specified as a model output array. +bool IsOutputArray(const Model& model, const string& array_name); + bool IsArrayConsumed(const Model& model, const string& name); int CountTrueOutputs(const Model& model, const Operator& op); @@ -147,6 +150,23 @@ void FixNoOrphanedArray(Model* model); // Fixes input/output arrays that may have issues during export or inference. void FixEdgeArrays(Model* model); +// Copies the contents of an array into another. +// Expects that the shape and data type match. +template +void CopyArrayBuffer(const Array& source_array, Array* target_array) { + int source_buffer_size = RequiredBufferSizeForShape(source_array.shape()); + int target_buffer_size = RequiredBufferSizeForShape(target_array->shape()); + CHECK_EQ(source_buffer_size, target_buffer_size) + << "Buffer sizes must match in element count"; + CHECK(source_array.data_type == target_array->data_type) + << "Data types must match"; + if (source_array.buffer) { + const auto& source_buffer = source_array.GetBuffer(); + auto& target_buffer = target_array->GetMutableBuffer(); + target_buffer.data = source_buffer.data; + } +} + // Inserts a no-op reshape operator between the source array and the target // array. This effectively just copies the data. void InsertCopyOperator(Model* model, const string& source_array_name, @@ -158,17 +178,6 @@ void CloneArray(Model* model, const string& source_array_name, void ResolveModelFlags(const ModelFlags& model_flags, Model* model); -template -void GetQuantizationParamsFromMinMax(const MinMax& minmax, - QuantizationParams* quantization_params) { - using Integer = DataType; - const double rmin = minmax.min; - const double rmax = minmax.max; - - *quantization_params = - ::tflite::ChooseQuantizationParams(rmin, rmax); -} - template T ConvertOperator(Operator* o, OperatorType type) { if (o != nullptr && o->type == type) { @@ -179,8 +188,6 @@ T ConvertOperator(Operator* o, OperatorType type) { } void CheckIsReadyForQuantization(const Model& model); -void UseDefaultMinMaxRangeValues(Model* model, double default_ranges_min, - double default_ranges_max); bool ReshapeIsEquivalentToTranspose(const Model& model, const TensorFlowReshapeOperator* op, diff --git a/tensorflow/contrib/lite/tools/BUILD b/tensorflow/contrib/lite/tools/BUILD index 44fde69a1e1536b8d2ecff16876248cfe66a9b8a..7b3569ea9c8b15959b15e8ba46cf44d159d5528c 100644 --- a/tensorflow/contrib/lite/tools/BUILD +++ b/tensorflow/contrib/lite/tools/BUILD @@ -78,6 +78,9 @@ cc_test( "//tensorflow/contrib/lite:testdata/test_model.bin", "//tensorflow/contrib/lite:testdata/test_model_broken.bin", ], + tags = [ + "tflite_not_portable_android", + ], deps = [ ":gen_op_registration", "@com_google_googletest//:gtest", diff --git a/tensorflow/contrib/lookup/BUILD b/tensorflow/contrib/lookup/BUILD index f616207d462954341dd0c4b2722471b50c06c917..e3928a82a2d453fdd36cb861ce178a776574269c 100644 --- a/tensorflow/contrib/lookup/BUILD +++ b/tensorflow/contrib/lookup/BUILD @@ -28,7 +28,7 @@ py_library( tf_py_test( name = "lookup_ops_test", - size = "small", + size = "medium", srcs = ["lookup_ops_test.py"], additional_deps = [ ":lookup_py", diff --git a/tensorflow/contrib/losses/python/losses/loss_ops.py b/tensorflow/contrib/losses/python/losses/loss_ops.py index 8c3a8afe7a0f6f5ad9ceae566288ba60be73d339..bdad34a665e47a4e060fcaddfffecfdc876a8fb0 100644 --- a/tensorflow/contrib/losses/python/losses/loss_ops.py +++ b/tensorflow/contrib/losses/python/losses/loss_ops.py @@ -29,6 +29,7 @@ from tensorflow.python.ops import nn from tensorflow.python.ops import nn_ops from tensorflow.python.util.deprecation import deprecated from tensorflow.python.util.deprecation import deprecated_args +from tensorflow.python.util.deprecation import deprecated_argument_lookup __all__ = [ "absolute_difference", "add_loss", "cosine_distance", @@ -651,11 +652,9 @@ def cosine_distance(predictions, ValueError: If `predictions` shape doesn't match `labels` shape, or `weights` is `None`. """ - if dim is not None: - if axis is not None: - raise ValueError("Cannot specify both 'axis' and 'dim'") - axis = dim - if axis is None and dim is None: + axis = deprecated_argument_lookup( + "axis", axis, "dim", dim) + if axis is None: raise ValueError("You must specify 'axis'.") with ops.name_scope(scope, "cosine_distance_loss", [predictions, labels, weights]) as scope: 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 2b9eee4ef7b418e2b90d388d2f165537b8660a9a..de76acb51ffe985162a66c617b266f47c5216b19 100644 --- a/tensorflow/contrib/losses/python/metric_learning/metric_loss_ops.py +++ b/tensorflow/contrib/losses/python/metric_learning/metric_loss_ops.py @@ -711,7 +711,7 @@ def _find_loss_augmented_facility_idx(pairwise_distances, labels, chosen_ids, candidate_scores, margin_multiplier * nmi_scores) argmax_index = math_ops.to_int32( - math_ops.argmax(candidate_scores, dimension=0)) + math_ops.argmax(candidate_scores, axis=0)) return candidate_ids[argmax_index] @@ -811,7 +811,7 @@ def update_medoid_per_cluster(pairwise_distances, pairwise_distances_subset, candidate_scores = math_ops.add(scores_fac, margin_multiplier * scores_margin) argmax_index = math_ops.to_int32( - math_ops.argmax(candidate_scores, dimension=0)) + math_ops.argmax(candidate_scores, axis=0)) best_medoid = math_ops.to_int32(cluster_member_ids[argmax_index]) chosen_ids = update_1d_tensor(chosen_ids, cluster_idx, best_medoid) diff --git a/tensorflow/contrib/makefile/Makefile b/tensorflow/contrib/makefile/Makefile index 05e8d9064bea748c935859f5f9b4c7e646f504cf..1a1ab54a53dd5866ca8357067846c002c5d5e9c1 100644 --- a/tensorflow/contrib/makefile/Makefile +++ b/tensorflow/contrib/makefile/Makefile @@ -89,6 +89,7 @@ HOST_INCLUDES := \ -I$(MAKEFILE_DIR)/downloads/gemmlowp \ -I$(MAKEFILE_DIR)/downloads/nsync/public \ -I$(MAKEFILE_DIR)/downloads/fft2d \ +-I$(MAKEFILE_DIR)/downloads/double_conversion \ -I$(HOST_GENDIR) ifeq ($(HAS_GEN_HOST_PROTOC),true) HOST_INCLUDES += -I$(MAKEFILE_DIR)/gen/protobuf-host/include @@ -125,7 +126,9 @@ PROTO_TEXT := $(HOST_BINDIR)proto_text # The list of dependencies is derived from the Bazel build file by running # the gen_file_lists.sh script on a system with a working Bazel setup. PROTO_TEXT_CC_FILES := $(shell cat $(MAKEFILE_DIR)/proto_text_cc_files.txt) -PROTO_TEXT_PB_CC_LIST := $(shell cat $(MAKEFILE_DIR)/proto_text_pb_cc_files.txt) +PROTO_TEXT_PB_CC_LIST := \ + $(shell cat $(MAKEFILE_DIR)/proto_text_pb_cc_files.txt) \ + $(wildcard tensorflow/contrib/makefile/downloads/double_conversion/double-conversion/*.cc) PROTO_TEXT_PB_H_LIST := $(shell cat $(MAKEFILE_DIR)/proto_text_pb_h_files.txt) # Locations of the intermediate files proto_text generates. @@ -171,6 +174,7 @@ INCLUDES := \ -I$(MAKEFILE_DIR)/downloads/gemmlowp \ -I$(MAKEFILE_DIR)/downloads/nsync/public \ -I$(MAKEFILE_DIR)/downloads/fft2d \ +-I$(MAKEFILE_DIR)/downloads/double_conversion \ -I$(PROTOGENDIR) \ -I$(PBTGENDIR) ifeq ($(HAS_GEN_HOST_PROTOC),true) @@ -326,6 +330,7 @@ $(MARCH_OPTION) \ -I$(MAKEFILE_DIR)/downloads/gemmlowp \ -I$(MAKEFILE_DIR)/downloads/nsync/public \ -I$(MAKEFILE_DIR)/downloads/fft2d \ +-I$(MAKEFILE_DIR)/downloads/double_conversion \ -I$(MAKEFILE_DIR)/gen/protobuf_android/$(ANDROID_ARCH)/include \ -I$(PROTOGENDIR) \ -I$(PBTGENDIR) @@ -603,6 +608,7 @@ $(wildcard tensorflow/core/platform/*/*.cc) \ $(wildcard tensorflow/core/platform/*/*/*.cc) \ $(wildcard tensorflow/core/util/*.cc) \ $(wildcard tensorflow/core/util/*/*.cc) \ +$(wildcard tensorflow/contrib/makefile/downloads/double_conversion/double-conversion/*.cc) \ tensorflow/core/util/version_info.cc # Remove duplicates (for version_info.cc) CORE_CC_ALL_SRCS := $(sort $(CORE_CC_ALL_SRCS)) diff --git a/tensorflow/contrib/makefile/download_dependencies.sh b/tensorflow/contrib/makefile/download_dependencies.sh index 8b415e6527f85a5a7844b9d4156fd39ecb1b637a..eff9081e35c285027c764c5bdbaf14f78bc5f512 100755 --- a/tensorflow/contrib/makefile/download_dependencies.sh +++ b/tensorflow/contrib/makefile/download_dependencies.sh @@ -27,12 +27,15 @@ if [ ! -f $BZL_FILE_PATH ]; then fi EIGEN_URL="$(grep -o 'http.*bitbucket.org/eigen/eigen/get/.*tar\.gz' "${BZL_FILE_PATH}" | grep -v mirror.bazel | head -n1)" -GEMMLOWP_URL="$(grep -o 'https://mirror.bazel.build/github.com/google/gemmlowp/.*zip' "${BZL_FILE_PATH}" | head -n1)" +# TODO (yongtang): Replace the following with 'https://mirror.bazel.build/github.com/google/gemmlowp/.*zip' once +# the archive has been propagated in mirror.bazel.build. +GEMMLOWP_URL="$(grep -o 'https://github.com/google/gemmlowp/.*zip' "${BZL_FILE_PATH}" | head -n1)" GOOGLETEST_URL="https://github.com/google/googletest/archive/release-1.8.0.tar.gz" NSYNC_URL="$(grep -o 'https://mirror.bazel.build/github.com/google/nsync/.*tar\.gz' "${BZL_FILE_PATH}" | head -n1)" PROTOBUF_URL="$(grep -o 'https://mirror.bazel.build/github.com/google/protobuf/.*tar\.gz' "${BZL_FILE_PATH}" | head -n1)" RE2_URL="$(grep -o 'https://mirror.bazel.build/github.com/google/re2/.*tar\.gz' "${BZL_FILE_PATH}" | head -n1)" -FFT2D_URL="$(grep -o 'http.*fft\.tgz' "${BZL_FILE_PATH}" | grep -v mirror.bazel | head -n1)" +FFT2D_URL="$(grep -o 'http.*fft\.tgz' "${BZL_FILE_PATH}" | grep -v bazel-mirror | head -n1)" +DOUBLE_CONVERSION_URL="$(grep -o "https.*google/double-conversion.*\.zip" "${BZL_FILE_PATH}" | head -n1)" ABSL_URL="$(grep -o 'https://github.com/abseil/abseil-cpp/.*tar.gz' "${BZL_FILE_PATH}" | head -n1)" CUB_URL="$(grep -o 'https.*cub/archive.*zip' "${BZL_FILE_PATH}" | grep -v mirror.bazel | head -n1)" @@ -87,6 +90,7 @@ download_and_extract "${NSYNC_URL}" "${DOWNLOADS_DIR}/nsync" download_and_extract "${PROTOBUF_URL}" "${DOWNLOADS_DIR}/protobuf" download_and_extract "${RE2_URL}" "${DOWNLOADS_DIR}/re2" download_and_extract "${FFT2D_URL}" "${DOWNLOADS_DIR}/fft2d" +download_and_extract "${DOUBLE_CONVERSION_URL}" "${DOWNLOADS_DIR}/double_conversion" download_and_extract "${ABSL_URL}" "${DOWNLOADS_DIR}/absl" download_and_extract "${CUB_URL}" "${DOWNLOADS_DIR}/cub/external/cub_archive" diff --git a/tensorflow/contrib/meta_graph_transform/meta_graph_transform.py b/tensorflow/contrib/meta_graph_transform/meta_graph_transform.py index ff88b4fa841673fc52b9f6fdc5ca43d30c44bbfd..f37a2593e2695ec55a10f69803cc7f446b0862df 100644 --- a/tensorflow/contrib/meta_graph_transform/meta_graph_transform.py +++ b/tensorflow/contrib/meta_graph_transform/meta_graph_transform.py @@ -30,7 +30,7 @@ from tensorflow.python.framework import importer as _importer from tensorflow.python.framework import ops as _ops from tensorflow.python.saved_model import constants as _saved_model_constants from tensorflow.python.training import saver as _saver_lib -from tensorflow.python.util import compat +from tensorflow.python.util import compat as _compat from tensorflow.tools import graph_transforms as _graph_transforms @@ -161,7 +161,7 @@ def _clean_save_and_restore(graph_def, op, removed_op_names): shapes = [] dtypes = [] for index, value in enumerate(name_op_value_tensor.string_val): - if not _is_removed(compat.as_str(value), removed_op_names): + if not _is_removed(_compat.as_str(value), removed_op_names): names.append(value) shapes.append(shape_op_value_tensor.string_val[index]) dtypes.append(op.attr['dtypes'].list.type[index]) @@ -348,7 +348,7 @@ def _freeze_graph_with_def_protos(input_graph_def, output_node_names, input_saver_def, input_checkpoint): """Converts all variables in a graph and checkpoint into constants. - During this process, we need to retain certain initialzer nodes (e.g. table + During this process, we need to retain certain initializer nodes (e.g. table initializer nodes). Instead of determining which dependencies of the shared initializer node (e.g. group_deps) to keep, we reconstruct the connections between the individual initializer nodes and @@ -651,7 +651,7 @@ def _is_removed_mentioned(s, removed_op_names): # /foo/bar. This regex ensures that we handle these two nodes # as separate entities. It matches on nodes having names in the form of # '/foo/bar_x' as well as nodes having names in the form of 'foo.' - s_names = _re.findall(r'((?:[\/]?[a-zA-Z0-9\_]*)*)', compat.as_str_any(s)) + s_names = _re.findall(r'((?:[\/]?[a-zA-Z0-9\_]*)*)', _compat.as_str_any(s)) for removed_op_name in removed_op_names: for s_name in s_names: if s_name.endswith(removed_op_name): @@ -737,9 +737,9 @@ def meta_graph_transform( for tag in tags: meta_graph_def.meta_info_def.tags.append(tag) - base_op_names = [compat.as_str(node.name) + base_op_names = [_compat.as_str(node.name) for node in base_meta_graph_def.graph_def.node] - retained_op_names = [compat.as_str(node.name) + retained_op_names = [_compat.as_str(node.name) for node in meta_graph_def.graph_def.node] removed_op_names = set(base_op_names) - set(retained_op_names) diff --git a/tensorflow/contrib/metrics/__init__.py b/tensorflow/contrib/metrics/__init__.py index de02dc8f457364450929776035829d86035d706b..5effea3596bb83a08e0a8627e411684262aef5f7 100644 --- a/tensorflow/contrib/metrics/__init__.py +++ b/tensorflow/contrib/metrics/__init__.py @@ -71,6 +71,7 @@ See the @{$python/contrib.metrics} guide. @@count @@precision_recall_at_equal_thresholds @@recall_at_precision +@@precision_at_recall """ from __future__ import absolute_import @@ -87,6 +88,7 @@ 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_at_recall from tensorflow.contrib.metrics.python.ops.metric_ops import precision_recall_at_equal_thresholds from tensorflow.contrib.metrics.python.ops.metric_ops import recall_at_precision from tensorflow.contrib.metrics.python.ops.metric_ops import sparse_recall_at_top_k diff --git a/tensorflow/contrib/metrics/python/ops/metric_ops.py b/tensorflow/contrib/metrics/python/ops/metric_ops.py index 088319a5572f346ebb3409f2176e8b3589791f5d..00a933e5e0c537033573b225d43581f74557b240 100644 --- a/tensorflow/contrib/metrics/python/ops/metric_ops.py +++ b/tensorflow/contrib/metrics/python/ops/metric_ops.py @@ -62,6 +62,7 @@ def _safe_div(numerator, denominator, name): 0, name=name) + @deprecated(None, 'Please switch to tf.metrics.true_positives. Note that the ' 'order of the labels and predictions arguments has been switched.') def streaming_true_positives(predictions, @@ -107,6 +108,7 @@ def streaming_true_positives(predictions, updates_collections=updates_collections, name=name) + @deprecated(None, 'Please switch to tf.metrics.true_negatives. Note that the ' 'order of the labels and predictions arguments has been switched.') def streaming_true_negatives(predictions, @@ -152,6 +154,7 @@ def streaming_true_negatives(predictions, updates_collections=updates_collections, name=name) + @deprecated(None, 'Please switch to tf.metrics.false_positives. Note that the ' 'order of the labels and predictions arguments has been switched.') def streaming_false_positives(predictions, @@ -197,6 +200,7 @@ def streaming_false_positives(predictions, updates_collections=updates_collections, name=name) + @deprecated(None, 'Please switch to tf.metrics.false_negatives. Note that the ' 'order of the labels and predictions arguments has been switched.') def streaming_false_negatives(predictions, @@ -241,6 +245,7 @@ def streaming_false_negatives(predictions, updates_collections=updates_collections, name=name) + @deprecated(None, 'Please switch to tf.metrics.mean') def streaming_mean(values, weights=None, @@ -290,6 +295,7 @@ def streaming_mean(values, updates_collections=updates_collections, name=name) + @deprecated(None, 'Please switch to tf.metrics.mean_tensor') def streaming_mean_tensor(values, weights=None, @@ -345,7 +351,7 @@ def streaming_mean_tensor(values, @deprecated(None, 'Please switch to tf.metrics.accuracy. Note that the order ' - 'of the labels and predictions arguments has been switched.') + 'of the labels and predictions arguments has been switched.') def streaming_accuracy(predictions, labels, weights=None, @@ -402,8 +408,9 @@ def streaming_accuracy(predictions, updates_collections=updates_collections, name=name) + @deprecated(None, 'Please switch to tf.metrics.precision. Note that the order ' - 'of the labels and predictions arguments has been switched.') + 'of the labels and predictions arguments has been switched.') def streaming_precision(predictions, labels, weights=None, @@ -459,8 +466,9 @@ def streaming_precision(predictions, updates_collections=updates_collections, name=name) + @deprecated(None, 'Please switch to tf.metrics.recall. Note that the order ' - 'of the labels and predictions arguments has been switched.') + 'of the labels and predictions arguments has been switched.') def streaming_recall(predictions, labels, weights=None, @@ -981,7 +989,7 @@ def streaming_curve_points(labels=None, @deprecated(None, 'Please switch to tf.metrics.auc. Note that the order of ' - 'the labels and predictions arguments has been switched.') + 'the labels and predictions arguments has been switched.') def streaming_auc(predictions, labels, weights=None, @@ -2580,6 +2588,121 @@ def recall_at_precision(labels, return recall, update_op +def precision_at_recall(labels, + predictions, + target_recall, + weights=None, + num_thresholds=200, + metrics_collections=None, + updates_collections=None, + name=None): + """Computes the precision at a given recall. + + This function creates variables to track the true positives, false positives, + true negatives, and false negatives at a set of thresholds. Among those + thresholds where recall is at least `target_recall`, precision is computed + at the threshold where recall is closest to `target_recall`. + + For estimation of the metric over a stream of data, the function creates an + `update_op` operation that updates these variables and returns the + precision at `target_recall`. `update_op` increments the counts of true + positives, false positives, true negatives, and false negatives with the + weight of each case found in the `predictions` and `labels`. + + If `weights` is `None`, weights default to 1. Use weights of 0 to mask values. + + For additional information about precision and recall, see + http://en.wikipedia.org/wiki/Precision_and_recall + + Args: + labels: The ground truth values, a `Tensor` whose dimensions must match + `predictions`. Will be cast to `bool`. + predictions: A floating point `Tensor` of arbitrary shape and whose values + are in the range `[0, 1]`. + target_recall: A scalar value in range `[0, 1]`. + 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 + be either `1`, or the same as the corresponding `labels` dimension). + num_thresholds: The number of thresholds to use for matching the given + recall. + metrics_collections: An optional list of collections to which `precision` + should be added. + updates_collections: An optional list of collections to which `update_op` + should be added. + name: An optional variable_scope name. + + Returns: + precision: A scalar `Tensor` representing the precision at the given + `target_recall` value. + update_op: An operation that increments the variables for tracking the + true positives, false positives, true negatives, and false negatives and + whose value matches `precision`. + + Raises: + ValueError: If `predictions` and `labels` have mismatched shapes, if + `weights` is not `None` and its shape doesn't match `predictions`, or if + `target_recall` is not between 0 and 1, or if either `metrics_collections` + or `updates_collections` are not a list or tuple. + RuntimeError: If eager execution is enabled. + """ + if context.executing_eagerly(): + raise RuntimeError('tf.metrics.precision_at_recall is not ' + 'supported when eager execution is enabled.') + + if target_recall < 0 or target_recall > 1: + raise ValueError('`target_recall` must be in the range [0, 1].') + + with variable_scope.variable_scope(name, 'precision_at_recall', + (predictions, labels, weights)): + kepsilon = 1e-7 # Used to avoid division by zero. + thresholds = [ + (i + 1) * 1.0 / (num_thresholds - 1) for i in range(num_thresholds - 2) + ] + thresholds = [0.0 - kepsilon] + thresholds + [1.0 + kepsilon] + + values, update_ops = _streaming_confusion_matrix_at_thresholds( + predictions, labels, thresholds, weights) + + def compute_precision_at_recall(tp, fp, fn, name): + """Computes the precision at a given recall. + + Args: + tp: True positives. + fp: False positives. + fn: False negatives. + name: A name for the operation. + + Returns: + The precision at the desired recall. + """ + recalls = math_ops.div(tp, tp + fn + kepsilon) + + # Because recall is monotone decreasing as a function of the threshold, + # the smallest recall exceeding target_recall occurs at the largest + # threshold where recall >= target_recall. + admissible_recalls = math_ops.cast( + math_ops.greater_equal(recalls, target_recall), dtypes.int64) + tf_index = math_ops.reduce_sum(admissible_recalls) - 1 + + # Now we have the threshold at which to compute precision: + return math_ops.div(tp[tf_index] + kepsilon, + tp[tf_index] + fp[tf_index] + kepsilon, + name) + + precision_value = compute_precision_at_recall( + values['tp'], values['fp'], values['fn'], 'value') + update_op = compute_precision_at_recall( + update_ops['tp'], update_ops['fp'], update_ops['fn'], 'update_op') + + if metrics_collections: + ops.add_to_collections(metrics_collections, precision_value) + + if updates_collections: + ops.add_to_collections(updates_collections, update_op) + + return precision_value, update_op + + def streaming_sparse_average_precision_at_k(predictions, labels, k, @@ -2711,7 +2834,9 @@ def streaming_sparse_average_precision_at_top_k(top_k_predictions, name=name) -@deprecated(None, 'Please switch to tf.metrics.mean.') +@deprecated(None, + 'Please switch to tf.metrics.mean_absolute_error. Note that the ' + 'order of the labels and predictions arguments has been switched.') def streaming_mean_absolute_error(predictions, labels, weights=None, @@ -2830,7 +2955,9 @@ def streaming_mean_relative_error(predictions, updates_collections=updates_collections, name=name) - +@deprecated(None, + 'Please switch to tf.metrics.mean_squared_error. Note that the ' + 'order of the labels and predictions arguments has been switched.') def streaming_mean_squared_error(predictions, labels, weights=None, @@ -2888,7 +3015,10 @@ def streaming_mean_squared_error(predictions, updates_collections=updates_collections, name=name) - +@deprecated( + None, + 'Please switch to tf.metrics.root_mean_squared_error. Note that the ' + 'order of the labels and predictions arguments has been switched.') def streaming_root_mean_squared_error(predictions, labels, weights=None, @@ -3228,7 +3358,7 @@ def streaming_mean_cosine_distance(predictions, radial_diffs = math_ops.reduce_sum( radial_diffs, reduction_indices=[ dim, - ], keep_dims=True) + ], keepdims=True) mean_distance, update_op = streaming_mean(radial_diffs, weights, None, None, name or 'mean_cosine_distance') mean_distance = math_ops.subtract(1.0, mean_distance) diff --git a/tensorflow/contrib/metrics/python/ops/metric_ops_test.py b/tensorflow/contrib/metrics/python/ops/metric_ops_test.py index 33eb655fb660f0ecdfe1c5ab870d7f17690ae3ff..76420db8bda39435bcc2be2fd3d8c3467d6753e2 100644 --- a/tensorflow/contrib/metrics/python/ops/metric_ops_test.py +++ b/tensorflow/contrib/metrics/python/ops/metric_ops_test.py @@ -3380,6 +3380,138 @@ class RecallAtPrecisionTest(test.TestCase): self.assertAlmostEqual(target_recall, recall.eval()) +class PrecisionAtRecallTest(test.TestCase): + + def setUp(self): + np.random.seed(1) + ops.reset_default_graph() + + def testVars(self): + metrics.precision_at_recall( + predictions=array_ops.ones((10, 1)), + labels=array_ops.ones((10, 1)), + target_recall=0.7) + _assert_metric_variables(self, + ('precision_at_recall/true_positives:0', + 'precision_at_recall/false_negatives:0', + 'precision_at_recall/false_positives:0', + 'precision_at_recall/true_negatives:0')) + + def testMetricsCollection(self): + my_collection_name = '__metrics__' + mean, _ = metrics.precision_at_recall( + predictions=array_ops.ones((10, 1)), + labels=array_ops.ones((10, 1)), + target_recall=0.7, + metrics_collections=[my_collection_name]) + self.assertListEqual(ops.get_collection(my_collection_name), [mean]) + + def testUpdatesCollection(self): + my_collection_name = '__updates__' + _, update_op = metrics.precision_at_recall( + predictions=array_ops.ones((10, 1)), + labels=array_ops.ones((10, 1)), + target_recall=0.7, + updates_collections=[my_collection_name]) + self.assertListEqual(ops.get_collection(my_collection_name), [update_op]) + + def testValueTensorIsIdempotent(self): + predictions = random_ops.random_uniform( + (10, 3), maxval=1, dtype=dtypes_lib.float32, seed=1) + labels = random_ops.random_uniform( + (10, 3), maxval=2, dtype=dtypes_lib.int64, seed=1) + precision, update_op = metrics.precision_at_recall( + labels, predictions, target_recall=0.7) + + with self.test_session() as sess: + sess.run(variables.local_variables_initializer()) + + # Run several updates. + for _ in range(10): + sess.run(update_op) + + # Then verify idempotency. + initial_precision = precision.eval() + for _ in range(10): + self.assertAlmostEqual(initial_precision, precision.eval(), places=5) + + def testAllCorrect(self): + inputs = np.random.randint(0, 2, size=(100, 1)) + + predictions = constant_op.constant(inputs, dtype=dtypes_lib.float32) + labels = constant_op.constant(inputs) + precision, update_op = metrics.precision_at_recall( + labels, predictions, target_recall=0.7) + + with self.test_session() as sess: + sess.run(variables.local_variables_initializer()) + self.assertEqual(1, sess.run(update_op)) + self.assertEqual(1, precision.eval()) + + def testAllIncorrect(self): + inputs = np.random.randint(0, 2, size=(100, 1)) + + predictions = constant_op.constant(inputs, dtype=dtypes_lib.float32) + labels = 1.0 - predictions + label_prior = math_ops.reduce_mean(labels) + precision, update_op = metrics.precision_at_recall( + labels, predictions, target_recall=0.2) + + with self.test_session() as sess: + sess.run(variables.local_variables_initializer()) + self.assertEqual(sess.run(label_prior), sess.run(update_op)) + self.assertEqual(sess.run(label_prior), precision.eval()) + + def testSomeCorrectHighRecall(self): + predictions_values = [0.1, 0.2, 0.5, 0.3, 0.0, 0.1, 0.45, 0.5, 0.8, 0.9] + labels_values = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1] + + predictions = constant_op.constant( + predictions_values, dtype=dtypes_lib.float32) + labels = constant_op.constant(labels_values) + precision, update_op = metrics.precision_at_recall( + labels, predictions, target_recall=0.8) + + with self.test_session() as sess: + sess.run(variables.local_variables_initializer()) + self.assertAlmostEqual(0.8, sess.run(update_op)) + self.assertAlmostEqual(0.8, precision.eval()) + + def testSomeCorrectLowRecall(self): + predictions_values = [0.1, 0.2, 0.7, 0.3, 0.0, 0.1, 0.45, 0.5, 0.6, 0.9] + labels_values = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1] + + predictions = constant_op.constant( + predictions_values, dtype=dtypes_lib.float32) + labels = constant_op.constant(labels_values) + precision, update_op = metrics.precision_at_recall( + labels, predictions, target_recall=0.4) + + with self.test_session() as sess: + sess.run(variables.local_variables_initializer()) + self.assertAlmostEqual(2.0/3, sess.run(update_op)) + self.assertAlmostEqual(2.0/3, precision.eval()) + + def testWeighted_multipleLabelDtypes(self): + for label_dtype in (dtypes_lib.bool, dtypes_lib.int32, dtypes_lib.float32): + predictions_values = [ + 0.0, 0.1, 0.2, 0.3, 0.4, 0.1, 0.22, 0.25, 0.31, 0.35] + labels_values = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1] + weights_values = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + + predictions = constant_op.constant( + predictions_values, dtype=dtypes_lib.float32) + labels = math_ops.cast(labels_values, dtype=label_dtype) + weights = constant_op.constant(weights_values) + precision, update_op = metrics.precision_at_recall( + labels, predictions, target_recall=0.8, weights=weights) + + with self.test_session() as sess: + sess.run(variables.local_variables_initializer()) + self.assertAlmostEqual(34.0/43, sess.run(update_op)) + self.assertAlmostEqual(34.0/43, precision.eval()) + + class StreamingFNRThresholdsTest(test.TestCase): def setUp(self): diff --git a/tensorflow/contrib/model_pruning/BUILD b/tensorflow/contrib/model_pruning/BUILD index f50575b2cf311e33f7b7c77488bc94b8d24c70ec..54bd39afacbec07f054f61b72eda0a3654858aa7 100644 --- a/tensorflow/contrib/model_pruning/BUILD +++ b/tensorflow/contrib/model_pruning/BUILD @@ -71,6 +71,17 @@ py_library( ], ) +py_library( + name = "pruning_utils", + srcs = ["python/pruning_utils.py"], + srcs_version = "PY2AND3", + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/python:platform", + "//third_party/py/numpy", + ], +) + py_library( name = "pruning", srcs = ["python/pruning.py"], @@ -78,9 +89,20 @@ py_library( visibility = ["//visibility:public"], deps = [ ":core_layers", + ":pruning_utils", "//tensorflow/contrib/training:training_py", "//tensorflow/python:platform", - "//third_party/py/numpy", + ], +) + +py_test( + name = "pruning_utils_test", + size = "small", + srcs = ["python/pruning_utils_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":pruning_utils", + "//tensorflow/python:client_testlib", ], ) diff --git a/tensorflow/contrib/model_pruning/README.md b/tensorflow/contrib/model_pruning/README.md index 52b659c69fdfc507e6259e928d79c65471f2f025..86f4fd6adf60d8fa54c13989bf4087e28f1e006f 100644 --- a/tensorflow/contrib/model_pruning/README.md +++ b/tensorflow/contrib/model_pruning/README.md @@ -45,7 +45,7 @@ The pruning library allows for specification of the following hyper parameters: | do_not_prune | list of strings | [""] | list of layers names that are not pruned | | threshold_decay | float | 0.9 | The decay factor to use for exponential decay of the thresholds | | pruning_frequency | integer | 10 | How often should the masks be updated? (in # of global_steps) | -| nbins | integer | 255 | Number of bins to use for histogram computation | +| nbins | integer | 256 | Number of bins to use for histogram computation | | block_height|integer | 1 | Number of rows in a block for block sparse matrices| | block_width |integer | 1 | Number of cols in a block for block sparse matrices| | block_pooling_function| string | AVG | The function to use to pool weight values in a block: average (AVG) or max (MAX)| diff --git a/tensorflow/contrib/model_pruning/python/pruning.py b/tensorflow/contrib/model_pruning/python/pruning.py index 5146a4a2de7806041991c04958de378b2d3dc810..ea6032e588cf398deaf497fb99087436ce1cb2e8 100644 --- a/tensorflow/contrib/model_pruning/python/pruning.py +++ b/tensorflow/contrib/model_pruning/python/pruning.py @@ -33,12 +33,14 @@ # Returns a list of all the weight tensors that have been masked get_weights() - The Pruning class uses a proto (defined in pruning.proto) to set up the - parameters for a pruning specification. Here's a typical usage: + The Pruning class uses a tf.hparams object to set up the + parameters for a model pruning. Here's a typical usage: - # Initialize a pruning spec from a proto - pruning_spec = '/tmp/pruning.pb' - p = Pruning(pruning_spec) + # Parse pruning hyperparameters + pruning_hparams = pruning.get_pruning_hparams().parse(FLAGS.pruning_hparams) + + # Create a pruning object using the pruning_hparams + p = pruning.Pruning(pruning_hparams) # Add mask update ops to the graph mask_update_op = p.conditional_mask_update_op() @@ -51,24 +53,20 @@ # An object of the pruning also accepts externally defined sparsity: sparsity = tf.Variable(0.5, name = "ConstantSparsity") - pruning_spec = '/tmp/pruning.pb' - p = Pruning(pruning_spec, sparsity=sparsity) - + p = pruning.Pruning(pruning_hparams, sparsity=sparsity) """ # pylint: disable=missing-docstring from __future__ import absolute_import from __future__ import division from __future__ import print_function -import numpy as np - +from tensorflow.contrib.model_pruning.python import pruning_utils from tensorflow.contrib.model_pruning.python.layers import core_layers as core from tensorflow.contrib.training.python.training import hparam +from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops 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 gen_array_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn_impl @@ -87,172 +85,18 @@ _WEIGHT_COLLECTION = core.WEIGHT_COLLECTION _MASKED_WEIGHT_NAME = core.MASKED_WEIGHT_NAME -def _weight_mask_variable(var, scope): - """Create a mask for the weights. - - This function adds a variable 'mask' to the graph. - - Args: - var: the weight variable that needs to be masked - scope: The variable scope of the variable var - - Returns: - the mask variable of the same size and shape as var, initialized to all 1s. - """ - with variable_scope.variable_scope(scope): - mask = variable_scope.get_variable( - 'mask', - var.get_shape(), - initializer=init_ops.ones_initializer(), - trainable=False, - dtype=var.dtype) - return mask - - -def _weight_threshold_variable(var, scope): - """Create a scalar threshold for the weights. - - This function adds a variable - 'threshold' to the graph. - - Args: - var: The weight variable that needs to be masked - scope: The variable scope of the variable var - - Returns: - a scalar threshold variable initialized to 0. - """ - with variable_scope.variable_scope(scope): - threshold = variable_scope.get_variable( - 'threshold', [], - initializer=init_ops.zeros_initializer(), - trainable=False, - dtype=var.dtype) - return threshold - - -def _kronecker_product(mat1, mat2): - """Computes the Kronecker product of two matrices mat1 and mat2. - - Args: - mat1: A matrix of size m x n - mat2: A matrix of size p x q - Returns: - Kronecker product of matrices mat1 and mat2 of size mp x nq - """ - - m1, n1 = mat1.get_shape().as_list() - mat1_rsh = array_ops.reshape(mat1, [m1, 1, n1, 1]) - m2, n2 = mat2.get_shape().as_list() - mat2_rsh = array_ops.reshape(mat2, [1, m2, 1, n2]) - return array_ops.reshape(mat1_rsh * mat2_rsh, [m1 * m2, n1 * n2]) - - -def _histogram(values, value_range, nbins=100, dtype=np.int32, name=None): - """Return histogram of values. - - Given the tensor `values`, this operation returns a rank 1 histogram counting - the number of entries in `values` that fell into every bin. The bins are - equal width and determined by the arguments `value_range` and `nbins`. - - Args: - values: Numeric `Tensor`. - value_range: Shape [2] `Tensor` of same `dtype` as `values`. - values <= value_range[0] will be mapped to hist[0], - values >= value_range[1] will be mapped to hist[-1]. - nbins: Scalar `int32 Tensor`. Number of histogram bins. - dtype: dtype for returned histogram. - name: A name for this operation (defaults to 'histogram'). - - Returns: - A 1-D `Tensor` holding histogram of values. - - """ - with ops.name_scope(name, 'histogram', [values, value_range, nbins]) as scope: - values = ops.convert_to_tensor(values, name='values') - values = gen_array_ops.reshape(values, [-1]) - value_range = ops.convert_to_tensor(value_range, name='value_range') - nbins = ops.convert_to_tensor(nbins, dtype=np.int32, name='nbins') - nbins_float = math_ops.cast(nbins, values.dtype) - - # Map tensor values that fall within value_range to [0, 1]. - scaled_values = math_ops.truediv( - values - value_range[0], - value_range[1] - value_range[0], - name='scaled_values') - - # map tensor values within the open interval value_range to {0,.., nbins-1}, - # values outside the open interval will be zero or less, or nbins or more. - indices = math_ops.floor(nbins_float * scaled_values, name='indices') - - # Clip edge cases (e.g. value = value_range[1]) or "outliers." - indices = math_ops.cast( - clip_ops.clip_by_value(indices, 0, nbins_float - 1), np.int32) - - return math_ops.unsorted_segment_sum( - array_ops.ones_like(indices, dtype=dtype), indices, nbins, name=scope) - - -def _determine_partitioned_axis(partitioned_variable): - partitioned_axis = 0 - concatenated_variable_shape = partitioned_variable.get_shape() - for partition in partitioned_variable: - partition_shape = partition.get_shape() - maybe_partitioned_axis = np.less(partition_shape, - concatenated_variable_shape) - # Sanity check: make sure number of partitioned axis == 1 - if np.count_nonzero(maybe_partitioned_axis) != 1: - raise ValueError('Number of partitioned axes %s not equal to 1' % - np.count_nonzero(maybe_partitioned_axis)) - partitioned_axis = np.where(maybe_partitioned_axis)[0][0] - return partitioned_axis - - -def _variable_assign(var, new_value): - return state_ops.assign(var, new_value, name=var.op.name + '_assign') - - -def _partitioned_variable_assign(partitioned_var, new_value): - """Assign op for partitioned variables. - - Args: - partitioned_var: A partitioned tensorflow variable - new_value: Value to be assigned to the variable var - - Returns: - A tensorflow op that groups the assign ops for each of the variable slices - """ - # Determine which axis was used to partition the variable. Currently - # tensorflow allows partitioning variable only along 1 axis. - axis = 0 if len(partitioned_var) == 1 else _determine_partitioned_axis( - partitioned_var) - - partition_sizes = np.array( - [partition.get_shape()[axis] for partition in partitioned_var]) - new_partitioned_values = array_ops.split( - new_value, - ops.convert_to_tensor(partition_sizes, dtype=np.int32), - axis=axis) - op_list = [] - for partition in partitioned_var: - op_list.append( - _variable_assign(partition, new_partitioned_values[len(op_list)])) - return control_flow_ops.group( - *op_list, name=partitioned_var.name + '_group_assign') - - def apply_mask(x, scope=''): """Apply mask to a given weight tensor. Args: x: Input weight tensor - scope: The current variable scope. Defaults to "" + scope: The current variable scope. Defaults to "". Returns: Tensor representing masked_weights """ - mask = _weight_mask_variable(x, scope) - threshold = _weight_threshold_variable(x, scope) + mask = pruning_utils.weight_mask_variable(x, scope) + threshold = pruning_utils.weight_threshold_variable(x, scope) # Add masked_weights in the weights namescope so as to make it easier # for the quantization library to add quant ops. masked_weights = math_ops.multiply(mask, x, _MASKED_WEIGHT_NAME) @@ -335,6 +179,8 @@ def get_pruning_hparams(): sparsity_function_exponent: float exponent = 1 is linearly varying sparsity between initial and final. exponent > 1 varies more slowly towards the end than the beginning + use_tpu: False + Indicates whether to use TPU We use the following sparsity function: @@ -357,7 +203,7 @@ def get_pruning_hparams(): do_not_prune=[''], threshold_decay=0.9, pruning_frequency=10, - nbins=255, + nbins=256, block_height=1, block_width=1, block_pooling_function='AVG', @@ -365,7 +211,8 @@ def get_pruning_hparams(): target_sparsity=0.5, sparsity_function_begin_step=0, sparsity_function_end_step=100, - sparsity_function_exponent=3) + sparsity_function_exponent=3, + use_tpu=False) class Pruning(object): @@ -414,7 +261,7 @@ class Pruning(object): if graph_global_step is None: graph_global_step = training_util.get_global_step() - return math_ops.cast(graph_global_step, np.int32) + return math_ops.cast(graph_global_step, dtypes.int32) def _setup_sparsity(self): begin_step = self._spec.sparsity_function_begin_step @@ -429,13 +276,13 @@ class Pruning(object): (begin_step, end_step)) with ops.name_scope(self._spec.name): - p = math_ops.minimum(1.0, - math_ops.maximum( - 0.0, - math_ops.div( - math_ops.cast(self._global_step - begin_step, - np.float32), - end_step - begin_step))) + p = math_ops.minimum( + 1.0, + math_ops.maximum( + 0.0, + math_ops.div( + math_ops.cast(self._global_step - begin_step, dtypes.float32), + end_step - begin_step))) sparsity = math_ops.add( math_ops.multiply(initial_sparsity - target_sparsity, math_ops.pow(1 - p, exponent)), @@ -445,17 +292,18 @@ class Pruning(object): return sparsity def _setup_last_update_step(self): - with variable_scope.variable_scope(self._spec.name) as scope: + with variable_scope.variable_scope( + self._spec.name, use_resource=self._spec.use_tpu) as scope: try: last_update_step = variable_scope.get_variable( 'last_mask_update_step', [], initializer=init_ops.zeros_initializer(), trainable=False, - dtype=np.int32) + dtype=dtypes.int32) except ValueError: scope.reuse_variables() last_update_step = variable_scope.get_variable( - 'last_mask_update_step', dtype=np.int32) + 'last_mask_update_step', dtype=dtypes.int32) return last_update_step def _exists_in_do_not_prune_list(self, tensor_name): @@ -497,18 +345,16 @@ class Pruning(object): with ops.name_scope(weights.op.name + '_pruning_ops'): abs_weights = math_ops.abs(weights) max_value = math_ops.reduce_max(abs_weights) - histogram = _histogram( - abs_weights, [0.0, max_value], - nbins=self._spec.nbins, - dtype=np.float32) + cdf_fn = pruning_utils.compute_cdf_from_histogram + if self._spec.use_tpu: + cdf_fn = pruning_utils.compute_cdf - cdf = math_ops.cumsum(histogram) - norm_cdf = math_ops.div(cdf, math_ops.reduce_sum(histogram)) + norm_cdf = cdf_fn(abs_weights, [0.0, max_value], nbins=self._spec.nbins) current_threshold = math_ops.multiply( math_ops.div( math_ops.reduce_sum( math_ops.cast( - math_ops.less(norm_cdf, self._sparsity), np.float32)), + math_ops.less(norm_cdf, self._sparsity), dtypes.float32)), float(self._spec.nbins)), max_value) smoothed_threshold = math_ops.add_n([ @@ -516,7 +362,7 @@ class Pruning(object): math_ops.multiply(threshold, self._spec.threshold_decay) ]) new_mask = math_ops.cast( - math_ops.greater(abs_weights, smoothed_threshold), np.float32) + math_ops.greater(abs_weights, smoothed_threshold), dtypes.float32) return smoothed_threshold, new_mask def _maybe_update_block_mask(self, weights, threshold): @@ -572,8 +418,8 @@ class Pruning(object): new_mask, [pooled_weights.get_shape()[1], pooled_weights.get_shape()[2]]) - updated_mask = _kronecker_product(reshaped_mask, - array_ops.ones(self._block_dim)) + updated_mask = pruning_utils.kronecker_product( + reshaped_mask, array_ops.ones(self._block_dim)) sliced_mask = array_ops.slice( updated_mask, [0, 0], [squeezed_weights.get_shape()[0], @@ -608,11 +454,12 @@ class Pruning(object): continue new_threshold, new_mask = self._maybe_update_block_mask(weight, threshold) - self._assign_ops.append(_variable_assign(threshold, new_threshold)) + self._assign_ops.append( + pruning_utils.variable_assign(threshold, new_threshold)) self._assign_ops.append( - _partitioned_variable_assign(mask, new_mask) - if is_partitioned else _variable_assign(mask, new_mask)) + pruning_utils.partitioned_variable_assign(mask, new_mask) + if is_partitioned else pruning_utils.variable_assign(mask, new_mask)) def mask_update_op(self): with ops.name_scope(self._spec.name): diff --git a/tensorflow/contrib/model_pruning/python/pruning_test.py b/tensorflow/contrib/model_pruning/python/pruning_test.py index 89e65713197afc6ed37346cb67a6e9be3fa9290f..f80b7c52c000f13b5ce98dd442ff21abfac37761 100644 --- a/tensorflow/contrib/model_pruning/python/pruning_test.py +++ b/tensorflow/contrib/model_pruning/python/pruning_test.py @@ -110,12 +110,12 @@ class PruningTest(test.TestCase): self.assertAllEqual(np.count_nonzero(masked_weights_val), 100) session.run(mask_update_op) masked_weights_val = masked_weights.eval() - self.assertAllEqual(np.count_nonzero(masked_weights_val), 51) + self.assertAllEqual(np.count_nonzero(masked_weights_val), 50) def _blockMasking(self, hparams, weights, expected_mask): threshold = variables.Variable(0.0, name="threshold") - sparsity = variables.Variable(0.51, name="sparsity") + sparsity = variables.Variable(0.5, name="sparsity") test_spec = ",".join(hparams) pruning_hparams = pruning.get_pruning_hparams().parse(test_spec) @@ -138,7 +138,8 @@ class PruningTest(test.TestCase): 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]] + expected_mask = [[0.0, 0.0, 0.0, 0.0], [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) @@ -155,7 +156,8 @@ class PruningTest(test.TestCase): 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]]] + expected_mask = [[[0.0, 0.0, 0.0, 0.0], [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) @@ -178,11 +180,12 @@ class PruningTest(test.TestCase): masked_weights_val = masked_weights.eval() session.run(mask_update_op) masked_weights_val = masked_weights.eval() - self.assertAllEqual(np.count_nonzero(masked_weights_val), 51) + self.assertAllEqual(np.count_nonzero(masked_weights_val), 50) def testConditionalMaskUpdate(self): param_list = [ - "pruning_frequency=2", "begin_pruning_step=1", "end_pruning_step=6" + "pruning_frequency=2", "begin_pruning_step=1", "end_pruning_step=6", + "nbins=100" ] test_spec = ",".join(param_list) pruning_hparams = pruning.get_pruning_hparams().parse(test_spec) diff --git a/tensorflow/contrib/model_pruning/python/pruning_utils.py b/tensorflow/contrib/model_pruning/python/pruning_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..56d3dcef20d1b1c34d6b04535e2b4dc7be7f7320 --- /dev/null +++ b/tensorflow/contrib/model_pruning/python/pruning_utils.py @@ -0,0 +1,269 @@ +# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Utility functions for adding pruning related ops to the graph. +""" +# pylint: disable=missing-docstring +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 +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 init_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import state_ops +from tensorflow.python.ops import variable_scope + +_NBINS = 256 + + +def weight_mask_variable(var, scope): + """Create a mask for the weights. + + This function adds a variable 'mask' to the graph. + + Args: + var: the weight variable that needs to be masked + scope: The variable scope of the variable var + + Returns: + the mask variable of the same size and shape as var, initialized to all 1s. + """ + with variable_scope.variable_scope(scope): + mask = variable_scope.get_variable( + 'mask', + var.get_shape(), + initializer=init_ops.ones_initializer(), + trainable=False, + dtype=var.dtype) + return mask + + +def weight_threshold_variable(var, scope): + """Create a scalar threshold for the weights. + + This function adds a variable + 'threshold' to the graph. + + Args: + var: The weight variable that needs to be masked + scope: The variable scope of the variable var + + Returns: + a scalar threshold variable initialized to 0. + """ + with variable_scope.variable_scope(scope): + threshold = variable_scope.get_variable( + 'threshold', [], + initializer=init_ops.zeros_initializer(), + trainable=False, + dtype=var.dtype) + return threshold + + +def kronecker_product(mat1, mat2): + """Computes the Kronecker product of two matrices mat1 and mat2. + + Args: + mat1: A matrix of size m x n + mat2: A matrix of size p x q + Returns: + Kronecker product of matrices mat1 and mat2 of size mp x nq + """ + + m1, n1 = mat1.get_shape().as_list() + mat1_rsh = array_ops.reshape(mat1, [m1, 1, n1, 1]) + m2, n2 = mat2.get_shape().as_list() + mat2_rsh = array_ops.reshape(mat2, [1, m2, 1, n2]) + return array_ops.reshape(mat1_rsh * mat2_rsh, [m1 * m2, n1 * n2]) + + +def _histogram(values, value_range, nbins=100, dtype=dtypes.int32, name=None): + """Return histogram of values. + + Given the tensor `values`, this operation returns a rank 1 histogram counting + the number of entries in `values` that fell into every bin. The bins are + equal width and determined by the arguments `value_range` and `nbins`. + + Args: + values: Numeric `Tensor`. + value_range: Shape [2] `Tensor` of same `dtype` as `values`. + values <= value_range[0] will be mapped to hist[0], + values >= value_range[1] will be mapped to hist[-1]. + nbins: Scalar `int32 Tensor`. Number of histogram bins. + dtype: dtype for returned histogram. + name: A name for this operation (defaults to 'histogram'). + + Returns: + A 1-D `Tensor` holding histogram of values. + + """ + with ops.name_scope(name, 'histogram', [values, value_range, nbins]) as scope: + values = ops.convert_to_tensor(values, name='values') + values = array_ops.reshape(values, [-1]) + value_range = ops.convert_to_tensor(value_range, name='value_range') + nbins_float = np.float32(nbins) + + # Map tensor values that fall within value_range to [0, 1]. + scaled_values = math_ops.truediv( + values - value_range[0], + value_range[1] - value_range[0], + name='scaled_values') + + # map tensor values within the open interval value_range to {0,.., nbins-1}, + # values outside the open interval will be zero or less, or nbins or more. + indices = math_ops.floor(nbins_float * scaled_values, name='indices') + + # Clip edge cases (e.g. value = value_range[1]) or "outliers." + indices = math_ops.cast( + clip_ops.clip_by_value(indices, 0, nbins_float - 1), dtypes.int32) + + return math_ops.unsorted_segment_sum( + array_ops.ones_like(indices, dtype=dtype), indices, nbins, name=scope) + + +def compute_cdf_from_histogram(values, value_range, **kwargs): + """Returns the normalized cumulative distribution of the given values tensor. + + Computes the histogram and uses tf.cumsum to arrive at cdf + + Args: + values: Numeric `Tensor`. + value_range: Shape [2] `Tensor` of same `dtype` as `values`. + **kwargs: keyword arguments: nbins, name + + Returns: + A 1-D `Tensor` holding normalized cdf of values. + + """ + nbins = kwargs.get('nbins', _NBINS) + name = kwargs.get('name', None) + with ops.name_scope(name, 'cdf', [values, value_range, nbins]): + histogram = _histogram( + values, value_range, dtype=dtypes.float32, nbins=nbins) + cdf = math_ops.cumsum(histogram) + return math_ops.div(cdf, math_ops.reduce_max(cdf)) + + +def compute_cdf(values, value_range, **kwargs): + """Returns the normalized cumulative distribution of the given values tensor. + + Uses tf.while_loop to directly compute the cdf of the values. Number of bins + for histogram is fixed at _NBINS=255 + + Args: + values: Numeric `Tensor`. + value_range: Shape [2] `Tensor` of same `dtype` as `values` + **kwargs: keyword arguments: name + + Returns: + A 1-D `Tensor` holding normalized cdf of values. + + """ + nbins = _NBINS + name = kwargs.get('name', None) + with ops.name_scope(name, 'cdf', [values, value_range, nbins]): + values = ops.convert_to_tensor(values, name='values') + value_range = ops.convert_to_tensor(value_range, name='value_range') + nbins_float = np.float32(nbins) + + # Map tensor values that fall within value_range to [0, 1]. + scaled_values = math_ops.truediv( + values - value_range[0], + value_range[1] - value_range[0], + name='scaled_values') + + # map tensor values within the open interval value_range to {0,.., nbins-1}, + # values outside the open interval will be zero or less, or nbins or more. + indices = math_ops.floor(nbins_float * scaled_values, name='indices') + + # Clip edge cases (e.g. value = value_range[1]) or "outliers." + indices = math_ops.cast( + clip_ops.clip_by_value(indices, 0, nbins_float - 1), dtypes.int32) + + cdf = array_ops.zeros(nbins) + i = constant_op.constant(0) + + def loop_cond(loop_count, _): + return math_ops.less(loop_count, nbins) + + def loop_body(loop_count, cdf): + temp = math_ops.reduce_sum( + math_ops.cast( + math_ops.less_equal(indices, loop_count), dtypes.float32)) + cdf = math_ops.add( + cdf, + array_ops.one_hot( + loop_count, depth=_NBINS, on_value=temp, off_value=0.0)) + return [loop_count + 1, cdf] + + _, cdf = control_flow_ops.while_loop( + loop_cond, loop_body, [i, cdf], maximum_iterations=nbins) + + return math_ops.div(cdf, math_ops.reduce_max(cdf)) + + +def determine_partitioned_axis(partitioned_variable): + partitioned_axis = 0 + concatenated_variable_shape = partitioned_variable.get_shape() + for partition in partitioned_variable: + partition_shape = partition.get_shape() + maybe_partitioned_axis = np.less(partition_shape, + concatenated_variable_shape) + # Sanity check: make sure number of partitioned axis == 1 + if np.count_nonzero(maybe_partitioned_axis) != 1: + raise ValueError('Number of partitioned axes %s not equal to 1' % + np.count_nonzero(maybe_partitioned_axis)) + partitioned_axis = np.where(maybe_partitioned_axis)[0][0] + return partitioned_axis + + +def variable_assign(var, new_value): + return state_ops.assign(var, new_value, name=var.op.name + '_assign') + + +def partitioned_variable_assign(partitioned_var, new_value): + """Assign op for partitioned variables. + + Args: + partitioned_var: A partitioned tensorflow variable + new_value: Value to be assigned to the variable var + + Returns: + A tensorflow op that groups the assign ops for each of the variable slices + """ + # Determine which axis was used to partition the variable. Currently + # tensorflow allows partitioning variable only along 1 axis. + axis = 0 if len(partitioned_var) == 1 else determine_partitioned_axis( + partitioned_var) + + partition_sizes = np.array( + [partition.get_shape()[axis] for partition in partitioned_var]) + new_partitioned_values = array_ops.split( + new_value, + ops.convert_to_tensor(partition_sizes, dtype=dtypes.int32), + axis=axis) + op_list = [] + for partition in partitioned_var: + op_list.append( + variable_assign(partition, new_partitioned_values[len(op_list)])) + return control_flow_ops.group( + *op_list, name=partitioned_var.name + '_group_assign') diff --git a/tensorflow/contrib/model_pruning/python/pruning_utils_test.py b/tensorflow/contrib/model_pruning/python/pruning_utils_test.py new file mode 100644 index 0000000000000000000000000000000000000000..10e1dd0a8eee88f357fbe60bf00f180c05f2c4d2 --- /dev/null +++ b/tensorflow/contrib/model_pruning/python/pruning_utils_test.py @@ -0,0 +1,86 @@ +# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for utility functions in pruning_utils.py.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.model_pruning.python import pruning_utils +from tensorflow.python.framework import constant_op +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 + + +class PruningUtilsTest(test.TestCase): + + def testHistogram(self): + width = 10 + height = 10 + nbins = 100 + expected_histogram = np.full(nbins, 1.0) + init = init_ops.constant_initializer(np.linspace(0.0, 1.0, width * height)) + weights = variable_scope.get_variable( + "weights", [width, height], initializer=init) + histogram = pruning_utils._histogram( + weights, [0, 1.0], nbins, dtype=np.float32) + with self.test_session(): + variables.global_variables_initializer().run() + computed_histogram = histogram.eval() + self.assertAllEqual(expected_histogram, computed_histogram) + + def testCDF(self): + nbins = 5 + weights = constant_op.constant([-1, 0, 1, 1.5, 2, 3, 4, 5, 10, 100]) + abs_weights = math_ops.abs(weights) + norm_cdf = pruning_utils.compute_cdf_from_histogram( + abs_weights, [0.0, 5.0], nbins=nbins) + expected_cdf = np.array([0.1, 0.4, 0.5, 0.6, 1.0], dtype=np.float32) + with self.test_session() as sess: + variables.global_variables_initializer().run() + norm_cdf_val = sess.run(norm_cdf) + self.assertAllEqual(len(norm_cdf_val), nbins) + self.assertAllEqual(expected_cdf, norm_cdf_val) + + def _compare_cdf(self, values): + abs_values = math_ops.abs(values) + max_value = math_ops.reduce_max(abs_values) + with self.test_session(): + variables.global_variables_initializer().run() + cdf_from_histogram = pruning_utils.compute_cdf_from_histogram( + abs_values, [0.0, max_value], nbins=pruning_utils._NBINS) + cdf = pruning_utils.compute_cdf(abs_values, [0.0, max_value]) + return cdf.eval(), cdf_from_histogram.eval() + + def testCDFEquivalence2D(self): + width = 100 + height = 100 + weights = variable_scope.get_variable("weights", shape=[width, height]) + cdf_val, cdf_from_histogram_val = self._compare_cdf(weights) + self.assertAllEqual(cdf_val, cdf_from_histogram_val) + + def testCDFEquivalence4D(self): + weights = variable_scope.get_variable("weights", shape=[5, 5, 128, 128]) + cdf_val, cdf_from_histogram_val = self._compare_cdf(weights) + self.assertAllEqual(cdf_val, cdf_from_histogram_val) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/nn/python/ops/sampling_ops.py b/tensorflow/contrib/nn/python/ops/sampling_ops.py index 63fc487dca69a4777821595a0366d0ae0b393ce2..e65925610c5f5125c2d2e92edc1cf708c54255d4 100644 --- a/tensorflow/contrib/nn/python/ops/sampling_ops.py +++ b/tensorflow/contrib/nn/python/ops/sampling_ops.py @@ -88,7 +88,7 @@ def _rank_resample(weights, biases, inputs, sampled_values, num_resampled, return math_ops.reduce_logsumexp( math_ops.matmul(embeddings, reweighted_inputs, transpose_b=True), axis=1, - keep_dims=False) + keepdims=False) # Calling this protected form of embedding_lookup allows co-locating # the logsumexp computation with the partitioned weights, which yields diff --git a/tensorflow/contrib/opt/BUILD b/tensorflow/contrib/opt/BUILD index c57c5e3f29f1e36fa2f36f8113cb208be6c6be3e..612ecc3e63891f4dabf97828fe75672dd7877a91 100644 --- a/tensorflow/contrib/opt/BUILD +++ b/tensorflow/contrib/opt/BUILD @@ -14,6 +14,7 @@ py_library( name = "opt_py", srcs = [ "__init__.py", + "python/training/adamax.py", "python/training/addsign.py", "python/training/drop_stale_gradient_optimizer.py", "python/training/elastic_average_optimizer.py", @@ -43,11 +44,27 @@ py_library( "//tensorflow/python:util", "//tensorflow/python:variable_scope", "//tensorflow/python:variables", + "//tensorflow/python/eager:context", "//third_party/py/numpy", "@six_archive//:six", ], ) +py_test( + name = "adamax_test", + srcs = ["python/training/adamax_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":opt_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:math_ops", + "//tensorflow/python:training", + "//third_party/py/numpy", + ], +) + py_test( name = "external_optimizer_test", srcs = ["python/training/external_optimizer_test.py"], diff --git a/tensorflow/contrib/opt/__init__.py b/tensorflow/contrib/opt/__init__.py index 6c1bb1adc096f5b8e6945ea1492727d16cf29e65..4c13c8e247185213b798eb733ddcf65a07a8f64d 100644 --- a/tensorflow/contrib/opt/__init__.py +++ b/tensorflow/contrib/opt/__init__.py @@ -19,6 +19,7 @@ from __future__ import division from __future__ import print_function # pylint: disable=wildcard-import +from tensorflow.contrib.opt.python.training.adamax import * from tensorflow.contrib.opt.python.training.addsign import * from tensorflow.contrib.opt.python.training.drop_stale_gradient_optimizer import * from tensorflow.contrib.opt.python.training.external_optimizer import * @@ -36,6 +37,7 @@ from tensorflow.python.util.all_util import remove_undocumented _allowed_symbols = [ + 'AdaMaxOptimizer', 'PowerSignOptimizer', 'AddSignOptimizer', 'DelayCompensatedGradientDescentOptimizer', diff --git a/tensorflow/contrib/opt/python/training/adamax.py b/tensorflow/contrib/opt/python/training/adamax.py new file mode 100644 index 0000000000000000000000000000000000000000..686bac0d8408cdb0d2937bbd1de8535e216f8bf0 --- /dev/null +++ b/tensorflow/contrib/opt/python/training/adamax.py @@ -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. +# ============================================================================== + +"""AdaMax for TensorFlow.""" +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 control_flow_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import resource_variable_ops +from tensorflow.python.ops import state_ops +from tensorflow.python.training import adam +from tensorflow.python.training import training_ops + + +class AdaMaxOptimizer(adam.AdamOptimizer): + """Optimizer that implements the AdaMax algorithm. + + Adamax is sometimes superior to adam, specially in models with embeddings, + see [Kingma et al., 2014](http://arxiv.org/abs/1412.6980) + ([pdf](http://arxiv.org/pdf/1412.6980.pdf)). + """ + + def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8, + use_locking=False, name="AdaMax"): + """Construct a new AdaMax optimizer. + + Initialization: + + ``` + m_0 <- 0 (Initialize initial 1st moment vector) + v_0 <- 0 (Initialize the exponentially weighted infinity norm) + t <- 0 (Initialize timestep) + ``` + + The update rule for `variable` with gradient `g` uses an optimization + described at the end of section 7.1 of the paper: + + ``` + t <- t + 1 + + m_t <- beta1 * m_{t-1} + (1 - beta1) * g + v_t <- max(beta2 * v_{t-1}, abs(g)) + variable <- variable - learning_rate / (1 - beta1^t) * m_t / (v_t + epsilon) + ``` + + Similar to AdamOptimizer, the epsilon is added for numerical stability + (especially to get rid of division by zero when v_t = 0). + + Contrast to AdamOptimizer, the sparse implementation of this algorithm + (used when the gradient is an IndexedSlices object, typically because of + `tf.gather` or an embedding lookup in the forward pass) only updates + variable slices and corresponding `m_t`, `v_t` terms when that part of + the variable was used in the forward pass. This means that the sparse + behavior is contrast to the dense behavior (similar to some momentum + implementations which ignore momentum unless a variable slice was actually + used). + + Args: + learning_rate: A Tensor or a floating point value. The learning rate. + beta1: A float value or a constant float tensor. + The exponential decay rate for the 1st moment estimates. + beta2: A float value or a constant float tensor. + The exponential decay rate for the exponentially weighted infinity norm. + epsilon: A small constant for numerical stability. + use_locking: If True use locks for update operations. + name: Optional name for the operations created when applying gradients. + Defaults to "AdaMax". + """ + super(AdaMaxOptimizer, self).__init__(learning_rate, beta1, beta2, + epsilon, use_locking, name) + + def _get_beta_accumulators(self): + if context.executing_eagerly(): + graph = None + else: + graph = ops.get_default_graph() + return self._get_non_slot_variable("beta1_power", graph=graph) + + def _create_slots(self, var_list): + # Create the beta1 accumulators on the same device as the first + # variable. Sort the var_list to make sure this device is consistent across + # workers (these need to go on the same PS, otherwise some updates are + # silently ignored). + first_var = min(var_list, key=lambda x: x.name) + self._create_non_slot_variable(initial_value=self._beta1, + name="beta1_power", + colocate_with=first_var) + + # Create slots for the first and second moments. + for v in var_list: + self._zeros_slot(v, "m", self._name) + self._zeros_slot(v, "v", self._name) + + def _apply_dense(self, grad, var): + m = self.get_slot(var, "m") + v = self.get_slot(var, "v") + beta1_power = self._get_beta_accumulators() + return training_ops.apply_ada_max( + var, m, v, + math_ops.cast(beta1_power, var.dtype.base_dtype), + math_ops.cast(self._lr_t, var.dtype.base_dtype), + math_ops.cast(self._beta1_t, var.dtype.base_dtype), + math_ops.cast(self._beta2_t, var.dtype.base_dtype), + math_ops.cast(self._epsilon_t, var.dtype.base_dtype), + grad, use_locking=self._use_locking).op + + def _resource_apply_dense(self, grad, var): + m = self.get_slot(var, "m") + v = self.get_slot(var, "v") + beta1_power = self._get_beta_accumulators() + return training_ops.resource_apply_ada_max( + var.handle, m.handle, v.handle, + math_ops.cast(beta1_power, grad.dtype.base_dtype), + math_ops.cast(self._lr_t, grad.dtype.base_dtype), + math_ops.cast(self._beta1_t, grad.dtype.base_dtype), + math_ops.cast(self._beta2_t, grad.dtype.base_dtype), + math_ops.cast(self._epsilon_t, grad.dtype.base_dtype), + grad, use_locking=self._use_locking) + + def _apply_sparse_shared(self, grad, var, indices, + scatter_add, scatter_update): + beta1_power = self._get_beta_accumulators() + beta1_power = math_ops.cast(beta1_power, var.dtype.base_dtype) + lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype) + beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype) + beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype) + epsilon_t = math_ops.cast(self._epsilon_t, var.dtype.base_dtype) + # m_t = beta1 * m + (1 - beta1) * g_t + m = self.get_slot(var, "m") + m_slice = array_ops.gather(m, indices) + m_t_slice = m_slice * beta1_t + grad * (1 - beta1_t) + with ops.control_dependencies([m_t_slice]): + m_t = scatter_update(m, indices, m_t_slice) + # u_t = max(beta2 * u, abs(g_t)) + v = self.get_slot(var, "v") + v_slice = array_ops.gather(v, indices) + v_t_slice = math_ops.maximum(v_slice * beta2_t, math_ops.abs(grad)) + with ops.control_dependencies([v_t_slice]): + v_t = scatter_update(v, indices, v_t_slice) + # theta_t = theta - lr / (1 - beta1^t) * m_t / u_t + var_slice = -lr_t / (1 - beta1_power) * (m_t_slice / + (v_t_slice + epsilon_t)) + with ops.control_dependencies([var_slice]): + var_update = scatter_add(var, indices, var_slice) + return control_flow_ops.group(*[var_update, m_t, v_t]) + + def _apply_sparse(self, grad, var): + return self._apply_sparse_shared( + grad.values, var, grad.indices, + lambda x, i, v: state_ops.scatter_add( # pylint: disable=g-long-lambda + x, i, v, use_locking=self._use_locking), + lambda x, i, v: state_ops.scatter_update( # pylint: disable=g-long-lambda + x, i, v, use_locking=self._use_locking)) + + def _resource_scatter_update(self, x, i, v): + with ops.control_dependencies( + [resource_variable_ops.resource_scatter_update( + x.handle, i, v)]): + return x.value() + + def _resource_apply_sparse(self, grad, var, indices): + return self._apply_sparse_shared( + grad, var, indices, + self._resource_scatter_add, self._resource_scatter_update) + + def _finish(self, update_ops, name_scope): + # Update the power accumulators. + with ops.control_dependencies(update_ops): + beta1_power = self._get_beta_accumulators() + with ops.colocate_with(beta1_power): + update_beta1 = beta1_power.assign( + beta1_power * self._beta1_t, use_locking=self._use_locking) + return control_flow_ops.group(*update_ops + [update_beta1], + name=name_scope) diff --git a/tensorflow/contrib/opt/python/training/adamax_test.py b/tensorflow/contrib/opt/python/training/adamax_test.py new file mode 100644 index 0000000000000000000000000000000000000000..bc92a7006f1a0a56adafc486a75afa94e965cb2c --- /dev/null +++ b/tensorflow/contrib/opt/python/training/adamax_test.py @@ -0,0 +1,348 @@ +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 AdaMax.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.opt.python.training import adamax +from tensorflow.python.client import session +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 test_util +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 + + +def adamax_update_numpy(param, + g_t, + t, + m, + v, + alpha=0.001, + beta1=0.9, + beta2=0.999, + epsilon=1e-8): + m_t = beta1 * m + (1 - beta1) * g_t + v_t = np.maximum(beta2 * v, np.abs(g_t)) + param_t = param - (alpha / (1 - beta1**t)) * (m_t / (v_t + epsilon)) + return param_t, m_t, v_t + + +def adamax_sparse_update_numpy(param, + indices, + g_t, + t, + m, + v, + alpha=0.001, + beta1=0.9, + beta2=0.999, + epsilon=1e-8): + m_t, v_t, param_t = np.copy(m), np.copy(v), np.copy(param) + m_t_slice = beta1 * m[indices] + (1 - beta1) * g_t + v_t_slice = np.maximum(beta2 * v[indices], np.abs(g_t)) + param_t_slice = param[indices] - ((alpha / (1 - beta1**t)) * + (m_t_slice / (v_t_slice + epsilon))) + m_t[indices] = m_t_slice + v_t[indices] = v_t_slice + param_t[indices] = param_t_slice + return param_t, m_t, v_t + + +class AdaMaxOptimizerTest(test.TestCase): + + def doTestSparse(self, use_resource=False): + for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: + with self.test_session(): + # Initialize variables for numpy implementation. + zero_slots = lambda: np.zeros((3), dtype=dtype.as_numpy_dtype) + m0, v0, m1, v1 = zero_slots(), zero_slots(), zero_slots(), zero_slots() + var0_np = np.array([1.0, 2.0, 3.0], dtype=dtype.as_numpy_dtype) + grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype) + var1_np = np.array([4.0, 5.0, 6.0], dtype=dtype.as_numpy_dtype) + grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype) + + if use_resource: + var0 = resource_variable_ops.ResourceVariable(var0_np) + var1 = resource_variable_ops.ResourceVariable(var1_np) + else: + var0 = variables.Variable(var0_np) + var1 = variables.Variable(var1_np) + grads0_np_indices = np.array([0, 1], dtype=np.int32) + grads0 = ops.IndexedSlices( + constant_op.constant(grads0_np), + constant_op.constant(grads0_np_indices), constant_op.constant([2])) + grads1_np_indices = np.array([2, 1], dtype=np.int32) + grads1 = ops.IndexedSlices( + constant_op.constant(grads1_np), + constant_op.constant(grads1_np_indices), constant_op.constant([2])) + opt = adamax.AdaMaxOptimizer() + update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) + variables.global_variables_initializer().run() + + # Fetch params to validate initial values + self.assertAllClose([1.0, 2.0, 3.0], var0.eval()) + self.assertAllClose([4.0, 5.0, 6.0], var1.eval()) + + beta1_power = opt._get_beta_accumulators() + + # Run 3 steps of AdaMax + for t in range(1, 4): + self.assertAllCloseAccordingToType(0.9**t, beta1_power.eval()) + update.run() + + var0_np, m0, v0 = adamax_sparse_update_numpy( + var0_np, grads0_np_indices, grads0_np, t, m0, v0) + var1_np, m1, v1 = adamax_sparse_update_numpy( + var1_np, grads1_np_indices, grads1_np, t, m1, v1) + + # Validate updated params + self.assertAllCloseAccordingToType(var0_np, var0.eval()) + self.assertAllCloseAccordingToType(var1_np, var1.eval()) + + def testSparse(self): + self.doTestSparse(use_resource=False) + + def testResourceSparse(self): + self.doTestSparse(use_resource=True) + + def testSparseDevicePlacement(self): + for index_dtype in [dtypes.int32, dtypes.int64]: + with self.test_session(force_gpu=test.is_gpu_available()): + # If a GPU is available, tests that all optimizer ops can be placed on + # it (i.e. they have GPU kernels). + var = variables.Variable([[1.0], [2.0]]) + indices = constant_op.constant([0, 1], dtype=index_dtype) + gathered_sum = math_ops.reduce_sum(array_ops.gather(var, indices)) + optimizer = adamax.AdaMaxOptimizer(3.0) + minimize_op = optimizer.minimize(gathered_sum) + variables.global_variables_initializer().run() + minimize_op.run() + + def testSparseRepeatedIndices(self): + for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: + with self.test_session(): + repeated_index_update_var = variables.Variable( + [[1.0], [2.0]], dtype=dtype) + aggregated_update_var = variables.Variable( + [[1.0], [2.0]], dtype=dtype) + grad_repeated_index = ops.IndexedSlices( + constant_op.constant( + [0.1, 0.1], shape=[2, 1], dtype=dtype), + constant_op.constant([1, 1]), + constant_op.constant([2, 1])) + grad_aggregated = ops.IndexedSlices( + constant_op.constant( + [0.2], shape=[1, 1], dtype=dtype), + constant_op.constant([1]), + constant_op.constant([2, 1])) + repeated_update = adamax.AdaMaxOptimizer().apply_gradients( + [(grad_repeated_index, repeated_index_update_var)]) + aggregated_update = adamax.AdaMaxOptimizer().apply_gradients( + [(grad_aggregated, aggregated_update_var)]) + variables.global_variables_initializer().run() + self.assertAllClose(aggregated_update_var.eval(), + repeated_index_update_var.eval()) + for _ in range(3): + repeated_update.run() + aggregated_update.run() + self.assertAllClose(aggregated_update_var.eval(), + repeated_index_update_var.eval()) + + def doTestBasic(self, use_resource=False): + for i, dtype in enumerate([dtypes.half, dtypes.float32, dtypes.float64]): + with self.test_session(graph=ops.Graph()): + # Initialize variables for numpy implementation. + m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0 + var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype) + grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype) + var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype) + grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype) + + if use_resource: + var0 = resource_variable_ops.ResourceVariable( + var0_np, name="var0_%d" % i) + var1 = resource_variable_ops.ResourceVariable( + var1_np, name="var1_%d" % i) + else: + var0 = variables.Variable(var0_np) + var1 = variables.Variable(var1_np) + grads0 = constant_op.constant(grads0_np) + grads1 = constant_op.constant(grads1_np) + + opt = adamax.AdaMaxOptimizer() + update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) + opt_variables = opt.variables() + beta1_power = opt._get_beta_accumulators() + self.assertTrue(beta1_power is not None) + self.assertIn(beta1_power, opt_variables) + + with ops.Graph().as_default(): + # Shouldn't return non-slot variables from other graphs. + self.assertEqual(0, len(opt.variables())) + + 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)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) + + beta1_power = opt._get_beta_accumulators() + + # Run 3 steps of AdaMax + for t in range(1, 4): + if not context.executing_eagerly(): + self.evaluate(update) + elif t > 1: + opt.apply_gradients(zip([grads0, grads1], [var0, var1])) + + self.assertAllCloseAccordingToType(0.9**(t + 1), + self.evaluate(beta1_power)) + + var0_np, m0, v0 = adamax_update_numpy(var0_np, grads0_np, t, m0, v0) + var1_np, m1, v1 = adamax_update_numpy(var1_np, grads1_np, t, m1, v1) + + # Validate updated params + self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0)) + self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1)) + if use_resource: + self.assertEqual("var0_%d/AdaMax:0" % (i,), + opt.get_slot(var=var0, name="m").name) + + def testBasic(self): + with self.test_session(): + self.doTestBasic(use_resource=False) + + @test_util.run_in_graph_and_eager_modes(reset_test=True) + def testResourceBasic(self): + self.doTestBasic(use_resource=True) + + def testTensorLearningRate(self): + for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: + with self.test_session(): + # Initialize variables for numpy implementation. + m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0 + var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype) + grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype) + var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype) + grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype) + + var0 = variables.Variable(var0_np) + var1 = variables.Variable(var1_np) + grads0 = constant_op.constant(grads0_np) + grads1 = constant_op.constant(grads1_np) + opt = adamax.AdaMaxOptimizer(constant_op.constant(0.001)) + update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) + variables.global_variables_initializer().run() + + # Fetch params to validate initial values + self.assertAllClose([1.0, 2.0], var0.eval()) + self.assertAllClose([3.0, 4.0], var1.eval()) + + beta1_power = opt._get_beta_accumulators() + + # Run 3 steps of AdaMax + for t in range(1, 4): + self.assertAllCloseAccordingToType(0.9**t, beta1_power.eval()) + update.run() + + var0_np, m0, v0 = adamax_update_numpy(var0_np, grads0_np, t, m0, v0) + var1_np, m1, v1 = adamax_update_numpy(var1_np, grads1_np, t, m1, v1) + + # Validate updated params + self.assertAllCloseAccordingToType(var0_np, var0.eval()) + self.assertAllCloseAccordingToType(var1_np, var1.eval()) + + def testSharing(self): + for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: + with self.test_session(): + # Initialize variables for numpy implementation. + m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0 + var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype) + grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype) + var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype) + grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype) + + var0 = variables.Variable(var0_np) + var1 = variables.Variable(var1_np) + grads0 = constant_op.constant(grads0_np) + grads1 = constant_op.constant(grads1_np) + opt = adamax.AdaMaxOptimizer() + update1 = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) + update2 = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) + variables.global_variables_initializer().run() + + beta1_power = opt._get_beta_accumulators() + + # Fetch params to validate initial values + self.assertAllClose([1.0, 2.0], var0.eval()) + self.assertAllClose([3.0, 4.0], var1.eval()) + + # Run 3 steps of intertwined AdaMax1 and AdaMax2. + for t in range(1, 4): + self.assertAllCloseAccordingToType(0.9**t, beta1_power.eval()) + if t % 2 == 0: + update1.run() + else: + update2.run() + + var0_np, m0, v0 = adamax_update_numpy(var0_np, grads0_np, t, m0, v0) + var1_np, m1, v1 = adamax_update_numpy(var1_np, grads1_np, t, m1, v1) + + # Validate updated params + self.assertAllCloseAccordingToType(var0_np, var0.eval()) + self.assertAllCloseAccordingToType(var1_np, var1.eval()) + + def testTwoSessions(self): + optimizer = adamax.AdaMaxOptimizer() + g = ops.Graph() + with g.as_default(): + with session.Session(): + var0 = variables.Variable(np.array([1.0, 2.0]), name="v0") + grads0 = constant_op.constant(np.array([0.1, 0.1])) + optimizer.apply_gradients([(grads0, var0)]) + + gg = ops.Graph() + with gg.as_default(): + with session.Session(): + var0 = variables.Variable(np.array([1.0, 2.0]), name="v0") + grads0 = constant_op.constant(np.array([0.1, 0.1])) + + # If the optimizer saves any state not keyed by graph the following line + # 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 = adamax.AdaMaxOptimizer(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(5, len(set(opt.variables()))) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/opt/python/training/lazy_adam_optimizer.py b/tensorflow/contrib/opt/python/training/lazy_adam_optimizer.py index aeca900bc8ff4c4cc26da490ce43dfec70fd9f11..72117c1e81a164b0517fabeaddec3ea5132af5a9 100644 --- a/tensorflow/contrib/opt/python/training/lazy_adam_optimizer.py +++ b/tensorflow/contrib/opt/python/training/lazy_adam_optimizer.py @@ -56,21 +56,21 @@ class LazyAdamOptimizer(adam.AdamOptimizer): epsilon_t = math_ops.cast(self._epsilon_t, var.dtype.base_dtype) lr = (lr_t * math_ops.sqrt(1 - beta2_power) / (1 - beta1_power)) - # m := beta1 * m + (1 - beta1) * g_t + # \\(m := beta1 * m + (1 - beta1) * g_t\\) m = self.get_slot(var, "m") m_t = state_ops.scatter_update(m, grad.indices, beta1_t * array_ops.gather(m, grad.indices) + (1 - beta1_t) * grad.values, use_locking=self._use_locking) - # v := beta2 * v + (1 - beta2) * (g_t * g_t) + # \\(v := beta2 * v + (1 - beta2) * (g_t * g_t)\\) v = self.get_slot(var, "v") v_t = state_ops.scatter_update(v, grad.indices, beta2_t * array_ops.gather(v, grad.indices) + (1 - beta2_t) * math_ops.square(grad.values), use_locking=self._use_locking) - # variable -= learning_rate * m_t / (epsilon_t + sqrt(v_t)) + # \\(variable -= learning_rate * m_t / (epsilon_t + sqrt(v_t))\\) m_t_slice = array_ops.gather(m_t, grad.indices) v_t_slice = array_ops.gather(v_t, grad.indices) denominator_slice = math_ops.sqrt(v_t_slice) + epsilon_t 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 85e3e8d3791f2331ed249c0b7f67a3dbde4fca08..ac04ad99110b016b62e091aa10c7f565e5093bc1 100644 --- a/tensorflow/contrib/opt/python/training/moving_average_optimizer_test.py +++ b/tensorflow/contrib/opt/python/training/moving_average_optimizer_test.py @@ -85,7 +85,7 @@ class MovingAverageOptimizerTest(test.TestCase): state_ops.assign_add(ema_var1, [4.0, 4.0]) ]) - # Test taht saver with missing ema variables will fail. + # Test that saver with missing ema variables will fail. with self.assertRaisesRegexp(ValueError, r'Variable to swap'): opt.swapping_saver(var_list=[var0]) @@ -123,7 +123,7 @@ class MovingAverageOptimizerTest(test.TestCase): 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. + # Restore back to previous state. train_saver.restore(sess, save_path) # If updates are parallel, this is not always true after the 1st step. diff --git a/tensorflow/contrib/optimizer_v2/BUILD b/tensorflow/contrib/optimizer_v2/BUILD index 26ea9135f57fb9fe95e61023bccb97d1d4f5ea1c..5225ecc14fef3cec9506eceb776805b74a87714e 100644 --- a/tensorflow/contrib/optimizer_v2/BUILD +++ b/tensorflow/contrib/optimizer_v2/BUILD @@ -48,6 +48,7 @@ py_library( srcs_version = "PY2AND3", deps = [ "//tensorflow/python:control_flow_ops", + "//tensorflow/python:distribute", "//tensorflow/python:framework", "//tensorflow/python:math_ops", "//tensorflow/python:resource_variable_ops", @@ -114,7 +115,6 @@ cuda_py_test( additional_deps = [ ":training", "@six_archive//:six", - "//tensorflow/contrib/eager/python:checkpointable_utils", "//tensorflow/python:constant_op", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", @@ -202,4 +202,5 @@ cuda_py_test( "//tensorflow/python:client_testlib", "//third_party/py/numpy", ], + tags = ["optonly"], ) diff --git a/tensorflow/contrib/optimizer_v2/adam.py b/tensorflow/contrib/optimizer_v2/adam.py index 42b7f92a76c1971e2a63722d769ee006c3f3210b..d538ad0fb02699ed8514f512208914f629a47436 100644 --- a/tensorflow/contrib/optimizer_v2/adam.py +++ b/tensorflow/contrib/optimizer_v2/adam.py @@ -40,23 +40,19 @@ class AdamOptimizer(optimizer_v2.OptimizerV2): Initialization: - ``` - m_0 <- 0 (Initialize initial 1st moment vector) - v_0 <- 0 (Initialize initial 2nd moment vector) - t <- 0 (Initialize timestep) - ``` + $$m_0 := 0 (Initialize initial 1st moment vector)$$ + $$v_0 := 0 (Initialize initial 2nd moment vector)$$ + $$t := 0 (Initialize timestep)$$ The update rule for `variable` with gradient `g` uses an optimization described at the end of section2 of the paper: - ``` - t <- t + 1 - lr_t <- learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t) + $$t := t + 1$$ + $$lr_t := \text{learning_rate} * \sqrt{(1 - beta_2^t) / (1 - beta_1^t)}$$ - m_t <- beta1 * m_{t-1} + (1 - beta1) * g - v_t <- beta2 * v_{t-1} + (1 - beta2) * g * g - variable <- variable - lr_t * m_t / (sqrt(v_t) + epsilon) - ``` + $$m_t := beta_1 * m_{t-1} + (1 - beta_1) * g$$ + $$v_t := beta_2 * v_{t-1} + (1 - beta_2) * g * g$$ + $$variable := variable - lr_t * m_t / (\sqrt{v_t} + \epsilon)$$ The default value of 1e-8 for epsilon might not be a good default in general. For example, when training an Inception network on ImageNet a diff --git a/tensorflow/contrib/optimizer_v2/checkpointable_utils_test.py b/tensorflow/contrib/optimizer_v2/checkpointable_utils_test.py index d219795aa1edf7c16e52e8e663c0b3cc167262ad..8ac9b581455f8f4c7af1a66432169ae179de1634 100644 --- a/tensorflow/contrib/optimizer_v2/checkpointable_utils_test.py +++ b/tensorflow/contrib/optimizer_v2/checkpointable_utils_test.py @@ -24,23 +24,24 @@ import os import six -from tensorflow.contrib.eager.python import checkpointable_utils from tensorflow.contrib.optimizer_v2 import adam 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 function 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.keras._impl.keras.layers import core from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import resource_variable_ops from tensorflow.python.ops import state_ops from tensorflow.python.ops import variable_scope from tensorflow.python.training import checkpointable +from tensorflow.python.training import checkpointable_utils from tensorflow.python.training import saver as core_saver from tensorflow.python.training import training_util @@ -372,6 +373,50 @@ class CheckpointingTests(test.TestCase): self.assertEqual(training_continuation + 1, self.evaluate(root.save_counter)) + # pylint: disable=cell-var-from-loop + @test_util.run_in_graph_and_eager_modes() + def testWithDefun(self): + num_training_steps = 2 + 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() + # Don't actually train so we can test variable values + optimizer = adam.AdamOptimizer(0.) + 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) + def train_fn(): + @function.defun + def _call_model(x): + return model(x) + with backprop.GradientTape() as tape: + loss = _call_model(constant_op.constant([[3.]])) + gradients = tape.gradient(loss, model.variables) + return optimizer.apply_gradients(zip(gradients, model.variables), + 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() + if training_continuation > 0: + status.assert_consumed() + self.assertAllClose([[42.]], self.evaluate(model.variables[0])) + else: + self.evaluate(model.variables[0].assign([[42.]])) + 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)) + # pylint: enable=cell-var-from-loop + def _get_checkpoint_name(self, name): root = checkpointable.Checkpointable() checkpointable_utils.add_variable( diff --git a/tensorflow/contrib/optimizer_v2/momentum_test.py b/tensorflow/contrib/optimizer_v2/momentum_test.py index f37eb48181d6bef195215b86f14f69d3df65a8ac..26724f66c2a1db1d01577b31b739af18f51d3976 100644 --- a/tensorflow/contrib/optimizer_v2/momentum_test.py +++ b/tensorflow/contrib/optimizer_v2/momentum_test.py @@ -237,7 +237,17 @@ class MomentumOptimizerTest(test.TestCase): @test_util.run_in_graph_and_eager_modes(reset_test=True) def testMinimizeSparseResourceVariable(self): for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype) + # This test invokes the ResourceSparseApplyMomentum operation, which + # did not have a registered GPU kernel as of April 2018. With graph + # execution, the placement algorithm notices this and automatically + # places the variable in CPU (host) memory. With eager execution, + # the variable would be placed in GPU memory if available, which + # would then conflict with the future invocation of the + # ResourceSparseApplyMomentum operation. + # To work around this discrepancy, for now we force the variable + # to be placed on CPU. + with ops.device("/cpu:0"): + var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype) # pylint: disable=cell-var-from-loop def loss(): @@ -256,7 +266,17 @@ class MomentumOptimizerTest(test.TestCase): @test_util.run_in_graph_and_eager_modes(reset_test=True) def testMinimizeWith2DIndiciesForEmbeddingLookup(self): - var0 = resource_variable_ops.ResourceVariable(array_ops.ones([2, 2])) + # This test invokes the ResourceSparseApplyMomentum operation, which + # did not have a registered GPU kernel as of April 2018. With graph + # execution, the placement algorithm notices this and automatically + # places the variable in CPU (host) memory. With eager execution, + # the variable would be placed in GPU memory if available, which + # would then conflict with the future invocation of the + # ResourceSparseApplyMomentum operation. + # To work around this discrepancy, for now we force the variable + # to be placed on CPU. + with ops.device("/cpu:0"): + var0 = resource_variable_ops.ResourceVariable(array_ops.ones([2, 2])) def loss(): return math_ops.reduce_sum(embedding_ops.embedding_lookup(var0, [[1]])) diff --git a/tensorflow/contrib/optimizer_v2/optimizer_v2.py b/tensorflow/contrib/optimizer_v2/optimizer_v2.py index 25d19578ea8c4f53019657ab85950a814d1a47b8..ce15db6f1ec067e5aeb6ddbc8939d2b773692269 100644 --- a/tensorflow/contrib/optimizer_v2/optimizer_v2.py +++ b/tensorflow/contrib/optimizer_v2/optimizer_v2.py @@ -579,7 +579,7 @@ class OptimizerV2(optimizer_v1.Optimizer): ### State - Internal methods apre passed a `state` argument with the correct + Internal methods are passed a `state` argument with the correct values to use for the slot and non-slot variables, and the hyper parameters. """ diff --git a/tensorflow/contrib/proto/python/kernel_tests/defaut_values.TestCase.pbtxt b/tensorflow/contrib/proto/python/kernel_tests/defaut_values.TestCase.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..4e316819077c7dbb28beefd4dc260568f26da680 --- /dev/null +++ b/tensorflow/contrib/proto/python/kernel_tests/defaut_values.TestCase.pbtxt @@ -0,0 +1,94 @@ +primitive { + # No fields specified, so we get all defaults +} +shape: 1 +sizes: 0 +field { + name: "double_default" + dtype: DT_DOUBLE + expected { double_value: 1.0 } +} +sizes: 0 +field { + name: "float_default" + dtype: DT_DOUBLE # Try casting the float field to double. + expected { double_value: 2.0 } +} +sizes: 0 +field { + name: "int64_default" + dtype: DT_INT64 + expected { int64_value: 3 } +} +sizes: 0 +field { + name: "uint64_default" + dtype: DT_INT64 + expected { int64_value: 4 } +} +sizes: 0 +field { + name: "int32_default" + dtype: DT_INT32 + expected { int32_value: 5 } +} +sizes: 0 +field { + name: "fixed64_default" + dtype: DT_INT64 + expected { int64_value: 6 } +} +sizes: 0 +field { + name: "fixed32_default" + dtype: DT_INT32 + expected { int32_value: 7 } +} +sizes: 0 +field { + name: "bool_default" + dtype: DT_BOOL + expected { bool_value: true } +} +sizes: 0 +field { + name: "string_default" + dtype: DT_STRING + expected { string_value: "a" } +} +sizes: 0 +field { + name: "bytes_default" + dtype: DT_STRING + expected { string_value: "a longer default string" } +} +sizes: 0 +field { + name: "uint32_default" + dtype: DT_INT32 + expected { int32_value: -1 } +} +sizes: 0 +field { + name: "sfixed32_default" + dtype: DT_INT32 + expected { int32_value: 10 } +} +sizes: 0 +field { + name: "sfixed64_default" + dtype: DT_INT64 + expected { int64_value: 11 } +} +sizes: 0 +field { + name: "sint32_default" + dtype: DT_INT32 + expected { int32_value: 12 } +} +sizes: 0 +field { + name: "sint64_default" + dtype: DT_INT64 + expected { int64_value: 13 } +} diff --git a/tensorflow/contrib/proto/python/kernel_tests/promote_unsigned.TestCase.pbtxt b/tensorflow/contrib/proto/python/kernel_tests/promote_unsigned.TestCase.pbtxt index db7555bf2dff58e894b1e84db63b6ec91e511a2a..bc07efc8f3038c6c540855c97b2254575e517ef3 100644 --- a/tensorflow/contrib/proto/python/kernel_tests/promote_unsigned.TestCase.pbtxt +++ b/tensorflow/contrib/proto/python/kernel_tests/promote_unsigned.TestCase.pbtxt @@ -4,7 +4,6 @@ primitive { } shape: 1 sizes: 1 -sizes: 1 field { name: "fixed32_value" dtype: DT_INT64 @@ -12,6 +11,7 @@ field { int64_value: 4294967295 } } +sizes: 1 field { name: "uint32_value" dtype: DT_INT64 @@ -19,3 +19,11 @@ field { int64_value: 4294967295 } } +sizes: 0 +field { + name: "uint32_default" + dtype: DT_INT64 + expected { + int64_value: 4294967295 # Comes from an explicitly-specified default + } +} diff --git a/tensorflow/contrib/proto/python/kernel_tests/test_example.proto b/tensorflow/contrib/proto/python/kernel_tests/test_example.proto index dc495034ffae69dde4057858d37bc6afa600cd88..a2c88e372bf7c6b7f14c5bb55776b66c4c06bcd4 100644 --- a/tensorflow/contrib/proto/python/kernel_tests/test_example.proto +++ b/tensorflow/contrib/proto/python/kernel_tests/test_example.proto @@ -72,6 +72,23 @@ message RepeatedPrimitiveValue { repeated sint32 sint32_value = 17; repeated sint64 sint64_value = 18; repeated PrimitiveValue message_value = 19; + + // Optional fields with explicitly-specified defaults. + optional double double_default = 20 [default = 1.0]; + optional float float_default = 21 [default = 2.0]; + optional int64 int64_default = 22 [default = 3]; + optional uint64 uint64_default = 23 [default = 4]; + optional int32 int32_default = 24 [default = 5]; + optional fixed64 fixed64_default = 25 [default = 6]; + optional fixed32 fixed32_default = 26 [default = 7]; + optional bool bool_default = 27 [default = true]; + optional string string_default = 28 [default = "a"]; + optional bytes bytes_default = 29 [default = "a longer default string"]; + optional uint32 uint32_default = 30 [default = 4294967295]; + optional sfixed32 sfixed32_default = 31 [default = 10]; + optional sfixed64 sfixed64_default = 32 [default = 11]; + optional sint32 sint32_default = 33 [default = 12]; + optional sint64 sint64_default = 34 [default = 13]; } // A PackedPrimitiveValue looks exactly the same as a RepeatedPrimitiveValue @@ -97,6 +114,22 @@ message PackedPrimitiveValue { repeated sint32 sint32_value = 17 [packed = true]; repeated sint64 sint64_value = 18 [packed = true]; repeated PrimitiveValue message_value = 19; + + optional double double_default = 20 [default = 1.0]; + optional float float_default = 21 [default = 2.0]; + optional int64 int64_default = 22 [default = 3]; + optional uint64 uint64_default = 23 [default = 4]; + optional int32 int32_default = 24 [default = 5]; + optional fixed64 fixed64_default = 25 [default = 6]; + optional fixed32 fixed32_default = 26 [default = 7]; + optional bool bool_default = 27 [default = true]; + optional string string_default = 28 [default = "a"]; + optional bytes bytes_default = 29 [default = "a longer default string"]; + optional uint32 uint32_default = 30 [default = 4294967295]; + optional sfixed32 sfixed32_default = 31 [default = 10]; + optional sfixed64 sfixed64_default = 32 [default = 11]; + optional sint32 sint32_default = 33 [default = 12]; + optional sint64 sint64_default = 34 [default = 13]; } message EnumValue { diff --git a/tensorflow/contrib/quantize/python/fold_batch_norms.py b/tensorflow/contrib/quantize/python/fold_batch_norms.py index 4a8f8a04cc521d9ee7885b4318814a6f15008eef..aa0ef643088ef36b84596d08f78c29594ceca2d6 100644 --- a/tensorflow/contrib/quantize/python/fold_batch_norms.py +++ b/tensorflow/contrib/quantize/python/fold_batch_norms.py @@ -545,7 +545,7 @@ def _GetBatchNormParams(graph, context, has_scaling): gamma_tensor = graph.get_tensor_by_name(op.name + ':0') if not has_scaling: - gamma_tensor = array_ops.ones(batch_mean_tensor.shape) + gamma_tensor = array_ops.ones(moving_mean_tensor.shape) return _BatchNormMatch( layer_op=None, diff --git a/tensorflow/contrib/rpc/python/kernel_tests/BUILD b/tensorflow/contrib/rpc/python/kernel_tests/BUILD index 2311c15a68c46090cec0f97bd950296506b0817e..f3e6731213fd0aa8f14845a7b1562001b134c7ac 100644 --- a/tensorflow/contrib/rpc/python/kernel_tests/BUILD +++ b/tensorflow/contrib/rpc/python/kernel_tests/BUILD @@ -28,6 +28,7 @@ py_library( py_library( name = "rpc_op_test_base", srcs = ["rpc_op_test_base.py"], + tags = ["notsan"], deps = [ ":test_example_proto_py", "//tensorflow/contrib/proto", diff --git a/tensorflow/contrib/saved_model/BUILD b/tensorflow/contrib/saved_model/BUILD index e431c464ef14e86faf30221ed06061f41da528fb..26fd4e2023806765ea4088f4c13a780ca7338bff 100644 --- a/tensorflow/contrib/saved_model/BUILD +++ b/tensorflow/contrib/saved_model/BUILD @@ -48,16 +48,14 @@ py_library( ], ) -py_test( - name = "reader_test", - size = "small", - srcs = ["python/saved_model/reader_test.py"], +py_library( + name = "reader", + srcs = ["python/saved_model/reader.py"], srcs_version = "PY2AND3", tags = ["no_windows"], # TODO: needs investigation on Windows - visibility = ["//visibility:private"], + visibility = ["//visibility:public"], deps = [ ":saved_model_py", - "//tensorflow/python:client_testlib", "//tensorflow/python:framework_for_generated_wrappers", "//tensorflow/python:lib", "//tensorflow/python:variables", @@ -66,6 +64,19 @@ py_test( ], ) +py_test( + name = "reader_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 = [ + ":reader", + "//tensorflow/python:client_testlib", + ], +) + py_test( name = "signature_def_utils_test", size = "small", diff --git a/tensorflow/contrib/seq2seq/BUILD b/tensorflow/contrib/seq2seq/BUILD index a62069a252155a8bd1c6251d9dd3a4564a81c295..1a1591d798f6f904e23987d9d7a60193c124c20e 100644 --- a/tensorflow/contrib/seq2seq/BUILD +++ b/tensorflow/contrib/seq2seq/BUILD @@ -3,9 +3,12 @@ licenses(["notice"]) # Apache 2.0 -exports_files(["LICENSE"]) +package(default_visibility = [ + "//learning/brain/google/xla/tests:__subpackages__", + "//tensorflow:__subpackages__", +]) -package(default_visibility = ["//tensorflow:__subpackages__"]) +exports_files(["LICENSE"]) load("//tensorflow:tensorflow.bzl", "cuda_py_test") load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") @@ -38,6 +41,7 @@ tf_custom_op_py_library( "//tensorflow/python:check_ops", "//tensorflow/python:clip_ops", "//tensorflow/python:control_flow_ops", + "//tensorflow/python:control_flow_util", "//tensorflow/python:embedding_ops", "//tensorflow/python:framework_for_generated_wrappers", "//tensorflow/python:functional_ops", 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 d508cf3f9db81aa7c3a1174ed13f2310b0595b04..cd162bae25aa1c1b6718b8e5b0b8687e5b80eab3 100644 --- a/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_test.py +++ b/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_test.py @@ -30,6 +30,7 @@ from tensorflow.contrib.seq2seq.python.ops import helper as helper_py from tensorflow.contrib.seq2seq.python.ops import basic_decoder from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes +from tensorflow.python.layers import core as layers_core from tensorflow.python.ops import array_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops @@ -110,7 +111,12 @@ class AttentionWrapperTest(test.TestCase): alignment_history=False, expected_final_alignment_history=None, attention_layer_size=6, + attention_layer=None, name=''): + attention_layer_sizes = ( + [attention_layer_size] if attention_layer_size is not None else None) + attention_layers = ( + [attention_layer] if attention_layer is not None else None) self._testWithMaybeMultiAttention( is_multi=False, create_attention_mechanisms=[create_attention_mechanism], @@ -119,7 +125,8 @@ class AttentionWrapperTest(test.TestCase): attention_mechanism_depths=[attention_mechanism_depth], alignment_history=alignment_history, expected_final_alignment_history=expected_final_alignment_history, - attention_layer_sizes=[attention_layer_size], + attention_layer_sizes=attention_layer_sizes, + attention_layers=attention_layers, name=name) def _testWithMaybeMultiAttention(self, @@ -131,6 +138,7 @@ class AttentionWrapperTest(test.TestCase): alignment_history=False, expected_final_alignment_history=None, attention_layer_sizes=None, + attention_layers=None, name=''): # Allow is_multi to be True with a single mechanism to enable test for # passing in a single mechanism in a list. @@ -144,12 +152,18 @@ class AttentionWrapperTest(test.TestCase): encoder_output_depth = 10 cell_depth = 9 - if attention_layer_sizes is None: - attention_depth = encoder_output_depth * len(create_attention_mechanisms) - else: + if attention_layer_sizes is not None: # Compute sum of attention_layer_sizes. Use encoder_output_depth if None. attention_depth = sum([attention_layer_size or encoder_output_depth for attention_layer_size in attention_layer_sizes]) + elif attention_layers is not None: + # Compute sum of attention_layers output depth. + attention_depth = sum( + attention_layer.compute_output_shape( + [batch_size, cell_depth + encoder_output_depth])[-1].value + for attention_layer in attention_layers) + else: + attention_depth = encoder_output_depth * len(create_attention_mechanisms) decoder_inputs = array_ops.placeholder_with_default( np.random.randn(batch_size, decoder_max_time, @@ -171,13 +185,20 @@ class AttentionWrapperTest(test.TestCase): with vs.variable_scope( 'root', initializer=init_ops.random_normal_initializer(stddev=0.01, seed=3)): + attention_layer_size = attention_layer_sizes + attention_layer = attention_layers + if not is_multi: + if attention_layer_size is not None: + attention_layer_size = attention_layer_size[0] + if attention_layer is not None: + attention_layer = attention_layer[0] cell = rnn_cell.LSTMCell(cell_depth) cell = wrapper.AttentionWrapper( cell, attention_mechanisms if is_multi else attention_mechanisms[0], - attention_layer_size=(attention_layer_sizes if is_multi - else attention_layer_sizes[0]), - alignment_history=alignment_history) + attention_layer_size=attention_layer_size, + alignment_history=alignment_history, + attention_layer=attention_layer) helper = helper_py.TrainingHelper(decoder_inputs, decoder_sequence_length) my_decoder = basic_decoder.BasicDecoder( @@ -260,6 +281,41 @@ class AttentionWrapperTest(test.TestCase): expected_final_alignment_history, final_alignment_history_info) + def testBahdanauNormalizedDType(self): + for dtype in [np.float16, np.float32, np.float64]: + num_units = 128 + encoder_outputs = array_ops.placeholder(dtype, shape=[64, None, 256]) + encoder_sequence_length = array_ops.placeholder(dtypes.int32, shape=[64]) + decoder_inputs = array_ops.placeholder(dtype, shape=[64, None, 128]) + decoder_sequence_length = array_ops.placeholder(dtypes.int32, shape=[64]) + batch_size = 64 + attention_mechanism = wrapper.BahdanauAttention( + num_units=num_units, + memory=encoder_outputs, + memory_sequence_length=encoder_sequence_length, + normalize=True, + dtype=dtype, + ) + cell = rnn_cell.LSTMCell(num_units) + cell = wrapper.AttentionWrapper(cell, attention_mechanism) + + helper = helper_py.TrainingHelper(decoder_inputs, + decoder_sequence_length) + my_decoder = basic_decoder.BasicDecoder( + cell=cell, + helper=helper, + initial_state=cell.zero_state( + dtype=dtype, batch_size=batch_size)) + + final_outputs, final_state, _ = decoder.dynamic_decode(my_decoder) + self.assertTrue( + isinstance(final_outputs, basic_decoder.BasicDecoderOutput)) + self.assertEqual(final_outputs.rnn_output.dtype, dtype) + self.assertTrue( + isinstance(final_state, wrapper.AttentionWrapperState)) + self.assertTrue( + isinstance(final_state.cell_state, rnn_cell.LSTMStateTuple)) + def testBahdanauNotNormalized(self): create_attention_mechanism = wrapper.BahdanauAttention @@ -355,11 +411,11 @@ class AttentionWrapperTest(test.TestCase): def testLuongScaledDType(self): # Test case for GitHub issue 18099 - for dtype in [np.float16, np.float32, np.float64]: + for dt in [np.float16, np.float32, np.float64]: num_units = 128 - encoder_outputs = array_ops.placeholder(dtype, shape=[64, None, 256]) + encoder_outputs = array_ops.placeholder(dt, shape=[64, None, 256]) encoder_sequence_length = array_ops.placeholder(dtypes.int32, shape=[64]) - decoder_inputs = array_ops.placeholder(dtype, shape=[64, None, 128]) + decoder_inputs = array_ops.placeholder(dt, shape=[64, None, 128]) decoder_sequence_length = array_ops.placeholder(dtypes.int32, shape=[64]) batch_size = 64 attention_mechanism = wrapper.LuongAttention( @@ -367,7 +423,7 @@ class AttentionWrapperTest(test.TestCase): memory=encoder_outputs, memory_sequence_length=encoder_sequence_length, scale=True, - dtype=dtype, + dtype=dt, ) cell = rnn_cell.LSTMCell(num_units) cell = wrapper.AttentionWrapper(cell, attention_mechanism) @@ -378,12 +434,12 @@ class AttentionWrapperTest(test.TestCase): cell=cell, helper=helper, initial_state=cell.zero_state( - dtype=dtype, batch_size=batch_size)) + dtype=dt, batch_size=batch_size)) final_outputs, final_state, _ = decoder.dynamic_decode(my_decoder) self.assertTrue( isinstance(final_outputs, basic_decoder.BasicDecoderOutput)) - self.assertEqual(final_outputs.rnn_output.dtype, dtype) + self.assertEqual(final_outputs.rnn_output.dtype, dt) self.assertTrue( isinstance(final_state, wrapper.AttentionWrapperState)) self.assertTrue( @@ -797,6 +853,48 @@ class AttentionWrapperTest(test.TestCase): expected_final_alignment_history=expected_final_alignment_history, name='testMultiAttention') + def testMultiAttentionWithLayerInstances(self): + create_attention_mechanisms = ( + wrapper.BahdanauAttention, wrapper.LuongAttention) + + expected_final_output = BasicDecoderOutput( + rnn_output=ResultSummary( + shape=(5, 3, 7), dtype=dtype('float32'), mean=0.0011709079), + sample_id=ResultSummary( + shape=(5, 3), dtype=dtype('int32'), mean=3.2000000000000002)) + expected_final_state = AttentionWrapperState( + cell_state=LSTMStateTuple( + c=ResultSummary( + shape=(5, 9), dtype=dtype('float32'), mean=-0.0038725811), + h=ResultSummary( + shape=(5, 9), dtype=dtype('float32'), mean=-0.0019329828)), + attention=ResultSummary( + shape=(5, 7), dtype=dtype('float32'), mean=0.001174294), + time=3, + alignments=( + ResultSummary(shape=(5, 8), dtype=dtype('float32'), mean=0.125), + ResultSummary(shape=(5, 8), dtype=dtype('float32'), mean=0.125)), + attention_state=( + ResultSummary(shape=(5, 8), dtype=dtype('float32'), mean=0.125), + ResultSummary(shape=(5, 8), dtype=dtype('float32'), mean=0.125)), + alignment_history=()) + + expected_final_alignment_history = ( + ResultSummary(shape=(3, 5, 8), dtype=dtype('float32'), mean=0.125), + ResultSummary(shape=(3, 5, 8), dtype=dtype('float32'), mean=0.125)) + + self._testWithMaybeMultiAttention( + True, + create_attention_mechanisms, + expected_final_output, + expected_final_state, + attention_mechanism_depths=[9, 9], + attention_layers=[layers_core.Dense(3, use_bias=False), + layers_core.Dense(4, use_bias=False)], + alignment_history=True, + expected_final_alignment_history=expected_final_alignment_history, + name='testMultiAttention') + def testLuongMonotonicHard(self): # Run attention mechanism with mode='hard', make sure probabilities are hard b, t, u, d = 10, 20, 30, 40 diff --git a/tensorflow/contrib/seq2seq/python/kernel_tests/decoder_test.py b/tensorflow/contrib/seq2seq/python/kernel_tests/decoder_test.py index ac830ae98e5f9878f0972fb851fd5edd1c1a8789..b549cbf568f254cbf18456145af751a8245dd379 100644 --- a/tensorflow/contrib/seq2seq/python/kernel_tests/decoder_test.py +++ b/tensorflow/contrib/seq2seq/python/kernel_tests/decoder_test.py @@ -92,14 +92,18 @@ class DynamicDecodeRNNTest(test.TestCase): # Mostly a smoke test time_steps = max_out + expected_length = sequence_length if maximum_iterations is not None: time_steps = min(max_out, maximum_iterations) + expected_length = [min(x, maximum_iterations) for x in expected_length] self.assertEqual( _t((batch_size, time_steps, cell_depth)), sess_results["final_outputs"].rnn_output.shape) self.assertEqual( _t((batch_size, time_steps)), sess_results["final_outputs"].sample_id.shape) + self.assertItemsEqual(expected_length, + sess_results["final_sequence_length"]) def testDynamicDecodeRNNBatchMajor(self): self._testDynamicDecodeRNN(time_major=False) diff --git a/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py b/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py index f0f143ddfcf17c0e471add804ac4920b02da68e0..1c9d179e3c55ad07fcf709f66028c91c20e8eea0 100644 --- a/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py +++ b/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py @@ -472,7 +472,8 @@ def _bahdanau_score(processed_query, keys, normalize): # Scalar used in weight normalization g = variable_scope.get_variable( "attention_g", dtype=dtype, - initializer=math.sqrt((1. / num_units))) + initializer=init_ops.constant_initializer(math.sqrt((1. / num_units))), + shape=()) # Bias added prior to the nonlinearity b = variable_scope.get_variable( "attention_b", [num_units], dtype=dtype, @@ -654,7 +655,7 @@ def monotonic_attention(p_choose_i, previous_attention, mode): shifted_1mp_choose_i = array_ops.concat( [array_ops.ones((batch_size, 1)), 1 - p_choose_i[:, :-1]], 1) # Compute attention distribution recursively as - # q[i] = (1 - p_choose_i[i])*q[i - 1] + previous_attention[i] + # q[i] = (1 - p_choose_i[i - 1])*q[i - 1] + previous_attention[i] # attention[i] = p_choose_i[i]*q[i] attention = p_choose_i*array_ops.transpose(functional_ops.scan( # Need to use reshape to remind TF of the shape between loop iterations @@ -1082,7 +1083,8 @@ class AttentionWrapper(rnn_cell_impl.RNNCell): cell_input_fn=None, output_attention=True, initial_cell_state=None, - name=None): + name=None, + attention_layer=None): """Construct the `AttentionWrapper`. **NOTE** If you are using the `BeamSearchDecoder` with a cell wrapped in @@ -1125,7 +1127,8 @@ class AttentionWrapper(rnn_cell_impl.RNNCell): (default), use the context as attention at each time step. Otherwise, feed the context and cell output into the attention layer to generate attention at each time step. If attention_mechanism is a list, - attention_layer_size must be a list of the same length. + attention_layer_size must be a list of the same length. If + attention_layer is set, this must be None. alignment_history: Python boolean, whether to store alignment history from all time steps in the final output state (currently stored as a time major `TensorArray` on which you must call `stack()`). @@ -1145,12 +1148,19 @@ class AttentionWrapper(rnn_cell_impl.RNNCell): does not match the batch size of `initial_cell_state`, proper behavior is not guaranteed. name: Name to use when creating ops. + attention_layer: A list of `tf.layers.Layer` instances or a + single `tf.layers.Layer` instance taking the context and cell output as + inputs to generate attention at each time step. If None (default), use + the context as attention at each time step. If attention_mechanism is a + list, attention_layer must be a list of the same length. If + attention_layers_size is set, this must be None. Raises: TypeError: `attention_layer_size` is not None and (`attention_mechanism` is a list but `attention_layer_size` is not; or vice versa). ValueError: if `attention_layer_size` is not None, `attention_mechanism` - is a list, and its length does not match that of `attention_layer_size`. + is a list, and its length does not match that of `attention_layer_size`; + if `attention_layer_size` and `attention_layer` are set simultaneously. """ super(AttentionWrapper, self).__init__(name=name) rnn_cell_impl.assert_like_rnncell("cell", cell) @@ -1181,6 +1191,10 @@ class AttentionWrapper(rnn_cell_impl.RNNCell): "cell_input_fn must be callable, saw type: %s" % type(cell_input_fn).__name__) + if attention_layer_size is not None and attention_layer is not None: + raise ValueError("Only one of attention_layer_size and attention_layer " + "should be set") + if attention_layer_size is not None: attention_layer_sizes = tuple( attention_layer_size @@ -1199,6 +1213,22 @@ class AttentionWrapper(rnn_cell_impl.RNNCell): dtype=attention_mechanisms[i].dtype) for i, attention_layer_size in enumerate(attention_layer_sizes)) self._attention_layer_size = sum(attention_layer_sizes) + elif attention_layer is not None: + self._attention_layers = tuple( + attention_layer + if isinstance(attention_layer, (list, tuple)) + else (attention_layer,)) + if len(self._attention_layers) != len(attention_mechanisms): + raise ValueError( + "If provided, attention_layer must contain exactly one " + "layer per attention_mechanism, saw: %d vs %d" + % (len(self._attention_layers), len(attention_mechanisms))) + self._attention_layer_size = sum( + layer.compute_output_shape( + [None, + cell.output_size + mechanism.values.shape[-1].value])[-1].value + for layer, mechanism in zip( + self._attention_layers, attention_mechanisms)) else: self._attention_layers = None self._attention_layer_size = sum( diff --git a/tensorflow/contrib/seq2seq/python/ops/decoder.py b/tensorflow/contrib/seq2seq/python/ops/decoder.py index 898493662d7594f9996400a9636378db3c6b4cd1..e69725ff8ab1ba4de880c914a6f5fdad5e54566d 100644 --- a/tensorflow/contrib/seq2seq/python/ops/decoder.py +++ b/tensorflow/contrib/seq2seq/python/ops/decoder.py @@ -28,6 +28,7 @@ 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 control_flow_ops +from tensorflow.python.ops import control_flow_util from tensorflow.python.ops import math_ops from tensorflow.python.ops import rnn from tensorflow.python.ops import rnn_cell_impl @@ -181,6 +182,15 @@ def dynamic_decode(decoder, raise TypeError("Expected decoder to be type Decoder, but saw: %s" % type(decoder)) + def _is_xla_tensor(tensor): + try: + op = tensor.op + except AttributeError: + return False + if control_flow_util.IsInXLAContext(op): + return True + return False + with variable_scope.variable_scope(scope, "decoder") as varscope: # Properly cache variable values inside the while_loop if varscope.caching_device is None: @@ -198,6 +208,11 @@ def dynamic_decode(decoder, decoder.output_dtype, decoder.batch_size) + is_xla = False + if any([_is_xla_tensor(i) for i in nest.flatten(initial_inputs)]): + is_xla = True + if is_xla and maximum_iterations is None: + raise ValueError("maximum_iterations is required for XLA compilation.") if maximum_iterations is not None: initial_finished = math_ops.logical_or( initial_finished, 0 >= maximum_iterations) @@ -215,11 +230,13 @@ def dynamic_decode(decoder, batch_size, name="batch_size")) return tensor_shape.TensorShape([batch_size]).concatenate(from_shape) + dynamic_size = maximum_iterations is None or not is_xla + def _create_ta(s, d): return tensor_array_ops.TensorArray( dtype=d, - size=0, - dynamic_size=True, + size=0 if dynamic_size else maximum_iterations, + dynamic_size=dynamic_size, element_shape=_shape(decoder.batch_size, s)) initial_outputs_ta = nest.map_structure(_create_ta, decoder.output_size, @@ -251,11 +268,8 @@ def dynamic_decode(decoder, next_finished = decoder_finished else: next_finished = math_ops.logical_or(decoder_finished, finished) - if maximum_iterations is not None: - next_finished = math_ops.logical_or( - next_finished, time + 1 >= maximum_iterations) next_sequence_lengths = array_ops.where( - math_ops.logical_and(math_ops.logical_not(finished), next_finished), + math_ops.logical_not(finished), array_ops.fill(array_ops.shape(sequence_lengths), time + 1), sequence_lengths) @@ -296,11 +310,16 @@ def dynamic_decode(decoder, res = control_flow_ops.while_loop( condition, body, - loop_vars=[ - initial_time, initial_outputs_ta, initial_state, initial_inputs, - initial_finished, initial_sequence_lengths, - ], + loop_vars=( + initial_time, + initial_outputs_ta, + initial_state, + initial_inputs, + initial_finished, + initial_sequence_lengths, + ), parallel_iterations=parallel_iterations, + maximum_iterations=maximum_iterations, swap_memory=swap_memory) final_outputs_ta = res[1] diff --git a/tensorflow/contrib/signal/python/kernel_tests/mel_ops_test.py b/tensorflow/contrib/signal/python/kernel_tests/mel_ops_test.py index 35c4b5bec172858b39dd4628a37e164efe87bdbf..345eb6cfaa67fd4cda6e7e3f01a1243bbf3c9fa1 100644 --- a/tensorflow/contrib/signal/python/kernel_tests/mel_ops_test.py +++ b/tensorflow/contrib/signal/python/kernel_tests/mel_ops_test.py @@ -24,6 +24,7 @@ from tensorflow.contrib.signal.python.kernel_tests import test_util from tensorflow.contrib.signal.python.ops import mel_ops from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops +from tensorflow.python.ops import array_ops from tensorflow.python.platform import test # mel spectrum constants and functions. @@ -173,6 +174,18 @@ class LinearToMelTest(test.TestCase): rewritten_graph = test_util.grappler_optimize(g, [mel_matrix]) self.assertEqual(1, len(rewritten_graph.node)) + def test_num_spectrogram_bins_dynamic(self): + with self.test_session(use_gpu=True): + num_spectrogram_bins = array_ops.placeholder(shape=(), + dtype=dtypes.int32) + mel_matrix_np = spectrogram_to_mel_matrix( + 20, 129, 8000.0, 125.0, 3800.0) + mel_matrix = mel_ops.linear_to_mel_weight_matrix( + 20, num_spectrogram_bins, 8000.0, 125.0, 3800.0) + self.assertAllClose( + mel_matrix_np, + mel_matrix.eval(feed_dict={num_spectrogram_bins: 129}), atol=3e-6) + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/signal/python/ops/mel_ops.py b/tensorflow/contrib/signal/python/ops/mel_ops.py index d1a36548d95cf44d2bf7e6108141aeb00853db04..1e84006116daa3f28c760037cb9eeafd53eaafb8 100644 --- a/tensorflow/contrib/signal/python/ops/mel_ops.py +++ b/tensorflow/contrib/signal/python/ops/mel_ops.py @@ -64,14 +64,11 @@ def _hertz_to_mel(frequencies_hertz, name=None): 1.0 + (frequencies_hertz / _MEL_BREAK_FREQUENCY_HERTZ)) -def _validate_arguments(num_mel_bins, num_spectrogram_bins, sample_rate, +def _validate_arguments(num_mel_bins, sample_rate, lower_edge_hertz, upper_edge_hertz, dtype): """Checks the inputs to linear_to_mel_weight_matrix.""" if num_mel_bins <= 0: raise ValueError('num_mel_bins must be positive. Got: %s' % num_mel_bins) - if num_spectrogram_bins <= 0: - raise ValueError('num_spectrogram_bins must be positive. Got: %s' % - num_spectrogram_bins) if sample_rate <= 0.0: raise ValueError('sample_rate must be positive. Got: %s' % sample_rate) if lower_edge_hertz < 0.0: @@ -122,9 +119,9 @@ def linear_to_mel_weight_matrix(num_mel_bins=20, Args: num_mel_bins: Python int. How many bands in the resulting mel spectrum. - num_spectrogram_bins: Python int. How many bins there are in the source - spectrogram data, which is understood to be `fft_size // 2 + 1`, i.e. the - spectrogram only contains the nonredundant FFT bins. + num_spectrogram_bins: An integer `Tensor`. How many bins there are in the + source spectrogram data, which is understood to be `fft_size // 2 + 1`, + i.e. the spectrogram only contains the nonredundant FFT bins. sample_rate: Python float. Samples per second of the input signal used to create the spectrogram. We need this to figure out the actual frequencies for each spectrogram bin, which dictates how they are mapped into the mel @@ -148,7 +145,10 @@ def linear_to_mel_weight_matrix(num_mel_bins=20, [mel]: https://en.wikipedia.org/wiki/Mel_scale """ with ops.name_scope(name, 'linear_to_mel_weight_matrix') as name: - _validate_arguments(num_mel_bins, num_spectrogram_bins, sample_rate, + # Note: As num_spectrogram_bins is passed to `math_ops.linspace` + # and the validation is already done in linspace (both in shape function + # and in kernel), there is no need to validate num_spectrogram_bins here. + _validate_arguments(num_mel_bins, sample_rate, lower_edge_hertz, upper_edge_hertz, dtype) # To preserve accuracy, we compute the matrix at float64 precision and then diff --git a/tensorflow/contrib/slim/README.md b/tensorflow/contrib/slim/README.md index 40f484fd78302163ba36142dec057478fe899189..746b95564237617359afe1791484809369c4a894 100644 --- a/tensorflow/contrib/slim/README.md +++ b/tensorflow/contrib/slim/README.md @@ -290,9 +290,9 @@ slim.stack(x, slim.conv2d, [(32, [3, 3]), (32, [1, 1]), (64, [3, 3]), (64, [1, 1 In addition to the types of scope mechanisms in TensorFlow ([name_scope](https://www.tensorflow.org/api_docs/python/tf/name_scope), -[variable_scope](https://www.tensorflow.org/api_docs/python/tf/variable_scope), +[variable_scope](https://www.tensorflow.org/api_docs/python/tf/variable_scope)), TF-Slim adds a new scoping mechanism called -[arg_scope](https://www.tensorflow.org/api_docs/python/tf/contrib/framework/arg_scope), +[arg_scope](https://www.tensorflow.org/api_docs/python/tf/contrib/framework/arg_scope). This new scope allows a user to specify one or more operations and a set of arguments which will be passed to each of the operations defined in the `arg_scope`. This functionality is best illustrated by example. Consider the @@ -761,8 +761,8 @@ parts: 3. Finalization: (optionally) perform any final operation to compute metric values. For example, computing means, mins, maxes, etc. -For example, to compute `mean_absolute_error`, two variables, a `count` and -`total` variable are *initialized* to zero. During *aggregation*, we observed +For example, to compute `mean_absolute_error`, two variables (`count` and +`total`) are *initialized* to zero. During *aggregation*, we observed some set of predictions and labels, compute their absolute differences and add the total to `total`. Each time we observe another value, `count` is incremented. Finally, during *finalization*, `total` is divided diff --git a/tensorflow/contrib/slim/python/slim/evaluation.py b/tensorflow/contrib/slim/python/slim/evaluation.py index 3caf4e02da3aa2d7e586c4e76807a11f84585ea6..5cfd5ee82e2a0fce33311a8783d2d4ceb031544d 100644 --- a/tensorflow/contrib/slim/python/slim/evaluation.py +++ b/tensorflow/contrib/slim/python/slim/evaluation.py @@ -230,6 +230,7 @@ def evaluation_loop(master, max_number_of_evaluations=None, session_config=None, timeout=None, + timeout_fn=None, hooks=None): """Runs TF-Slim's Evaluation Loop. @@ -261,6 +262,9 @@ def evaluation_loop(master, configure the `Session`. If left as `None`, the default will be used. timeout: The maximum amount of time to wait between checkpoints. If left as `None`, then the process will wait indefinitely. + timeout_fn: Optional function to call after a timeout. If the function + returns True, then it means that no new checkpoints will be generated and + the iterator will exit. The function is called with no arguments. hooks: A list of additional `SessionRunHook` objects to pass during repeated evaluations. @@ -298,4 +302,5 @@ def evaluation_loop(master, hooks=all_hooks, config=session_config, max_number_of_evaluations=max_number_of_evaluations, - timeout=timeout) + timeout=timeout, + timeout_fn=timeout_fn) diff --git a/tensorflow/contrib/slim/python/slim/evaluation_test.py b/tensorflow/contrib/slim/python/slim/evaluation_test.py index c24bd048512daaae116e732ac437f7c9b6f6d7fc..94fc12ca814721acf62f16b72ffa50473043cc8b 100644 --- a/tensorflow/contrib/slim/python/slim/evaluation_test.py +++ b/tensorflow/contrib/slim/python/slim/evaluation_test.py @@ -177,6 +177,17 @@ class EvaluationTest(test.TestCase): # The timeout kicked in. self.assertLess(end, start + 1.1) + def testTimeoutFnOnEvaluationLoop(self): + # We require a mutable object (e.g. list but not an int) to maintain state + # across calls of a nested function. + timeout_fn_calls = [0] + def _TimeoutFn(): + timeout_fn_calls[0] += 1 + return timeout_fn_calls[0] >= 3 + # Need not do any evaluation, but should just call timeout_fn repeatedly. + evaluation.evaluation_loop('', '', '', timeout=0, timeout_fn=_TimeoutFn) + self.assertEqual(timeout_fn_calls[0], 3) + def testMonitorCheckpointsLoopTimeout(self): ret = list( evaluation_lib.checkpoints_iterator( diff --git a/tensorflow/contrib/slim/python/slim/learning.py b/tensorflow/contrib/slim/python/slim/learning.py index 6a200de1ea172b4ccb38c0f5d889566ccaeef893..8a2c74742a8ebbfdca702943ee2f631531c7b2ca 100644 --- a/tensorflow/contrib/slim/python/slim/learning.py +++ b/tensorflow/contrib/slim/python/slim/learning.py @@ -389,7 +389,7 @@ def create_train_op(total_loss, total_loss: A `Tensor` representing the total loss. optimizer: A tf.Optimizer to use for computing the gradients. global_step: A `Tensor` representing the global step variable. If left as - `_USE_GLOBAL_STEP`, then slim.variables.global_step() is used. + `_USE_GLOBAL_STEP`, then tf.contrib.framework.global_step() is used. update_ops: An optional list of updates to execute. If `update_ops` is `None`, then the update ops are set to the contents of the `tf.GraphKeys.UPDATE_OPS` collection. If `update_ops` is not `None`, but @@ -578,7 +578,8 @@ def train(train_op, is_chief: Specifies whether or not the training is being run by the primary replica during replica training. global_step: The `Tensor` representing the global step. If left as `None`, - then slim.variables.get_or_create_global_step() is used. + then training_util.get_or_create_global_step(), that is, + tf.contrib.framework.global_step() is used. number_of_steps: The max number of gradient steps to take during training, as measured by 'global_step': training will stop if global_step is greater than 'number_of_steps'. If the value is left as None, training diff --git a/tensorflow/contrib/slim/python/slim/nets/resnet_v1.py b/tensorflow/contrib/slim/python/slim/nets/resnet_v1.py index 235a595de49f956e1df740fd821936c80eefaa55..11c4214176a8e3d69065066bb5ac4d668da10574 100644 --- a/tensorflow/contrib/slim/python/slim/nets/resnet_v1.py +++ b/tensorflow/contrib/slim/python/slim/nets/resnet_v1.py @@ -207,7 +207,7 @@ def resnet_v1(inputs, net = resnet_utils.stack_blocks_dense(net, blocks, output_stride) if global_pool: # Global average pooling. - net = math_ops.reduce_mean(net, [1, 2], name='pool5', keep_dims=True) + net = math_ops.reduce_mean(net, [1, 2], name='pool5', keepdims=True) if num_classes is not None: net = layers.conv2d( net, diff --git a/tensorflow/contrib/slim/python/slim/nets/resnet_v2.py b/tensorflow/contrib/slim/python/slim/nets/resnet_v2.py index 61665c9c8ba7817377a16bf3f2673447cab0518e..19e0538dd1e272b2bcaada3c83944e9c9c1f9eef 100644 --- a/tensorflow/contrib/slim/python/slim/nets/resnet_v2.py +++ b/tensorflow/contrib/slim/python/slim/nets/resnet_v2.py @@ -221,7 +221,7 @@ def resnet_v2(inputs, net, activation_fn=nn_ops.relu, scope='postnorm') if global_pool: # Global average pooling. - net = math_ops.reduce_mean(net, [1, 2], name='pool5', keep_dims=True) + net = math_ops.reduce_mean(net, [1, 2], name='pool5', keepdims=True) if num_classes is not None: net = layers_lib.conv2d( net, diff --git a/tensorflow/contrib/slim/python/slim/summaries.py b/tensorflow/contrib/slim/python/slim/summaries.py index 358359d6ebeea209fe83f7282f47db8be63747ac..a7dc3f6723a0d1a55dd3e8dce006e6acce083e6f 100644 --- a/tensorflow/contrib/slim/python/slim/summaries.py +++ b/tensorflow/contrib/slim/python/slim/summaries.py @@ -144,7 +144,7 @@ def add_zero_fraction_summary(tensor, name=None, prefix=None, A scalar `Tensor` of type `string` whose contents are the serialized `Summary` protocol buffer. """ - name = _get_summary_name(tensor, name, prefix, 'Fraction of Zero Values') + name = _get_summary_name(tensor, name, prefix, 'Fraction_of_Zero_Values') tensor = nn.zero_fraction(tensor) return add_scalar_summary(tensor, name, print_summary=print_summary) diff --git a/tensorflow/contrib/stateless/__init__.py b/tensorflow/contrib/stateless/__init__.py index ca937546f50df46b7e5b1144dcbdc380cb04ca9b..0cca40f071c889773736ce009b32ba17728041ce 100644 --- a/tensorflow/contrib/stateless/__init__.py +++ b/tensorflow/contrib/stateless/__init__.py @@ -22,6 +22,7 @@ WARNING: These ops are in contrib, and are not stable. They should be consistent across multiple runs on the same hardware, but only for the same version of the code. +@@stateless_multinomial @@stateless_random_uniform @@stateless_random_normal @@stateless_truncated_normal @@ -37,6 +38,7 @@ from tensorflow.contrib.stateless.gen_stateless_random_ops import * from tensorflow.python.framework import ops from tensorflow.python.util.all_util import remove_undocumented +ops.NotDifferentiable("StatelessMultinomial") ops.NotDifferentiable("StatelessRandomNormal") ops.NotDifferentiable("StatelessRandomUniform") ops.NotDifferentiable("StatelessTruncatedNormal") diff --git a/tensorflow/contrib/stateless/python/kernel_tests/stateless_random_ops_test.py b/tensorflow/contrib/stateless/python/kernel_tests/stateless_random_ops_test.py index bea6341cfdcf7d56f255bec275b7861228e44e12..d724a5c014d2f9f5f6e3a6704341bcb8c429ae06 100644 --- a/tensorflow/contrib/stateless/python/kernel_tests/stateless_random_ops_test.py +++ b/tensorflow/contrib/stateless/python/kernel_tests/stateless_random_ops_test.py @@ -96,6 +96,52 @@ class StatelessOpsTest(test.TestCase): for s1, v1 in values: self.assertEqual(s0 == s1, np.all(v0 == v1)) + def testMatchStatefulMultinomial(self): + # Stateless ops should be the same as stateful ops on the first call + # after seed scrambling. + key = 0x3ec8f720, 0x02461e29 + num_samples = 4 + for logits_dtype in np.float16, np.float32, np.float64: + for output_dtype in dtypes.int32, dtypes.int64: + for seed in (7, 17), (11, 5), (2, 3): + preseed = invert_philox(key, + (seed[0], 0, seed[1], 0)).astype(np.uint64) + preseed = preseed[::2] | preseed[1::2] << 32 + random_seed.set_random_seed(seed[0]) + with self.test_session(use_gpu=True): + for logits in ([[0.1, 0.25, 0.5, 0.15]], [[0.5, 0.5], [0.8, 0.2], + [0.25, 0.75]]): + logits_t = constant_op.constant(logits, dtype=logits_dtype) + stateful = random_ops.multinomial( + logits_t, + num_samples, + seed=seed[1], + output_dtype=output_dtype) + pure = stateless.stateless_multinomial( + logits_t, + num_samples, + seed=preseed, + output_dtype=output_dtype) + self.assertAllEqual(stateful.eval(), pure.eval()) + + def testDeterminismMultinomial(self): + # Stateless values should be equal iff the seeds are equal (roughly) + num_samples = 10 + with self.test_session(use_gpu=True): + for seed_type in [dtypes.int32, dtypes.int64]: + seed_t = array_ops.placeholder(seed_type, shape=[2]) + seeds = [(x, y) for x in range(5) for y in range(5)] * 3 + for logits in ([[0.1, 0.25, 0.5, 0.15]], [[0.5, 0.5], [0.8, 0.2], + [0.25, 0.75]]): + pure = stateless.stateless_multinomial( + logits, num_samples, seed=seed_t) + values = [ + (seed, pure.eval(feed_dict={seed_t: seed})) for seed in seeds + ] + for s0, v0 in values: + for s1, v1 in values: + self.assertEqual(s0 == s1, np.all(v0 == v1)) + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/summary/BUILD b/tensorflow/contrib/summary/BUILD index fda1367b156c86f385f31cc41c5fca747cf8668d..f88b03ec4c2b1f250091594ea12d7d1862029fa2 100644 --- a/tensorflow/contrib/summary/BUILD +++ b/tensorflow/contrib/summary/BUILD @@ -15,7 +15,6 @@ py_test( srcs = ["summary_ops_test.py"], srcs_version = "PY2AND3", deps = [ - ":summary_ops", ":summary_test_util", "//tensorflow/python:array_ops", "//tensorflow/python:errors", @@ -23,6 +22,7 @@ py_test( "//tensorflow/python:framework_test_lib", "//tensorflow/python:platform", "//tensorflow/python:state_ops", + "//tensorflow/python:summary_ops_v2", "//tensorflow/python:training", "//tensorflow/python/eager:function", "//tensorflow/python/eager:test", @@ -35,7 +35,6 @@ py_test( srcs = ["summary_ops_graph_test.py"], srcs_version = "PY2AND3", deps = [ - ":summary_ops", ":summary_test_util", "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", @@ -44,31 +43,9 @@ py_test( "//tensorflow/python:control_flow_ops", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", + "//tensorflow/python:summary_ops_v2", "//tensorflow/python:training", - "@six_archive//:six", - ], -) - -py_library( - name = "summary_ops", - srcs = ["summary_ops.py"], - srcs_version = "PY2AND3", - visibility = ["//tensorflow:internal"], - deps = [ - "//tensorflow/core:protos_all_py", - "//tensorflow/python:array_ops", - "//tensorflow/python:constant_op", - "//tensorflow/python:control_flow_ops", - "//tensorflow/python:dtypes", - "//tensorflow/python:framework_ops", - "//tensorflow/python:layers_base", - "//tensorflow/python:math_ops", - "//tensorflow/python:resource_variable_ops", - "//tensorflow/python:summary_op_util", - "//tensorflow/python:summary_ops_gen", - "//tensorflow/python:training", - "//tensorflow/python:util", - "//tensorflow/python/eager:context", + "//tensorflow/python:variables", "@six_archive//:six", ], ) @@ -79,7 +56,7 @@ py_library( srcs_version = "PY2AND3", visibility = ["//tensorflow:internal"], deps = [ - ":summary_ops", + "//tensorflow/python:summary_ops_v2", ], ) @@ -92,8 +69,10 @@ py_library( visibility = ["//tensorflow:internal"], deps = [ "//tensorflow/core:protos_all_py", + "//tensorflow/python:framework_test_lib", "//tensorflow/python:lib", "//tensorflow/python:platform", + "//tensorflow/python:summary_ops_v2", "@org_sqlite//:python", ], ) diff --git a/tensorflow/contrib/summary/summary.py b/tensorflow/contrib/summary/summary.py index 2d6d7ea6a3eff2562ba8def4117e3aa6f818b6fd..99ced53e1167ec5486d0b75cff81ffbf857c2be7 100644 --- a/tensorflow/contrib/summary/summary.py +++ b/tensorflow/contrib/summary/summary.py @@ -61,23 +61,23 @@ from __future__ import division from __future__ import print_function # pylint: disable=unused-import -from tensorflow.contrib.summary.summary_ops import all_summary_ops -from tensorflow.contrib.summary.summary_ops import always_record_summaries -from tensorflow.contrib.summary.summary_ops import audio -from tensorflow.contrib.summary.summary_ops import create_db_writer -from tensorflow.contrib.summary.summary_ops import create_file_writer -from tensorflow.contrib.summary.summary_ops import create_summary_file_writer -from tensorflow.contrib.summary.summary_ops import eval_dir -from tensorflow.contrib.summary.summary_ops import flush -from tensorflow.contrib.summary.summary_ops import generic -from tensorflow.contrib.summary.summary_ops import graph -from tensorflow.contrib.summary.summary_ops import histogram -from tensorflow.contrib.summary.summary_ops import image -from tensorflow.contrib.summary.summary_ops import import_event -from tensorflow.contrib.summary.summary_ops import initialize -from tensorflow.contrib.summary.summary_ops import never_record_summaries -from tensorflow.contrib.summary.summary_ops import record_summaries_every_n_global_steps -from tensorflow.contrib.summary.summary_ops import scalar -from tensorflow.contrib.summary.summary_ops import should_record_summaries -from tensorflow.contrib.summary.summary_ops import summary_writer_initializer_op -from tensorflow.contrib.summary.summary_ops import SummaryWriter +from tensorflow.python.ops.summary_ops_v2 import all_summary_ops +from tensorflow.python.ops.summary_ops_v2 import always_record_summaries +from tensorflow.python.ops.summary_ops_v2 import audio +from tensorflow.python.ops.summary_ops_v2 import create_db_writer +from tensorflow.python.ops.summary_ops_v2 import create_file_writer +from tensorflow.python.ops.summary_ops_v2 import create_summary_file_writer +from tensorflow.python.ops.summary_ops_v2 import eval_dir +from tensorflow.python.ops.summary_ops_v2 import flush +from tensorflow.python.ops.summary_ops_v2 import generic +from tensorflow.python.ops.summary_ops_v2 import graph +from tensorflow.python.ops.summary_ops_v2 import histogram +from tensorflow.python.ops.summary_ops_v2 import image +from tensorflow.python.ops.summary_ops_v2 import import_event +from tensorflow.python.ops.summary_ops_v2 import initialize +from tensorflow.python.ops.summary_ops_v2 import never_record_summaries +from tensorflow.python.ops.summary_ops_v2 import record_summaries_every_n_global_steps +from tensorflow.python.ops.summary_ops_v2 import scalar +from tensorflow.python.ops.summary_ops_v2 import should_record_summaries +from tensorflow.python.ops.summary_ops_v2 import summary_writer_initializer_op +from tensorflow.python.ops.summary_ops_v2 import SummaryWriter diff --git a/tensorflow/contrib/summary/summary_ops_graph_test.py b/tensorflow/contrib/summary/summary_ops_graph_test.py index 3aba04540eba12092d884cca10e23546eb91c91d..ae8336daaf8ea9113716b90b6ea9be9de7303596 100644 --- a/tensorflow/contrib/summary/summary_ops_graph_test.py +++ b/tensorflow/contrib/summary/summary_ops_graph_test.py @@ -16,27 +16,220 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import os import tempfile +import time import six -from tensorflow.contrib.summary import summary_ops from tensorflow.contrib.summary import summary_test_util from tensorflow.core.framework import graph_pb2 from tensorflow.core.framework import node_def_pb2 from tensorflow.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 control_flow_ops from tensorflow.python.ops import state_ops +from tensorflow.python.ops import summary_ops_v2 as summary_ops +from tensorflow.python.ops import variables +from tensorflow.python.platform import gfile from tensorflow.python.platform import test from tensorflow.python.training import training_util get_all = summary_test_util.get_all -class DbTest(summary_test_util.SummaryDbTest): +class GraphFileTest(test_util.TensorFlowTestCase): + + def testSummaryOps(self): + logdir = self.get_temp_dir() + writer = summary_ops.create_file_writer(logdir, max_queue=0) + with writer.as_default(), summary_ops.always_record_summaries(): + summary_ops.generic('tensor', 1, step=1) + summary_ops.scalar('scalar', 2.0, step=1) + summary_ops.histogram('histogram', [1.0], step=1) + summary_ops.image('image', [[[[1.0]]]], step=1) + summary_ops.audio('audio', [[1.0]], 1.0, 1, step=1) + with self.test_session() as sess: + sess.run(summary_ops.summary_writer_initializer_op()) + sess.run(summary_ops.all_summary_ops()) + # The working condition of the ops is tested in the C++ test so we just + # test here that we're calling them correctly. + self.assertTrue(gfile.Exists(logdir)) + + def testSummaryName(self): + logdir = self.get_temp_dir() + writer = summary_ops.create_file_writer(logdir, max_queue=0) + with writer.as_default(), summary_ops.always_record_summaries(): + summary_ops.scalar('scalar', 2.0, step=1) + with self.test_session() as sess: + sess.run(summary_ops.summary_writer_initializer_op()) + sess.run(summary_ops.all_summary_ops()) + events = summary_test_util.events_from_logdir(logdir) + self.assertEqual(2, len(events)) + self.assertEqual('scalar', events[1].summary.value[0].tag) + + def testSummaryNameScope(self): + logdir = self.get_temp_dir() + writer = summary_ops.create_file_writer(logdir, max_queue=0) + with writer.as_default(), summary_ops.always_record_summaries(): + with ops.name_scope('scope'): + summary_ops.scalar('scalar', 2.0, step=1) + with self.test_session() as sess: + sess.run(summary_ops.summary_writer_initializer_op()) + sess.run(summary_ops.all_summary_ops()) + events = summary_test_util.events_from_logdir(logdir) + self.assertEqual(2, len(events)) + self.assertEqual('scope/scalar', events[1].summary.value[0].tag) + + def testSummaryGlobalStep(self): + training_util.get_or_create_global_step() + logdir = self.get_temp_dir() + writer = summary_ops.create_file_writer(logdir, max_queue=0) + with writer.as_default(), summary_ops.always_record_summaries(): + summary_ops.scalar('scalar', 2.0) + with self.test_session() as sess: + sess.run(variables.global_variables_initializer()) + sess.run(summary_ops.summary_writer_initializer_op()) + step, _ = sess.run( + [training_util.get_global_step(), summary_ops.all_summary_ops()]) + events = summary_test_util.events_from_logdir(logdir) + self.assertEqual(2, len(events)) + self.assertEqual(step, events[1].step) + + def testMaxQueue(self): + logdir = self.get_temp_dir() + writer = summary_ops.create_file_writer( + logdir, max_queue=1, flush_millis=999999) + with writer.as_default(), summary_ops.always_record_summaries(): + summary_ops.scalar('scalar', 2.0, step=1) + with self.test_session() as sess: + sess.run(summary_ops.summary_writer_initializer_op()) + get_total = lambda: len(summary_test_util.events_from_logdir(logdir)) + # Note: First tf.Event is always file_version. + self.assertEqual(1, get_total()) + sess.run(summary_ops.all_summary_ops()) + self.assertEqual(1, get_total()) + # Should flush after second summary since max_queue = 1 + sess.run(summary_ops.all_summary_ops()) + self.assertEqual(3, get_total()) + + def testFlushFunction(self): + logdir = self.get_temp_dir() + writer = summary_ops.create_file_writer( + logdir, max_queue=999999, flush_millis=999999) + with writer.as_default(), summary_ops.always_record_summaries(): + summary_ops.scalar('scalar', 2.0, step=1) + flush_op = summary_ops.flush() + with self.test_session() as sess: + sess.run(summary_ops.summary_writer_initializer_op()) + get_total = lambda: len(summary_test_util.events_from_logdir(logdir)) + # Note: First tf.Event is always file_version. + self.assertEqual(1, get_total()) + sess.run(summary_ops.all_summary_ops()) + self.assertEqual(1, get_total()) + sess.run(flush_op) + self.assertEqual(2, get_total()) + # Test "writer" parameter + sess.run(summary_ops.all_summary_ops()) + sess.run(summary_ops.flush(writer=writer)) + self.assertEqual(3, get_total()) + sess.run(summary_ops.all_summary_ops()) + sess.run(summary_ops.flush(writer=writer._resource)) # pylint:disable=protected-access + self.assertEqual(4, get_total()) + + def testSharedName(self): + logdir = self.get_temp_dir() + with summary_ops.always_record_summaries(): + # Create with default shared name (should match logdir) + writer1 = summary_ops.create_file_writer(logdir) + with writer1.as_default(): + summary_ops.scalar('one', 1.0, step=1) + # Create with explicit logdir shared name (should be same resource/file) + shared_name = 'logdir:' + logdir + writer2 = summary_ops.create_file_writer(logdir, name=shared_name) + with writer2.as_default(): + summary_ops.scalar('two', 2.0, step=2) + # Create with different shared name (should be separate resource/file) + writer3 = summary_ops.create_file_writer(logdir, name='other') + with writer3.as_default(): + summary_ops.scalar('three', 3.0, step=3) + + with self.test_session() as sess: + # Run init ops across writers sequentially to avoid race condition. + # TODO(nickfelt): fix race condition in resource manager lookup or create + sess.run(writer1.init()) + sess.run(writer2.init()) + time.sleep(1.1) # Ensure filename has a different timestamp + sess.run(writer3.init()) + sess.run(summary_ops.all_summary_ops()) + sess.run([writer1.flush(), writer2.flush(), writer3.flush()]) + + event_files = iter(sorted(gfile.Glob(os.path.join(logdir, '*tfevents*')))) + + # First file has tags "one" and "two" + events = summary_test_util.events_from_file(next(event_files)) + self.assertEqual('brain.Event:2', events[0].file_version) + tags = [e.summary.value[0].tag for e in events[1:]] + self.assertItemsEqual(['one', 'two'], tags) + + # Second file has tag "three" + events = summary_test_util.events_from_file(next(event_files)) + self.assertEqual('brain.Event:2', events[0].file_version) + tags = [e.summary.value[0].tag for e in events[1:]] + self.assertItemsEqual(['three'], tags) + + # No more files + self.assertRaises(StopIteration, lambda: next(event_files)) + + def testWriterInitAndClose(self): + logdir = self.get_temp_dir() + with summary_ops.always_record_summaries(): + writer = summary_ops.create_file_writer( + logdir, max_queue=100, flush_millis=1000000) + with writer.as_default(): + summary_ops.scalar('one', 1.0, step=1) + with self.test_session() as sess: + sess.run(summary_ops.summary_writer_initializer_op()) + get_total = lambda: len(summary_test_util.events_from_logdir(logdir)) + self.assertEqual(1, get_total()) # file_version Event + # Running init() again while writer is open has no effect + sess.run(writer.init()) + self.assertEqual(1, get_total()) + sess.run(summary_ops.all_summary_ops()) + self.assertEqual(1, get_total()) + # Running close() should do an implicit flush + sess.run(writer.close()) + self.assertEqual(2, get_total()) + # Running init() on a closed writer should start a new file + time.sleep(1.1) # Ensure filename has a different timestamp + sess.run(writer.init()) + sess.run(summary_ops.all_summary_ops()) + sess.run(writer.close()) + files = sorted(gfile.Glob(os.path.join(logdir, '*tfevents*'))) + self.assertEqual(2, len(files)) + self.assertEqual(2, len(summary_test_util.events_from_file(files[1]))) + + def testWriterFlush(self): + logdir = self.get_temp_dir() + with summary_ops.always_record_summaries(): + writer = summary_ops.create_file_writer( + logdir, max_queue=100, flush_millis=1000000) + with writer.as_default(): + summary_ops.scalar('one', 1.0, step=1) + with self.test_session() as sess: + sess.run(summary_ops.summary_writer_initializer_op()) + get_total = lambda: len(summary_test_util.events_from_logdir(logdir)) + self.assertEqual(1, get_total()) # file_version Event + sess.run(summary_ops.all_summary_ops()) + self.assertEqual(1, get_total()) + sess.run(writer.flush()) + self.assertEqual(2, get_total()) + + +class GraphDbTest(summary_test_util.SummaryDbTest): def testGraphPassedToGraph_isForbiddenForThineOwnSafety(self): with self.assertRaises(TypeError): diff --git a/tensorflow/contrib/summary/summary_ops_test.py b/tensorflow/contrib/summary/summary_ops_test.py index c756f8b27055f9cf86a311e485d97745a3c7a95b..f1ef218e74bbd225071324a8269fdfeb5de0e038 100644 --- a/tensorflow/contrib/summary/summary_ops_test.py +++ b/tensorflow/contrib/summary/summary_ops_test.py @@ -16,12 +16,13 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import os import tempfile +import time import numpy as np import six -from tensorflow.contrib.summary import summary_ops from tensorflow.contrib.summary import summary_test_util from tensorflow.core.framework import graph_pb2 from tensorflow.core.framework import node_def_pb2 @@ -33,6 +34,7 @@ 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 +from tensorflow.python.ops import summary_ops_v2 as summary_ops from tensorflow.python.platform import gfile from tensorflow.python.training import training_util @@ -57,7 +59,7 @@ _NUMPY_NUMERIC_TYPES = { } -class TargetTest(test_util.TensorFlowTestCase): +class EagerFileTest(test_util.TensorFlowTestCase): def testShouldRecordSummary(self): self.assertFalse(summary_ops.should_record_summaries()) @@ -138,21 +140,22 @@ class TargetTest(test_util.TensorFlowTestCase): def testMaxQueue(self): logs = tempfile.mkdtemp() with summary_ops.create_file_writer( - logs, max_queue=2, flush_millis=999999, + logs, max_queue=1, flush_millis=999999, name='lol').as_default(), summary_ops.always_record_summaries(): get_total = lambda: len(summary_test_util.events_from_logdir(logs)) # Note: First tf.Event is always file_version. self.assertEqual(1, get_total()) summary_ops.scalar('scalar', 2.0, step=1) self.assertEqual(1, get_total()) + # Should flush after second summary since max_queue = 1 summary_ops.scalar('scalar', 2.0, step=2) self.assertEqual(3, get_total()) - def testFlush(self): + def testFlushFunction(self): logs = tempfile.mkdtemp() - with summary_ops.create_file_writer( - logs, max_queue=999999, flush_millis=999999, - name='lol').as_default(), summary_ops.always_record_summaries(): + writer = summary_ops.create_file_writer( + logs, max_queue=999999, flush_millis=999999, name='lol') + with writer.as_default(), summary_ops.always_record_summaries(): get_total = lambda: len(summary_test_util.events_from_logdir(logs)) # Note: First tf.Event is always file_version. self.assertEqual(1, get_total()) @@ -161,9 +164,103 @@ class TargetTest(test_util.TensorFlowTestCase): self.assertEqual(1, get_total()) summary_ops.flush() self.assertEqual(3, get_total()) + # Test "writer" parameter + summary_ops.scalar('scalar', 2.0, step=3) + summary_ops.flush(writer=writer) + self.assertEqual(4, get_total()) + summary_ops.scalar('scalar', 2.0, step=4) + summary_ops.flush(writer=writer._resource) # pylint:disable=protected-access + self.assertEqual(5, get_total()) + + def testSharedName(self): + logdir = self.get_temp_dir() + with summary_ops.always_record_summaries(): + # Create with default shared name (should match logdir) + writer1 = summary_ops.create_file_writer(logdir) + with writer1.as_default(): + summary_ops.scalar('one', 1.0, step=1) + summary_ops.flush() + # Create with explicit logdir shared name (should be same resource/file) + shared_name = 'logdir:' + logdir + writer2 = summary_ops.create_file_writer(logdir, name=shared_name) + with writer2.as_default(): + summary_ops.scalar('two', 2.0, step=2) + summary_ops.flush() + # Create with different shared name (should be separate resource/file) + time.sleep(1.1) # Ensure filename has a different timestamp + writer3 = summary_ops.create_file_writer(logdir, name='other') + with writer3.as_default(): + summary_ops.scalar('three', 3.0, step=3) + summary_ops.flush() + + event_files = iter(sorted(gfile.Glob(os.path.join(logdir, '*tfevents*')))) + + # First file has tags "one" and "two" + events = iter(summary_test_util.events_from_file(next(event_files))) + self.assertEqual('brain.Event:2', next(events).file_version) + self.assertEqual('one', next(events).summary.value[0].tag) + self.assertEqual('two', next(events).summary.value[0].tag) + self.assertRaises(StopIteration, lambda: next(events)) + + # Second file has tag "three" + events = iter(summary_test_util.events_from_file(next(event_files))) + self.assertEqual('brain.Event:2', next(events).file_version) + self.assertEqual('three', next(events).summary.value[0].tag) + self.assertRaises(StopIteration, lambda: next(events)) + + # No more files + self.assertRaises(StopIteration, lambda: next(event_files)) + + def testWriterInitAndClose(self): + logdir = self.get_temp_dir() + get_total = lambda: len(summary_test_util.events_from_logdir(logdir)) + with summary_ops.always_record_summaries(): + writer = summary_ops.create_file_writer( + logdir, max_queue=100, flush_millis=1000000) + self.assertEqual(1, get_total()) # file_version Event + # Calling init() again while writer is open has no effect + writer.init() + self.assertEqual(1, get_total()) + try: + # Not using .as_default() to avoid implicit flush when exiting + writer.set_as_default() + summary_ops.scalar('one', 1.0, step=1) + self.assertEqual(1, get_total()) + # Calling .close() should do an implicit flush + writer.close() + self.assertEqual(2, get_total()) + # Calling init() on a closed writer should start a new file + time.sleep(1.1) # Ensure filename has a different timestamp + writer.init() + files = sorted(gfile.Glob(os.path.join(logdir, '*tfevents*'))) + self.assertEqual(2, len(files)) + get_total = lambda: len(summary_test_util.events_from_file(files[1])) + self.assertEqual(1, get_total()) # file_version Event + summary_ops.scalar('two', 2.0, step=2) + writer.close() + self.assertEqual(2, get_total()) + finally: + # Clean up by resetting default writer + summary_ops.create_file_writer(None).set_as_default() + + def testWriterFlush(self): + logdir = self.get_temp_dir() + get_total = lambda: len(summary_test_util.events_from_logdir(logdir)) + with summary_ops.always_record_summaries(): + writer = summary_ops.create_file_writer( + logdir, max_queue=100, flush_millis=1000000) + self.assertEqual(1, get_total()) # file_version Event + with writer.as_default(): + summary_ops.scalar('one', 1.0, step=1) + self.assertEqual(1, get_total()) + writer.flush() + self.assertEqual(2, get_total()) + summary_ops.scalar('two', 2.0, step=2) + # Exiting the "as_default()" should do an implicit flush of the "two" tag + self.assertEqual(3, get_total()) -class DbTest(summary_test_util.SummaryDbTest): +class EagerDbTest(summary_test_util.SummaryDbTest): def testIntegerSummaries(self): step = training_util.create_global_step() diff --git a/tensorflow/contrib/summary/summary_test_internal.py b/tensorflow/contrib/summary/summary_test_internal.py deleted file mode 100644 index d0d3384735fb1eb1a048c7aa6da0037ee9fc6936..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/summary/summary_test_internal.py +++ /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. -# ============================================================================== -"""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 8506c4be9c4ca8305b62da17c7246e6e18313bd3..b4ae43302cb22ad17c04050eb84433c470757bf1 100644 --- a/tensorflow/contrib/summary/summary_test_util.py +++ b/tensorflow/contrib/summary/summary_test_util.py @@ -24,10 +24,10 @@ import os import sqlite3 -from tensorflow.contrib.summary import summary_ops from tensorflow.core.util import event_pb2 from tensorflow.python.framework import test_util from tensorflow.python.lib.io import tf_record +from tensorflow.python.ops import summary_ops_v2 as summary_ops from tensorflow.python.platform import gfile diff --git a/tensorflow/contrib/tensor_forest/client/eval_metrics.py b/tensorflow/contrib/tensor_forest/client/eval_metrics.py index 90033015ebc5e44ea70fbf2bc9735d0aeb4ec27d..e893e1d1c836cc7feef15757dde79d0db362cbaf 100644 --- a/tensorflow/contrib/tensor_forest/client/eval_metrics.py +++ b/tensorflow/contrib/tensor_forest/client/eval_metrics.py @@ -37,7 +37,7 @@ def _top_k_generator(k): def _top_k(probabilities, targets): targets = math_ops.to_int32(targets) if targets.get_shape().ndims > 1: - targets = array_ops.squeeze(targets, squeeze_dims=[1]) + targets = array_ops.squeeze(targets, axis=[1]) return metric_ops.streaming_mean(nn.in_top_k(probabilities, targets, k)) return _top_k @@ -57,7 +57,7 @@ def _r2(probabilities, targets, weights=None): def _squeeze_and_onehot(targets, depth): - targets = array_ops.squeeze(targets, squeeze_dims=[1]) + targets = array_ops.squeeze(targets, axis=[1]) return array_ops.one_hot(math_ops.to_int32(targets), depth) diff --git a/tensorflow/contrib/tensor_forest/client/random_forest.py b/tensorflow/contrib/tensor_forest/client/random_forest.py index 4abcc20ed334e706c8ae59e2127dfd6f4e152361..35e8c92aba325d9115c7ee566363a1625e6e76fc 100644 --- a/tensorflow/contrib/tensor_forest/client/random_forest.py +++ b/tensorflow/contrib/tensor_forest/client/random_forest.py @@ -399,7 +399,7 @@ def get_combined_model_fn(model_fns): training ops: tf.group them. loss: average them. predictions: concat probabilities such that predictions[*][0-C1] are the - probablities for output 1 (where C1 is the number of classes in output 1), + probabilities for output 1 (where C1 is the number of classes in output 1), predictions[*][C1-(C1+C2)] are the probabilities for output 2 (where C2 is the number of classes in output 2), etc. Also stack predictions such that predictions[i][j] is the class prediction for example i and output j. diff --git a/tensorflow/contrib/tensor_forest/hybrid/core/ops/hard_routing_function_op.cc b/tensorflow/contrib/tensor_forest/hybrid/core/ops/hard_routing_function_op.cc index cf0db788a419f64ed891df8aa097fa8826f6de91..06bfe871fdff29fcdac1a3a50500b94dbd2998d7 100644 --- a/tensorflow/contrib/tensor_forest/hybrid/core/ops/hard_routing_function_op.cc +++ b/tensorflow/contrib/tensor_forest/hybrid/core/ops/hard_routing_function_op.cc @@ -80,7 +80,7 @@ REGISTER_OP("HardRoutingFunction") regression model that translates from node features to probabilities. - path_probility: `path_probability[i]` gives the probability of reaching each + path_probability: `path_probability[i]` gives the probability of reaching each node in `path[i]`. path: `path[i][j]` gives the jth node in the path taken by the ith data instance. diff --git a/tensorflow/contrib/tensor_forest/hybrid/core/ops/stochastic_hard_routing_function_op.cc b/tensorflow/contrib/tensor_forest/hybrid/core/ops/stochastic_hard_routing_function_op.cc index c9df09bfda44e665ed013da383e1e9a2c665c454..1a055756c084016b3862ad131937e4d400f26d55 100644 --- a/tensorflow/contrib/tensor_forest/hybrid/core/ops/stochastic_hard_routing_function_op.cc +++ b/tensorflow/contrib/tensor_forest/hybrid/core/ops/stochastic_hard_routing_function_op.cc @@ -85,7 +85,7 @@ REGISTER_OP("StochasticHardRoutingFunction") regression model that translates from node features to probabilities. - path_probility: `path_probability[i]` gives the probability of reaching each + path_probability: `path_probability[i]` gives the probability of reaching each node in `path[i]`. path: `path[i][j]` gives the jth node in the path taken by the ith data instance. diff --git a/tensorflow/contrib/tensor_forest/hybrid/core/ops/stochastic_hard_routing_gradient_op.cc b/tensorflow/contrib/tensor_forest/hybrid/core/ops/stochastic_hard_routing_gradient_op.cc index b0d8b832b5437db7a4b3026e80ae99d0391d7f7a..7d092bbc24d5c6611fed431c04200b2f950887f7 100644 --- a/tensorflow/contrib/tensor_forest/hybrid/core/ops/stochastic_hard_routing_gradient_op.cc +++ b/tensorflow/contrib/tensor_forest/hybrid/core/ops/stochastic_hard_routing_gradient_op.cc @@ -81,7 +81,7 @@ REGISTER_OP("StochasticHardRoutingGradient") tree_biases: `tree_biases[i]` gives the bias of the logistic regression model that translates from node features to probabilities. - path_probility: `path_probability[i]` gives the probability of reaching each + path_probability: `path_probability[i]` gives the probability of reaching each node in `path[i]`. path: `path[i][j]` gives the jth node in the path taken by the ith data instance. diff --git a/tensorflow/contrib/tensor_forest/hybrid/python/layers/fully_connected.py b/tensorflow/contrib/tensor_forest/hybrid/python/layers/fully_connected.py index ff3ab21eaa9a4aa823f2ae7d3dd39674abea3d2a..745a5b1caf2fe348f1b276ccc245aa2ef350a62e 100644 --- a/tensorflow/contrib/tensor_forest/hybrid/python/layers/fully_connected.py +++ b/tensorflow/contrib/tensor_forest/hybrid/python/layers/fully_connected.py @@ -55,7 +55,7 @@ class ManyToOneLayer(hybrid_layer.HybridLayer): # There is always one activation per instance by definition, so squeeze # away the extra dimension. - return array_ops.squeeze(nn_activations, squeeze_dims=[1]) + return array_ops.squeeze(nn_activations, axis=[1]) class FlattenedFullyConnectedLayer(hybrid_layer.HybridLayer): diff --git a/tensorflow/contrib/tensor_forest/kernels/tree_utils.cc b/tensorflow/contrib/tensor_forest/kernels/tree_utils.cc index 44997ec5d6d5fdb9aab52ab7a50f46a731bfda66..cefcc960510293c15569391091e5c1b4611c8835 100644 --- a/tensorflow/contrib/tensor_forest/kernels/tree_utils.cc +++ b/tensorflow/contrib/tensor_forest/kernels/tree_utils.cc @@ -421,7 +421,7 @@ double getChebyshevEpsilon(const std::vector& mu1, const std::vector& mu2) { // Math time!! // We are trying to minimize d = |mu1 - x|^2 + |mu2 - y|^2 over the surface. - // Using Langrange multipliers, we get + // Using Lagrange multipliers, we get // partial d / partial x = -2 mu1 + 2 x = lambda_1 1 + 2 lambda_3 x // partial d / partial y = -2 mu2 + 2 y = lambda_2 1 - 2 lambda_3 y // or @@ -485,7 +485,7 @@ double getChebyshevEpsilon(const std::vector& mu1, } double sdiscrim = sqrt(discrim); - // TODO(thomaswc): Analyze whetever one of these is always closer. + // TODO(thomaswc): Analyze whatever one of these is always closer. double v1 = (-b + sdiscrim) / (2 * a); double v2 = (-b - sdiscrim) / (2 * a); double dist1 = getDistanceFromLambda3(v1, mu1, mu2); diff --git a/tensorflow/contrib/tensor_forest/kernels/tree_utils.h b/tensorflow/contrib/tensor_forest/kernels/tree_utils.h index edbac6700677633cbd4d41f7040b4859ca599c4a..03aab1b61ee58a647edb24f6b97e517a411e996c 100644 --- a/tensorflow/contrib/tensor_forest/kernels/tree_utils.h +++ b/tensorflow/contrib/tensor_forest/kernels/tree_utils.h @@ -123,7 +123,7 @@ bool BestSplitDominatesRegression(const Tensor& total_sums, const Tensor& split_squares, int32 accumulator); -// Performs booststrap_samples bootstrap samples of the best split's class +// Performs bootstrap_samples bootstrap samples of the best split's class // counts and the second best splits's class counts, and returns true if at // least dominate_fraction of the time, the former has a better (lower) // Gini impurity. Does not take over ownership of *rand. diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/decision-tree-resource.h b/tensorflow/contrib/tensor_forest/kernels/v4/decision-tree-resource.h index 328af28725af016e90b30ae2d303ffba15c81c1f..d3edb43733761a906c6e5bf8b65f76e3e1ae56fc 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/decision-tree-resource.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/decision-tree-resource.h @@ -60,7 +60,7 @@ class DecisionTreeResource : public ResourceBase { mutex* get_mutex() { return &mu_; } // Return the TreeNode for the leaf that the example ends up at according - // to decsion_tree_. Also fill in that leaf's depth if it isn't nullptr. + // to decision_tree_. Also fill in that leaf's depth if it isn't nullptr. int32 TraverseTree(const std::unique_ptr& input_data, int example, int32* depth, TreePath* path) const; diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/decision_node_evaluator.h b/tensorflow/contrib/tensor_forest/kernels/v4/decision_node_evaluator.h index bf2b2aaa3c8f433ab4fc145217857112f7a0a579..3db351c328c73beb94d6994aa503e3e2c4c06390 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/decision_node_evaluator.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/decision_node_evaluator.h @@ -60,7 +60,7 @@ class InequalityDecisionNodeEvaluator : public BinaryDecisionNodeEvaluator { bool include_equals_; }; -// Evalutor for splits with multiple weighted features. +// Evaluator for splits with multiple weighted features. class ObliqueInequalityDecisionNodeEvaluator : public BinaryDecisionNodeEvaluator { public: diff --git a/tensorflow/contrib/tensor_forest/ops/model_ops.cc b/tensorflow/contrib/tensor_forest/ops/model_ops.cc index 3099cccdf8b1b7b1d49d12ec454f9a0dc7325e9b..98124d519c761907ed0e62ee2a396ef84716a911 100644 --- a/tensorflow/contrib/tensor_forest/ops/model_ops.cc +++ b/tensorflow/contrib/tensor_forest/ops/model_ops.cc @@ -165,7 +165,7 @@ tree_handle: The handle to the tree. leaf_ids: `leaf_ids[i]` is the leaf id for input i. input_labels: The training batch's labels as a 1 or 2-d tensor. 'input_labels[i][j]' gives the j-th label/target for the i-th input. -input_weights: The training batch's eample weights as a 1-d tensor. +input_weights: The training batch's weights as a 1-d tensor. 'input_weights[i]' gives the weight for the i-th input. )doc"); diff --git a/tensorflow/contrib/tensor_forest/ops/stats_ops.cc b/tensorflow/contrib/tensor_forest/ops/stats_ops.cc index e8b5c5d8a6efae83b9e408ed3d05ac72d848ef7f..5be581aaec4cab342ef3fd49fa0294e5e702ba1c 100644 --- a/tensorflow/contrib/tensor_forest/ops/stats_ops.cc +++ b/tensorflow/contrib/tensor_forest/ops/stats_ops.cc @@ -75,7 +75,7 @@ REGISTER_OP("GrowTreeV4") .Attr("params: string") .Input("tree_handle: resource") .Input("stats_handle: resource") - .Input("finshed_nodes: int32") + .Input("finished_nodes: int32") .SetShapeFn(tensorflow::shape_inference::NoOutputs) .Doc(R"doc( Grows the tree for finished nodes and allocates waiting nodes. @@ -83,7 +83,7 @@ Grows the tree for finished nodes and allocates waiting nodes. params: A serialized TensorForestParams proto. tree_handle: The handle to the tree. stats_handle: The handle to the stats. -finshed_nodes: A 1-d Tensor of finished node ids from ProcessInput. +finished_nodes: A 1-d Tensor of finished node ids from ProcessInput. )doc"); REGISTER_OP("ProcessInputV4") @@ -119,7 +119,7 @@ sparse_input_values: The values tensor from the SparseTensor input. sparse_input_shape: The shape tensor from the SparseTensor input. input_labels: The training batch's labels as a 1 or 2-d tensor. 'input_labels[i][j]' gives the j-th label/target for the i-th input. -input_weights: The training batch's eample weights as a 1-d tensor. +input_weights: The training batch's weights as a 1-d tensor. 'input_weights[i]' gives the weight for the i-th input. finished_nodes: A 1-d tensor of node ids that have finished and are ready to grow. diff --git a/tensorflow/contrib/tensor_forest/python/tensor_forest.py b/tensorflow/contrib/tensor_forest/python/tensor_forest.py index 3650b5d52fe8a1b87a239d41ecfa3de677fffc72..7a35a70bbe3112e0649cefd8116cc50565978da5 100644 --- a/tensorflow/contrib/tensor_forest/python/tensor_forest.py +++ b/tensorflow/contrib/tensor_forest/python/tensor_forest.py @@ -212,7 +212,7 @@ class ForestHParams(object): self.regression = getattr(self, 'regression', False) # Num_outputs is the actual number of outputs (a single prediction for - # classification, a N-dimenensional point for regression). + # classification, a N-dimensional point for regression). self.num_outputs = self.num_classes if self.regression else 1 # Add an extra column to classes for storing counts, which is needed for @@ -445,7 +445,7 @@ class RandomForestGraphs(object): mask = math_ops.less( r, array_ops.ones_like(r) * self.params.bagging_fraction) gather_indices = array_ops.squeeze( - array_ops.where(mask), squeeze_dims=[1]) + array_ops.where(mask), axis=[1]) # TODO(thomaswc): Calculate out-of-bag data and labels, and store # them for use in calculating statistics later. tree_data = array_ops.gather(processed_dense_features, gather_indices) diff --git a/tensorflow/contrib/tensorboard/db/summary_file_writer.cc b/tensorflow/contrib/tensorboard/db/summary_file_writer.cc index 85b3e7231bcb433e9510522597c03c5f764f06cf..3f24f58f03aac2ba6d368d7eccf8731f611a81b4 100644 --- a/tensorflow/contrib/tensorboard/db/summary_file_writer.cc +++ b/tensorflow/contrib/tensorboard/db/summary_file_writer.cc @@ -132,7 +132,7 @@ class SummaryFileWriter : public SummaryWriterInterface { Status WriteEvent(std::unique_ptr event) override { mutex_lock ml(mu_); queue_.emplace_back(std::move(event)); - if (queue_.size() >= max_queue_ || + if (queue_.size() > max_queue_ || env_->NowMicros() - last_flush_ > 1000 * flush_millis_) { return InternalFlush(); } diff --git a/tensorflow/contrib/tensorrt/BUILD b/tensorflow/contrib/tensorrt/BUILD index fd3582e175ee91ba7222d41fd3f834c522b5e28d..f80b4f1b112dcfb44244cff56d6915af9f455caa 100644 --- a/tensorflow/contrib/tensorrt/BUILD +++ b/tensorflow/contrib/tensorrt/BUILD @@ -11,6 +11,7 @@ exports_files(["LICENSE"]) load( "//tensorflow:tensorflow.bzl", + "py_test", "tf_cc_test", "tf_copts", "tf_cuda_library", @@ -139,6 +140,7 @@ tf_custom_op_py_library( ]), srcs_version = "PY2AND3", deps = [ + "//tensorflow/contrib/util:util_py", "//tensorflow/python:framework_for_generated_wrappers", "//tensorflow/python:resources", ], @@ -173,6 +175,7 @@ py_library( srcs_version = "PY2AND3", deps = [ ":wrap_conversion", + "//tensorflow/python:tf_optimizer", ], ) @@ -272,3 +275,19 @@ tf_cc_test( "//tensorflow/core:test_main", ], ) + +py_test( + name = "tf_trt_integration_test", + srcs = ["test/tf_trt_integration_test.py"], + main = "test/tf_trt_integration_test.py", + srcs_version = "PY2AND3", + tags = [ + "manual", + "notap", + ], + deps = [ + ":init_py", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_test_lib", + ], +) diff --git a/tensorflow/contrib/tensorrt/README.md b/tensorflow/contrib/tensorrt/README.md index 6eafc1754ca5102c8adf04f00e33dc2f8ff970f6..687dee07e1327d50fabc4e14c25a357ae6c959e7 100644 --- a/tensorflow/contrib/tensorrt/README.md +++ b/tensorflow/contrib/tensorrt/README.md @@ -1,59 +1,29 @@ # 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. +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. +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 script +After the installation of tensorflow package, TensorRT transformation will be +available. An example use can be found in test/test_tftrt.py script ## Installing TensorRT 3.0.4 -In order to make use of TensorRT integration, you will need a local installation of TensorRT 3.0.4 from the [NVIDIA Developer website](https://developer.nvidia.com/tensorrt). Due to compiler compatibility, you will need to download and install the TensorRT 3.0.4 tarball for _Ubuntu 14.04_, i.e., **_TensorRT-3.0.4.Ubuntu-14.04.5.x86_64.cuda-9.0.cudnn7.0-tar.gz_**, even if you are using Ubuntu 16.04 or later. - -### Preparing TensorRT installation - -Once you have downloaded TensorRT-3.0.4.Ubuntu-14.04.5.x86_64.cuda-9.0.cudnn7.0-tar.gz, you will need to unpack it to an installation directory, which will be referred to as . Please replace with the full path of actual installation directory you choose in commands below. - -```shell -cd && tar -zxf /path/to/TensorRT-3.0.4.Ubuntu-14.04.5.x86_64.cuda-9.0.cudnn7.0-tar.gz -``` - -After unpacking the binaries, you have several options to use them: - -#### To run TensorFlow as a user without superuser privileges - -For a regular user without any sudo rights, you should add TensorRT to your `$LD_LIBRARY_PATH`: - - ```shell - export LD_LIBRARY_PATH=/TensorRT-3.0.4/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}} - ``` - -Then you are ready to use TensorFlow-TensorRT integration. `$LD_LIBRARY_PATH` must contain the path to TensorRT installation for TensorFlow-TensorRT integration to work. If you are using a VirtualEnv-like setup, you can add the command above to your `bin/activate` script or to your `.bashrc` script. - -#### To run TensorFlow as a superuser - - When running as a superuser, such as in a container or via sudo, the `$LD_LIBRARY_PATH` approach above may not work. The following is preferred when the user has superuser privileges: - - ```shell - echo "/TensorRT-3.0.4/lib" | sudo tee /etc/ld.so.conf.d/tensorrt304.conf && sudo ldconfig - ``` - - Please ensure that any existing deb package installation of TensorRT is removed before following these instructions to avoid package conflicts. \ No newline at end of file +In order to make use of TensorRT integration, you will need a local installation +of TensorRT 3.0.4 from the [NVIDIA Developer website](https://developer.nvidia.com/tensorrt). +Installation instructions for compatibility with TensorFlow are provided on the +[TensorFlow Installation page](https://www.tensorflow.org/install/install_linux#nvidia_requirements_to_run_tensorflow_with_gpu_support). diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc index 567b4af88df65b8dd83c85a8430370bff611e917..b81ae9dc3eeed6f7b7c6eeac0186700bdd692245 100644 --- a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc +++ b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc @@ -444,8 +444,8 @@ class Converter { * remove this and annotate the edge as a control dependency. ************************************************************************/ // skip control nodes - if (input_name[0] == '^' ) continue; - string name = input_name; + if (input_name[0] == '^') continue; + string name = input_name; auto first = name.find_first_of(':'); if (first != string::npos && first + 2 == name.size() && name[first + 1] == '0') @@ -2511,7 +2511,7 @@ tensorflow::Status ConvertSubGraphToTensorRTNodeDef( std::vector input_names; std::vector input_dtypes; for (const std::pair& input : s.input_inds) { - VLOG(2) << "parsing input. Node id= " << input.first ; + 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); diff --git a/tensorflow/contrib/tensorrt/test/tf_trt_integration_test.py b/tensorflow/contrib/tensorrt/test/tf_trt_integration_test.py new file mode 100644 index 0000000000000000000000000000000000000000..7a4732876286a9484bc607242ae19a31941313db --- /dev/null +++ b/tensorflow/contrib/tensorrt/test/tf_trt_integration_test.py @@ -0,0 +1,156 @@ +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 warnings +import numpy as np + +from tensorflow.contrib import tensorrt as trt +from tensorflow.core.protobuf import config_pb2 as cpb2 +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.framework import test_util +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 +from tensorflow.python.platform import googletest + + +@test_util.with_c_api +class IntegrationTest(test_util.TensorFlowTestCase): + """Class to test Tensorflow-TensorRT integration.""" + + def setUp(self): + """Setup method.""" + super(IntegrationTest, self).setUp() + warnings.simplefilter("always") + inp_dims = (100, 24, 24, 2) + self._input = np.random.random_sample(inp_dims) + self._original_graph = self.get_simple_graph_def() + self._gpu_options = cpb2.GPUOptions( + per_process_gpu_memory_fraction=0.50) + self._config = cpb2.ConfigProto(gpu_options=self._gpu_options) + self._reference = self.run_graph(self._original_graph, self._input) + + def get_simple_graph_def(self): + """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(self, gdef, dumm_inp): + """Run given graphdef once.""" + 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 self.test_session( + graph=g, config=self._config, use_gpu=True, + force_gpu=True) 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(self, gdef, dumm_inp): + """Run given calibration graph multiple times.""" + 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] + # run over real calibration data here, we are mimicking a calibration + # set of 30 different batches. Use as much calibration data as you want + with self.test_session( + graph=g, config=self._config, use_gpu=True, + force_gpu=True) as sess: + for _ in range(30): + val = sess.run(out, {inp: dumm_inp}) + return val + + def get_trt_graph(self, mode): + """Return trt converted graph.""" + if mode in ["FP32", "FP16", "INT8"]: + return trt.create_inference_graph( + input_graph_def=self._original_graph, + outputs=["output"], + max_batch_size=self._input.shape[0], + max_workspace_size_bytes=1 << 25, + precision_mode=mode, # TRT Engine precision "FP32","FP16" or "INT8" + minimum_segment_size=2 # minimum number of nodes in an engine + ) + return None + + def testFP32(self): + """Test FP32 conversion. Results should be identical to native case.""" + trt_graph = self.get_trt_graph("FP32") + result = self.run_graph(trt_graph, self._input) + self.assertAllEqual(self._reference, result) + result1 = self.run_graph(trt_graph, self._input) + self.assertAllEqual(result1, result) + + def testFP16(self): + """Test FP16 conversion. Results may be different from native case.""" + trt_graph = self.get_trt_graph("FP16") + result = self.run_graph(trt_graph, self._input) + self.assertAllClose(self._reference, result, rtol=1.e-03) + result1 = self.run_graph(trt_graph, self._input) + self.assertAllEqual(result1, result) + + def testINT8(self): + """Test INT8 conversion. Results may be different from native case.""" + calib_graph = self.get_trt_graph("INT8") + result = self.run_calibration(calib_graph, self._input) + self.assertAllEqual(self._reference, result) + int8_graph = trt.calib_graph_to_infer_graph(calib_graph) + result = self.run_graph(int8_graph, self._input) + self.assertAllClose(self._reference, result, rtol=1.e-03) + result1 = self.run_graph(int8_graph, self._input) + self.assertAllEqual(result1, result) + + +if __name__ == "__main__": + googletest.main() diff --git a/tensorflow/contrib/timeseries/examples/BUILD b/tensorflow/contrib/timeseries/examples/BUILD index 32e948a009741b126e21a64473ac2d020a25a7af..355303acf6ddf866ecf18815b394fcea8488d67d 100644 --- a/tensorflow/contrib/timeseries/examples/BUILD +++ b/tensorflow/contrib/timeseries/examples/BUILD @@ -8,14 +8,22 @@ licenses(["notice"]) # Apache 2.0 exports_files(["LICENSE"]) +config_setting( + name = "empty_condition", + values = {"define": "UNUSED=unused"}, +) + py_binary( name = "predict", srcs = ["predict.py"], srcs_version = "PY2AND3", tags = ["no_pip"], - deps = [ - "//tensorflow:tensorflow_py", + deps = select({ + ":empty_condition": [], + "//conditions:default": [], + }) + [ "//third_party/py/numpy", + "//tensorflow:tensorflow_py", ], ) @@ -41,9 +49,12 @@ py_binary( data = ["data/changepoints.csv"], srcs_version = "PY2AND3", tags = ["no_pip"], - deps = [ - "//tensorflow:tensorflow_py", + deps = select({ + ":empty_condition": [], + "//conditions:default": [], + }) + [ "//third_party/py/numpy", + "//tensorflow:tensorflow_py", ], ) @@ -64,9 +75,12 @@ py_binary( data = ["data/multivariate_level.csv"], srcs_version = "PY2AND3", tags = ["no_pip"], - deps = [ - "//tensorflow:tensorflow_py", + deps = select({ + ":empty_condition": [], + "//conditions:default": [], + }) + [ "//third_party/py/numpy", + "//tensorflow:tensorflow_py", ], ) @@ -89,11 +103,14 @@ py_binary( data = ["data/multivariate_periods.csv"], srcs_version = "PY2AND3", tags = ["no_pip"], - deps = [ + deps = select({ + ":empty_condition": [], + "//conditions:default": [], + }) + [ + "//third_party/py/numpy", "//tensorflow:tensorflow_py", "//tensorflow/contrib/timeseries/python/timeseries:estimators", "//tensorflow/contrib/timeseries/python/timeseries:model", - "//third_party/py/numpy", ], ) diff --git a/tensorflow/contrib/timeseries/examples/lstm_test.py b/tensorflow/contrib/timeseries/examples/lstm_test.py index ca56e38ca079f71b38cf29605a295a50929945e8..c58e24e6d9748868791d21b0ff4ec28ca2f646c3 100644 --- a/tensorflow/contrib/timeseries/examples/lstm_test.py +++ b/tensorflow/contrib/timeseries/examples/lstm_test.py @@ -36,17 +36,14 @@ 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=2, 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) self.assertAllEqual([200, 5], predicted_values.shape) - self.assertGreater( - predicted_values[100, 4] - - predicted_values[115, 4], # Amplitude of fifth component - 0.2) - + # TODO(allenl): Make the model deterministic so you can check something + # substantive. if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/timeseries/python/timeseries/math_utils.py b/tensorflow/contrib/timeseries/python/timeseries/math_utils.py index 26793c80bfbb3c9394e81a5bbfae360deb95ca58..9b593fecbb3fbc3b8b57848462c85dff4c3b7577 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/math_utils.py +++ b/tensorflow/contrib/timeseries/python/timeseries/math_utils.py @@ -60,7 +60,7 @@ def clip_covariance( # TODO(allenl): Smarter scaling here so that correlations are preserved when # fiddling with diagonal elements. diagonal = array_ops.matrix_diag_part(covariance_matrix) - maximum = math_ops.reduce_max(diagonal, axis=-1, keep_dims=True) + maximum = math_ops.reduce_max(diagonal, axis=-1, keepdims=True) new_diagonal = gen_math_ops.maximum( diagonal, maximum / maximum_variance_ratio) return array_ops.matrix_set_diag( diff --git a/tensorflow/contrib/timeseries/python/timeseries/state_management_test.py b/tensorflow/contrib/timeseries/python/timeseries/state_management_test.py index d5dce30fda0353bd70f44ec567ac91acce1e9394..5f7e3da2db6da26f50aad9d500959238063a3e3c 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/state_management_test.py +++ b/tensorflow/contrib/timeseries/python/timeseries/state_management_test.py @@ -78,7 +78,7 @@ class StubTimeSeriesModel(model.TimeSeriesModel): batch_end_values = array_ops.squeeze( array_ops.slice(values, [0, array_ops.shape(times)[1] - 1, 0], [-1, 1, -1]), - squeeze_dims=[1, 2]) + axis=[1, 2]) # A pretty odd but easy to think about loss: L1 loss on the batch end # values. loss = math_ops.reduce_sum( diff --git a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/kalman_filter.py b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/kalman_filter.py index 1fcd3e391b63c2362d6187da9556e2c71836dbaa..a614386121e000961bf8b32625a28e1251654320 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/kalman_filter.py +++ b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/kalman_filter.py @@ -170,7 +170,7 @@ class KalmanFilter(object): math_ops.matmul( transition_matrices, prior_state[..., None]), - squeeze_dims=[-1]) + axis=[-1]) return advanced_state def predict_state_var( @@ -254,7 +254,7 @@ class KalmanFilter(object): kalman_gain_transposed, array_ops.expand_dims(residual, -1), adjoint_a=True), - squeeze_dims=[-1]) + axis=[-1]) gain_obs = math_ops.matmul( kalman_gain_transposed, observation_model, adjoint_a=True) identity_extradim = linalg_ops.eye( @@ -332,7 +332,7 @@ class KalmanFilter(object): array_ops.expand_dims(state_mean, 1), observation_model, adjoint_b=True), - squeeze_dims=[1]) + axis=[1]) observed_var = math_ops.matmul( math_ops.matmul(observation_model, state_var), observation_model, diff --git a/tensorflow/contrib/tpu/BUILD b/tensorflow/contrib/tpu/BUILD index 2f4a76720d2fa1a982716a16ab831b70416f3bbc..9646d15486ef618f206936ce55a5eb6ca0387e41 100644 --- a/tensorflow/contrib/tpu/BUILD +++ b/tensorflow/contrib/tpu/BUILD @@ -46,7 +46,6 @@ py_library( deps = [ ":tpu_lib", ":tpu_py", - "//tensorflow/contrib/summary:summary_ops", "//tensorflow/contrib/training:training_py", "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", @@ -57,6 +56,7 @@ py_library( "//tensorflow/python:platform", "//tensorflow/python:state_ops", "//tensorflow/python:summary", + "//tensorflow/python:summary_ops_v2", "//tensorflow/python:training", "//tensorflow/python:variable_scope", "//tensorflow/python:variables", @@ -198,7 +198,8 @@ py_library( ], srcs_version = "PY2AND3", deps = [ - "//tensorflow/contrib/data/python/ops:transformation_ops", + "//tensorflow/contrib/data/python/ops:batching", + "//tensorflow/contrib/data/python/ops:interleave_ops", "//tensorflow/python:dtypes", "//tensorflow/python:function", "//tensorflow/python:functional_ops", diff --git a/tensorflow/contrib/tpu/__init__.py b/tensorflow/contrib/tpu/__init__.py index bb60f3e2d771b351058322b908dfe68df4abed30..dc9066855990f372c28dc481959117daa4c2da97 100644 --- a/tensorflow/contrib/tpu/__init__.py +++ b/tensorflow/contrib/tpu/__init__.py @@ -43,6 +43,7 @@ @@TPUEstimator @@TPUEstimatorSpec @@RunConfig +@@InputPipelineConfig @@TPUConfig """ diff --git a/tensorflow/contrib/tpu/ops/outfeed_ops.cc b/tensorflow/contrib/tpu/ops/outfeed_ops.cc index 5900c61a38726551391c212f92b9b9eacd4a465b..b05c76ca64fbaedc205ab06cc31616787ccc84b8 100644 --- a/tensorflow/contrib/tpu/ops/outfeed_ops.cc +++ b/tensorflow/contrib/tpu/ops/outfeed_ops.cc @@ -26,6 +26,7 @@ REGISTER_OP("OutfeedEnqueue") .Input("input: dtype") .Attr("dtype: type") .SetIsStateful() + .SetShapeFn(shape_inference::NoOutputs) .Doc(R"doc( An op which emits a single Tensor value from an XLA computation. @@ -36,6 +37,7 @@ REGISTER_OP("OutfeedEnqueueTuple") .Input("inputs: dtypes") .Attr("dtypes: list(type)") .SetIsStateful() + .SetShapeFn(shape_inference::NoOutputs) .Doc(R"doc( An op which emits multiple Tensor values from an XLA computation. diff --git a/tensorflow/contrib/tpu/profiler/BUILD b/tensorflow/contrib/tpu/profiler/BUILD index 1c32993e8e546a17b8b3c289a306ad8f8388c345..dbf1ab6bbf0ddc7429d8e19279451eb862981e0c 100644 --- a/tensorflow/contrib/tpu/profiler/BUILD +++ b/tensorflow/contrib/tpu/profiler/BUILD @@ -46,6 +46,7 @@ tf_cc_binary( visibility = ["//visibility:public"], deps = [ ":dump_tpu_profile", + ":tpu_profiler_analysis_proto_cc", ":tpu_profiler_proto_cc", ":version", "//tensorflow/core:framework_internal", diff --git a/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc b/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc index 6b198dbc16e544686e35c1ffee8a7f4d3955dafc..816897499b7a49365060c026d60b977990f3ecdc 100644 --- a/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc +++ b/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc @@ -26,6 +26,7 @@ limitations under the License. #include "tensorflow/contrib/tpu/profiler/dump_tpu_profile.h" #include "tensorflow/contrib/tpu/profiler/tpu_profiler.grpc.pb.h" +#include "tensorflow/contrib/tpu/profiler/tpu_profiler_analysis.grpc.pb.h" #include "tensorflow/contrib/tpu/profiler/version.h" #include "tensorflow/core/distributed_runtime/rpc/grpc_util.h" #include "tensorflow/core/lib/core/errors.h" @@ -40,6 +41,7 @@ namespace tensorflow { namespace tpu { namespace { +using ::tensorflow::TPUProfileAnalysis; using ::tensorflow::TPUProfiler; constexpr uint64 kMaxEvents = 1000000; @@ -64,11 +66,10 @@ Status ValidateHostPortPair(const string& host_port) { return Status::OK(); } -// Returns whether the returned trace is empty. -// Failure are handled by CHECK, i.e. abort() -bool Profile(const string& service_addr, const string& logdir, int duration_ms, - const string& repository_root, const string& session_id, - const ProfileOptions& opts) { +ProfileRequest PopulateProfileRequest(int duration_ms, + const string& repository_root, + const string& session_id, + const ProfileOptions& opts) { ProfileRequest request; request.set_duration_ms(duration_ms); request.set_max_events(kMaxEvents); @@ -83,6 +84,17 @@ bool Profile(const string& service_addr, const string& logdir, int duration_ms, *request.mutable_opts() = opts; std::cout << "Limiting the number of trace events to " << kMaxEvents << std::endl; + return request; +} + +// Returns whether the returned trace is empty. +// Failure are handled by CHECK, i.e. abort() +bool Profile(const string& service_addr, const string& logdir, int duration_ms, + const string& repository_root, const string& session_id, + const ProfileOptions& opts) { + ProfileRequest request = + PopulateProfileRequest(duration_ms, repository_root, session_id, opts); + ::grpc::ClientContext context; ::grpc::ChannelArguments channel_args; // TODO(ioeric): use `SetMaxReceiveMessageSize` instead once it's available. @@ -120,7 +132,36 @@ bool NewSession(const string& service_addr, const std::vector& hostnames, int duration_ms, const string& repository_root, const string& session_id, const ProfileOptions& opts) { - return true; + NewProfileSessionRequest new_session_request; + *new_session_request.mutable_request() = + PopulateProfileRequest(duration_ms, repository_root, session_id, opts); + new_session_request.set_repository_root(repository_root); + new_session_request.set_session_id(session_id); + for (const auto& hostname : hostnames) { + new_session_request.add_hosts(hostname); + } + + ::grpc::ClientContext context; + ::grpc::ChannelArguments channel_args; + // TODO(qiuminxu): use `NewHostPortGrpcChannel` instead once their + // `ValidateHostPortPair` checks for empty host string case. + channel_args.SetMaxReceiveMessageSize(std::numeric_limits::max()); + // TODO(jiesun): GRPC support following relevant naming scheme: + // 1. dns:///host:port + // 2. ipv4:host:port or ipv6:[host]:port + // We might need to change the prefix which depends on what TPU name resolver + // will give us. + std::unique_ptr stub = + TPUProfileAnalysis::NewStub(::grpc::CreateCustomChannel( + "dns:///" + service_addr, ::grpc::InsecureChannelCredentials(), + channel_args)); + NewProfileSessionResponse new_session_response; + TF_QCHECK_OK(FromGrpcStatus( + stub->NewSession(&context, new_session_request, &new_session_response))); + + std::cout << "Profile session succeed for host(s):" + << str_util::Join(hostnames, ",") << std::endl; + return new_session_response.empty_trace(); } } // namespace diff --git a/tensorflow/contrib/tpu/profiler/dump_tpu_profile.cc b/tensorflow/contrib/tpu/profiler/dump_tpu_profile.cc index ae508583f848a94d9a52d8663af96d85d8fff74c..5e85a967ad4ea373e213fa90c3640e9ab1f92d25 100644 --- a/tensorflow/contrib/tpu/profiler/dump_tpu_profile.cc +++ b/tensorflow/contrib/tpu/profiler/dump_tpu_profile.cc @@ -64,7 +64,8 @@ Status WriteGzippedDataToFile(const string& filename, const string& data) { Status DumpTraceToLogDirectory(StringPiece run_dir, const string& host_prefix, const string& encoded_trace, std::ostream* os) { - string proto_path = JoinPath(run_dir, kProtoTraceFileName); + string proto_path = + JoinPath(run_dir, StrCat(host_prefix, kProtoTraceFileName)); TF_RETURN_IF_ERROR( WriteStringToFile(Env::Default(), proto_path, encoded_trace)); LOG(INFO) << "Dumped raw-proto trace data to " << proto_path; @@ -127,6 +128,7 @@ Status WriteTensorboardTPUProfile(const string& logdir, const string& run, // Dumps profile data to /plugins/profile//. string host_prefix = host.empty() ? "" : StrCat(host, "."); string profile_run_dir = JoinPath(logdir, kProfilePluginDirectory, run); + *os << "Creating directory: " << profile_run_dir; TF_RETURN_IF_ERROR(Env::Default()->RecursivelyCreateDir(profile_run_dir)); // Ignore computation_graph for now. 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 0b78cf8695091daf797bcb80586397e7ab1c6284..508c7a842fb82ec080082d7e7f02f8d2f2a79447 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 @@ -37,12 +37,17 @@ flags.DEFINE_string( '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.') + 'specify either this flag or --service_addr.') # 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( + 'workers_list', None, 'The list of worker TPUs that we are about to profile' + ' e.g. 10.0.1.2, 10.0.1.3. You can specify this flag with --tpu_name or ' + '--service_addr to profile a subset of tpu nodes. You can also use only' + '--tpu_name and leave this flag unspecified to profile all the tpus.') flags.DEFINE_string('logdir', None, 'Path of TensorBoard log directory e.g. /tmp/tb_log, ' 'gs://tb_bucket') @@ -56,18 +61,25 @@ flags.DEFINE_boolean('include_dataset_ops', True, FLAGS = flags.FLAGS EXECUTABLE = 'data/capture_tpu_profile' +JOB_NAME = 'worker' +def get_workers_list(cluster_resolver): + cluster_spec = cluster_resolver.cluster_spec() + task_indices = cluster_spec.task_indices(JOB_NAME) + workers_list = [cluster_spec.task_address(JOB_NAME, i).split(':')[0] + for i in task_indices] + return ','.join(workers_list) def run_main(): tf.app.run(main) - def main(unused_argv=None): 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.') + tpu_cluster_resolver = None 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 ' @@ -82,6 +94,12 @@ def main(unused_argv=None): service_addr = tpu_cluster_resolver.get_master() service_addr = service_addr.replace('grpc://', '').replace(':8470', ':8466') + workers_list = "" + if FLAGS.workers_list is not None: + workers_list = FLAGS.workers_list + elif tpu_cluster_resolver is not None: + workers_list = get_workers_list(tpu_cluster_resolver) + if not FLAGS.logdir: sys.exit('logdir must be provided.') executable_path = os.path.join(os.path.dirname(__file__), EXECUTABLE) @@ -89,6 +107,7 @@ def main(unused_argv=None): cmd = [executable_path] cmd.append('--logdir=' + logdir) cmd.append('--service_addr=' + service_addr) + cmd.append('--workers_list=' + workers_list) 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()) diff --git a/tensorflow/contrib/tpu/profiler/pip_package/setup.py b/tensorflow/contrib/tpu/profiler/pip_package/setup.py index 8d99835b64152629c66607e6792495eb36319eb8..ebd478fd02295108b9d2454963eb06165828b523 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.6.0-rc1' +_VERSION = '1.6.0' CONSOLE_SCRIPTS = [ 'capture_tpu_profile=cloud_tpu_profiler.main:run_main', diff --git a/tensorflow/contrib/tpu/profiler/tf_op_stats.proto b/tensorflow/contrib/tpu/profiler/tf_op_stats.proto index 2a158756279b5be7f818950b7f7b70571b5bc38e..63955d18068fc9d3b3ca1a657a3fd526edf10e6f 100644 --- a/tensorflow/contrib/tpu/profiler/tf_op_stats.proto +++ b/tensorflow/contrib/tpu/profiler/tf_op_stats.proto @@ -66,6 +66,10 @@ message OpMetricsDbResult { // The total of the difference between the start times of two // consecutive infeed-enqueues (per host) in picoseconds. optional uint64 total_host_infeed_enq_start_timestamp_ps_diff = 3; + // The total device time in microseconds. + optional double total_device_time_in_us = 4; + // The total host time in microseconds. + optional double total_host_time_in_us = 5; } // Result proto for StepInfo. diff --git a/tensorflow/contrib/tpu/profiler/tpu_profiler.proto b/tensorflow/contrib/tpu/profiler/tpu_profiler.proto index 8505c4bc69b9444519d0bc5d23bab093b8a57163..7be694e866729c58efae4ccf7932dd929c03ed91 100644 --- a/tensorflow/contrib/tpu/profiler/tpu_profiler.proto +++ b/tensorflow/contrib/tpu/profiler/tpu_profiler.proto @@ -96,5 +96,10 @@ message ProfileResponse { // Data payload for each required tools. repeated ProfileToolData tool_data = 6; - // next-field: 7 + + // When we write profiling data directly to repository directory, we need a + // way to figure out whether the captured trace is empty (due to idle TPU). + bool empty_trace = 7; + + // next-field: 8 } diff --git a/tensorflow/contrib/tpu/profiler/tpu_profiler_analysis.proto b/tensorflow/contrib/tpu/profiler/tpu_profiler_analysis.proto index a4fc8d4e879eb85522f35663c9c628ecd5ef562c..8b0bbde98e6a1dee8ade789328f3ba0624049562 100644 --- a/tensorflow/contrib/tpu/profiler/tpu_profiler_analysis.proto +++ b/tensorflow/contrib/tpu/profiler/tpu_profiler_analysis.proto @@ -7,13 +7,15 @@ message NewProfileSessionRequest { ProfileRequest request = 1; string repository_root = 2; repeated string hosts = 3; + string session_id = 4; } message NewProfileSessionResponse { // Auxiliary error_message. string error_message = 1; - // If success, return session identifier for future reference. - string session_id = 2; + + // Whether all hosts had returned a empty trace. + bool empty_trace = 2; } message EnumProfileSessionsAndToolsRequest { diff --git a/tensorflow/contrib/tpu/profiler/version.h b/tensorflow/contrib/tpu/profiler/version.h index dc6a934891138018d32d511750120453bdf290cf..618479e1a6ccf26a4103ea1f182b662d7d9998da 100644 --- a/tensorflow/contrib/tpu/profiler/version.h +++ b/tensorflow/contrib/tpu/profiler/version.h @@ -16,6 +16,6 @@ limitations under the License. #ifndef TENSORFLOW_CONTRIB_TPU_PROFILER_VERSION_H_ #define TENSORFLOW_CONTRIB_TPU_PROFILER_VERSION_H_ -#define TPU_PROFILER_VERSION "1.5.0" +#define TPU_PROFILER_VERSION "1.6.0" #endif // TENSORFLOW_CONTRIB_TPU_PROFILER_VERSION_H_ diff --git a/tensorflow/contrib/tpu/python/tpu/datasets.py b/tensorflow/contrib/tpu/python/tpu/datasets.py index 465c668fd8b42f150892f8e4b52de76c6fe13fa9..2e472a2805f98b15505f56af403aa6223e28c667 100644 --- a/tensorflow/contrib/tpu/python/tpu/datasets.py +++ b/tensorflow/contrib/tpu/python/tpu/datasets.py @@ -170,7 +170,7 @@ def StreamingFilesDataset(files, args=[source_handle], Tout=[dtypes.string], f=LoadingFunc, - target='/job:%s/replica:0/task:0/cpu:0' % file_reader_job) + target='/job:%s/replica:0/task:0/cpu:0' % file_reader_job)[0] with ops.device('/job:%s' % worker_job): output_dataset = dataset_ops.Dataset.range(2).repeat().map( diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_config.py b/tensorflow/contrib/tpu/python/tpu/tpu_config.py index cc1a7fd801506e3f0b758c4848205f1c375403d2..6d7331e3c79ade9c12c15de79f550cf3973c4e6c 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_config.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_config.py @@ -210,8 +210,9 @@ class RunConfig(run_config_lib.RunConfig): raise ValueError( 'You cannot provide a ClusterResolver and ' 'session_config.cluster_def.') - self._session_config.cluster_def.CopyFrom( - self._cluster_spec.as_cluster_def()) + if self._cluster_spec: + self._session_config.cluster_def.CopyFrom( + self._cluster_spec.as_cluster_def()) @property def evaluation_master(self): diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py index 1332108d04c3a41cdf61fdcd95d6dad4c26d04d8..98eb0e240f0666b2d4a1b5135faef383e49b7468 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py @@ -30,7 +30,6 @@ import six from six.moves import queue as Queue # pylint: disable=redefined-builtin from six.moves import xrange # pylint: disable=redefined-builtin -from tensorflow.contrib.summary import summary_ops as contrib_summary from tensorflow.contrib.tpu.python.ops import tpu_ops from tensorflow.contrib.tpu.python.tpu import tpu from tensorflow.contrib.tpu.python.tpu import tpu_config @@ -57,6 +56,7 @@ from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import resource_variable_ops from tensorflow.python.ops import state_ops +from tensorflow.python.ops import summary_ops_v2 as contrib_summary from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables from tensorflow.python.platform import tf_logging as logging @@ -2054,6 +2054,16 @@ class TPUEstimator(estimator_lib.Estimator): }, every_n_secs=30) ] + input_hooks + chief_hooks = [] + if (self._config.save_checkpoints_secs or + self._config.save_checkpoints_steps): + chief_hooks.append( + training.CheckpointSaverHook( + self.model_dir, + save_secs=self._config.save_checkpoints_secs, + save_steps=self._config.save_checkpoints_steps, + steps_per_run=self._config.tpu_config.iterations_per_loop, + scaffold=scaffold)) summary.scalar(model_fn_lib.LOSS_METRIC_KEY, loss) with ops.control_dependencies([loss]): update_ops = _sync_variables_ops() @@ -2067,6 +2077,7 @@ class TPUEstimator(estimator_lib.Estimator): return model_fn_lib.EstimatorSpec( mode, loss=loss, + training_chief_hooks=chief_hooks, training_hooks=hooks, train_op=train_op, scaffold=scaffold) diff --git a/tensorflow/contrib/training/BUILD b/tensorflow/contrib/training/BUILD index 4d2bfd3e434e60b3fac408931688e8e486b7e494..5de55b5f7f2a41ac6edd27e5a102e565f33df12c 100644 --- a/tensorflow/contrib/training/BUILD +++ b/tensorflow/contrib/training/BUILD @@ -60,6 +60,7 @@ py_library( "//tensorflow/python:util", "//tensorflow/python:variable_scope", "//tensorflow/python:variables", + "//tensorflow/python/data", "//tensorflow/python/estimator:inputs_queues", "//third_party/py/numpy", "@six_archive//:six", diff --git a/tensorflow/contrib/training/__init__.py b/tensorflow/contrib/training/__init__.py index da2de3e421b841937e4125168ea1ecea066ff841..edd71fb2502cf6c965a97485e074d20f876fd504 100644 --- a/tensorflow/contrib/training/__init__.py +++ b/tensorflow/contrib/training/__init__.py @@ -57,6 +57,8 @@ from tensorflow.contrib.training.python.training.hparam import * from tensorflow.contrib.training.python.training.resample import * from tensorflow.contrib.training.python.training.sampling_ops import * from tensorflow.contrib.training.python.training.sequence_queueing_state_saver import * +from tensorflow.contrib.training.python.training.tensor_queue_dataset import enqueue_in_queue_dataset +from tensorflow.contrib.training.python.training.tensor_queue_dataset import prepend_from_queue_and_padded_batch_dataset from tensorflow.contrib.training.python.training.training import add_gradients_summaries from tensorflow.contrib.training.python.training.training import clip_gradient_norms from tensorflow.contrib.training.python.training.training import clip_gradient_norms_fn @@ -75,6 +77,7 @@ _allowed_symbols = [ 'FeedingQueueRunner', 'get_or_create_eval_step', 'StopAfterNEvalsHook', 'SummaryAtEndHook', 'wait_for_new_checkpoint', 'add_gradients_summaries', 'clip_gradient_norms', 'clip_gradient_norms_fn', 'create_train_op', - 'multiply_gradients', 'train'] + 'multiply_gradients', 'enqueue_in_queue_dataset', + 'prepend_from_queue_and_padded_batch_dataset', 'train'] remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/training/python/training/evaluation.py b/tensorflow/contrib/training/python/training/evaluation.py index 4bb53e867811b27dc95857cfdfe936dd2e3b5c6e..f7fd66d33fc0c329db7daaf87373385156d84217 100644 --- a/tensorflow/contrib/training/python/training/evaluation.py +++ b/tensorflow/contrib/training/python/training/evaluation.py @@ -138,7 +138,6 @@ from __future__ import print_function import time -from tensorflow.contrib.framework.python.ops import variables from tensorflow.python.ops import state_ops from tensorflow.python.platform import tf_logging as logging from tensorflow.python.summary import summary @@ -298,7 +297,7 @@ class SummaryAtEndHook(session_run_hook.SessionRunHook): def begin(self): if self._replace_summary_op: self._summary_op = summary.merge_all() - self._global_step = variables.get_or_create_global_step() + self._global_step = training_util.get_or_create_global_step() def after_create_session(self, session, coord): if self._summary_writer is None and self._log_dir: diff --git a/tensorflow/contrib/training/python/training/hparam.py b/tensorflow/contrib/training/python/training/hparam.py index 185f70a86d00fc4bd195fdccc2504515c8c42a89..6c59b68053cfc6c1aebfca149bfba583d645a1e7 100644 --- a/tensorflow/contrib/training/python/training/hparam.py +++ b/tensorflow/contrib/training/python/training/hparam.py @@ -315,7 +315,7 @@ class HParams(object): Hyperparameters have type, which is inferred from the type of their value passed at construction type. The currently supported types are: integer, - float, string, and list of integer, float, or string. + float, boolean, string, and list of integer, float, boolean, or string. You can override hyperparameter values by calling the [`parse()`](#HParams.parse) method, passing a string of comma separated diff --git a/tensorflow/contrib/training/python/training/resample.py b/tensorflow/contrib/training/python/training/resample.py index b16159bc16bb5f3ccd9b5cc7f4af64d9c24b22af..7b8332b1d672dfdfb233e24c0e4f760e49382bc0 100644 --- a/tensorflow/contrib/training/python/training/resample.py +++ b/tensorflow/contrib/training/python/training/resample.py @@ -77,7 +77,7 @@ def resample_at_rate(inputs, rates, scope=None, seed=None, back_prop=False): Args: inputs: A list of tensors, each of which has a shape of `[batch_size, ...]` - rates: A tensor of shape `[batch_size]` contiaining the resampling rates + rates: A tensor of shape `[batch_size]` containing the resampling rates for each input. scope: Scope for the op. seed: Random seed to use. diff --git a/tensorflow/contrib/training/python/training/sampling_ops.py b/tensorflow/contrib/training/python/training/sampling_ops.py index ba888f87dc8c12788c1160f7199c0b762be186c1..7140f2a46d57f0f3b76ff4f1ea9d0d73808405c8 100644 --- a/tensorflow/contrib/training/python/training/sampling_ops.py +++ b/tensorflow/contrib/training/python/training/sampling_ops.py @@ -123,7 +123,7 @@ def rejection_sample(tensors, batch_size=batch_size, num_threads=queue_threads) - # Queues return a single tensor if the list of enqued tensors is one. Since + # Queues return a single tensor if the list of enqueued tensors is one. Since # we want the type to always be the same, always return a list. if isinstance(minibatch, ops.Tensor): minibatch = [minibatch] @@ -312,7 +312,7 @@ def _verify_input(tensor_list, labels, probs_list): """Verify that batched inputs are well-formed.""" checked_probs_list = [] for probs in probs_list: - # Since number of classes shouldn't change at runtime, probalities shape + # Since number of classes shouldn't change at runtime, probabilities shape # should be fully defined. probs.get_shape().assert_is_fully_defined() @@ -407,7 +407,7 @@ def _calculate_acceptance_probabilities(init_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]``` """ # Make list of t_i / p_i. diff --git a/tensorflow/contrib/training/python/training/sequence_queueing_state_saver.py b/tensorflow/contrib/training/python/training/sequence_queueing_state_saver.py index 99d486b1833ccaa81a873d457a8edb06f3d9c7a5..39d75a080604e3a7ae93391652d4c03be9857218 100644 --- a/tensorflow/contrib/training/python/training/sequence_queueing_state_saver.py +++ b/tensorflow/contrib/training/python/training/sequence_queueing_state_saver.py @@ -876,7 +876,7 @@ class SequenceQueueingStateSaver(object): ]): self._length = array_ops.identity(self._length) - # Only create barrier; enqueu and dequeue operations happen when you + # Only create barrier; enqueue and dequeue operations happen when you # access prefetch_op and next_batch. self._create_barrier() self._scope = scope @@ -1637,7 +1637,7 @@ def _move_sparse_tensor_out_context(input_context, input_sequences, num_unroll): For `key, value` pairs in `input_context` with `SparseTensor` `value` removes them from `input_context` and transforms the `value` into a sequence and - then adding `key`, transformed `value` into `input_seuqences`. + then adding `key`, transformed `value` into `input_sequences`. The transformation is done by adding a new first dimension of `value_length` equal to that of the other values in input_sequences` and tiling the `value` every `num_unroll` steps. diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD index c5ca421ced2434c5fae3682b0c52bfe1ef9883c5..24ea732877576b3faaa1da713f6a9bf617301511 100644 --- a/tensorflow/core/BUILD +++ b/tensorflow/core/BUILD @@ -149,6 +149,7 @@ load( "//third_party/mkl:build_defs.bzl", "if_mkl", ) +load("@io_bazel_rules_closure//closure:defs.bzl", "closure_proto_library") exports_files(["ops/ops.pbtxt"]) @@ -159,6 +160,7 @@ exports_files(["ops/ops.pbtxt"]) # # Note that some protos are in neither additional_core_proto_srcs nor this # filegroup; e.g. ones with individual proto_library targets. +# LINT.IfChange CORE_PROTO_SRCS = [ "example/example.proto", "example/feature.proto", @@ -200,6 +202,7 @@ CORE_PROTO_SRCS = [ "util/memmapped_file_system.proto", "util/saved_tensor_slice.proto", ] +# LINT.ThenChange(//tensorflow/core/android_proto_config.asciipb) # Protos which are not needed on mobile builds, but should be included in # protos_all. @@ -208,6 +211,7 @@ CORE_PROTO_SRCS = [ # ones with individual proto_library targets. ADDITIONAL_CORE_PROTO_SRCS = [ "example/example_parser_configuration.proto", + "protobuf/checkpointable_object_graph.proto", "protobuf/control_flow.proto", # TODO(ebrevdo): Re-enable once CriticalSection is in core. # "protobuf/critical_section.proto", @@ -244,6 +248,21 @@ tf_nano_proto_library( deps = [":protos_all_cc"], ) +proto_library( + name = "example_protos", + srcs = [ + "example/example.proto", + "example/feature.proto", + ], + visibility = ["//visibility:public"], +) + +closure_proto_library( + name = "example_protos_closure", + visibility = ["//visibility:public"], + deps = [":example_protos"], +) + exports_files([ "framework/types.proto", ]) @@ -334,7 +353,9 @@ cc_library( "lib/bfloat16/bfloat16.h", ] + tf_additional_proto_hdrs() + glob(tf_env_time_hdrs()), copts = tf_copts(), - deps = tf_lib_proto_parsing_deps(), + deps = tf_lib_proto_parsing_deps() + [ + "@double_conversion//:double-conversion", + ], ) # This build rule (along with :lib_internal, :framework, and @@ -544,6 +565,7 @@ tf_cuda_library( "framework/selective_registration.h", "framework/session_state.h", "framework/shape_inference.h", + "framework/stats_aggregator.h", "framework/tensor.h", "framework/tensor_shape.h", "framework/tensor_slice.h", @@ -1052,6 +1074,7 @@ cc_library( hdrs = [ "common_runtime/function_testlib.h", "common_runtime/kernel_benchmark_testlib.h", + "common_runtime/test_collective_executor_mgr.h", "framework/fake_input.h", "framework/function_testlib.h", "framework/shape_inference_testutil.h", @@ -1226,6 +1249,7 @@ cc_library( deps = [ ":protos_all_cc_impl", "//third_party/eigen3", + "@double_conversion//:double-conversion", "@nsync//:nsync_cpp", "@protobuf_archive//:protobuf", ], @@ -1265,6 +1289,7 @@ cc_library( deps = [ ":protos_all_cc_impl", "//third_party/eigen3", + "@double_conversion//:double-conversion", "@nsync//:nsync_cpp", "@protobuf_archive//:protobuf", ], @@ -1328,6 +1353,7 @@ cc_library( deps = [ ":protos_all_cc_impl", "//third_party/eigen3", + "@double_conversion//:double-conversion", "@nsync//:nsync_cpp", "@protobuf_archive//:protobuf", ], @@ -1350,6 +1376,7 @@ cc_library( deps = [ ":protos_all_cc_impl", "//third_party/eigen3", + "@double_conversion//:double-conversion", "@nsync//:nsync_cpp", "@protobuf_archive//:protobuf", ], @@ -1746,6 +1773,7 @@ cc_library( "//tensorflow/core/platform/default/build_config:platformlib", "@snappy", "@zlib_archive//:zlib", + "@double_conversion//:double-conversion", "@protobuf_archive//:protobuf", ] + tf_protos_all_impl() + tf_protos_grappler_impl(), ) @@ -2249,7 +2277,9 @@ tf_cuda_library( CORE_CPU_LIB_HEADERS = CORE_CPU_BASE_HDRS + [ "common_runtime/allocator_retry.h", + "common_runtime/base_collective_executor.h", "common_runtime/bfc_allocator.h", + "common_runtime/broadcaster.h", "common_runtime/buf_rendezvous.h", "common_runtime/build_graph_options.h", "common_runtime/collective_executor_mgr.h", @@ -2277,6 +2307,7 @@ CORE_CPU_LIB_HEADERS = CORE_CPU_BASE_HDRS + [ "common_runtime/renamed_device.h", "common_runtime/rendezvous_mgr.h", "common_runtime/rendezvous_util.h", + "common_runtime/ring_reducer.h", "common_runtime/scoped_allocator.h", "common_runtime/scoped_allocator_mgr.h", "common_runtime/session_factory.h", @@ -2294,7 +2325,9 @@ tf_cuda_library( srcs = [ "common_runtime/accumulate_n_optimizer.cc", "common_runtime/allocator_retry.cc", + "common_runtime/base_collective_executor.cc", "common_runtime/bfc_allocator.cc", + "common_runtime/broadcaster.cc", "common_runtime/buf_rendezvous.cc", "common_runtime/build_graph_options.cc", "common_runtime/collective_executor_mgr.cc", @@ -2324,6 +2357,7 @@ tf_cuda_library( "common_runtime/renamed_device.cc", "common_runtime/rendezvous_mgr.cc", "common_runtime/rendezvous_util.cc", + "common_runtime/ring_reducer.cc", "common_runtime/scoped_allocator.cc", "common_runtime/scoped_allocator_mgr.cc", "common_runtime/session.cc", @@ -2751,7 +2785,6 @@ tf_cc_tests( "lib/monitoring/sampler_test.cc", "lib/random/distribution_sampler_test.cc", "lib/random/philox_random_test.cc", - "lib/random/random_distributions_test.cc", "lib/random/random_test.cc", "lib/random/simple_philox_test.cc", "lib/strings/base64_test.cc", @@ -2781,6 +2814,21 @@ tf_cc_tests( ], ) +tf_cc_test( + name = "lib_random_random_distributions_test", + srcs = ["lib/random/random_distributions_test.cc"], + tags = ["optonly"], + deps = [ + ":lib", + ":lib_internal", + ":lib_test_internal", + ":protos_all_cc", + ":test", + ":test_main", + "//third_party/eigen3", + ], +) + tf_cc_test( name = "platform_env_test", size = "small", @@ -3089,6 +3137,62 @@ tf_cc_test( ], ) +tf_cc_tests_gpu( + name = "ring_reducer_test", + size = "medium", + srcs = [ + "common_runtime/ring_reducer_test.cc", + ], + linkstatic = tf_kernel_tests_linkstatic(), + tags = tf_cuda_tests_tags(), + deps = [ + ":all_kernels", + ":core", + ":core_cpu", + ":core_cpu_internal", + ":direct_session_internal", + ":framework", + ":framework_internal", + ":gpu_runtime", + ":lib", + ":lib_internal", + ":ops", + ":protos_all_cc", + ":protos_test_cc", + ":test", + ":test_main", + ":testlib", + ], +) + +tf_cc_tests_gpu( + name = "broadcaster_test", + size = "small", + srcs = [ + "common_runtime/broadcaster_test.cc", + ], + linkstatic = tf_kernel_tests_linkstatic(), + tags = tf_cuda_tests_tags(), + deps = [ + ":all_kernels", + ":core", + ":core_cpu", + ":core_cpu_internal", + ":direct_session_internal", + ":framework", + ":framework_internal", + ":gpu_runtime", + ":lib", + ":lib_internal", + ":ops", + ":protos_all_cc", + ":protos_test_cc", + ":test", + ":test_main", + ":testlib", + ], +) + tf_cc_test_mkl( name = "mkl_runtime_tests", size = "small", diff --git a/tensorflow/core/api_def/base_api/api_def_ApplyAdaMax.pbtxt b/tensorflow/core/api_def/base_api/api_def_ApplyAdaMax.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..145d05de59a16ed67af4b978e72593ae88da3d05 --- /dev/null +++ b/tensorflow/core/api_def/base_api/api_def_ApplyAdaMax.pbtxt @@ -0,0 +1,78 @@ +op { + graph_op_name: "ApplyAdaMax" + visibility: HIDDEN + in_arg { + name: "var" + description: <>> x = tf.constant([1, 2, 3]) +>>> y = tf.broadcast_to(x, [3, 3]) +>>> sess.run(y) +array([[1, 2, 3], + [1, 2, 3], + [1, 2, 3]], dtype=int32) +``` +In the above example, the input Tensor with the shape of `[1, 3]` +is broadcasted to output Tensor with shape of `[3, 3]`. +END +} diff --git a/tensorflow/core/api_def/base_api/api_def_DatasetToTFRecord.pbtxt b/tensorflow/core/api_def/base_api/api_def_DatasetToTFRecord.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..e1b8a9abdd2bec0fda690f96d266569b2fb2fcab --- /dev/null +++ b/tensorflow/core/api_def/base_api/api_def_DatasetToTFRecord.pbtxt @@ -0,0 +1,24 @@ +op { + graph_op_name: "DatasetToTFRecord" + visibility: HIDDEN + in_arg { + name: "input_dataset" + description: < [[0, 0, 0, 0, 0, 0] [0, 0, 2, 2, 0, 0] [0, 0, 0, 0, 0, 0]] ``` + END } diff --git a/tensorflow/core/api_def/base_api/api_def_QuantizeV2.pbtxt b/tensorflow/core/api_def/base_api/api_def_QuantizeV2.pbtxt index b9e75caf02b3b557e632a8137cf4780178b8ad8a..37ac10dddb7fc583d270fd6a75af6ab82f9206a5 100644 --- a/tensorflow/core/api_def/base_api/api_def_QuantizeV2.pbtxt +++ b/tensorflow/core/api_def/base_api/api_def_QuantizeV2.pbtxt @@ -44,6 +44,7 @@ In 'MIN_COMBINED' mode, each value of the tensor will undergo the following: out[i] = (in[i] - min_range) * range(T) / (max_range - min_range) if T == qint8, out[i] -= (range(T) + 1) / 2.0 ``` + here `range(T) = numeric_limits::max() - numeric_limits::min()` *MIN_COMBINED Mode Example* @@ -87,6 +88,7 @@ choosing to elide the lowest possible value for symmetry (e.g., output range is We first find the range of values in our tensor. The range we use is always centered on 0, so we find m such that + ```c++ m = max(abs(input_min), abs(input_max)) ``` @@ -95,6 +97,7 @@ Our input tensor range is then `[-m, m]`. Next, we choose our fixed-point quantization buckets, `[min_fixed, max_fixed]`. If T is signed, this is + ``` num_bits = sizeof(T) * 8 [min_fixed, max_fixed] = @@ -102,16 +105,19 @@ If T is signed, this is ``` Otherwise, if T is unsigned, the fixed-point range is + ``` [min_fixed, max_fixed] = [0, (1 << num_bits) - 1] ``` From this we compute our scaling factor, s: + ```c++ s = (max_fixed - min_fixed) / (2 * m) ``` Now we can quantize the elements of our tensor: + ```c++ result = round(input * s) ``` diff --git a/tensorflow/core/api_def/base_api/api_def_ResourceApplyAdaMax.pbtxt b/tensorflow/core/api_def/base_api/api_def_ResourceApplyAdaMax.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..a3f2188ba50723cafab560f417da3ebf793a8661 --- /dev/null +++ b/tensorflow/core/api_def/base_api/api_def_ResourceApplyAdaMax.pbtxt @@ -0,0 +1,72 @@ +op { + graph_op_name: "ResourceApplyAdaMax" + visibility: HIDDEN + in_arg { + name: "var" + description: < +## Build the C or Java libraries + +The instructions above are tailored to building the TensorFlow Python packages. + +If you're interested in building the libraries for the TensorFlow C API, do the +following: + +1. Follow the steps up to [Configure the installation](#ConfigureInstallation) +2. Build the C libraries following instructions in the [README](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/lib_package/README.md). + +If you're interested inv building the libraries for the TensorFlow Java API, +do the following: + +1. Follow the steps up to [Configure the installation](#ConfigureInstallation) +2. Build the Java library following instructions in the [README](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/lib_package/README.md). diff --git a/tensorflow/docs_src/mobile/android_build.md b/tensorflow/docs_src/mobile/android_build.md index 08a5fbe41c87c88399682208c38bf7a892d8fc1a..c35530061dcaf2a4a894dcdf54fd794907d98162 100644 --- a/tensorflow/docs_src/mobile/android_build.md +++ b/tensorflow/docs_src/mobile/android_build.md @@ -51,7 +51,8 @@ If you haven't already, do the following two things: // set to 'bazel', 'cmake', 'makefile', 'none' def nativeBuildSystem = 'none' -4. Click the Run button (the green arrow) or use **Run -> Run 'android'** from the top menu. +4. Click the *Run* button (the green arrow) or select *Run > Run 'android'* from the + top menu. You may need to rebuild the project using *Build > Rebuild Project*. If it asks you to use Instant Run, click **Proceed Without Instant Run**. diff --git a/tensorflow/docs_src/performance/performance_guide.md b/tensorflow/docs_src/performance/performance_guide.md index 580a899ac4e4f5c3d97ce023f25083168fe00d01..b1796cf9b2d0bf7459e70ab542b6e6fcb203667a 100644 --- a/tensorflow/docs_src/performance/performance_guide.md +++ b/tensorflow/docs_src/performance/performance_guide.md @@ -475,7 +475,7 @@ optimizations. ### TensorFlow with Intel® MKL DNN Intel® has added optimizations to TensorFlow for Intel® Xeon® and Intel® Xeon -Phiâ„¢ though the use of Intel® Math Kernel Library for Deep Neural Networks +Phiâ„¢ through the use of the Intel® Math Kernel Library for Deep Neural Networks (Intel® MKL-DNN) optimized primitives. The optimizations also provide speedups for the consumer line of processors, e.g. i5 and i7 Intel processors. The Intel published paper @@ -581,9 +581,9 @@ Each variable that impacts performance is discussed below. for optimal settings. * **intra_op_parallelism_threads**: Setting this equal to the number of - physical cores is recommended. Setting the value to 0, which is the default - and will result in the value being set to the number of logical cores, is an - option to try for some architectures. This value and `OMP_NUM_THREADS` + physical cores is recommended. Setting the value to 0, which is the default, + results in the value being set to the number of logical cores - this is an + alternate option to try for some architectures. This value and `OMP_NUM_THREADS` should be equal. * **inter_op_parallelism_threads**: Setting this equal to the number of diff --git a/tensorflow/docs_src/performance/quantization.md b/tensorflow/docs_src/performance/quantization.md index 411889cb1c616130f809e6228cc692ba3f951d48..2fea02d861d314cc61f2ba20475bf08ebea8fb5f 100644 --- a/tensorflow/docs_src/performance/quantization.md +++ b/tensorflow/docs_src/performance/quantization.md @@ -110,7 +110,7 @@ we've added a separate rewrite for the *eval graph*: ``` # Build eval model -logits = tf.nn.softmax_cross_entropy_with_logits(...) +logits = tf.nn.softmax_cross_entropy_with_logits_v2(...) # Call the eval rewrite which rewrites the graph in-place with # FakeQuantization nodes and fold batchnorm for eval. diff --git a/tensorflow/docs_src/performance/xla/operation_semantics.md b/tensorflow/docs_src/performance/xla/operation_semantics.md index 217ab596b72bc263ae5dda377a8faab8a39b0a3c..8373a1219dab0cfe57b3b01ad0d4cb00ef72093e 100644 --- a/tensorflow/docs_src/performance/xla/operation_semantics.md +++ b/tensorflow/docs_src/performance/xla/operation_semantics.md @@ -854,12 +854,13 @@ calculation of 'start_indices') is currently implementation-defined. | `operand` | `ComputationDataHandle` | N dimensional array of type T | | `update` | `ComputationDataHandle` | N dimensional array of type T | : : : containing the slice update. : -: : : Each dimension of update shape : +: : : Each dimension of update shape : : : : must be strictly greater than : : : : zero, and start + update must be : -: : : less than operand size for each : -: : : dimension to avoid generating : -: : : out-of-bounds update indices. : +: : : less than or equal to the operand: +: : : size for each dimension to avoid : +: : : generating out-of-bounds update : +: : : indices. : | `start_indices` | `ComputationDataHandle` | Rank 1 array of N integers | : : : containing the starting indices : : : : of the slice for each dimension. : @@ -1416,12 +1417,12 @@ Applies a reduction function to an array. | `dimensions` | `int64` array | unordered array of dimensions to | : : : reduce : -Conceptually, this operation reduces one or more dimensions in the input array -into scalars. The rank of the result array is `rank(operand) - len(dimensions)`. -`init_value` is the initial value used for every reduction and may also be -inserted anywhere during computation if the back-end chooses to do so. So in -most cases `init_value` should be an identity of the reduction function (for -example, 0 for addition). +This operation reduces one or more dimensions of the input array into scalars. +The rank of the returned array is `rank(operand) - len(dimensions)`. +`init_value` is the initial value used for every reduction and may be inserted +anywhere during computation by the back-end. In most cases, `init_value` is an +identity of the reduction function (for example, 0 for addition). The applied +`computation` is always passed the `init_value` on the left-hand side. The evaluation order of the reduction function is arbitrary and may be non-deterministic. Therefore, the reduction function should not be overly @@ -1441,8 +1442,7 @@ could be computed as but there are also many other possibilities, e.g. -`f(init_value, f(f(10, f(init_value, 11)), f(f(init_value, 12), f(13, -init_value))))` +`f(init_value, f(f(10, f(init_value, 11)), f(f(init_value, 12), f(init_value, 13))))` The following is a rough pseudo-code example of how reduction could be implemented, using summation as the reduction computation with an initial value @@ -1560,7 +1560,9 @@ See also Applies a reduction function to all elements in each window of the input multi-dimensional array, producing an output multi-dimensional array with the same number of elements as the number of valid positions of the window. A -pooling layer can be expressed as a `ReduceWindow`. +pooling layer can be expressed as a `ReduceWindow`. Similar to +[`Reduce`](#reduce), the applied `computation` is always passed the `init_value` +on the left-hand side. `ReduceWindow(operand, init_value, computation, window_dimensions, window_strides, padding)` diff --git a/tensorflow/docs_src/programmers_guide/datasets.md b/tensorflow/docs_src/programmers_guide/datasets.md index 9ccdbde627e6b2415835f7c0771eca1afa04f7f8..67be41b1a68bc72913d859f83ccd74b529350a4c 100644 --- a/tensorflow/docs_src/programmers_guide/datasets.md +++ b/tensorflow/docs_src/programmers_guide/datasets.md @@ -540,7 +540,7 @@ batched into a fixed size. # to a fixed shape. def _parse_function(filename, label): image_string = tf.read_file(filename) - image_decoded = tf.image.decode_image(image_string) + image_decoded = tf.image.decode_jpeg(image_string) image_resized = tf.image.resize_images(image_decoded, [28, 28]) return image_resized, label diff --git a/tensorflow/docs_src/programmers_guide/debugger.md b/tensorflow/docs_src/programmers_guide/debugger.md index f5a0eb0a2000a5c35f7e3641e6552d40629305a6..f7817b06d4c8bd8fe6ff8c3ffae2db856611cb8e 100644 --- a/tensorflow/docs_src/programmers_guide/debugger.md +++ b/tensorflow/docs_src/programmers_guide/debugger.md @@ -400,7 +400,7 @@ diff = -(y_ * tf.log(y)) to the built-in, numerically-stable implementation of softmax cross-entropy: ```python -diff = tf.losses.sparse_softmax_cross_entropy(labels=y_, logits=logits) +diff = tf.losses.softmax_cross_entropy(labels=y_, logits=logits) ``` Rerun with the `--debug` flag as follows: diff --git a/tensorflow/docs_src/programmers_guide/faq.md b/tensorflow/docs_src/programmers_guide/faq.md index 392ac6f7f12532c3efce5bec1917691f55c7bee5..51c1a1e032baae7eff334da785fc5ffa2438e0ca 100644 --- a/tensorflow/docs_src/programmers_guide/faq.md +++ b/tensorflow/docs_src/programmers_guide/faq.md @@ -121,7 +121,7 @@ dimensions: devices, which makes it possible to speed up @{$deep_cnn$CIFAR-10 training using multiple GPUs}. * The Session API allows multiple concurrent steps (i.e. calls to - @{tf.Session.run} in parallel. This + @{tf.Session.run} in parallel). This enables the runtime to get higher throughput, if a single step does not use all of the resources in your computer. diff --git a/tensorflow/docs_src/programmers_guide/graphs.md b/tensorflow/docs_src/programmers_guide/graphs.md index e69b717432e6a8fab0085eb419dcbc0991cd9d28..f0dd8def17fd6dfed241167a5ebb5be678152c16 100644 --- a/tensorflow/docs_src/programmers_guide/graphs.md +++ b/tensorflow/docs_src/programmers_guide/graphs.md @@ -96,7 +96,7 @@ to all API functions in the same context. For example: (See @{$programmers_guide/variables} for more information about variables.) * Calling @{tf.train.Optimizer.minimize} will add operations and tensors to the - default graph that calculate gradients, and return a @{tf.Operation} that, + default graph that calculates gradients, and return a @{tf.Operation} that, when run, will apply those gradients to a set of variables. Most programs rely solely on the default graph. However, @@ -210,7 +210,7 @@ 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}, +If you are deploying TensorFlow in a @{$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"`): @@ -362,7 +362,7 @@ operations that are needed to compute the result. @{tf.Session.run} requires you to specify a list of **fetches**, which determine the return values, and may be a @{tf.Operation}, a @{tf.Tensor}, or -a [tensor-like type](#tensor-like-objects) such as @{tf.Variable}. These fetches +a [tensor-like type](#tensor-like_objects) such as @{tf.Variable}. These fetches determine what **subgraph** of the overall @{tf.Graph} must be executed to produce the result: this is the subgraph that contains all operations named in the fetch list, plus all operations whose outputs are used to compute the value @@ -505,7 +505,7 @@ multiple graphs in the same process. As noted above, TensorFlow provides a "default graph" that is implicitly passed to all API functions in the same context. For many applications, a single graph is sufficient. However, TensorFlow also provides methods for manipulating -the default graph, which can be useful in more advanced used cases. For example: +the default graph, which can be useful in more advanced use cases. For example: * A @{tf.Graph} defines the namespace for @{tf.Operation} objects: each operation in a single graph must have a unique name. TensorFlow will diff --git a/tensorflow/docs_src/programmers_guide/index.md b/tensorflow/docs_src/programmers_guide/index.md index 017db0e8cb4d239fa4b6be6a5f9d6b0c582a82c2..648d001bd3535fe3dcc460c9ebdb6e6a997dc332 100644 --- a/tensorflow/docs_src/programmers_guide/index.md +++ b/tensorflow/docs_src/programmers_guide/index.md @@ -5,7 +5,7 @@ works. The units are as follows: ## High Level APIs - * @{$programmers_guide/eager}, which is the easiest way to use tensorflow. + * @{$programmers_guide/eager}, which is the easiest way to use TensorFlow. * @{$programmers_guide/estimators}, which introduces a high-level TensorFlow API that greatly simplifies ML programming. * @{$programmers_guide/datasets}, which explains how to diff --git a/tensorflow/docs_src/programmers_guide/saved_model.md b/tensorflow/docs_src/programmers_guide/saved_model.md index 55ee42dd6405db6bd34b064d71deaeb94839b0fa..c6ef87c54a3bc37dbfc0553232a8e3d30f8ee2f6 100644 --- a/tensorflow/docs_src/programmers_guide/saved_model.md +++ b/tensorflow/docs_src/programmers_guide/saved_model.md @@ -485,31 +485,7 @@ portion of the signature. That is, when writing a 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. - -### 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, - strip_default_attrs=True) -``` - -This method builds a new graph by first calling the -`serving_input_receiver_fn()` to obtain feature `Tensor`s, and then calling -this `Estimator`'s `model_fn()` to generate the model graph based on those -features. It starts a fresh `Session`, and, by default, restores the most recent -checkpoint into it. (A different checkpoint may be passed, if needed.) -Finally it creates a time-stamped export directory below the given -`export_dir_base` (i.e., `export_dir_base/`), and writes a -SavedModel into it containing a single `MetaGraphDef` saved from this -Session. - -> Note: It is your responsibility to garbage-collect old exports. -> Otherwise, successive exports will accumulate under `export_dir_base`. - + ### Specify the outputs of a custom model When writing a custom `model_fn`, you must populate the `export_outputs` element @@ -541,6 +517,30 @@ using [`signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY`](https://www.tens indicating which `SignatureDef` will be served when an inference request does not specify one. + +### 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, + strip_default_attrs=True) +``` + +This method builds a new graph by first calling the +`serving_input_receiver_fn()` to obtain feature `Tensor`s, and then calling +this `Estimator`'s `model_fn()` to generate the model graph based on those +features. It starts a fresh `Session`, and, by default, restores the most recent +checkpoint into it. (A different checkpoint may be passed, if needed.) +Finally it creates a time-stamped export directory below the given +`export_dir_base` (i.e., `export_dir_base/`), and writes a +SavedModel into it containing a single `MetaGraphDef` saved from this +Session. + +> Note: It is your responsibility to garbage-collect old exports. +> Otherwise, successive exports will accumulate under `export_dir_base`. ### Serve the exported model locally diff --git a/tensorflow/docs_src/programmers_guide/using_tpu.md b/tensorflow/docs_src/programmers_guide/using_tpu.md index cb0d86fc4c54ac690f13c93ebbd10805c7738c62..5e3e49d43402cd76f8b7062483259df4598bd8ff 100644 --- a/tensorflow/docs_src/programmers_guide/using_tpu.md +++ b/tensorflow/docs_src/programmers_guide/using_tpu.md @@ -280,8 +280,8 @@ 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 +free from the 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 diff --git a/tensorflow/docs_src/tutorials/audio_recognition.md b/tensorflow/docs_src/tutorials/audio_recognition.md index 7d79f433c41b42a268816d8277ea69b0d62a04f3..372ab47df7df309ab926836ca19f34d2d0d38915 100644 --- a/tensorflow/docs_src/tutorials/audio_recognition.md +++ b/tensorflow/docs_src/tutorials/audio_recognition.md @@ -280,7 +280,7 @@ tool: ``` bazel run tensorflow/examples/wav_to_spectrogram:wav_to_spectrogram -- \ --input_wav=/tmp/speech_dataset/happy/ab00c4b2_nohash_0.wav \ ---output_png=/tmp/spectrogram.png +--output_image=/tmp/spectrogram.png ``` If you open up `/tmp/spectrogram.png` you should see something like this: diff --git a/tensorflow/docs_src/tutorials/image_retraining.md b/tensorflow/docs_src/tutorials/image_retraining.md index 93d7c86e42aa90d145d27b56edc0abfec7034686..27784eef9cdb5c6f8b9af44b3fc3f876cda39d13 100644 --- a/tensorflow/docs_src/tutorials/image_retraining.md +++ b/tensorflow/docs_src/tutorials/image_retraining.md @@ -1,404 +1,4 @@ # How to Retrain Inception's Final Layer for New Categories -Modern object recognition models have millions of parameters and can take weeks -to fully train. Transfer learning is a technique that shortcuts a lot of this -work by taking a fully-trained model for a set of categories like ImageNet, and -retrains from the existing weights for new classes. In this example we'll be -retraining the final layer from scratch, while leaving all the others untouched. -For more information on the approach you can see -[this paper on Decaf](https://arxiv.org/pdf/1310.1531v1.pdf). - -Though it's not as good as a full training run, this is surprisingly effective -for many applications, and can be run in as little as thirty minutes on a -laptop, without requiring a GPU. This tutorial will show you how to run the -example script on your own images, and will explain some of the options you have -to help control the training process. - -Note: A version of this tutorial is also available -[as a codelab](https://codelabs.developers.google.com/codelabs/tensorflow-for-poets/#0). - -Before you start, you must @{$install$install tensorflow}. - -[TOC] - -## Training on Flowers - -![Daisies by Kelly Sikkema](https://www.tensorflow.org/images/daisies.jpg) - -[Image by Kelly Sikkema](https://www.flickr.com/photos/95072945@N05/9922116524/) - -Before you start any training, you'll need a set of images to teach the network -about the new classes you want to recognize. There's a later section that -explains how to prepare your own images, but to make it easy we've created an -archive of creative-commons licensed flower photos to use initially. To get the -set of flower photos, run these commands: - -```sh -cd ~ -curl -O http://download.tensorflow.org/example_images/flower_photos.tgz -tar xzf flower_photos.tgz -``` - -Once you have the images, you can clone the tensorflow repository using the -following command (these examples are not included in the installation): - -```sh -git clone https://github.com/tensorflow/tensorflow -``` - -Then checkout the version of the tensorflow repository matching your -installation and this tutorial as follows: - -``` sh -cd tensorflow -git checkout {version} -``` - -In the simplest cases the retrainer can then be run like this: - -```sh -python tensorflow/examples/image_retraining/retrain.py --image_dir ~/flower_photos -``` - -The script has many other options. You can get a full listing with: - -```sh -python tensorflow/examples/image_retraining/retrain.py -h -``` - -This script loads the pre-trained Inception v3 model, removes the old top layer, -and trains a new one on the flower photos you've downloaded. None of the flower -species were in the original ImageNet classes the full network was trained on. -The magic of transfer learning is that lower layers that have been trained to -distinguish between some objects can be reused for many recognition tasks -without any alteration. - -## Bottlenecks - -The script can take thirty minutes or more to complete, depending on the speed -of your machine. The first phase analyzes all the images on disk and calculates -the bottleneck values for each of them. 'Bottleneck' is an informal term we -often use for the layer just before the final output layer that actually does -the classification. This penultimate layer has been trained to output a set of -values that's good enough for the classifier to use to distinguish between all -the classes it's been asked to recognize. That means it has to be a meaningful -and compact summary of the images, since it has to contain enough information -for the classifier to make a good choice in a very small set of values. The -reason our final layer retraining can work on new classes is that it turns out -the kind of information needed to distinguish between all the 1,000 classes in -ImageNet is often also useful to distinguish between new kinds of objects. - -Because every image is reused multiple times during training and calculating -each bottleneck takes a significant amount of time, it speeds things up to -cache these bottleneck values on disk so they don't have to be repeatedly -recalculated. By default they're stored in the `/tmp/bottleneck` directory, and -if you rerun the script they'll be reused so you don't have to wait for this -part again. - -## Training - -Once the bottlenecks are complete, the actual training of the top layer of the -network begins. You'll see a series of step outputs, each one showing training -accuracy, validation accuracy, and the cross entropy. The training accuracy -shows what percent of the images used in the current training batch were -labeled with the correct class. The validation accuracy is the precision on a -randomly-selected group of images from a different set. The key difference is -that the training accuracy is based on images that the network has been able -to learn from so the network can overfit to the noise in the training data. A -true measure of the performance of the network is to measure its performance on -a data set not contained in the training data -- this is measured by the -validation accuracy. If the train accuracy is high but the validation accuracy -remains low, that means the network is overfitting and memorizing particular -features in the training images that aren't helpful more generally. Cross -entropy is a loss function which gives a glimpse into how well the learning -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 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 -through the back-propagation process. As the process continues you should see -the reported accuracy improve, and after all the steps are done, a final test -accuracy evaluation is run on a set of images kept separate from the training -and validation pictures. This test evaluation is the best estimate of how the -trained model will perform on the classification task. You should see an -accuracy value of between 90% and 95%, though the exact value will vary from run -to run since there's randomness in the training process. This number is based on -the percent of the images in the test set that are given the correct label -after the model is fully trained. - -## Visualizing the Retraining with TensorBoard - -The script includes TensorBoard summaries that make it easier to understand, debug, and optimize the retraining. For example, you can visualize the graph and statistics, such as how the weights or accuracy varied during training. - -To launch TensorBoard, run this command during or after retraining: - -```sh -tensorboard --logdir /tmp/retrain_logs -``` - -Once TensorBoard is running, navigate your web browser to `localhost:6006` to view the TensorBoard. - -The script will log TensorBoard summaries to `/tmp/retrain_logs` by default. You can change the directory with the `--summaries_dir` flag. - -The [TensorBoard's GitHub](https://github.com/tensorflow/tensorboard) has a lot more information on TensorBoard usage, including tips & tricks, and debugging information. - -## Using the Retrained Model - -The script will write out a version of the Inception v3 network with a final -layer retrained to your categories to /tmp/output_graph.pb, and a text file -containing the labels to /tmp/output_labels.txt. These are both in a format that -the @{$image_recognition$C++ and Python image classification examples} -can read in, so you can start using your new model immediately. Since you've -replaced the top layer, you will need to specify the new name in the script, for -example with the flag `--output_layer=final_result` if you're using label_image. - -Here's an example of how to run the label_image example with your -retrained graphs: - -```sh -python tensorflow/examples/label_image/label_image.py \ ---graph=/tmp/output_graph.pb --labels=/tmp/output_labels.txt \ ---input_layer=Mul \ ---output_layer=final_result \ ---input_mean=128 --input_std=128 \ ---image=$HOME/flower_photos/daisy/21652746_cc379e0eea_m.jpg -``` - -You should see a list of flower labels, in most cases with daisy on top -(though each retrained model may be slightly different). You can replace the -`--image` parameter with your own images to try those out. - -If you'd like to use the retrained model in your own Python program, then the -above -[`label_image` script](https://www.tensorflow.org/code/tensorflow/examples/label_image/label_image.py) -is a reasonable starting point. The `label_image` -directory also contains C++ code which you can use as a template to integrate -tensorflow with your own applications. - -If you find the default Inception v3 model is too large or slow for your -application, take a look at the [Other Model Architectures section](/tutorials/image_retraining#other_model_architectures) -below for options to speed up and slim down your network. - -## Training on Your Own Categories - -If you've managed to get the script working on the flower example images, you -can start looking at teaching it to recognize categories you care about instead. -In theory all you'll need to do is point it at a set of sub-folders, each named -after one of your categories and containing only images from that category. If -you do that and pass the root folder of the subdirectories as the argument to -`--image_dir`, the script should train just like it did for the flowers. - -Here's what the folder structure of the flowers archive looks like, to give you -and example of the kind of layout the script is looking for: - -![Folder Structure](https://www.tensorflow.org/images/folder_structure.png) - -In practice it may take some work to get the accuracy you want. I'll try to -guide you through some of the common problems you might encounter below. - -## Creating a Set of Training Images - -The first place to start is by looking at the images you've gathered, since the -most common issues we see with training come from the data that's being fed in. - -For training to work well, you should gather at least a hundred photos of each -kind of object you want to recognize. The more you can gather, the better the -accuracy of your trained model is likely to be. You also need to make sure that -the photos are a good representation of what your application will actually -encounter. For example, if you take all your photos indoors against a blank wall -and your users are trying to recognize objects outdoors, you probably won't see -good results when you deploy. - -Another pitfall to avoid is that the learning process will pick up on anything -that the labeled images have in common with each other, and if you're not -careful that might be something that's not useful. For example if you photograph -one kind of object in a blue room, and another in a green one, then the model -will end up basing its prediction on the background color, not the features of -the object you actually care about. To avoid this, try to take pictures in as -wide a variety of situations as you can, at different times, and with different -devices. If you want to know more about this problem, you can read about the -classic (and possibly apocryphal) -[tank recognition problem](https://www.jefftk.com/p/detecting-tanks). - -You may also want to think about the categories you use. It might be worth -splitting big categories that cover a lot of different physical forms into -smaller ones that are more visually distinct. For example instead of 'vehicle' -you might use 'car', 'motorbike', and 'truck'. It's also worth thinking about -whether you have a 'closed world' or an 'open world' problem. In a closed world, -the only things you'll ever be asked to categorize are the classes of object you -know about. This might apply to a plant recognition app where you know the user -is likely to be taking a picture of a flower, so all you have to do is decide -which species. By contrast a roaming robot might see all sorts of different -things through its camera as it wanders around the world. In that case you'd -want the classifier to report if it wasn't sure what it was seeing. This can be -hard to do well, but often if you collect a large number of typical 'background' -photos with no relevant objects in them, you can add them to an extra 'unknown' -class in your image folders. - -It's also worth checking to make sure that all of your images are labeled -correctly. Often user-generated tags are unreliable for our purposes, for -example using #daisy for pictures of a person named Daisy. If you go through -your images and weed out any mistakes it can do wonders for your overall -accuracy. - -## Training Steps - -If you're happy with your images, you can take a look at improving your results -by altering the details of the learning process. The simplest one to try is -`--how_many_training_steps`. This defaults to 4,000, but if you increase it to -8,000 it will train for twice as long. The rate of improvement in the accuracy -slows the longer you train for, and at some point will stop altogether, but you -can experiment to see when you hit that limit for your model. - -## Distortions - -A common way of improving the results of image training is by deforming, -cropping, or brightening the training inputs in random ways. This has the -advantage of expanding the effective size of the training data thanks to all the -possible variations of the same images, and tends to help the network learn to -cope with all the distortions that will occur in real-life uses of the -classifier. The biggest disadvantage of enabling these distortions in our script -is that the bottleneck caching is no longer useful, since input images are never -reused exactly. This means the training process takes a lot longer, so I -recommend trying this as a way of fine-tuning your model once you've got one -that you're reasonably happy with. - -You enable these distortions by passing `--random_crop`, `--random_scale` and -`--random_brightness` to the script. These are all percentage values that -control how much of each of the distortions is applied to each image. It's -reasonable to start with values of 5 or 10 for each of them and then experiment -to see which of them help with your application. `--flip_left_right` will -randomly mirror half of the images horizontally, which makes sense as long as -those inversions are likely to happen in your application. For example it -wouldn't be a good idea if you were trying to recognize letters, since flipping -them destroys their meaning. - -## Hyper-parameters - -There are several other parameters you can try adjusting to see if they help -your results. The `--learning_rate` controls the magnitude of the updates to the -final layer during training. Intuitively if this is smaller then the learning -will take longer, but it can end up helping the overall precision. That's not -always the case though, so you need to experiment carefully to see what works -for your case. The `--train_batch_size` controls how many images are examined -during one training step, and because the learning rate is applied per batch -you'll need to reduce it if you have larger batches to get the same overall -effect. - -## Training, Validation, and Testing Sets - -One of the things the script does under the hood when you point it at a folder -of images is divide them up into three different sets. The largest is usually -the training set, which are all the images fed into the network during training, -with the results used to update the model's weights. You might wonder why we -don't use all the images for training? A big potential problem when we're doing -machine learning is that our model may just be memorizing irrelevant details of -the training images to come up with the right answers. For example, you could -imagine a network remembering a pattern in the background of each photo it was -shown, and using that to match labels with objects. It could produce good -results on all the images it's seen before during training, but then fail on new -images because it's not learned general characteristics of the objects, just -memorized unimportant details of the training images. - -This problem is known as overfitting, and to avoid it we keep some of our data -out of the training process, so that the model can't memorize them. We then use -those images as a check to make sure that overfitting isn't occurring, since if -we see good accuracy on them it's a good sign the network isn't overfitting. The -usual split is to put 80% of the images into the main training set, keep 10% -aside to run as validation frequently during training, and then have a final 10% -that are used less often as a testing set to predict the real-world performance -of the classifier. These ratios can be controlled using the -`--testing_percentage` and `--validation_percentage` flags. In general -you should be able to leave these values at their defaults, since you won't -usually find any advantage to training to adjusting them. - -Note that the script uses the image filenames (rather than a completely random -function) to divide the images among the training, validation, and test sets. -This is done to ensure that images don't get moved between training and testing -sets on different runs, since that could be a problem if images that had been -used for training a model were subsequently used in a validation set. - -You might notice that the validation accuracy fluctuates among iterations. Much -of this fluctuation arises from the fact that a random subset of the validation -set is chosen for each validation accuracy measurement. The fluctuations can be -greatly reduced, at the cost of some increase in training time, by choosing -`--validation_batch_size=-1`, which uses the entire validation set for each -accuracy computation. - -Once training is complete, you may find it insightful to examine misclassified -images in the test set. This can be done by adding the flag -`--print_misclassified_test_images`. This may help you get a feeling for which -types of images were most confusing for the model, and which categories were -most difficult to distinguish. For instance, you might discover that some -subtype of a particular category, or some unusual photo angle, is particularly -difficult to identify, which may encourage you to add more training images of -that subtype. Oftentimes, examining misclassified images can also point to -errors in the input data set, such as mislabeled, low-quality, or ambiguous -images. However, one should generally avoid point-fixing individual errors in -the test set, since they are likely to merely reflect more general problems in -the (much larger) training set. - -## Other Model Architectures - -By default the script uses a pretrained version of the Inception v3 model -architecture. This is a good place to start because it provides high accuracy -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 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 (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 -``` - -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 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 -float range into the 'input' tensor. Typically 24-bit images are in the range -[0,255], and you must convert them to the [-1,1] float range expected by the -model with the formula `(image - 128.)/128.`. - -The default arguments for the `label_image` script are set for Inception V3. -To use it with a MobileNet, specify the above normalization parameters as -`input_mean` and `input_std` on the command line. You also must specify the -image size that your model expects, as follows: - -```sh -python tensorflow/examples/label_image/label_image.py \ ---graph=/tmp/output_graph.pb --labels=/tmp/output_labels.txt \ ---input_layer=input \ ---output_layer=final_result \ ---input_height=224 --input_width=224 \ ---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). +**NOTE: This tutorial has moved to** +https://github.com/tensorflow/hub/tree/master/docs/tutorials/image_retraining.md diff --git a/tensorflow/docs_src/tutorials/layers.md b/tensorflow/docs_src/tutorials/layers.md index cadaec391d8970faf5847c9b9e39bccb31f885ed..37cd2bb1397deaddec9f7194d8c5093145e08644 100644 --- a/tensorflow/docs_src/tutorials/layers.md +++ b/tensorflow/docs_src/tutorials/layers.md @@ -192,8 +192,7 @@ dive deeper into the `tf.layers` code used to create each layer, as well as how 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). +skip ahead to ["Training and Evaluating the CNN MNIST Classifier"](#train_eval_mnist). ### Input Layer @@ -536,8 +535,9 @@ 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 Estimators 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 Estimations in tf.estimator"} tutorial. + ### Add evaluation metrics @@ -552,7 +552,8 @@ return tf.estimator.EstimatorSpec( mode=mode, loss=loss, eval_metric_ops=eval_metric_ops) ``` -## Training and Evaluating the CNN MNIST Classifier {#training_and_evaluating_the_cnn_mnist_classifier} + +## Training and Evaluating the CNN MNIST Classifier We've coded our MNIST CNN model function; now we're ready to train and evaluate it. @@ -612,9 +613,9 @@ following to `main()`: ```python # Set up logging for predictions - tensors_to_log = {"probabilities": "softmax_tensor"} - logging_hook = tf.train.LoggingTensorHook( - tensors=tensors_to_log, every_n_iter=50) +tensors_to_log = {"probabilities": "softmax_tensor"} +logging_hook = tf.train.LoggingTensorHook( + tensors=tensors_to_log, every_n_iter=50) ``` We store a dict of the tensors we want to log in `tensors_to_log`. Each key is a diff --git a/tensorflow/examples/android/BUILD b/tensorflow/examples/android/BUILD index a088d7cf2f05c81cf2e60cb5aa8de79957a30de2..aa594a63c6ad5ab7129e452e7a6345114b994231 100644 --- a/tensorflow/examples/android/BUILD +++ b/tensorflow/examples/android/BUILD @@ -76,7 +76,6 @@ android_binary( custom_package = "org.tensorflow.demo", inline_constants = 1, manifest = "AndroidManifest.xml", - manifest_merger = "legacy", resource_files = glob(["res/**"]), tags = [ "manual", diff --git a/tensorflow/examples/image_retraining/BUILD b/tensorflow/examples/image_retraining/BUILD deleted file mode 100644 index ecd79a3b004d0ca9f50d2a6f140dbc353efe30cb..0000000000000000000000000000000000000000 --- a/tensorflow/examples/image_retraining/BUILD +++ /dev/null @@ -1,51 +0,0 @@ -# Description: -# Transfer learning example for TensorFlow. - -licenses(["notice"]) # Apache 2.0 - -exports_files(["LICENSE"]) - -load("//tensorflow:tensorflow.bzl", "py_test") - -py_binary( - name = "retrain", - srcs = [ - "retrain.py", - ], - srcs_version = "PY2AND3", - visibility = ["//tensorflow:__subpackages__"], - deps = [ - "//tensorflow:tensorflow_py", - "//tensorflow/python:framework", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:graph_util", - "//tensorflow/python:platform", - "//tensorflow/python:util", - "//third_party/py/numpy", - ], -) - -py_test( - name = "retrain_test", - size = "small", - srcs = [ - "retrain.py", - "retrain_test.py", - ], - data = [ - ":data/labels.txt", - "//tensorflow/examples/label_image:data/grace_hopper.jpg", - ], - srcs_version = "PY2AND3", - deps = [ - ":retrain", - "//tensorflow:tensorflow_py", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:graph_util", - "//tensorflow/python:platform", - "//tensorflow/python:platform_test", - "//tensorflow/python:tensor_shape", - "//tensorflow/python:util", - "//third_party/py/numpy", - ], -) diff --git a/tensorflow/examples/image_retraining/README.md b/tensorflow/examples/image_retraining/README.md index 8a49525c6eff003f2c7acb592f213285e627eb51..3f0b3d12682b81e7c13ed6c2d32149746d506cd3 100644 --- a/tensorflow/examples/image_retraining/README.md +++ b/tensorflow/examples/image_retraining/README.md @@ -1,12 +1,15 @@ -retrain.py is an example script that shows how one can adapt a pretrained -network for other classification problems. A detailed overview of this script -can be found at: -https://codelabs.developers.google.com/codelabs/tensorflow-for-poets/#0 +**NOTE: This code has moved to** +https://github.com/tensorflow/hub/tree/master/examples/image_retraining -The script also shows how one can train layers -with quantized weights and activations instead of taking a pre-trained floating -point model and then quantizing weights and activations. -The output graphdef produced by this script is compatible with the TensorFlow -Lite Optimizing Converter and can be converted to TFLite format. +retrain.py is an example script that shows how one can adapt a pretrained +network for other classification problems (including use with TFLite and +quantization). +As of TensorFlow 1.7, it is recommended to use a pretrained network from +TensorFlow Hub, using the new version of this example found in the location +above, as explained in TensorFlow's revised [image retraining +tutorial](https://www.tensorflow.org/tutorials/image_retraining). +Older versions of this example (using frozen GraphDefs instead of +TensorFlow Hub modules) are available in the release branches of +TensorFlow versions up to and including 1.7. diff --git a/tensorflow/examples/image_retraining/data/labels.txt b/tensorflow/examples/image_retraining/data/labels.txt deleted file mode 100644 index bc1131ac4591ca1bdb840695b55f79a6feb95db3..0000000000000000000000000000000000000000 --- a/tensorflow/examples/image_retraining/data/labels.txt +++ /dev/null @@ -1,3 +0,0 @@ -Runner-up -Winner -Loser diff --git a/tensorflow/examples/image_retraining/retrain.py b/tensorflow/examples/image_retraining/retrain.py deleted file mode 100644 index fcc191250fe8c9d80e788b6d345b041c7ea22f2f..0000000000000000000000000000000000000000 --- a/tensorflow/examples/image_retraining/retrain.py +++ /dev/null @@ -1,1487 +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. -# ============================================================================== -r"""Simple transfer learning with Inception v3 or Mobilenet models. - -With support for TensorBoard. - -This example shows how to take a Inception v3 or Mobilenet model trained on -ImageNet images, and train a new top layer that can recognize other classes of -images. - -The top layer receives as input a 2048-dimensional vector (1001-dimensional for -Mobilenet) for each image. We train a softmax layer on top of this -representation. Assuming the softmax layer contains N labels, this corresponds -to learning N + 2048*N (or 1001*N) model parameters corresponding to the -learned biases and weights. - -Here's an example, which assumes you have a folder containing class-named -subfolders, each full of images for each label. The example folder flower_photos -should have a structure like this: - -~/flower_photos/daisy/photo1.jpg -~/flower_photos/daisy/photo2.jpg -... -~/flower_photos/rose/anotherphoto77.jpg -... -~/flower_photos/sunflower/somepicture.jpg - -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 \ - --image_dir ~/flower_photos -``` - -Or, if you have a pip installation of tensorflow, `retrain.py` can be run -without bazel: - -```bash -python tensorflow/examples/image_retraining/retrain.py \ - --image_dir ~/flower_photos -``` - -You can replace the image_dir argument with any folder containing subfolders of -images. The label for each image is taken from the name of the subfolder it's -in. - -This produces a new model file that can be loaded and run by any TensorFlow -program, for example the label_image sample code. - -By default this script will use the high accuracy, but comparatively large and -slow Inception v3 model architecture. It's recommended that you start with this -to validate that you have gathered good training data, but if you want to deploy -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 mobilenet, instrumented for quantization: - -```bash -python tensorflow/examples/image_retraining/retrain.py \ - --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 -'224', '192', '160', or '128', with smaller sizes running faster. See -https://research.googleblog.com/2017/06/mobilenets-open-source-models-for.html -for more information on Mobilenet. - -To use with TensorBoard: - -By default, this script will log summaries to /tmp/retrain_logs directory - -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 -from __future__ import print_function - -import argparse -from datetime import datetime -import hashlib -import os.path -import random -import re -import sys -import tarfile - -import numpy as np -from six.moves import urllib -import tensorflow as tf - -from tensorflow.python.framework import graph_util -from tensorflow.python.framework import tensor_shape -from tensorflow.python.platform import gfile -from tensorflow.python.util import compat - -FLAGS = None - -# These are all parameters that are tied to the particular model architecture -# we're using for Inception v3. These include things like tensor names and their -# sizes. If you want to adapt this script to work with another model, you will -# 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. - - Analyzes the sub folders in the image directory, splits them into stable - training, testing, and validation sets, and returns a data structure - describing the lists of images for each label and their paths. - - Args: - image_dir: String path to a folder containing subfolders of images. - testing_percentage: Integer percentage of the images to reserve for tests. - validation_percentage: Integer percentage of images reserved for validation. - - Returns: - A dictionary containing an entry for each label subfolder, with images split - into training, testing, and validation sets within each label. - """ - if not gfile.Exists(image_dir): - tf.logging.error("Image directory '" + image_dir + "' not found.") - return None - result = {} - sub_dirs = [x[0] for x in gfile.Walk(image_dir)] - # The root directory comes first, so skip it. - is_root_dir = True - for sub_dir in sub_dirs: - if is_root_dir: - is_root_dir = False - continue - extensions = ['jpg', 'jpeg', 'JPG', 'JPEG'] - file_list = [] - dir_name = os.path.basename(sub_dir) - if dir_name == image_dir: - continue - tf.logging.info("Looking for images in '" + dir_name + "'") - for extension in extensions: - file_glob = os.path.join(image_dir, dir_name, '*.' + extension) - file_list.extend(gfile.Glob(file_glob)) - if not file_list: - tf.logging.warning('No files found') - continue - if len(file_list) < 20: - tf.logging.warning( - 'WARNING: Folder has less than 20 images, which may cause issues.') - elif len(file_list) > MAX_NUM_IMAGES_PER_CLASS: - tf.logging.warning( - 'WARNING: Folder {} has more than {} images. Some images will ' - 'never be selected.'.format(dir_name, MAX_NUM_IMAGES_PER_CLASS)) - label_name = re.sub(r'[^a-z0-9]+', ' ', dir_name.lower()) - training_images = [] - testing_images = [] - validation_images = [] - for file_name in file_list: - base_name = os.path.basename(file_name) - # We want to ignore anything after '_nohash_' in the file name when - # deciding which set to put an image in, the data set creator has a way of - # grouping photos that are close variations of each other. For example - # this is used in the plant disease data set to group multiple pictures of - # the same leaf. - hash_name = re.sub(r'_nohash_.*$', '', file_name) - # This looks a bit magical, but we need to decide whether this file should - # go into the training, testing, or validation sets, and we want to keep - # existing files in the same set even if more files are subsequently - # added. - # To do that, we need a stable way of deciding based on just the file name - # itself, so we do a hash of that and then use that to generate a - # probability value that we use to assign it. - hash_name_hashed = hashlib.sha1(compat.as_bytes(hash_name)).hexdigest() - percentage_hash = ((int(hash_name_hashed, 16) % - (MAX_NUM_IMAGES_PER_CLASS + 1)) * - (100.0 / MAX_NUM_IMAGES_PER_CLASS)) - if percentage_hash < validation_percentage: - validation_images.append(base_name) - elif percentage_hash < (testing_percentage + validation_percentage): - testing_images.append(base_name) - else: - training_images.append(base_name) - result[label_name] = { - 'dir': dir_name, - 'training': training_images, - 'testing': testing_images, - 'validation': validation_images, - } - return result - - -def get_image_path(image_lists, label_name, index, image_dir, category): - """"Returns a path to an image for a label at the given index. - - Args: - image_lists: Dictionary of training images for each label. - label_name: Label string we want to get an image for. - index: Int offset of the image we want. This will be moduloed by the - available number of images for the label, so it can be arbitrarily large. - image_dir: Root folder string of the subfolders containing the training - images. - category: Name string of set to pull images from - training, testing, or - validation. - - Returns: - File system path string to an image that meets the requested parameters. - - """ - if label_name not in image_lists: - tf.logging.fatal('Label does not exist %s.', label_name) - label_lists = image_lists[label_name] - if category not in label_lists: - tf.logging.fatal('Category does not exist %s.', category) - category_list = label_lists[category] - if not category_list: - tf.logging.fatal('Label %s has no images in the category %s.', - label_name, category) - mod_index = index % len(category_list) - base_name = category_list[mod_index] - sub_dir = label_lists['dir'] - full_path = os.path.join(image_dir, sub_dir, base_name) - return full_path - - -def get_bottleneck_path(image_lists, label_name, index, bottleneck_dir, - category, architecture): - """"Returns a path to a bottleneck file for a label at the given index. - - Args: - image_lists: Dictionary of training images for each label. - label_name: Label string we want to get an image for. - index: Integer offset of the image we want. This will be moduloed by the - available number of images for the label, so it can be arbitrarily large. - bottleneck_dir: Folder string holding cached files of bottleneck values. - category: Name string of set to pull images from - training, testing, or - validation. - architecture: The name of the model architecture. - - Returns: - File system path string to an image that meets the requested parameters. - """ - return get_image_path(image_lists, label_name, index, bottleneck_dir, - category) + '_' + architecture + '.txt' - - -def create_model_graph(model_info): - """"Creates a graph from saved GraphDef file and returns a Graph object. - - Args: - model_info: Dictionary containing information about the model architecture. - - Returns: - Graph holding the trained Inception network, and various tensors we'll be - manipulating. - """ - with tf.Graph().as_default() as graph: - model_path = os.path.join(FLAGS.model_dir, model_info['model_file_name']) - print('Model path: ', model_path) - with gfile.FastGFile(model_path, 'rb') as f: - graph_def = tf.GraphDef() - graph_def.ParseFromString(f.read()) - bottleneck_tensor, resized_input_tensor = (tf.import_graph_def( - graph_def, - name='', - return_elements=[ - model_info['bottleneck_tensor_name'], - model_info['resized_input_tensor_name'], - ])) - return graph, bottleneck_tensor, resized_input_tensor - - -def run_bottleneck_on_image(sess, image_data, image_data_tensor, - decoded_image_tensor, resized_input_tensor, - bottleneck_tensor): - """Runs inference on an image to extract the 'bottleneck' summary layer. - - Args: - sess: Current active TensorFlow Session. - image_data: String of raw JPEG data. - image_data_tensor: Input data layer in the graph. - decoded_image_tensor: Output of initial image resizing and preprocessing. - resized_input_tensor: The input node of the recognition graph. - bottleneck_tensor: Layer before the final softmax. - - Returns: - Numpy array of bottleneck values. - """ - # First decode the JPEG image, resize it, and rescale the pixel values. - resized_input_values = sess.run(decoded_image_tensor, - {image_data_tensor: image_data}) - # Then run it through the recognition network. - bottleneck_values = sess.run(bottleneck_tensor, - {resized_input_tensor: resized_input_values}) - bottleneck_values = np.squeeze(bottleneck_values) - return bottleneck_values - - -def maybe_download_and_extract(data_url): - """Download and extract model tar file. - - If the pretrained model we're using doesn't already exist, this function - downloads it from the TensorFlow.org website and unpacks it into a directory. - - Args: - data_url: Web location of the tar file containing the pretrained model. - """ - dest_directory = FLAGS.model_dir - if not os.path.exists(dest_directory): - os.makedirs(dest_directory) - filename = data_url.split('/')[-1] - filepath = os.path.join(dest_directory, filename) - if not os.path.exists(filepath): - - def _progress(count, block_size, total_size): - sys.stdout.write('\r>> Downloading %s %.1f%%' % - (filename, - float(count * block_size) / float(total_size) * 100.0)) - sys.stdout.flush() - - filepath, _ = urllib.request.urlretrieve(data_url, filepath, _progress) - print() - statinfo = os.stat(filepath) - 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: - print('Not extracting or downloading files, model already present in disk') - - -def ensure_dir_exists(dir_name): - """Makes sure the folder exists on disk. - - Args: - dir_name: Path string to the folder we want to create. - """ - if not os.path.exists(dir_name): - os.makedirs(dir_name) - - -bottleneck_path_2_bottleneck_values = {} - - -def create_bottleneck_file(bottleneck_path, image_lists, label_name, index, - image_dir, category, sess, jpeg_data_tensor, - decoded_image_tensor, resized_input_tensor, - bottleneck_tensor): - """Create a single bottleneck file.""" - tf.logging.info('Creating bottleneck at ' + bottleneck_path) - image_path = get_image_path(image_lists, label_name, index, - image_dir, category) - if not gfile.Exists(image_path): - tf.logging.fatal('File does not exist %s', image_path) - image_data = gfile.FastGFile(image_path, 'rb').read() - try: - bottleneck_values = run_bottleneck_on_image( - sess, image_data, jpeg_data_tensor, decoded_image_tensor, - resized_input_tensor, bottleneck_tensor) - except Exception as e: - raise RuntimeError('Error during processing file %s (%s)' % (image_path, - str(e))) - bottleneck_string = ','.join(str(x) for x in bottleneck_values) - with open(bottleneck_path, 'w') as bottleneck_file: - bottleneck_file.write(bottleneck_string) - - -def get_or_create_bottleneck(sess, image_lists, label_name, index, image_dir, - category, bottleneck_dir, jpeg_data_tensor, - decoded_image_tensor, resized_input_tensor, - bottleneck_tensor, architecture): - """Retrieves or calculates bottleneck values for an image. - - If a cached version of the bottleneck data exists on-disk, return that, - otherwise calculate the data and save it to disk for future use. - - Args: - sess: The current active TensorFlow Session. - image_lists: Dictionary of training images for each label. - label_name: Label string we want to get an image for. - index: Integer offset of the image we want. This will be modulo-ed by the - available number of images for the label, so it can be arbitrarily large. - image_dir: Root folder string of the subfolders containing the training - images. - category: Name string of which set to pull images from - training, testing, - or validation. - bottleneck_dir: Folder string holding cached files of bottleneck values. - jpeg_data_tensor: The tensor to feed loaded jpeg data into. - decoded_image_tensor: The output of decoding and resizing the image. - resized_input_tensor: The input node of the recognition graph. - bottleneck_tensor: The output tensor for the bottleneck values. - architecture: The name of the model architecture. - - Returns: - Numpy array of values produced by the bottleneck layer for the image. - """ - label_lists = image_lists[label_name] - sub_dir = label_lists['dir'] - sub_dir_path = os.path.join(bottleneck_dir, sub_dir) - ensure_dir_exists(sub_dir_path) - bottleneck_path = get_bottleneck_path(image_lists, label_name, index, - bottleneck_dir, category, architecture) - if not os.path.exists(bottleneck_path): - create_bottleneck_file(bottleneck_path, image_lists, label_name, index, - image_dir, category, sess, jpeg_data_tensor, - decoded_image_tensor, resized_input_tensor, - bottleneck_tensor) - with open(bottleneck_path, 'r') as bottleneck_file: - bottleneck_string = bottleneck_file.read() - did_hit_error = False - try: - bottleneck_values = [float(x) for x in bottleneck_string.split(',')] - except ValueError: - tf.logging.warning('Invalid float found, recreating bottleneck') - did_hit_error = True - if did_hit_error: - create_bottleneck_file(bottleneck_path, image_lists, label_name, index, - image_dir, category, sess, jpeg_data_tensor, - decoded_image_tensor, resized_input_tensor, - bottleneck_tensor) - with open(bottleneck_path, 'r') as bottleneck_file: - bottleneck_string = bottleneck_file.read() - # Allow exceptions to propagate here, since they shouldn't happen after a - # fresh creation - bottleneck_values = [float(x) for x in bottleneck_string.split(',')] - return bottleneck_values - - -def cache_bottlenecks(sess, image_lists, image_dir, bottleneck_dir, - jpeg_data_tensor, decoded_image_tensor, - resized_input_tensor, bottleneck_tensor, architecture): - """Ensures all the training, testing, and validation bottlenecks are cached. - - Because we're likely to read the same image multiple times (if there are no - distortions applied during training) it can speed things up a lot if we - calculate the bottleneck layer values once for each image during - preprocessing, and then just read those cached values repeatedly during - training. Here we go through all the images we've found, calculate those - values, and save them off. - - Args: - sess: The current active TensorFlow Session. - image_lists: Dictionary of training images for each label. - image_dir: Root folder string of the subfolders containing the training - images. - bottleneck_dir: Folder string holding cached files of bottleneck values. - jpeg_data_tensor: Input tensor for jpeg data from file. - decoded_image_tensor: The output of decoding and resizing the image. - resized_input_tensor: The input node of the recognition graph. - bottleneck_tensor: The penultimate output layer of the graph. - architecture: The name of the model architecture. - - Returns: - Nothing. - """ - how_many_bottlenecks = 0 - ensure_dir_exists(bottleneck_dir) - for label_name, label_lists in image_lists.items(): - for category in ['training', 'testing', 'validation']: - category_list = label_lists[category] - for index, unused_base_name in enumerate(category_list): - get_or_create_bottleneck( - sess, image_lists, label_name, index, image_dir, category, - bottleneck_dir, jpeg_data_tensor, decoded_image_tensor, - resized_input_tensor, bottleneck_tensor, architecture) - - how_many_bottlenecks += 1 - if how_many_bottlenecks % 100 == 0: - tf.logging.info( - str(how_many_bottlenecks) + ' bottleneck files created.') - - -def get_random_cached_bottlenecks(sess, image_lists, how_many, category, - bottleneck_dir, image_dir, jpeg_data_tensor, - decoded_image_tensor, resized_input_tensor, - bottleneck_tensor, architecture): - """Retrieves bottleneck values for cached images. - - If no distortions are being applied, this function can retrieve the cached - bottleneck values directly from disk for images. It picks a random set of - images from the specified category. - - Args: - sess: Current TensorFlow Session. - image_lists: Dictionary of training images for each label. - how_many: If positive, a random sample of this size will be chosen. - If negative, all bottlenecks will be retrieved. - category: Name string of which set to pull from - training, testing, or - validation. - bottleneck_dir: Folder string holding cached files of bottleneck values. - image_dir: Root folder string of the subfolders containing the training - images. - jpeg_data_tensor: The layer to feed jpeg image data into. - decoded_image_tensor: The output of decoding and resizing the image. - resized_input_tensor: The input node of the recognition graph. - bottleneck_tensor: The bottleneck output layer of the CNN graph. - architecture: The name of the model architecture. - - Returns: - List of bottleneck arrays, their corresponding ground truths, and the - relevant filenames. - """ - class_count = len(image_lists.keys()) - bottlenecks = [] - ground_truths = [] - filenames = [] - if how_many >= 0: - # Retrieve a random sample of bottlenecks. - for unused_i in range(how_many): - label_index = random.randrange(class_count) - label_name = list(image_lists.keys())[label_index] - image_index = random.randrange(MAX_NUM_IMAGES_PER_CLASS + 1) - image_name = get_image_path(image_lists, label_name, image_index, - image_dir, category) - bottleneck = get_or_create_bottleneck( - sess, image_lists, label_name, image_index, image_dir, category, - bottleneck_dir, jpeg_data_tensor, decoded_image_tensor, - resized_input_tensor, bottleneck_tensor, architecture) - bottlenecks.append(bottleneck) - ground_truths.append(label_index) - filenames.append(image_name) - else: - # Retrieve all bottlenecks. - for label_index, label_name in enumerate(image_lists.keys()): - for image_index, image_name in enumerate( - image_lists[label_name][category]): - image_name = get_image_path(image_lists, label_name, image_index, - image_dir, category) - bottleneck = get_or_create_bottleneck( - sess, image_lists, label_name, image_index, image_dir, category, - bottleneck_dir, jpeg_data_tensor, decoded_image_tensor, - resized_input_tensor, bottleneck_tensor, architecture) - bottlenecks.append(bottleneck) - ground_truths.append(label_index) - filenames.append(image_name) - return bottlenecks, ground_truths, filenames - - -def get_random_distorted_bottlenecks( - sess, image_lists, how_many, category, image_dir, input_jpeg_tensor, - distorted_image, resized_input_tensor, bottleneck_tensor): - """Retrieves bottleneck values for training images, after distortions. - - If we're training with distortions like crops, scales, or flips, we have to - recalculate the full model for every image, and so we can't use cached - bottleneck values. Instead we find random images for the requested category, - run them through the distortion graph, and then the full graph to get the - bottleneck results for each. - - Args: - sess: Current TensorFlow Session. - image_lists: Dictionary of training images for each label. - how_many: The integer number of bottleneck values to return. - category: Name string of which set of images to fetch - training, testing, - or validation. - image_dir: Root folder string of the subfolders containing the training - images. - input_jpeg_tensor: The input layer we feed the image data to. - distorted_image: The output node of the distortion graph. - resized_input_tensor: The input node of the recognition graph. - bottleneck_tensor: The bottleneck output layer of the CNN graph. - - Returns: - List of bottleneck arrays and their corresponding ground truths. - """ - class_count = len(image_lists.keys()) - bottlenecks = [] - ground_truths = [] - for unused_i in range(how_many): - label_index = random.randrange(class_count) - label_name = list(image_lists.keys())[label_index] - image_index = random.randrange(MAX_NUM_IMAGES_PER_CLASS + 1) - image_path = get_image_path(image_lists, label_name, image_index, image_dir, - category) - if not gfile.Exists(image_path): - tf.logging.fatal('File does not exist %s', image_path) - jpeg_data = gfile.FastGFile(image_path, 'rb').read() - # Note that we materialize the distorted_image_data as a numpy array before - # sending running inference on the image. This involves 2 memory copies and - # might be optimized in other implementations. - distorted_image_data = sess.run(distorted_image, - {input_jpeg_tensor: jpeg_data}) - bottleneck_values = sess.run(bottleneck_tensor, - {resized_input_tensor: distorted_image_data}) - bottleneck_values = np.squeeze(bottleneck_values) - bottlenecks.append(bottleneck_values) - ground_truths.append(label_index) - return bottlenecks, ground_truths - - -def should_distort_images(flip_left_right, random_crop, random_scale, - random_brightness): - """Whether any distortions are enabled, from the input flags. - - Args: - flip_left_right: Boolean whether to randomly mirror images horizontally. - random_crop: Integer percentage setting the total margin used around the - crop box. - random_scale: Integer percentage of how much to vary the scale by. - random_brightness: Integer range to randomly multiply the pixel values by. - - Returns: - Boolean value indicating whether any distortions should be applied. - """ - return (flip_left_right or (random_crop != 0) or (random_scale != 0) or - (random_brightness != 0)) - - -def add_input_distortions(flip_left_right, random_crop, random_scale, - random_brightness, input_width, input_height, - input_depth, input_mean, input_std): - """Creates the operations to apply the specified distortions. - - During training it can help to improve the results if we run the images - through simple distortions like crops, scales, and flips. These reflect the - kind of variations we expect in the real world, and so can help train the - model to cope with natural data more effectively. Here we take the supplied - parameters and construct a network of operations to apply them to an image. - - Cropping - ~~~~~~~~ - - Cropping is done by placing a bounding box at a random position in the full - image. The cropping parameter controls the size of that box relative to the - input image. If it's zero, then the box is the same size as the input and no - cropping is performed. If the value is 50%, then the crop box will be half the - width and height of the input. In a diagram it looks like this: - - < width > - +---------------------+ - | | - | width - crop% | - | < > | - | +------+ | - | | | | - | | | | - | | | | - | +------+ | - | | - | | - +---------------------+ - - Scaling - ~~~~~~~ - - Scaling is a lot like cropping, except that the bounding box is always - centered and its size varies randomly within the given range. For example if - the scale percentage is zero, then the bounding box is the same size as the - input and no scaling is applied. If it's 50%, then the bounding box will be in - a random range between half the width and height and full size. - - Args: - flip_left_right: Boolean whether to randomly mirror images horizontally. - random_crop: Integer percentage setting the total margin used around the - crop box. - random_scale: Integer percentage of how much to vary the scale by. - random_brightness: Integer range to randomly multiply the pixel values by. - graph. - input_width: Horizontal size of expected input image to model. - input_height: Vertical size of expected input image to model. - input_depth: How many channels the expected input image should have. - input_mean: Pixel value that should be zero in the image for the graph. - input_std: How much to divide the pixel values by before recognition. - - Returns: - The jpeg input layer and the distorted result tensor. - """ - - jpeg_data = tf.placeholder(tf.string, name='DistortJPGInput') - decoded_image = tf.image.decode_jpeg(jpeg_data, channels=input_depth) - decoded_image_as_float = tf.cast(decoded_image, dtype=tf.float32) - decoded_image_4d = tf.expand_dims(decoded_image_as_float, 0) - margin_scale = 1.0 + (random_crop / 100.0) - resize_scale = 1.0 + (random_scale / 100.0) - margin_scale_value = tf.constant(margin_scale) - resize_scale_value = tf.random_uniform(tensor_shape.scalar(), - minval=1.0, - maxval=resize_scale) - scale_value = tf.multiply(margin_scale_value, resize_scale_value) - precrop_width = tf.multiply(scale_value, input_width) - precrop_height = tf.multiply(scale_value, input_height) - precrop_shape = tf.stack([precrop_height, precrop_width]) - precrop_shape_as_int = tf.cast(precrop_shape, dtype=tf.int32) - precropped_image = tf.image.resize_bilinear(decoded_image_4d, - precrop_shape_as_int) - precropped_image_3d = tf.squeeze(precropped_image, squeeze_dims=[0]) - cropped_image = tf.random_crop(precropped_image_3d, - [input_height, input_width, input_depth]) - if flip_left_right: - flipped_image = tf.image.random_flip_left_right(cropped_image) - else: - flipped_image = cropped_image - brightness_min = 1.0 - (random_brightness / 100.0) - brightness_max = 1.0 + (random_brightness / 100.0) - brightness_value = tf.random_uniform(tensor_shape.scalar(), - minval=brightness_min, - maxval=brightness_max) - brightened_image = tf.multiply(flipped_image, brightness_value) - offset_image = tf.subtract(brightened_image, input_mean) - mul_image = tf.multiply(offset_image, 1.0 / input_std) - distort_result = tf.expand_dims(mul_image, 0, name='DistortResult') - return jpeg_data, distort_result - - -def variable_summaries(var): - """Attach a lot of summaries to a Tensor (for TensorBoard visualization).""" - with tf.name_scope('summaries'): - mean = tf.reduce_mean(var) - tf.summary.scalar('mean', mean) - with tf.name_scope('stddev'): - stddev = tf.sqrt(tf.reduce_mean(tf.square(var - mean))) - tf.summary.scalar('stddev', stddev) - tf.summary.scalar('max', tf.reduce_max(var)) - tf.summary.scalar('min', tf.reduce_min(var)) - tf.summary.histogram('histogram', var) - - -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 - weights, and then sets up all the gradients for the backward pass. - - The set up for the softmax and fully-connected layers is based on: - https://www.tensorflow.org/versions/master/tutorials/mnist/beginners/index.html - - Args: - class_count: Integer of how many categories of things we're trying to - recognize. - final_tensor_name: Name string for the new final node that produces results. - 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 - 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 - bottleneck input and ground truth input. - """ - with tf.name_scope('input'): - bottleneck_input = tf.placeholder_with_default( - bottleneck_tensor, - shape=[None, bottleneck_tensor_size], - name='BottleneckInputPlaceholder') - - ground_truth_input = tf.placeholder( - tf.int64, [None], name='GroundTruthInput') - - # 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') - variable_summaries(layer_weights) - - with tf.name_scope('biases'): - layer_biases = tf.Variable(tf.zeros([class_count]), name='final_biases') - variable_summaries(layer_biases) - - with tf.name_scope('Wx_plus_b'): - 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) - - tf.summary.scalar('cross_entropy', cross_entropy_mean) - - with tf.name_scope('train'): - optimizer = tf.train.GradientDescentOptimizer(FLAGS.learning_rate) - train_step = optimizer.minimize(cross_entropy_mean) - - return (train_step, cross_entropy_mean, bottleneck_input, ground_truth_input, - final_tensor) - - -def add_evaluation_step(result_tensor, ground_truth_tensor): - """Inserts the operations we need to evaluate the accuracy of our results. - - Args: - result_tensor: The new final node that produces results. - ground_truth_tensor: The node we feed ground truth data - into. - - Returns: - Tuple of (evaluation step, prediction). - """ - with tf.name_scope('accuracy'): - with tf.name_scope('correct_prediction'): - prediction = tf.argmax(result_tensor, 1) - correct_prediction = tf.equal(prediction, ground_truth_tensor) - with tf.name_scope('accuracy'): - evaluation_step = tf.reduce_mean(tf.cast(correct_prediction, tf.float32)) - tf.summary.scalar('accuracy', evaluation_step) - return evaluation_step, prediction - - -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. - """ - 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)) - - (sess, bottleneck_input, ground_truth_input, evaluation_step, - prediction) = build_eval_session(model_info, class_count) - - 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()) - - -def prepare_file_system(): - # Setup the directory we'll write summaries to for TensorBoard - if tf.gfile.Exists(FLAGS.summaries_dir): - tf.gfile.DeleteRecursively(FLAGS.summaries_dir) - tf.gfile.MakeDirs(FLAGS.summaries_dir) - if FLAGS.intermediate_store_frequency > 0: - ensure_dir_exists(FLAGS.intermediate_output_graphs_dir) - return - - -def create_model_info(architecture): - """Given the name of a model architecture, returns information about it. - - There are different base image recognition pretrained models that can be - retrained using transfer learning, and this function translates from the name - of a model to the attributes that are needed to download and train with it. - - Args: - architecture: Name of a model architecture. - - Returns: - Dictionary of information about the model, or None if the name isn't - recognized - - Raises: - ValueError: If architecture name is unknown. - """ - architecture = architecture.lower() - is_quantized = False - if architecture == 'inception_v3': - # pylint: disable=line-too-long - data_url = 'http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz' - # pylint: enable=line-too-long - bottleneck_tensor_name = 'pool_3/_reshape:0' - bottleneck_tensor_size = 2048 - input_width = 299 - input_height = 299 - input_depth = 3 - resized_input_tensor_name = 'Mul:0' - model_file_name = 'classify_image_graph_def.pb' - input_mean = 128 - input_std = 128 - elif architecture.startswith('mobilenet_'): - parts = architecture.split('_') - if len(parts) != 3 and len(parts) != 4: - tf.logging.error("Couldn't understand architecture name '%s'", - architecture) - return None - version_string = parts[1] - if (version_string != '1.0' and version_string != '0.75' and - version_string != '0.5' and version_string != '0.25'): - tf.logging.error( - """"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 - size_string != '160' and size_string != '128'): - tf.logging.error( - """The Mobilenet input size should be '224', '192', '160', or '128', - but found '%s' for architecture '%s'""", - size_string, architecture) - return None - if len(parts) == 3: - is_quantized = False - else: - 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: - 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 - input_mean = 127.5 - input_std = 127.5 - else: - tf.logging.error("Couldn't understand architecture name '%s'", architecture) - raise ValueError('Unknown architecture', architecture) - - return { - 'data_url': data_url, - 'bottleneck_tensor_name': bottleneck_tensor_name, - 'bottleneck_tensor_size': bottleneck_tensor_size, - 'input_width': input_width, - 'input_height': input_height, - 'input_depth': input_depth, - 'resized_input_tensor_name': resized_input_tensor_name, - 'model_file_name': model_file_name, - 'input_mean': input_mean, - 'input_std': input_std, - 'quantize_layer': is_quantized, - } - - -def add_jpeg_decoding(input_width, input_height, input_depth, input_mean, - input_std): - """Adds operations that perform JPEG decoding and resizing to the graph.. - - Args: - input_width: Desired width of the image fed into the recognizer graph. - input_height: Desired width of the image fed into the recognizer graph. - input_depth: Desired channels of the image fed into the recognizer graph. - input_mean: Pixel value that should be zero in the image for the graph. - input_std: How much to divide the pixel values by before recognition. - - Returns: - Tensors for the node to feed JPEG data into, and the output of the - preprocessing steps. - """ - jpeg_data = tf.placeholder(tf.string, name='DecodeJPGInput') - decoded_image = tf.image.decode_jpeg(jpeg_data, channels=input_depth) - decoded_image_as_float = tf.cast(decoded_image, dtype=tf.float32) - decoded_image_4d = tf.expand_dims(decoded_image_as_float, 0) - resize_shape = tf.stack([input_height, input_width]) - resize_shape_as_int = tf.cast(resize_shape, dtype=tf.int32) - resized_image = tf.image.resize_bilinear(decoded_image_4d, - resize_shape_as_int) - offset_image = tf.subtract(resized_image, input_mean) - mul_image = tf.multiply(offset_image, 1.0 / input_std) - 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 - tf.logging.set_verbosity(tf.logging.INFO) - - # Prepare necessary directories that can be used during training - prepare_file_system() - - # Gather information about the model architecture we'll be using. - model_info = create_model_info(FLAGS.architecture) - if not model_info: - tf.logging.error('Did not recognize architecture flag') - return -1 - - # 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) - class_count = len(image_lists.keys()) - if class_count == 0: - tf.logging.error('No valid folders of images found at ' + FLAGS.image_dir) - return -1 - if class_count == 1: - tf.logging.error('Only one valid folder of images found at ' + - FLAGS.image_dir + - ' - multiple classes are needed for classification.') - return -1 - - # See if the command-line flags mean we're applying any distortions. - do_distort_images = should_distort_images( - 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( - model_info['input_width'], model_info['input_height'], - model_info['input_depth'], model_info['input_mean'], - model_info['input_std']) - - if do_distort_images: - # We will be applying distortions, so setup the operations we'll need. - (distorted_jpeg_data_tensor, - distorted_image_tensor) = add_input_distortions( - FLAGS.flip_left_right, FLAGS.random_crop, FLAGS.random_scale, - FLAGS.random_brightness, model_info['input_width'], - model_info['input_height'], model_info['input_depth'], - model_info['input_mean'], model_info['input_std']) - else: - # We'll make sure we've calculated the 'bottleneck' image summaries and - # cached them on disk. - cache_bottlenecks(sess, image_lists, FLAGS.image_dir, - FLAGS.bottleneck_dir, jpeg_data_tensor, - decoded_image_tensor, resized_image_tensor, - bottleneck_tensor, FLAGS.architecture) - - # Create the operations we need to evaluate the accuracy of our new layer. - 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() - train_writer = tf.summary.FileWriter(FLAGS.summaries_dir + '/train', - sess.graph) - - 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) - - # Run the training for as many cycles as requested on the command line. - for i in range(FLAGS.how_many_training_steps): - # Get a batch of input bottleneck values, either calculated fresh every - # time with distortions applied, or from the cache stored on disk. - if do_distort_images: - (train_bottlenecks, - train_ground_truth) = get_random_distorted_bottlenecks( - sess, image_lists, FLAGS.train_batch_size, 'training', - FLAGS.image_dir, distorted_jpeg_data_tensor, - distorted_image_tensor, resized_image_tensor, bottleneck_tensor) - else: - (train_bottlenecks, - train_ground_truth, _) = get_random_cached_bottlenecks( - sess, image_lists, FLAGS.train_batch_size, 'training', - FLAGS.bottleneck_dir, FLAGS.image_dir, jpeg_data_tensor, - decoded_image_tensor, resized_image_tensor, bottleneck_tensor, - FLAGS.architecture) - # Feed the bottlenecks and ground truth into the graph, and run a training - # step. Capture training summaries for TensorBoard with the `merged` op. - train_summary, _ = sess.run( - [merged, train_step], - feed_dict={bottleneck_input: train_bottlenecks, - ground_truth_input: train_ground_truth}) - train_writer.add_summary(train_summary, i) - - # Every so often, print out how well the graph is training. - is_last_step = (i + 1 == FLAGS.how_many_training_steps) - if (i % FLAGS.eval_step_interval) == 0 or is_last_step: - train_accuracy, cross_entropy_value = sess.run( - [evaluation_step, cross_entropy], - feed_dict={bottleneck_input: train_bottlenecks, - ground_truth_input: train_ground_truth}) - tf.logging.info('%s: Step %d: Train accuracy = %.1f%%' % - (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', - FLAGS.bottleneck_dir, FLAGS.image_dir, jpeg_data_tensor, - decoded_image_tensor, resized_image_tensor, bottleneck_tensor, - FLAGS.architecture)) - # Run a validation step and capture training summaries for TensorBoard - # with the `merged` op. - validation_summary, validation_accuracy = sess.run( - [merged, evaluation_step], - feed_dict={bottleneck_input: validation_bottlenecks, - ground_truth_input: validation_ground_truth}) - validation_writer.add_summary(validation_summary, i) - tf.logging.info('%s: Step %d: Validation accuracy = %.1f%% (N=%d)' % - (datetime.now(), i, validation_accuracy * 100, - len(validation_bottlenecks))) - - # Store intermediate results - intermediate_frequency = FLAGS.intermediate_store_frequency - - 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(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. - 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(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() - parser.add_argument( - '--image_dir', - type=str, - default='', - help='Path to folders of labeled images.' - ) - parser.add_argument( - '--output_graph', - type=str, - default='/tmp/output_graph.pb', - help='Where to save the trained graph.' - ) - parser.add_argument( - '--intermediate_output_graphs_dir', - type=str, - default='/tmp/intermediate_graph/', - help='Where to save the intermediate graphs.' - ) - parser.add_argument( - '--intermediate_store_frequency', - type=int, - default=0, - help="""\ - How many steps to store intermediate graph. If "0" then will not - store.\ - """ - ) - parser.add_argument( - '--output_labels', - type=str, - default='/tmp/output_labels.txt', - help='Where to save the trained graph\'s labels.' - ) - parser.add_argument( - '--summaries_dir', - type=str, - default='/tmp/retrain_logs', - help='Where to save summary logs for TensorBoard.' - ) - parser.add_argument( - '--how_many_training_steps', - type=int, - default=4000, - help='How many training steps to run before ending.' - ) - parser.add_argument( - '--learning_rate', - type=float, - default=0.01, - help='How large a learning rate to use when training.' - ) - parser.add_argument( - '--testing_percentage', - type=int, - default=10, - help='What percentage of images to use as a test set.' - ) - parser.add_argument( - '--validation_percentage', - type=int, - default=10, - help='What percentage of images to use as a validation set.' - ) - parser.add_argument( - '--eval_step_interval', - type=int, - default=10, - help='How often to evaluate the training results.' - ) - parser.add_argument( - '--train_batch_size', - type=int, - default=100, - help='How many images to train on at a time.' - ) - parser.add_argument( - '--test_batch_size', - type=int, - default=-1, - help="""\ - How many images to test on. This test set is only used once, to evaluate - the final accuracy of the model after training completes. - A value of -1 causes the entire test set to be used, which leads to more - stable results across runs.\ - """ - ) - parser.add_argument( - '--validation_batch_size', - type=int, - default=100, - help="""\ - How many images to use in an evaluation batch. This validation set is - used much more often than the test set, and is an early indicator of how - accurate the model is during training. - A value of -1 causes the entire validation set to be used, which leads to - more stable results across training iterations, but may be slower on large - training sets.\ - """ - ) - parser.add_argument( - '--print_misclassified_test_images', - default=False, - help="""\ - Whether to print out a list of all misclassified test images.\ - """, - action='store_true' - ) - parser.add_argument( - '--model_dir', - type=str, - default='/tmp/imagenet', - help="""\ - Path to classify_image_graph_def.pb, - imagenet_synset_to_human_label_map.txt, and - imagenet_2012_challenge_label_map_proto.pbtxt.\ - """ - ) - parser.add_argument( - '--bottleneck_dir', - type=str, - default='/tmp/bottleneck', - help='Path to cache bottleneck layer values as files.' - ) - parser.add_argument( - '--final_tensor_name', - type=str, - default='final_result', - help="""\ - The name of the output classification layer in the retrained graph.\ - """ - ) - parser.add_argument( - '--flip_left_right', - default=False, - help="""\ - Whether to randomly flip half of the training images horizontally.\ - """, - action='store_true' - ) - parser.add_argument( - '--random_crop', - type=int, - default=0, - help="""\ - A percentage determining how much of a margin to randomly crop off the - training images.\ - """ - ) - parser.add_argument( - '--random_scale', - type=int, - default=0, - help="""\ - A percentage determining how much to randomly scale up the size of the - training images by.\ - """ - ) - parser.add_argument( - '--random_brightness', - type=int, - default=0, - help="""\ - A percentage determining how much to randomly multiply the training image - input pixels up or down by.\ - """ - ) - parser.add_argument( - '--architecture', - type=str, - default='inception_v3', - help="""\ - Which model architecture to use. 'inception_v3' is the most accurate, but - also the slowest. For faster or smaller models, chose a MobileNet with the - 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 - 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 deleted file mode 100644 index fb7324c58ac1be60baad840207f31a61ec6182be..0000000000000000000000000000000000000000 --- a/tensorflow/examples/image_retraining/retrain_test.py +++ /dev/null @@ -1,148 +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=g-bad-import-order,unused-import -"""Tests the graph freezing tool.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import tensorflow as tf -import os - -from tensorflow.examples.image_retraining import retrain -from tensorflow.python.framework import test_util - - -class ImageRetrainingTest(test_util.TensorFlowTestCase): - - def dummyImageLists(self): - return {'label_one': {'dir': 'somedir', 'training': ['image_one.jpg', - 'image_two.jpg'], - 'testing': ['image_three.jpg', 'image_four.jpg'], - 'validation': ['image_five.jpg', 'image_six.jpg']}, - 'label_two': {'dir': 'otherdir', 'training': ['image_one.jpg', - 'image_two.jpg'], - 'testing': ['image_three.jpg', 'image_four.jpg'], - 'validation': ['image_five.jpg', 'image_six.jpg']}} - - def testGetImagePath(self): - image_lists = self.dummyImageLists() - self.assertEqual('image_dir/somedir/image_one.jpg', retrain.get_image_path( - image_lists, 'label_one', 0, 'image_dir', 'training')) - self.assertEqual('image_dir/otherdir/image_four.jpg', - retrain.get_image_path(image_lists, 'label_two', 1, - 'image_dir', 'testing')) - - def testGetBottleneckPath(self): - image_lists = self.dummyImageLists() - self.assertEqual('bottleneck_dir/somedir/image_five.jpg_imagenet_v3.txt', - retrain.get_bottleneck_path( - image_lists, 'label_one', 0, 'bottleneck_dir', - 'validation', 'imagenet_v3')) - - def testShouldDistortImage(self): - self.assertEqual(False, retrain.should_distort_images(False, 0, 0, 0)) - self.assertEqual(True, retrain.should_distort_images(True, 0, 0, 0)) - self.assertEqual(True, retrain.should_distort_images(False, 10, 0, 0)) - self.assertEqual(True, retrain.should_distort_images(False, 0, 1, 0)) - self.assertEqual(True, retrain.should_distort_images(False, 0, 0, 50)) - - def testAddInputDistortions(self): - with tf.Graph().as_default(): - with tf.Session() as sess: - retrain.add_input_distortions(True, 10, 10, 10, 299, 299, 3, 128, 128) - self.assertIsNotNone(sess.graph.get_tensor_by_name('DistortJPGInput:0')) - self.assertIsNotNone(sess.graph.get_tensor_by_name('DistortResult:0')) - - @tf.test.mock.patch.object(retrain, 'FLAGS', learning_rate=0.01) - 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_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 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, 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(): - final = tf.placeholder(tf.float32, [1], name='final') - gt = tf.placeholder(tf.int64, [1], name='gt') - self.assertIsNotNone(retrain.add_evaluation_step(final, gt)) - - def testAddJpegDecoding(self): - with tf.Graph().as_default(): - jpeg_data, mul_image = retrain.add_jpeg_decoding(10, 10, 3, 0, 255) - self.assertIsNotNone(jpeg_data) - self.assertIsNotNone(mul_image) - - def testCreateModelInfo(self): - did_raise_value_error = False - try: - retrain.create_model_info('no_such_model_name') - except ValueError: - did_raise_value_error = True - self.assertTrue(did_raise_value_error) - model_info = retrain.create_model_info('inception_v3') - self.assertIsNotNone(model_info) - self.assertEqual(299, model_info['input_width']) - - def testCreateModelInfoQuantized(self): - # Test for mobilenet_quantized - model_info = retrain.create_model_info('mobilenet_1.0_224') - self.assertIsNotNone(model_info) - self.assertEqual(224, model_info['input_width']) - - -if __name__ == '__main__': - tf.test.main() diff --git a/tensorflow/examples/speech_commands/input_data.py b/tensorflow/examples/speech_commands/input_data.py index e7db9cddf02daf9a32d3ed859ee9bd35b2cae838..63dd18457fea42acb09058b9ddd4623d72d1fd04 100644 --- a/tensorflow/examples/speech_commands/input_data.py +++ b/tensorflow/examples/speech_commands/input_data.py @@ -457,7 +457,7 @@ class AudioProcessor(object): self.time_shift_offset_placeholder_: time_shift_offset, } # Choose a section of background noise to mix in. - if use_background: + if use_background or sample['label'] == SILENCE_LABEL: background_index = np.random.randint(len(self.background_data)) background_samples = self.background_data[background_index] background_offset = np.random.randint( @@ -465,7 +465,9 @@ class AudioProcessor(object): background_clipped = background_samples[background_offset:( background_offset + desired_samples)] background_reshaped = background_clipped.reshape([desired_samples, 1]) - if np.random.uniform(0, 1) < background_frequency: + if sample['label'] == SILENCE_LABEL: + background_volume = np.random.uniform(0, 1) + elif np.random.uniform(0, 1) < background_frequency: background_volume = np.random.uniform(0, background_volume_range) else: background_volume = 0 diff --git a/tensorflow/examples/image_retraining/__init__.py b/tensorflow/examples/tutorials/estimators/__init__.py similarity index 100% rename from tensorflow/examples/image_retraining/__init__.py rename to tensorflow/examples/tutorials/estimators/__init__.py diff --git a/tensorflow/examples/tutorials/input_fn/__init__.py b/tensorflow/examples/tutorials/input_fn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tensorflow/examples/tutorials/layers/__init__.py b/tensorflow/examples/tutorials/layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tensorflow/examples/tutorials/mnist/BUILD b/tensorflow/examples/tutorials/mnist/BUILD index aa1b2ec2db34f3cb0350bfde88a1598ed71456de..d7bc6a5a7d1e4cd3927c7c5067ccc22993885994 100644 --- a/tensorflow/examples/tutorials/mnist/BUILD +++ b/tensorflow/examples/tutorials/mnist/BUILD @@ -51,6 +51,7 @@ py_binary( "fully_connected_feed.py", ], srcs_version = "PY2AND3", + tags = ["optonly"], deps = [ ":input_data", ":mnist", diff --git a/tensorflow/examples/tutorials/monitors/__init__.py b/tensorflow/examples/tutorials/monitors/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tensorflow/examples/tutorials/monitors/iris_monitors.py b/tensorflow/examples/tutorials/monitors/iris_monitors.py index 850d105f7b1b33fadd40bc6a6cab3d08c0da3734..a2b7fe60237da0604f74f31c0a09951f708e908b 100644 --- a/tensorflow/examples/tutorials/monitors/iris_monitors.py +++ b/tensorflow/examples/tutorials/monitors/iris_monitors.py @@ -32,9 +32,9 @@ IRIS_TEST = os.path.join(os.path.dirname(__file__), "iris_test.csv") def main(unused_argv): # Load datasets. training_set = tf.contrib.learn.datasets.base.load_csv_with_header( - filename=IRIS_TRAINING, target_dtype=np.int, features_dtype=np.float) + filename=IRIS_TRAINING, target_dtype=np.int, features_dtype=np.float32) test_set = tf.contrib.learn.datasets.base.load_csv_with_header( - filename=IRIS_TEST, target_dtype=np.int, features_dtype=np.float) + filename=IRIS_TEST, target_dtype=np.int, features_dtype=np.float32) validation_metrics = { "accuracy": @@ -83,7 +83,7 @@ def main(unused_argv): # Classify two new flower samples. new_samples = np.array( - [[6.4, 3.2, 4.5, 1.5], [5.8, 3.1, 5.0, 1.7]], dtype=float) + [[6.4, 3.2, 4.5, 1.5], [5.8, 3.1, 5.0, 1.7]], dtype=np.float32) y = list(classifier.predict(new_samples)) print("Predictions: {}".format(str(y))) diff --git a/tensorflow/examples/tutorials/word2vec/word2vec_basic.py b/tensorflow/examples/tutorials/word2vec/word2vec_basic.py index 14ae7fbf35836ad7f5d56101ae0fc33a3f3fb9ba..b09ee9976897fcab2e90fdc17e8030532080aca8 100644 --- a/tensorflow/examples/tutorials/word2vec/word2vec_basic.py +++ b/tensorflow/examples/tutorials/word2vec/word2vec_basic.py @@ -224,7 +224,7 @@ with graph.as_default(): optimizer = tf.train.GradientDescentOptimizer(1.0).minimize(loss) # Compute the cosine similarity between minibatch examples and all embeddings. - norm = tf.sqrt(tf.reduce_sum(tf.square(embeddings), 1, keep_dims=True)) + norm = tf.sqrt(tf.reduce_sum(tf.square(embeddings), 1, keepdims=True)) normalized_embeddings = embeddings / norm valid_embeddings = tf.nn.embedding_lookup(normalized_embeddings, valid_dataset) diff --git a/tensorflow/go/op/wrappers.go b/tensorflow/go/op/wrappers.go index 3d261c9d0ad7c74350e09d23de3584dd28276b7f..c31ca8b67a10707e4e145243b6f1a1992601b4b1 100644 --- a/tensorflow/go/op/wrappers.go +++ b/tensorflow/go/op/wrappers.go @@ -43,7 +43,7 @@ type FakeQuantWithMinMaxVarsPerChannelGradientAttr func(optionalAttr) // FakeQuantWithMinMaxVarsPerChannelGradientNumBits sets the optional num_bits attribute to value. // -// value: The bitwidth of the quantization; between 2 and 8, inclusive. +// value: The bitwidth of the quantization; between 2 and 16, inclusive. // If not specified, defaults to 8 func FakeQuantWithMinMaxVarsPerChannelGradientNumBits(value int64) FakeQuantWithMinMaxVarsPerChannelGradientAttr { return func(m optionalAttr) { @@ -124,7 +124,7 @@ func FakeQuantWithMinMaxVarsPerChannelNarrowRange(value bool) FakeQuantWithMinMa // `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. +// `num_bits` is the bitwidth of the quantization; between 2 and 16, inclusive. // // This operation has a gradient and thus allows for training `min` and `max` // values. @@ -305,7 +305,7 @@ func FakeQuantWithMinMaxArgsNarrowRange(value bool) FakeQuantWithMinMaxArgsAttr // `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. +// `num_bits` is the bitwidth of the quantization; between 2 and 16, 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) { @@ -401,6 +401,9 @@ func FakeQuantWithMinMaxArgs(scope *Scope, inputs tf.Output, optional ...FakeQua // [[5, 5, 5, 5], [6, 6, 6, 6], [7, 7, 7, 7], [8, 8, 8, 8]], // [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]] // +// Note that on CPU, if an out of bound index is found, an error is returned. +// On GPU, if an out of bound index is found, the index is ignored. +// // Arguments: // indices: Index tensor. // updates: Updates to scatter into output. @@ -1845,6 +1848,93 @@ func ReverseSequence(scope *Scope, input tf.Output, seq_lengths tf.Output, seq_d return op.Output(0) } +// UniqueWithCountsV2Attr is an optional argument to UniqueWithCountsV2. +type UniqueWithCountsV2Attr func(optionalAttr) + +// UniqueWithCountsV2OutIdx sets the optional out_idx attribute to value. +// If not specified, defaults to DT_INT32 +func UniqueWithCountsV2OutIdx(value tf.DataType) UniqueWithCountsV2Attr { + return func(m optionalAttr) { + m["out_idx"] = value + } +} + +// Finds unique elements along an axis of a tensor. +// +// This operation either returns a tensor `y` containing unique elements +// along the `axis` of a tensor. The returned unique elements is sorted +// in the same order as they occur along `axis` in `x`. +// This operation also returns a tensor `idx` and a tensor `count` +// that are the same size as the number of the elements in `x` along the +// `axis` dimension. The `idx` contains the index in the unique output `y` +// and the `count` contains the count in the unique output `y`. +// In other words, for an `1-D` tensor `x` with `axis = None: +// +// `y[idx[i]] = x[i] for i in [0, 1,...,rank(x) - 1]` +// +// For example: +// +// ``` +// # tensor 'x' is [1, 1, 2, 4, 4, 4, 7, 8, 8] +// y, idx, count = unique_with_counts(x) +// y ==> [1, 2, 4, 7, 8] +// idx ==> [0, 0, 1, 2, 2, 2, 3, 4, 4] +// count ==> [2, 1, 3, 1, 2] +// ``` +// +// For an `2-D` tensor `x` with `axis = 0`: +// +// ``` +// # tensor 'x' is [[1, 0, 0], +// # [1, 0, 0], +// # [2, 0, 0]] +// y, idx, count = unique_with_counts(x, axis=0) +// y ==> [[1, 0, 0], +// [2, 0, 0]] +// idx ==> [0, 0, 1] +// count ==> [2, 1] +// ``` +// +// For an `2-D` tensor `x` with `axis = 1`: +// +// ``` +// # tensor 'x' is [[1, 0, 0], +// # [1, 0, 0], +// # [2, 0, 0]] +// y, idx, count = unique_with_counts(x, axis=1) +// y ==> [[1, 0], +// [1, 0], +// [2, 0]] +// idx ==> [0, 1, 1] +// count ==> [1, 2] +// ``` +// +// Arguments: +// x: A `Tensor`. +// axis: A `Tensor` of type `int32` (default: None). The axis of the Tensor to +// find the unique elements. +// +// Returns A `Tensor`. Unique elements along the `axis` of `Tensor` x.A 1-D Tensor. Has the same type as x that contains the index of each +// value of x in the output y.A 1-D Tensor. The count of each value of x in the output y. +func UniqueWithCountsV2(scope *Scope, x tf.Output, axis tf.Output, optional ...UniqueWithCountsV2Attr) (y tf.Output, idx tf.Output, count tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "UniqueWithCountsV2", + Input: []tf.Input{ + x, axis, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0), op.Output(1), op.Output(2) +} + // UniqueWithCountsAttr is an optional argument to UniqueWithCounts. type UniqueWithCountsAttr func(optionalAttr) @@ -1910,12 +2000,15 @@ func UniqueV2OutIdx(value tf.DataType) UniqueV2Attr { } } -// Finds unique elements in a 1-D tensor. +// Finds unique elements along an axis of a tensor. // -// This operation returns a tensor `y` containing all of the unique elements of `x` -// sorted in the same order that they occur in `x`. This operation also returns a -// tensor `idx` the same size as `x` that contains the index of each value of `x` -// in the unique output `y`. In other words: +// This operation either returns a tensor `y` containing unique elements +// along the `axis` of a tensor. The returned unique elements is sorted +// in the same order as they occur along `axis` in `x`. +// This operation also returns a tensor `idx` that is the same size as +// the number of the elements in `x` along the `axis` dimension. It +// contains the index in the unique output `y`. +// In other words, for an `1-D` tensor `x` with `axis = None: // // `y[idx[i]] = x[i] for i in [0, 1,...,rank(x) - 1]` // @@ -1928,9 +2021,34 @@ func UniqueV2OutIdx(value tf.DataType) UniqueV2Attr { // idx ==> [0, 0, 1, 2, 2, 2, 3, 4, 4] // ``` // +// For an `2-D` tensor `x` with `axis = 0`: +// +// ``` +// # tensor 'x' is [[1, 0, 0], +// # [1, 0, 0], +// # [2, 0, 0]] +// y, idx = unique(x, axis=0) +// y ==> [[1, 0, 0], +// [2, 0, 0]] +// idx ==> [0, 0, 1] +// ``` +// +// For an `2-D` tensor `x` with `axis = 1`: +// +// ``` +// # tensor 'x' is [[1, 0, 0], +// # [1, 0, 0], +// # [2, 0, 0]] +// y, idx = unique(x, axis=1) +// y ==> [[1, 0], +// [1, 0], +// [2, 0]] +// idx ==> [0, 1, 1] +// ``` +// // Arguments: // x: A `Tensor`. -// axis: A `Tensor` of type `int64` (default: 0). The axis of the Tensor to +// axis: A `Tensor` of type `int32` (default: None). The axis of the Tensor to // find the unique elements. // // Returns A `Tensor`. Unique elements along the `axis` of `Tensor` x.A 1-D Tensor. Has the same type as x that contains the index of each @@ -2217,6 +2335,35 @@ func ResourceSparseApplyMomentum(scope *Scope, var_ tf.Output, accum tf.Output, return scope.AddOperation(opspec) } +// Clips tensor values to a specified min and max. +// +// Given a tensor `t`, this operation returns a tensor of the same type and +// shape as `t` with its values clipped to `clip_value_min` and `clip_value_max`. +// Any values less than `clip_value_min` are set to `clip_value_min`. Any values +// greater than `clip_value_max` are set to `clip_value_max`. +// +// Arguments: +// t: A `Tensor`. +// clip_value_min: A 0-D (scalar) `Tensor`, or a `Tensor` with the same shape +// as `t`. The minimum value to clip by. +// clip_value_max: A 0-D (scalar) `Tensor`, or a `Tensor` with the same shape +// as `t`. The maximum value to clip by. +// +// Returns A clipped `Tensor` with the same shape as input 't'. +func ClipByValue(scope *Scope, t tf.Output, clip_value_min tf.Output, clip_value_max tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "ClipByValue", + Input: []tf.Input{ + t, clip_value_min, clip_value_max, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Creates a sequence of numbers. // // This operation creates a sequence of numbers that begins at `start` and @@ -2277,7 +2424,7 @@ func SparseSegmentSqrtNGrad(scope *Scope, grad tf.Output, indices tf.Output, seg // Computes the mean along sparse segments of a tensor. // -// Read @{$math_ops#segmentation$the section on segmentation} for an explanation of +// 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 @@ -2332,7 +2479,7 @@ func StackPopV2(scope *Scope, handle tf.Output, elem_type tf.DataType) (elem tf. // Like `SparseSegmentSum`, but allows missing ids in `segment_ids`. If an id is // misisng, the `output` tensor at that position will be zeroed. // -// Read @{$math_ops#segmentation$the section on segmentation} for an explanation of +// Read @{$math_ops#Segmentation$the section on segmentation} for an explanation of // segments. // // For example: @@ -2505,42 +2652,9 @@ func SparseToDense(scope *Scope, sparse_indices tf.Output, output_shape tf.Outpu return op.Output(0) } -// Counts the number of occurrences of each value in an integer array. -// -// Outputs a vector with length `size` and the same dtype as `weights`. If -// `weights` are empty, then index `i` stores the number of times the value `i` is -// counted in `arr`. If `weights` are non-empty, then index `i` stores the sum of -// the value in `weights` at each index where the corresponding value in `arr` is -// `i`. -// -// Values in `arr` outside of the range [0, size) are ignored. -// -// Arguments: -// arr: int32 `Tensor`. -// size: non-negative int32 scalar `Tensor`. -// weights: is an int32, int64, float32, or float64 `Tensor` with the same -// shape as `arr`, or a length-0 `Tensor`, in which case it acts as all weights -// equal to 1. -// -// Returns 1D `Tensor` with length equal to `size`. The counts or summed weights for -// each value in the range [0, size). -func Bincount(scope *Scope, arr tf.Output, size tf.Output, weights tf.Output) (bins tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "Bincount", - Input: []tf.Input{ - arr, size, weights, - }, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - // Computes the sum along sparse segments of a tensor. // -// Read @{$math_ops#segmentation$the section on segmentation} for an explanation of +// 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 @@ -2605,6 +2719,44 @@ func Sinh(scope *Scope, x tf.Output) (y tf.Output) { return op.Output(0) } +// Computes the minimum 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 minimum such that: +// +// \\(output_i = \min_j data_j\\) where min is over `j` such +// that `segment_ids[j] == i`. +// +// If the minimum is empty for a given segment ID `i`, it outputs the largest +// possible value for the specific numeric type, +// `output[i] = numeric_limits::max()`. +// +// Arguments: +// +// segment_ids: A 1-D tensor whose rank is equal to the rank of `data`'s +// first dimension. +// +// +// Returns Has same shape as data, except for dimension 0 which +// has size `num_segments`. +func UnsortedSegmentMin(scope *Scope, data tf.Output, segment_ids tf.Output, num_segments tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "UnsortedSegmentMin", + Input: []tf.Input{ + data, segment_ids, num_segments, + }, + } + op := scope.AddOperation(opspec) + 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 { @@ -2622,7 +2774,7 @@ func Relu6(scope *Scope, features tf.Output) (activations tf.Output) { // Computes the sum along segments of a tensor. // -// Read @{$math_ops#segmentation$the section on segmentation} for an explanation of +// Read @{$math_ops#Segmentation$the section on segmentation} for an explanation of // segments. // // Computes a tensor such that @@ -2953,6 +3105,68 @@ func Betainc(scope *Scope, a tf.Output, b tf.Output, x tf.Output) (z tf.Output) return op.Output(0) } +// Return a tensor with the same shape and contents as the input tensor or value. +func Identity(scope *Scope, input tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "Identity", + Input: []tf.Input{ + input, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Computes arctangent of `y/x` element-wise, respecting signs of the arguments. +// +// This is the angle \( \theta \in [-\pi, \pi] \) such that +// \[ x = r \cos(\theta) \] +// and +// \[ y = r \sin(\theta) \] +// where \(r = \sqrt(x^2 + y^2) \). +func Atan2(scope *Scope, y tf.Output, x tf.Output) (z tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "Atan2", + Input: []tf.Input{ + y, x, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Creates a dataset that passes a sliding window over `input_dataset`. +// +// Arguments: +// +// window_size: A scalar representing the number of elements in the +// sliding window. +// stride: A scalar representing the steps moving the sliding window +// forward in one iteration. It must be in `[1, window_size)`. +// +// +func SlideDataset(scope *Scope, input_dataset tf.Output, window_size tf.Output, stride 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: "SlideDataset", + Input: []tf.Input{ + input_dataset, window_size, stride, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Computes the sum along sparse segments of a tensor divided by the sqrt of N. // // N is the size of the segment being reduced. @@ -2960,7 +3174,7 @@ func Betainc(scope *Scope, a tf.Output, b tf.Output, x tf.Output) (z tf.Output) // Like `SparseSegmentSqrtN`, but allows missing ids in `segment_ids`. If an id is // misisng, the `output` tensor at that position will be zeroed. // -// Read @{$math_ops#segmentation$the section on segmentation} for an explanation of +// Read @{$math_ops#Segmentation$the section on segmentation} for an explanation of // segments. // // Arguments: @@ -3266,20 +3480,21 @@ func NthElement(scope *Scope, input tf.Output, n tf.Output, optional ...NthEleme return op.Output(0) } -// Computes the Max along segments of a tensor. +// Computes the maximum along segments of a tensor. // -// Read @{$math_ops#segmentation$the section on segmentation} for an explanation of +// Read @{$math_ops#Segmentation$the section on segmentation} for an explanation of // segments. // -// This operator is similar to the [unsorted segment sum operator](../../../api_docs/python/math_ops.md#UnsortedSegmentSum). -// Instead of computing the sum over segments, it computes the maximum -// such that: +// This operator is similar to the unsorted segment sum operator found +// [(here)](../../../api_docs/python/math_ops.md#UnsortedSegmentSum). +// Instead of computing the sum over segments, it computes the maximum such that: // // \\(output_i = \max_j data_j\\) where max is over `j` such // that `segment_ids[j] == i`. // -// If the maximum is empty for a given segment ID `i`, it outputs the smallest possible value for specific numeric type, -// `output[i] = numeric_limits::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()`. // //
// @@ -3689,7 +3904,7 @@ func CholeskyGrad(scope *Scope, l tf.Output, grad tf.Output) (output tf.Output) // Like `SparseSegmentMean`, but allows missing ids in `segment_ids`. If an id is // misisng, the `output` tensor at that position will be zeroed. // -// Read @{$math_ops#segmentation$the section on segmentation} for an explanation of +// Read @{$math_ops#Segmentation$the section on segmentation} for an explanation of // segments. // // Arguments: @@ -3791,9 +4006,8 @@ type ResizeBicubicAttr func(optionalAttr) // ResizeBicubicAlignCorners sets the optional align_corners attribute to value. // -// value: If true, rescale input by (new_height - 1) / (height - 1), which -// exactly aligns the 4 corners of images and resized images. If false, rescale -// by new_height / height. Treat similarly the width dimension. +// value: If true, the centers of the 4 corner pixels of the input and output tensors are +// aligned, preserving the values at the corner pixels. Defaults to false. // If not specified, defaults to false func ResizeBicubicAlignCorners(value bool) ResizeBicubicAttr { return func(m optionalAttr) { @@ -4204,6 +4418,26 @@ func LogicalAnd(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) { return op.Output(0) } +// Checks whether a tree ensemble has been initialized. +// +// Arguments: +// tree_ensemble_handle: Handle to the tree ensemble resouce. +// +// Returns output boolean on whether it is initialized or not. +func IsBoostedTreesEnsembleInitialized(scope *Scope, tree_ensemble_handle tf.Output) (is_initialized tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "IsBoostedTreesEnsembleInitialized", + Input: []tf.Input{ + tree_ensemble_handle, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Cast x of type SrcT to y of DstT. func Cast(scope *Scope, x tf.Output, DstT tf.DataType) (y tf.Output) { if scope.Err() != nil { @@ -4878,6 +5112,23 @@ func SoftsignGrad(scope *Scope, gradients tf.Output, features tf.Output) (backpr return op.Output(0) } +// Provides the time since epoch in seconds. +// +// Returns the timestamp as a `float64` for seconds since the Unix epoch. +// +// Note: the timestamp is computed when the op is executed, not when it is added +// to the graph. +func Timestamp(scope *Scope) (ts tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "Timestamp", + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // BatchMatMulAttr is an optional argument to BatchMatMul. type BatchMatMulAttr func(optionalAttr) @@ -4965,34 +5216,98 @@ func IsNan(scope *Scope, x tf.Output) (y tf.Output) { return op.Output(0) } -// Computes rectified linear gradients for a Relu operation. -// -// Arguments: -// gradients: The backpropagated gradients to the corresponding Relu operation. -// features: The features passed as input to the corresponding Relu operation, OR -// the outputs of that operation (both work equivalently). +// Identity op for gradient debugging. // -// Returns `gradients * (features > 0)`. -func ReluGrad(scope *Scope, gradients tf.Output, features tf.Output) (backprops tf.Output) { +// This op is hidden from public in Python. It is used by TensorFlow Debugger to +// register gradient tensors for gradient debugging. +// This op operates on non-reference-type tensors. +func DebugGradientIdentity(scope *Scope, input tf.Output) (output tf.Output) { if scope.Err() != nil { return } opspec := tf.OpSpec{ - Type: "ReluGrad", + Type: "DebugGradientIdentity", Input: []tf.Input{ - gradients, features, + input, }, } op := scope.AddOperation(opspec) return op.Output(0) } -// Computes the gradient of morphological 2-D dilation with respect to the input. +// ResourceSparseApplyAdadeltaAttr is an optional argument to ResourceSparseApplyAdadelta. +type ResourceSparseApplyAdadeltaAttr func(optionalAttr) + +// ResourceSparseApplyAdadeltaUseLocking sets the optional use_locking attribute to value. // -// Arguments: -// input: 4-D with shape `[batch, in_height, in_width, depth]`. -// filter: 3-D with shape `[filter_height, filter_width, depth]`. -// out_backprop: 4-D with shape `[batch, out_height, out_width, depth]`. +// 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 ResourceSparseApplyAdadeltaUseLocking(value bool) ResourceSparseApplyAdadeltaAttr { + return func(m optionalAttr) { + m["use_locking"] = value + } +} + +// var: Should be from a Variable(). +// +// Arguments: +// +// accum: Should be from a Variable(). +// accum_update: : Should be from a Variable(). +// lr: Learning rate. Must be a scalar. +// rho: Decay factor. Must be a scalar. +// epsilon: Constant factor. Must be a scalar. +// grad: The gradient. +// indices: A vector of indices into the first dimension of var and accum. +// +// Returns the created operation. +func ResourceSparseApplyAdadelta(scope *Scope, var_ tf.Output, accum tf.Output, accum_update tf.Output, lr tf.Output, rho tf.Output, epsilon tf.Output, grad tf.Output, indices tf.Output, optional ...ResourceSparseApplyAdadeltaAttr) (o *tf.Operation) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "ResourceSparseApplyAdadelta", + Input: []tf.Input{ + var_, accum, accum_update, lr, rho, epsilon, grad, indices, + }, + Attrs: attrs, + } + return scope.AddOperation(opspec) +} + +// Computes rectified linear gradients for a Relu operation. +// +// Arguments: +// gradients: The backpropagated gradients to the corresponding Relu operation. +// features: The features passed as input to the corresponding Relu operation, OR +// the outputs of that operation (both work equivalently). +// +// Returns `gradients * (features > 0)`. +func ReluGrad(scope *Scope, gradients tf.Output, features tf.Output) (backprops tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "ReluGrad", + Input: []tf.Input{ + gradients, features, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Computes the gradient of morphological 2-D dilation with respect to the input. +// +// Arguments: +// input: 4-D with shape `[batch, in_height, in_width, depth]`. +// filter: 3-D with shape `[filter_height, filter_width, depth]`. +// out_backprop: 4-D with shape `[batch, out_height, out_width, depth]`. // strides: 1-D of length 4. The stride of the sliding window for each dimension of // the input tensor. Must be: `[1, stride_height, stride_width, 1]`. // rates: 1-D of length 4. The input stride for atrous morphological dilation. @@ -5284,6 +5599,51 @@ func MaxPoolGradWithArgmax(scope *Scope, input tf.Output, grad tf.Output, argmax return op.Output(0) } +// MutexV2Attr is an optional argument to MutexV2. +type MutexV2Attr func(optionalAttr) + +// MutexV2Container sets the optional container attribute to value. +// +// value: If non-empty, this variable is placed in the given container. +// Otherwise, a default container is used. +// If not specified, defaults to "" +func MutexV2Container(value string) MutexV2Attr { + return func(m optionalAttr) { + m["container"] = value + } +} + +// MutexV2SharedName sets the optional shared_name attribute to value. +// +// value: If non-empty, this variable is named in the given bucket +// with this shared_name. Otherwise, the node name is used instead. +// If not specified, defaults to "" +func MutexV2SharedName(value string) MutexV2Attr { + return func(m optionalAttr) { + m["shared_name"] = value + } +} + +// Creates a Mutex resource that can be locked by `MutexLock`. +// +// Returns The mutex resource. +func MutexV2(scope *Scope, optional ...MutexV2Attr) (resource tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "MutexV2", + + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // AvgPool3DAttr is an optional argument to AvgPool3D. type AvgPool3DAttr func(optionalAttr) @@ -5557,6 +5917,17 @@ func Sqrt(scope *Scope, x tf.Output) (y tf.Output) { return op.Output(0) } +// Conv3DBackpropFilterAttr is an optional argument to Conv3DBackpropFilter. +type Conv3DBackpropFilterAttr func(optionalAttr) + +// Conv3DBackpropFilterDilations sets the optional dilations attribute to value. +// If not specified, defaults to +func Conv3DBackpropFilterDilations(value []int64) Conv3DBackpropFilterAttr { + return func(m optionalAttr) { + m["dilations"] = value + } +} + // Computes the gradients of 3-D convolution with respect to the filter. // // DEPRECATED at GraphDef version 10: Use Conv3DBackpropFilterV2 @@ -5570,11 +5941,14 @@ func Sqrt(scope *Scope, x tf.Output) (y tf.Output) { // strides: 1-D tensor of length 5. The stride of the sliding window for each // dimension of `input`. Must have `strides[0] = strides[4] = 1`. // padding: The type of padding algorithm to use. -func Conv3DBackpropFilter(scope *Scope, input tf.Output, filter tf.Output, out_backprop tf.Output, strides []int64, padding string) (output tf.Output) { +func Conv3DBackpropFilter(scope *Scope, input tf.Output, filter tf.Output, out_backprop tf.Output, strides []int64, padding string, optional ...Conv3DBackpropFilterAttr) (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: "Conv3DBackpropFilter", Input: []tf.Input{ @@ -6567,6 +6941,85 @@ func OneHot(scope *Scope, indices tf.Output, depth tf.Output, on_value tf.Output return op.Output(0) } +// Transforms a vector of brain.Example protos (as strings) into typed tensors. +// +// Arguments: +// serialized: A vector containing a batch of binary serialized Example protos. +// names: A vector containing the names of the serialized protos. +// May contain, for example, table key (descriptive) names for the +// corresponding serialized protos. These are purely useful for debugging +// purposes, and the presence of values here has no effect on the output. +// May also be an empty vector if no names are available. +// If non-empty, this vector must be the same length as "serialized". +// sparse_keys: A list of Nsparse string Tensors (scalars). +// The keys expected in the Examples' features associated with sparse values. +// dense_keys: A list of Ndense string Tensors (scalars). +// The keys expected in the Examples' features associated with dense values. +// dense_defaults: A list of Ndense Tensors (some may be empty). +// dense_defaults[j] provides default values +// when the example's feature_map lacks dense_key[j]. If an empty Tensor is +// provided for dense_defaults[j], then the Feature dense_keys[j] is required. +// The input type is inferred from dense_defaults[j], even when it's empty. +// If dense_defaults[j] is not empty, and dense_shapes[j] is fully defined, +// then the shape of dense_defaults[j] must match that of dense_shapes[j]. +// If dense_shapes[j] has an undefined major dimension (variable strides dense +// feature), dense_defaults[j] must contain a single element: +// the padding element. +// sparse_types: A list of Nsparse types; the data types of data in each Feature +// given in sparse_keys. +// Currently the ParseExample supports DT_FLOAT (FloatList), +// DT_INT64 (Int64List), and DT_STRING (BytesList). +// dense_shapes: A list of Ndense shapes; the shapes of data in each Feature +// given in dense_keys. +// The number of elements in the Feature corresponding to dense_key[j] +// must always equal dense_shapes[j].NumEntries(). +// If dense_shapes[j] == (D0, D1, ..., DN) then the shape of output +// Tensor dense_values[j] will be (|serialized|, D0, D1, ..., DN): +// The dense outputs are just the inputs row-stacked by batch. +// This works for dense_shapes[j] = (-1, D1, ..., DN). In this case +// the shape of the output Tensor dense_values[j] will be +// (|serialized|, M, D1, .., DN), where M is the maximum number of blocks +// of elements of length D1 * .... * DN, across all minibatch entries +// in the input. Any minibatch entry with less than M blocks of elements of +// length D1 * ... * DN will be padded with the corresponding default_value +// scalar element along the second dimension. +func ParseExample(scope *Scope, serialized tf.Output, names tf.Output, sparse_keys []tf.Output, dense_keys []tf.Output, dense_defaults []tf.Output, sparse_types []tf.DataType, dense_shapes []tf.Shape) (sparse_indices []tf.Output, sparse_values []tf.Output, sparse_shapes []tf.Output, dense_values []tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"sparse_types": sparse_types, "dense_shapes": dense_shapes} + opspec := tf.OpSpec{ + Type: "ParseExample", + Input: []tf.Input{ + serialized, names, tf.OutputList(sparse_keys), tf.OutputList(dense_keys), tf.OutputList(dense_defaults), + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + if scope.Err() != nil { + return + } + var idx int + var err error + if sparse_indices, idx, err = makeOutputList(op, idx, "sparse_indices"); err != nil { + scope.UpdateErr("ParseExample", err) + return + } + if sparse_values, idx, err = makeOutputList(op, idx, "sparse_values"); err != nil { + scope.UpdateErr("ParseExample", err) + return + } + if sparse_shapes, idx, err = makeOutputList(op, idx, "sparse_shapes"); err != nil { + scope.UpdateErr("ParseExample", err) + return + } + if dense_values, idx, err = makeOutputList(op, idx, "dense_values"); err != nil { + scope.UpdateErr("ParseExample", err) + return + } + return sparse_indices, sparse_values, sparse_shapes, dense_values +} + // Real-valued fast Fourier transform. // // Computes the 1-dimensional discrete Fourier transform of a real-valued signal @@ -6959,6 +7412,44 @@ func PopulationCount(scope *Scope, x tf.Output) (y tf.Output) { return op.Output(0) } +// Broadcasts a tensor value to one or more other devices. +func CollectiveBcastSend(scope *Scope, input tf.Output, group_size int64, group_key int64, instance_key int64, shape tf.Shape) (data tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"group_size": group_size, "group_key": group_key, "instance_key": instance_key, "shape": shape} + opspec := tf.OpSpec{ + Type: "CollectiveBcastSend", + Input: []tf.Input{ + input, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Makes a copy of `x`. +// +// Arguments: +// x: The source tensor of type `T`. +// +// Returns y: A `Tensor` of type `T`. A copy of `x`. Guaranteed that `y` +// is not an alias of `x`. +func DeepCopy(scope *Scope, x tf.Output) (y tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "DeepCopy", + Input: []tf.Input{ + x, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Split a `SparseTensor` into `num_split` tensors along one dimension. // // If the `shape[split_dim]` is not an integer multiple of `num_split`. Slices @@ -7123,22 +7614,6 @@ func ResourceSparseApplyFtrlV2(scope *Scope, var_ tf.Output, accum tf.Output, li return scope.AddOperation(opspec) } -// Associates the given iterator with the given statistics aggregator. -// -// Returns the created operation. -func IteratorSetStatsAggregator(scope *Scope, iterator_handle tf.Output, stats_aggregator_handle tf.Output) (o *tf.Operation) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "IteratorSetStatsAggregator", - Input: []tf.Input{ - iterator_handle, stats_aggregator_handle, - }, - } - return scope.AddOperation(opspec) -} - // DataFormatVecPermuteAttr is an optional argument to DataFormatVecPermute. type DataFormatVecPermuteAttr func(optionalAttr) @@ -7232,6 +7707,46 @@ func Tan(scope *Scope, x tf.Output) (y tf.Output) { return op.Output(0) } +// Updates the tree ensemble by either adding a layer to the last tree being grown +// +// or by starting a new tree. +// +// Arguments: +// tree_ensemble_handle: Handle to the ensemble variable. +// feature_ids: Rank 1 tensor with ids for each feature. This is the real id of +// the feature that will be used in the split. +// node_ids: List of rank 1 tensors representing the nodes for which this feature +// has a split. +// gains: List of rank 1 tensors representing the gains for each of the feature's +// split. +// thresholds: List of rank 1 tensors representing the thesholds for each of the +// feature's split. +// left_node_contribs: List of rank 2 tensors with left leaf contribs for each of +// the feature's splits. Will be added to the previous node values to constitute +// the values of the left nodes. +// right_node_contribs: List of rank 2 tensors with right leaf contribs for each +// of the feature's splits. Will be added to the previous node values to constitute +// the values of the right nodes. +// max_depth: Max depth of the tree to build. +// learning_rate: shrinkage const for each new tree. +// pruning_mode: 0-No pruning, 1-Pre-pruning, 2-Post-pruning. +// +// Returns the created operation. +func BoostedTreesUpdateEnsemble(scope *Scope, tree_ensemble_handle tf.Output, feature_ids tf.Output, node_ids []tf.Output, gains []tf.Output, thresholds []tf.Output, left_node_contribs []tf.Output, right_node_contribs []tf.Output, max_depth tf.Output, learning_rate tf.Output, pruning_mode int64) (o *tf.Operation) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"pruning_mode": pruning_mode} + opspec := tf.OpSpec{ + Type: "BoostedTreesUpdateEnsemble", + Input: []tf.Input{ + tree_ensemble_handle, feature_ids, tf.OutputList(node_ids), tf.OutputList(gains), tf.OutputList(thresholds), tf.OutputList(left_node_contribs), tf.OutputList(right_node_contribs), max_depth, learning_rate, + }, + Attrs: attrs, + } + return scope.AddOperation(opspec) +} + // ResourceSparseApplyFtrlAttr is an optional argument to ResourceSparseApplyFtrl. type ResourceSparseApplyFtrlAttr func(optionalAttr) @@ -7309,7 +7824,7 @@ func IsInf(scope *Scope, x tf.Output) (y tf.Output) { // // N is the size of the segment being reduced. // -// Read @{$math_ops#segmentation$the section on segmentation} for an explanation of +// Read @{$math_ops#Segmentation$the section on segmentation} for an explanation of // segments. // // Arguments: @@ -7333,15 +7848,38 @@ func SparseSegmentSqrtN(scope *Scope, data tf.Output, indices tf.Output, segment return op.Output(0) } -// StatelessTruncatedNormalAttr is an optional argument to StatelessTruncatedNormal. -type StatelessTruncatedNormalAttr func(optionalAttr) - -// StatelessTruncatedNormalDtype sets the optional dtype attribute to value. +// Adds up a `SparseTensor` and a dense `Tensor`, producing a dense `Tensor`. // -// value: The type of the output. -// If not specified, defaults to DT_FLOAT -func StatelessTruncatedNormalDtype(value tf.DataType) StatelessTruncatedNormalAttr { - return func(m optionalAttr) { +// This Op does not require `a_indices` be sorted in standard lexicographic order. +// +// Arguments: +// a_indices: 2-D. The `indices` of the `SparseTensor`, with shape `[nnz, ndims]`. +// a_values: 1-D. The `values` of the `SparseTensor`, with shape `[nnz]`. +// a_shape: 1-D. The `shape` of the `SparseTensor`, with shape `[ndims]`. +// b: `ndims`-D Tensor. With shape `a_shape`. +func SparseTensorDenseAdd(scope *Scope, a_indices tf.Output, a_values tf.Output, a_shape tf.Output, b tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "SparseTensorDenseAdd", + Input: []tf.Input{ + a_indices, a_values, a_shape, b, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// StatelessTruncatedNormalAttr is an optional argument to StatelessTruncatedNormal. +type StatelessTruncatedNormalAttr func(optionalAttr) + +// StatelessTruncatedNormalDtype sets the optional dtype attribute to value. +// +// value: The type of the output. +// If not specified, defaults to DT_FLOAT +func StatelessTruncatedNormalDtype(value tf.DataType) StatelessTruncatedNormalAttr { + return func(m optionalAttr) { m["dtype"] = value } } @@ -7519,6 +8057,47 @@ func Complex(scope *Scope, real tf.Output, imag tf.Output, optional ...ComplexAt return op.Output(0) } +// Divides sparse updates into the variable referenced by `resource`. +// +// This operation computes +// +// # Scalar indices +// ref[indices, ...] /= updates[...] +// +// # Vector indices (for each i) +// ref[indices[i], ...] /= updates[i, ...] +// +// # High rank indices (for each i, ..., j) +// ref[indices[i, ..., j], ...] /= updates[i, ..., j, ...] +// +// Duplicate entries are handled correctly: if multiple `indices` reference +// the same location, their contributions multiply. +// +// Requires `updates.shape = indices.shape + ref.shape[1:]` or `updates.shape = []`. +// +//
+// +//
+// +// Arguments: +// resource: Should be from a `Variable` node. +// indices: A tensor of indices into the first dimension of `ref`. +// updates: A tensor of updated values to add to `ref`. +// +// Returns the created operation. +func ResourceScatterDiv(scope *Scope, resource tf.Output, indices tf.Output, updates tf.Output) (o *tf.Operation) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "ResourceScatterDiv", + Input: []tf.Input{ + resource, indices, updates, + }, + } + return scope.AddOperation(opspec) +} + // StatelessRandomNormalAttr is an optional argument to StatelessRandomNormal. type StatelessRandomNormalAttr func(optionalAttr) @@ -7562,6 +8141,47 @@ func StatelessRandomNormal(scope *Scope, shape tf.Output, seed tf.Output, option return op.Output(0) } +// Reduces sparse updates into the variable referenced by `resource` using the `min` operation. +// +// This operation computes +// +// # Scalar indices +// ref[indices, ...] = min(ref[indices, ...], updates[...]) +// +// # Vector indices (for each i) +// ref[indices[i], ...] = min(ref[indices[i], ...], updates[i, ...]) +// +// # High rank indices (for each i, ..., j) +// ref[indices[i, ..., j], ...] = min(ref[indices[i, ..., j], ...], updates[i, ..., j, ...]) +// +// Duplicate entries are handled correctly: if multiple `indices` reference +// the same location, their contributions are combined. +// +// Requires `updates.shape = indices.shape + ref.shape[1:]` or `updates.shape = []`. +// +//
+// +//
+// +// Arguments: +// resource: Should be from a `Variable` node. +// indices: A tensor of indices into the first dimension of `ref`. +// updates: A tensor of updated values to add to `ref`. +// +// Returns the created operation. +func ResourceScatterMin(scope *Scope, resource tf.Output, indices tf.Output, updates tf.Output) (o *tf.Operation) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "ResourceScatterMin", + Input: []tf.Input{ + resource, indices, updates, + }, + } + return scope.AddOperation(opspec) +} + // Reshapes a quantized tensor as per the Reshape op. // // ``` @@ -7771,9 +8391,8 @@ type ResizeBilinearAttr func(optionalAttr) // ResizeBilinearAlignCorners sets the optional align_corners attribute to value. // -// value: If true, rescale input by (new_height - 1) / (height - 1), which -// exactly aligns the 4 corners of images and resized images. If false, rescale -// by new_height / height. Treat similarly the width dimension. +// value: If true, the centers of the 4 corner pixels of the input and output tensors are +// aligned, preserving the values at the corner pixels. Defaults to false. // If not specified, defaults to false func ResizeBilinearAlignCorners(value bool) ResizeBilinearAttr { return func(m optionalAttr) { @@ -7826,6 +8445,26 @@ func Softsign(scope *Scope, features tf.Output) (activations tf.Output) { return op.Output(0) } +// Creates a TensorList which, when stacked, has the value of `tensor`. +// +// Each tensor in the result list corresponds to one row of the input tensor. +// +// tensor: The input tensor. +// output_handle: The list. +func TensorListFromTensor(scope *Scope, tensor tf.Output, element_shape tf.Output) (output_handle tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "TensorListFromTensor", + Input: []tf.Input{ + tensor, element_shape, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // GenerateVocabRemappingAttr is an optional argument to GenerateVocabRemapping. type GenerateVocabRemappingAttr func(optionalAttr) @@ -7933,6 +8572,30 @@ func ResourceScatterUpdate(scope *Scope, resource tf.Output, indices tf.Output, return scope.AddOperation(opspec) } +// Creates and returns an empty tensor list. +// +// All list elements must be tensors of dtype element_dtype and shape compatible +// with element_shape. +// +// handle: an empty tensor list. +// element_dtype: the type of elements in the list. +// element_shape: a shape compatible with that of elements in the list. +func EmptyTensorList(scope *Scope, element_shape tf.Output, element_dtype tf.DataType) (handle tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"element_dtype": element_dtype} + opspec := tf.OpSpec{ + Type: "EmptyTensorList", + Input: []tf.Input{ + element_shape, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // AvgPoolGradAttr is an optional argument to AvgPoolGrad. type AvgPoolGradAttr func(optionalAttr) @@ -8414,98 +9077,92 @@ func StringToHashBucketStrong(scope *Scope, input tf.Output, num_buckets int64, 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: +// RegexReplaceAttr is an optional argument to RegexReplace. +type RegexReplaceAttr func(optionalAttr) + +// RegexReplaceReplaceGlobal sets the optional replace_global attribute to value. // -// (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. +// value: If True, the replacement is global, otherwise the replacement +// is done only on the first match. +// If not specified, defaults to true +func RegexReplaceReplaceGlobal(value bool) RegexReplaceAttr { + return func(m optionalAttr) { + m["replace_global"] = value + } +} + +// Replaces the match of pattern in input with rewrite. // -// Hence, the `SparseTensor` result has exactly the same non-zero indices and -// shape. +// It follows the re2 syntax (https://github.com/google/re2/wiki/Syntax) // // 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. +// input: The text to be processed. +// pattern: The regular expression to match the input. +// rewrite: The rewrite to be applied to the matched expresion. // -// 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) { +// Returns The text after applying pattern and rewrite. +func RegexReplace(scope *Scope, input tf.Output, pattern tf.Output, rewrite tf.Output, optional ...RegexReplaceAttr) (output tf.Output) { if scope.Err() != nil { return } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } opspec := tf.OpSpec{ - Type: "SparseSoftmax", + Type: "RegexReplace", Input: []tf.Input{ - sp_indices, sp_values, sp_shape, + input, pattern, rewrite, }, + Attrs: attrs, } op := scope.AddOperation(opspec) return op.Output(0) } -// 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] -// ``` +// Computes numerical negative value element-wise. // -// See `dynamic_stitch` for an example on how to merge partitions back. +// I.e., \\(y = -x\\). +func Neg(scope *Scope, x tf.Output) (y tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "Neg", + Input: []tf.Input{ + x, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Execute a sub graph on a remote processor. // -//
-// -//
+// The graph specifications(such as graph itself, input tensors and output names) +// are stored as a serialized protocol buffer of RemoteFusedGraphExecuteInfo +// as serialized_remote_fused_graph_execute_info. +// The specifications will be passed to a dedicated registered +// remote fused graph executor. The executor will send the graph specifications +// to a remote processor and execute that graph. The execution results +// will be passed to consumer nodes as outputs of this node. // // Arguments: +// inputs: Arbitrary number of tensors with arbitrary data types // -// 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) { +// serialized_remote_fused_graph_execute_info: Serialized protocol buffer +// of RemoteFusedGraphExecuteInfo which contains graph specifications. +// +// Returns Arbitrary number of tensors with arbitrary data types +func RemoteFusedGraphExecute(scope *Scope, inputs []tf.Output, Toutputs []tf.DataType, serialized_remote_fused_graph_execute_info string) (outputs []tf.Output) { if scope.Err() != nil { return } - attrs := map[string]interface{}{"num_partitions": num_partitions} + attrs := map[string]interface{}{"Toutputs": Toutputs, "serialized_remote_fused_graph_execute_info": serialized_remote_fused_graph_execute_info} opspec := tf.OpSpec{ - Type: "DynamicPartition", + Type: "RemoteFusedGraphExecute", Input: []tf.Input{ - data, partitions, + tf.OutputList(inputs), }, Attrs: attrs, } @@ -8516,119 +9173,117 @@ func DynamicPartition(scope *Scope, data tf.Output, partitions tf.Output, num_pa var idx int var err error if outputs, idx, err = makeOutputList(op, idx, "outputs"); err != nil { - scope.UpdateErr("DynamicPartition", err) + scope.UpdateErr("RemoteFusedGraphExecute", err) return } return outputs } -// ResourceApplyAdagradAttr is an optional argument to ResourceApplyAdagrad. -type ResourceApplyAdagradAttr func(optionalAttr) +// MaxPool3DGradGradAttr is an optional argument to MaxPool3DGradGrad. +type MaxPool3DGradGradAttr func(optionalAttr) -// ResourceApplyAdagradUseLocking sets the optional use_locking attribute to value. +// MaxPool3DGradGradDataFormat sets the optional data_format 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 { +// value: The data format of the input and output data. With the +// default format "NDHWC", the data is stored in the order of: +// [batch, in_depth, in_height, in_width, in_channels]. +// Alternatively, the format could be "NCDHW", the data storage order is: +// [batch, in_channels, in_depth, in_height, in_width]. +// If not specified, defaults to "NDHWC" +func MaxPool3DGradGradDataFormat(value string) MaxPool3DGradGradAttr { return func(m optionalAttr) { - m["use_locking"] = value + m["data_format"] = value } } -// Update '*var' according to the adagrad scheme. -// -// accum += grad * grad -// var -= lr * grad * (1 / sqrt(accum)) +// Computes second-order gradients of the maxpooling function. // // Arguments: -// var_: Should be from a Variable(). -// accum: Should be from a Variable(). -// lr: Scaling factor. Must be a scalar. -// grad: The gradient. +// orig_input: The original input tensor. +// orig_output: The original output tensor. +// grad: Output backprop of shape `[batch, depth, rows, cols, channels]`. +// ksize: 1-D tensor of length 5. The size of the window for each dimension of +// the input tensor. Must have `ksize[0] = ksize[4] = 1`. +// strides: 1-D tensor of length 5. The stride of the sliding window for each +// dimension of `input`. Must have `strides[0] = strides[4] = 1`. +// padding: The type of padding algorithm to use. // -// 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) { +// Returns Gradients of gradients w.r.t. the input to `max_pool`. +func MaxPool3DGradGrad(scope *Scope, orig_input tf.Output, orig_output tf.Output, grad tf.Output, ksize []int64, strides []int64, padding string, optional ...MaxPool3DGradGradAttr) (output tf.Output) { if scope.Err() != nil { return } - attrs := map[string]interface{}{} + attrs := map[string]interface{}{"ksize": ksize, "strides": strides, "padding": padding} for _, a := range optional { a(attrs) } opspec := tf.OpSpec{ - Type: "ResourceApplyAdagrad", + Type: "MaxPool3DGradGrad", Input: []tf.Input{ - var_, accum, lr, grad, + orig_input, orig_output, grad, }, Attrs: attrs, } - return scope.AddOperation(opspec) -} - -// Return the shape of s0 op s1 with broadcast. -// -// Given `s0` and `s1`, tensors that represent shapes, compute `r0`, the -// broadcasted shape. `s0`, `s1` and `r0` are all integer vectors. -func BroadcastArgs(scope *Scope, s0 tf.Output, s1 tf.Output) (r0 tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "BroadcastArgs", - Input: []tf.Input{ - s0, s1, - }, - } op := scope.AddOperation(opspec) return op.Output(0) } -// DataFormatDimMapAttr is an optional argument to DataFormatDimMap. -type DataFormatDimMapAttr func(optionalAttr) +// Conv3DBackpropFilterV2Attr is an optional argument to Conv3DBackpropFilterV2. +type Conv3DBackpropFilterV2Attr func(optionalAttr) -// DataFormatDimMapSrcFormat sets the optional src_format attribute to value. +// Conv3DBackpropFilterV2DataFormat sets the optional data_format attribute to value. // -// value: source data format. -// If not specified, defaults to "NHWC" -func DataFormatDimMapSrcFormat(value string) DataFormatDimMapAttr { +// value: The data format of the input and output data. With the +// default format "NDHWC", the data is stored in the order of: +// [batch, in_depth, in_height, in_width, in_channels]. +// Alternatively, the format could be "NCDHW", the data storage order is: +// [batch, in_channels, in_depth, in_height, in_width]. +// If not specified, defaults to "NDHWC" +func Conv3DBackpropFilterV2DataFormat(value string) Conv3DBackpropFilterV2Attr { return func(m optionalAttr) { - m["src_format"] = value + m["data_format"] = value } } -// DataFormatDimMapDstFormat sets the optional dst_format attribute to value. +// Conv3DBackpropFilterV2Dilations sets the optional dilations attribute to value. // -// value: destination data format. -// If not specified, defaults to "NCHW" -func DataFormatDimMapDstFormat(value string) DataFormatDimMapAttr { +// value: 1-D tensor of length 5. The dilation factor for each dimension of +// `input`. If set to k > 1, there will be k-1 skipped cells between each +// filter element on that dimension. The dimension order is determined by the +// value of `data_format`, see above for details. Dilations in the batch and +// depth dimensions must be 1. +// If not specified, defaults to +func Conv3DBackpropFilterV2Dilations(value []int64) Conv3DBackpropFilterV2Attr { return func(m optionalAttr) { - m["dst_format"] = value + m["dilations"] = value } } -// Returns the dimension index in the destination data format given the one in -// -// the source data format. +// Computes the gradients of 3-D convolution with respect to the filter. // // Arguments: -// x: A Tensor with each element as a dimension index in source data format. -// Must be in the range [-4, 4). -// -// Returns A Tensor with each element as a dimension index in destination data format. -func DataFormatDimMap(scope *Scope, x tf.Output, optional ...DataFormatDimMapAttr) (y tf.Output) { +// input: Shape `[batch, depth, rows, cols, in_channels]`. +// filter_sizes: An integer vector representing the tensor shape of `filter`, +// where `filter` is a 5-D +// `[filter_depth, filter_height, filter_width, in_channels, out_channels]` +// tensor. +// out_backprop: Backprop signal of shape `[batch, out_depth, out_rows, out_cols, +// out_channels]`. +// strides: 1-D tensor of length 5. The stride of the sliding window for each +// dimension of `input`. Must have `strides[0] = strides[4] = 1`. +// padding: The type of padding algorithm to use. +func Conv3DBackpropFilterV2(scope *Scope, input tf.Output, filter_sizes tf.Output, out_backprop tf.Output, strides []int64, padding string, optional ...Conv3DBackpropFilterV2Attr) (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: "DataFormatDimMap", + Type: "Conv3DBackpropFilterV2", Input: []tf.Input{ - x, + input, filter_sizes, out_backprop, }, Attrs: attrs, } @@ -8636,38 +9291,38 @@ func DataFormatDimMap(scope *Scope, x tf.Output, optional ...DataFormatDimMapAtt return op.Output(0) } -// ResourceApplyPowerSignAttr is an optional argument to ResourceApplyPowerSign. -type ResourceApplyPowerSignAttr func(optionalAttr) +// FakeQuantWithMinMaxVarsAttr is an optional argument to FakeQuantWithMinMaxVars. +type FakeQuantWithMinMaxVarsAttr func(optionalAttr) -// ResourceApplyPowerSignUseLocking sets the optional use_locking attribute to value. -// -// value: If `True`, updating of the var and m tensors is -// protected by a lock; otherwise the behavior is undefined, but may exhibit less -// contention. +// 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 ResourceApplyPowerSignUseLocking(value bool) ResourceApplyPowerSignAttr { +func FakeQuantWithMinMaxVarsNarrowRange(value bool) FakeQuantWithMinMaxVarsAttr { return func(m optionalAttr) { - m["use_locking"] = value + m["narrow_range"] = value } } -// Update '*var' according to the AddSign update. +// Fake-quantize the 'inputs' tensor of type float via global float scalars `min` // -// m_t <- beta1 * m_{t-1} + (1 - beta1) * g -// update <- exp(logbase * sign_decay * sign(g) * sign(m_t)) * g -// variable <- variable - lr_t * update +// and `max` to 'outputs' tensor of same shape as `inputs`. // -// Arguments: -// var_: Should be from a Variable(). -// m: Should be from a Variable(). -// lr: Scaling factor. Must be a scalar. -// logbase: Must be a scalar. -// sign_decay: Must be a scalar. -// beta: 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 16, inclusive. // -// Returns the created operation. -func ResourceApplyPowerSign(scope *Scope, var_ tf.Output, m tf.Output, lr tf.Output, logbase tf.Output, sign_decay tf.Output, beta tf.Output, grad tf.Output, optional ...ResourceApplyPowerSignAttr) (o *tf.Operation) { +// 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 } @@ -8676,225 +9331,152 @@ func ResourceApplyPowerSign(scope *Scope, var_ tf.Output, m tf.Output, lr tf.Out a(attrs) } opspec := tf.OpSpec{ - Type: "ResourceApplyPowerSign", + Type: "FakeQuantWithMinMaxVars", Input: []tf.Input{ - var_, m, lr, logbase, sign_decay, beta, grad, + inputs, min, max, }, Attrs: attrs, } - return scope.AddOperation(opspec) + op := scope.AddOperation(opspec) + return op.Output(0) } -// Computes the mean along segments of a tensor. +// Applies softmax to a batched N-D `SparseTensor`. // -// Read @{$math_ops#segmentation$the section on segmentation} for an explanation of -// segments. +// The inputs represent an N-D SparseTensor with logical shape `[..., B, C]` +// (where `N >= 2`), and with indices sorted in the canonical lexicographic order. // -// Computes a tensor such that -// \\(output_i = \frac{\sum_j data_j}{N}\\) where `mean` is -// over `j` such that `segment_ids[j] == i` and `N` is the total number of -// values summed. +// 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: // -// If the mean is empty for a given segment ID `i`, `output[i] = 0`. +// (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. // -// segment_ids: A 1-D tensor whose rank is equal to the rank of `data`'s -// first dimension. Values should be sorted and can be repeated. -// -// Returns Has same shape as data, except for dimension 0 which -// has size `k`, the number of segments. -func SegmentMean(scope *Scope, data tf.Output, segment_ids tf.Output) (output tf.Output) { +// 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: "SegmentMean", + Type: "SparseSoftmax", Input: []tf.Input{ - data, segment_ids, + sp_indices, sp_values, sp_shape, }, } op := scope.AddOperation(opspec) return op.Output(0) } -// ResourceSparseApplyCenteredRMSPropAttr is an optional argument to ResourceSparseApplyCenteredRMSProp. -type ResourceSparseApplyCenteredRMSPropAttr func(optionalAttr) - -// ResourceSparseApplyCenteredRMSPropUseLocking sets the optional use_locking attribute to value. -// -// value: If `True`, updating of the var, mg, ms, and mom tensors is -// protected by a lock; otherwise the behavior is undefined, but may exhibit less -// contention. -// If not specified, defaults to false -func ResourceSparseApplyCenteredRMSPropUseLocking(value bool) ResourceSparseApplyCenteredRMSPropAttr { - return func(m optionalAttr) { - m["use_locking"] = value - } -} - -// Update '*var' according to the centered RMSProp algorithm. +// Partitions `data` into `num_partitions` tensors using indices from `partitions`. // -// The centered RMSProp algorithm uses an estimate of the centered second moment -// (i.e., the variance) for normalization, as opposed to regular RMSProp, which -// uses the (uncentered) second moment. This often helps with training, but is -// slightly more expensive in terms of computation and memory. +// 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, // -// Note that in dense implementation of this algorithm, mg, ms, and mom will -// update even if the grad is zero, but in this sparse implementation, mg, ms, -// and mom will not update in iterations during which the grad is zero. +// ```python +// outputs[i].shape = [sum(partitions == i)] + data.shape[partitions.ndim:] // -// mean_square = decay * mean_square + (1-decay) * gradient ** 2 -// mean_grad = decay * mean_grad + (1-decay) * gradient -// Delta = learning_rate * gradient / sqrt(mean_square + epsilon - mean_grad ** 2) +// outputs[i] = pack([data[js, ...] for js if partitions[js] == i]) +// ``` // -// ms <- rho * ms_{t-1} + (1-rho) * grad * grad -// mom <- momentum * mom_{t-1} + lr * grad / sqrt(ms + epsilon) -// var <- var - mom +// `data.shape` must start with `partitions.shape`. // -// Arguments: -// var_: Should be from a Variable(). -// mg: Should be from a Variable(). -// ms: Should be from a Variable(). -// mom: Should be from a Variable(). -// lr: Scaling factor. Must be a scalar. -// rho: Decay rate. Must be a scalar. +// For example: // -// epsilon: Ridge term. Must be a scalar. -// grad: The gradient. -// indices: A vector of indices into the first dimension of var, ms and mom. +// ```python +// # Scalar partitions. +// partitions = 1 +// num_partitions = 2 +// data = [10, 20] +// outputs[0] = [] # Empty with shape [0, 2] +// outputs[1] = [[10, 20]] // -// Returns the created operation. -func ResourceSparseApplyCenteredRMSProp(scope *Scope, var_ tf.Output, mg tf.Output, ms tf.Output, mom tf.Output, lr tf.Output, rho tf.Output, momentum tf.Output, epsilon tf.Output, grad tf.Output, indices tf.Output, optional ...ResourceSparseApplyCenteredRMSPropAttr) (o *tf.Operation) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "ResourceSparseApplyCenteredRMSProp", - Input: []tf.Input{ - var_, mg, ms, mom, lr, rho, momentum, epsilon, grad, indices, - }, - Attrs: attrs, - } - return scope.AddOperation(opspec) -} - -// Creates a dataset that batches `batch_size` elements from `input_dataset`. +// # 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] +// ``` // -// Arguments: +// See `dynamic_stitch` for an example on how to merge partitions back. // -// batch_size: A scalar representing the number of elements to accumulate in a -// batch. +//
+// +//
// +// Arguments: // -func BatchDataset(scope *Scope, input_dataset tf.Output, batch_size tf.Output, output_types []tf.DataType, output_shapes []tf.Shape) (handle 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{}{"output_types": output_types, "output_shapes": output_shapes} + attrs := map[string]interface{}{"num_partitions": num_partitions} opspec := tf.OpSpec{ - Type: "BatchDataset", + Type: "DynamicPartition", Input: []tf.Input{ - input_dataset, batch_size, + data, partitions, }, Attrs: attrs, } op := scope.AddOperation(opspec) - return op.Output(0) -} - -// Inverse fast Fourier transform. -// -// Computes the inverse 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 inverse 1D Fourier transform. -// -// @compatibility(numpy) -// Equivalent to np.fft.ifft -// @end_compatibility -func IFFT(scope *Scope, input tf.Output) (output tf.Output) { if scope.Err() != nil { return } - opspec := tf.OpSpec{ - Type: "IFFT", - Input: []tf.Input{ - input, - }, - } - 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 { + var idx int + var err error + if outputs, idx, err = makeOutputList(op, idx, "outputs"); err != nil { + scope.UpdateErr("DynamicPartition", err) return } - opspec := tf.OpSpec{ - Type: "LinSpace", - Input: []tf.Input{ - start, stop, num, - }, - } - op := scope.AddOperation(opspec) - return op.Output(0) + 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 } @@ -8903,75 +9485,66 @@ 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, } return scope.AddOperation(opspec) } -// LRNAttr is an optional argument to LRN. -type LRNAttr func(optionalAttr) - -// LRNDepthRadius sets the optional depth_radius attribute to value. +// Return the shape of s0 op s1 with broadcast. // -// value: 0-D. Half-width of the 1-D normalization window. -// If not specified, defaults to 5 -func LRNDepthRadius(value int64) LRNAttr { - return func(m optionalAttr) { - m["depth_radius"] = value +// Given `s0` and `s1`, tensors that represent shapes, compute `r0`, the +// broadcasted shape. `s0`, `s1` and `r0` are all integer vectors. +func BroadcastArgs(scope *Scope, s0 tf.Output, s1 tf.Output) (r0 tf.Output) { + if scope.Err() != nil { + return } -} - -// LRNBias sets the optional bias attribute to value. -// -// value: An offset (usually positive to avoid dividing by 0). -// If not specified, defaults to 1 -func LRNBias(value float32) LRNAttr { - return func(m optionalAttr) { - m["bias"] = value + opspec := tf.OpSpec{ + Type: "BroadcastArgs", + Input: []tf.Input{ + s0, s1, + }, } + op := scope.AddOperation(opspec) + return op.Output(0) } -// LRNAlpha sets the optional alpha attribute to value. +// DataFormatDimMapAttr is an optional argument to DataFormatDimMap. +type DataFormatDimMapAttr func(optionalAttr) + +// DataFormatDimMapSrcFormat sets the optional src_format attribute to value. // -// value: A scale factor, usually positive. -// If not specified, defaults to 1 -func LRNAlpha(value float32) LRNAttr { +// value: source data format. +// If not specified, defaults to "NHWC" +func DataFormatDimMapSrcFormat(value string) DataFormatDimMapAttr { return func(m optionalAttr) { - m["alpha"] = value + m["src_format"] = value } } -// LRNBeta sets the optional beta attribute to value. +// DataFormatDimMapDstFormat sets the optional dst_format attribute to value. // -// value: An exponent. -// If not specified, defaults to 0.5 -func LRNBeta(value float32) LRNAttr { +// value: destination data format. +// If not specified, defaults to "NCHW" +func DataFormatDimMapDstFormat(value string) DataFormatDimMapAttr { return func(m optionalAttr) { - m["beta"] = value + m["dst_format"] = value } } -// Local Response Normalization. -// -// The 4-D `input` tensor is treated as a 3-D array of 1-D vectors (along the last -// dimension), and each vector is normalized independently. Within a given vector, -// each component is divided by the weighted, squared sum of inputs within -// `depth_radius`. In detail, -// -// sqr_sum[a, b, c, d] = -// sum(input[a, b, c, d - depth_radius : d + depth_radius + 1] ** 2) -// output = input / (bias + alpha * sqr_sum) ** beta +// Returns the dimension index in the destination data format given the one in // -// For details, see [Krizhevsky et al., ImageNet classification with deep -// convolutional neural networks (NIPS 2012)](http://papers.nips.cc/paper/4824-imagenet-classification-with-deep-convolutional-neural-networks). +// the source data format. // // Arguments: -// input: 4-D. -func LRN(scope *Scope, input tf.Output, optional ...LRNAttr) (output tf.Output) { +// x: A Tensor with each element as a dimension index in source data format. +// Must be in the range [-4, 4). +// +// Returns A Tensor with each element as a dimension index in destination data format. +func DataFormatDimMap(scope *Scope, x tf.Output, optional ...DataFormatDimMapAttr) (y tf.Output) { if scope.Err() != nil { return } @@ -8980,26 +9553,9 @@ func LRN(scope *Scope, input tf.Output, optional ...LRNAttr) (output tf.Output) a(attrs) } opspec := tf.OpSpec{ - Type: "LRN", - Input: []tf.Input{ - input, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// Creates a dataset that zips together `input_datasets`. -func ZipDataset(scope *Scope, input_datasets []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: "ZipDataset", + Type: "DataFormatDimMap", Input: []tf.Input{ - tf.OutputList(input_datasets), + x, }, Attrs: attrs, } @@ -9007,36 +9563,38 @@ func ZipDataset(scope *Scope, input_datasets []tf.Output, output_types []tf.Data return op.Output(0) } -// ResourceSparseApplyAdagradAttr is an optional argument to ResourceSparseApplyAdagrad. -type ResourceSparseApplyAdagradAttr func(optionalAttr) +// ResourceApplyPowerSignAttr is an optional argument to ResourceApplyPowerSign. +type ResourceApplyPowerSignAttr func(optionalAttr) -// ResourceSparseApplyAdagradUseLocking sets the optional use_locking attribute to value. +// ResourceApplyPowerSignUseLocking 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 +// value: If `True`, updating of the var and m tensors is +// protected by a lock; otherwise the behavior is undefined, but may exhibit less // contention. // If not specified, defaults to false -func ResourceSparseApplyAdagradUseLocking(value bool) ResourceSparseApplyAdagradAttr { +func ResourceApplyPowerSignUseLocking(value bool) ResourceApplyPowerSignAttr { return func(m optionalAttr) { m["use_locking"] = value } } -// Update relevant entries in '*var' and '*accum' according to the adagrad scheme. +// Update '*var' according to the AddSign update. // -// That is for rows we have grad for, we update var and accum as follows: -// accum += grad * grad -// var -= lr * grad * (1 / sqrt(accum)) +// m_t <- beta1 * m_{t-1} + (1 - beta1) * g +// update <- exp(logbase * sign_decay * sign(g) * sign(m_t)) * g +// variable <- variable - lr_t * update // // Arguments: // var_: Should be from a Variable(). -// accum: Should be from a Variable(). -// lr: Learning rate. Must be a scalar. -// grad: The gradient. -// indices: A vector of indices into the first dimension of var and accum. -// +// m: Should be from a Variable(). +// lr: Scaling factor. Must be a scalar. +// logbase: Must be a scalar. +// sign_decay: Must be a scalar. +// beta: Must be a scalar. +// grad: The gradient. +// // Returns the created operation. -func ResourceSparseApplyAdagrad(scope *Scope, var_ tf.Output, accum tf.Output, lr tf.Output, grad tf.Output, indices tf.Output, optional ...ResourceSparseApplyAdagradAttr) (o *tf.Operation) { +func ResourceApplyPowerSign(scope *Scope, var_ tf.Output, m tf.Output, lr tf.Output, logbase tf.Output, sign_decay tf.Output, beta tf.Output, grad tf.Output, optional ...ResourceApplyPowerSignAttr) (o *tf.Operation) { if scope.Err() != nil { return } @@ -9045,87 +9603,161 @@ func ResourceSparseApplyAdagrad(scope *Scope, var_ tf.Output, accum tf.Output, l a(attrs) } opspec := tf.OpSpec{ - Type: "ResourceSparseApplyAdagrad", + Type: "ResourceApplyPowerSign", Input: []tf.Input{ - var_, accum, lr, grad, indices, + var_, m, lr, logbase, sign_decay, beta, grad, }, Attrs: attrs, } return scope.AddOperation(opspec) } -// 2D real-valued fast Fourier transform. +// Locks a mutex resource. The output is the lock. So long as the lock tensor // -// Computes the 2-dimensional discrete Fourier transform of a real-valued signal -// over the inner-most 2 dimensions of `input`. +// is alive, any other request to use `MutexLock` with this mutex will wait. // -// Since the DFT of a real signal is Hermitian-symmetric, `RFFT2D` only returns the -// `fft_length / 2 + 1` unique components of the FFT for the inner-most dimension -// of `output`: the zero-frequency term, followed by the `fft_length / 2` -// positive-frequency terms. +// This is particularly useful for creating a critical section when used in +// conjunction with `MutexLockIdentity`: // -// Along each axis `RFFT2D` is computed on, if `fft_length` is smaller than the -// corresponding dimension of `input`, the dimension is cropped. If it is larger, -// the dimension is padded with zeros. +// ```python +// +// mutex = mutex_v2( +// shared_name=handle_name, container=container, name=name) +// +// def execute_in_critical_section(fn, *args, **kwargs): +// lock = gen_resource_variable_ops.mutex_lock(mutex) +// +// with ops.control_dependencies([lock]): +// r = fn(*args, **kwargs) +// +// with ops.control_dependencies(nest.flatten(r)): +// with ops.colocate_with(mutex): +// ensure_lock_exists = mutex_lock_identity(lock) +// +// # Make sure that if any element of r is accessed, all of +// # them are executed together. +// r = nest.map_structure(tf.identity, r) +// +// with ops.control_dependencies([ensure_lock_exists]): +// return nest.map_structure(tf.identity, r) +// ``` +// +// While `fn` is running in the critical section, no other functions which wish to +// use this critical section may run. +// +// Often the use case is that two executions of the same graph, in parallel, +// wish to run `fn`; and we wish to ensure that only one of them executes +// at a time. This is especially important if `fn` modifies one or more +// variables at a time. +// +// It is also useful if two separate functions must share a resource, but we +// wish to ensure the usage is exclusive. // // Arguments: -// input: A float32 tensor. -// fft_length: An int32 tensor of shape [2]. The FFT length for each dimension. +// mutex: The mutex resource to lock. // -// Returns A complex64 tensor of the same rank as `input`. The inner-most 2 -// dimensions of `input` are replaced with their 2D Fourier transform. The -// inner-most dimension contains `fft_length / 2 + 1` unique frequency -// components. +// Returns A tensor that keeps a shared pointer to a lock on the mutex; +// when the Tensor is destroyed, the use count on the shared pointer is decreased +// by 1. When it reaches 0, the lock is released. +func MutexLock(scope *Scope, mutex tf.Output) (mutex_lock tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "MutexLock", + Input: []tf.Input{ + mutex, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Computes the mean along segments of a tensor. // -// @compatibility(numpy) -// Equivalent to np.fft.rfft2 -// @end_compatibility -func RFFT2D(scope *Scope, input tf.Output, fft_length tf.Output) (output tf.Output) { +// Read @{$math_ops#Segmentation$the section on segmentation} for an explanation of +// segments. +// +// Computes a tensor such that +// \\(output_i = \frac{\sum_j data_j}{N}\\) where `mean` is +// over `j` such that `segment_ids[j] == i` and `N` is the total number of +// values summed. +// +// If the mean is empty for a given segment ID `i`, `output[i] = 0`. +// +//
+// +//
+// +// Arguments: +// +// segment_ids: A 1-D tensor whose rank is equal to the rank of `data`'s +// first dimension. Values should be sorted and can be repeated. +// +// Returns Has same shape as data, except for dimension 0 which +// has size `k`, the number of segments. +func SegmentMean(scope *Scope, data tf.Output, segment_ids tf.Output) (output tf.Output) { if scope.Err() != nil { return } opspec := tf.OpSpec{ - Type: "RFFT2D", + Type: "SegmentMean", Input: []tf.Input{ - input, fft_length, + data, segment_ids, }, } op := scope.AddOperation(opspec) return op.Output(0) } -// ResizeAreaAttr is an optional argument to ResizeArea. -type ResizeAreaAttr func(optionalAttr) +// ResourceSparseApplyCenteredRMSPropAttr is an optional argument to ResourceSparseApplyCenteredRMSProp. +type ResourceSparseApplyCenteredRMSPropAttr func(optionalAttr) -// ResizeAreaAlignCorners sets the optional align_corners attribute to value. +// ResourceSparseApplyCenteredRMSPropUseLocking sets the optional use_locking attribute to value. // -// value: If true, rescale input by (new_height - 1) / (height - 1), which -// exactly aligns the 4 corners of images and resized images. If false, rescale -// by new_height / height. Treat similarly the width dimension. +// value: If `True`, updating of the var, mg, ms, and mom tensors is +// protected by a lock; otherwise the behavior is undefined, but may exhibit less +// contention. // If not specified, defaults to false -func ResizeAreaAlignCorners(value bool) ResizeAreaAttr { +func ResourceSparseApplyCenteredRMSPropUseLocking(value bool) ResourceSparseApplyCenteredRMSPropAttr { return func(m optionalAttr) { - m["align_corners"] = value + m["use_locking"] = value } } -// Resize `images` to `size` using area interpolation. +// Update '*var' according to the centered RMSProp algorithm. // -// Input images can be of different types but output images are always float. +// The centered RMSProp algorithm uses an estimate of the centered second moment +// (i.e., the variance) for normalization, as opposed to regular RMSProp, which +// uses the (uncentered) second moment. This often helps with training, but is +// slightly more expensive in terms of computation and memory. // -// Each output pixel is computed by first transforming the pixel's footprint into -// the input tensor and then averaging the pixels that intersect the footprint. An -// input pixel's contribution to the average is weighted by the fraction of its -// area that intersects the footprint. This is the same as OpenCV's INTER_AREA. +// Note that in dense implementation of this algorithm, mg, ms, and mom will +// update even if the grad is zero, but in this sparse implementation, mg, ms, +// and mom will not update in iterations during which the grad is zero. +// +// mean_square = decay * mean_square + (1-decay) * gradient ** 2 +// mean_grad = decay * mean_grad + (1-decay) * gradient +// Delta = learning_rate * gradient / sqrt(mean_square + epsilon - mean_grad ** 2) +// +// ms <- rho * ms_{t-1} + (1-rho) * grad * grad +// mom <- momentum * mom_{t-1} + lr * grad / sqrt(ms + epsilon) +// var <- var - mom // // Arguments: -// images: 4-D with shape `[batch, height, width, channels]`. -// size: = A 1-D int32 Tensor of 2 elements: `new_height, new_width`. The -// new size for the images. +// var_: Should be from a Variable(). +// mg: Should be from a Variable(). +// ms: Should be from a Variable(). +// mom: Should be from a Variable(). +// lr: Scaling factor. Must be a scalar. +// rho: Decay rate. Must be a scalar. // -// Returns 4-D with shape -// `[batch, new_height, new_width, channels]`. -func ResizeArea(scope *Scope, images tf.Output, size tf.Output, optional ...ResizeAreaAttr) (resized_images tf.Output) { +// epsilon: Ridge term. Must be a scalar. +// grad: The gradient. +// indices: A vector of indices into the first dimension of var, ms and mom. +// +// Returns the created operation. +func ResourceSparseApplyCenteredRMSProp(scope *Scope, var_ tf.Output, mg tf.Output, ms tf.Output, mom tf.Output, lr tf.Output, rho tf.Output, momentum tf.Output, epsilon tf.Output, grad tf.Output, indices tf.Output, optional ...ResourceSparseApplyCenteredRMSPropAttr) (o *tf.Operation) { if scope.Err() != nil { return } @@ -9134,9 +9766,32 @@ func ResizeArea(scope *Scope, images tf.Output, size tf.Output, optional ...Resi a(attrs) } opspec := tf.OpSpec{ - Type: "ResizeArea", + Type: "ResourceSparseApplyCenteredRMSProp", Input: []tf.Input{ - images, size, + var_, mg, ms, mom, lr, rho, momentum, epsilon, grad, indices, + }, + Attrs: attrs, + } + return scope.AddOperation(opspec) +} + +// Creates a dataset that batches `batch_size` elements from `input_dataset`. +// +// Arguments: +// +// batch_size: A scalar representing the number of elements to accumulate in a +// batch. +// +// +func BatchDataset(scope *Scope, input_dataset tf.Output, batch_size 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: "BatchDataset", + Input: []tf.Input{ + input_dataset, batch_size, }, Attrs: attrs, } @@ -9144,91 +9799,90 @@ func ResizeArea(scope *Scope, images tf.Output, size tf.Output, optional ...Resi return op.Output(0) } -// Pads a tensor with zeros. -// -// This operation pads a `input` with zeros according to the `paddings` you -// specify. `paddings` is an integer tensor with shape `[Dn, 2]`, where n is the -// rank of `input`. For each dimension D of `input`, `paddings[D, 0]` indicates -// how many zeros to add before the contents of `input` in that dimension, and -// `paddings[D, 1]` indicates how many zeros to add after the contents of `input` -// in that dimension. +// Inverse fast Fourier transform. // -// The padded size of each dimension D of the output is: +// Computes the inverse 1-dimensional discrete Fourier transform over the +// inner-most dimension of `input`. // -// `paddings(D, 0) + input.dim_size(D) + paddings(D, 1)` +// Arguments: +// input: A complex64 tensor. // -// For example: +// Returns A complex64 tensor of the same shape as `input`. The inner-most +// dimension of `input` is replaced with its inverse 1D Fourier transform. // -// ``` -// # 't' is [[1, 1], [2, 2]] -// # 'paddings' is [[1, 1], [2, 2]] -// # rank of 't' is 2 -// pad(t, paddings) ==> [[0, 0, 0, 0, 0, 0] -// [0, 0, 1, 1, 0, 0] -// [0, 0, 2, 2, 0, 0] -// [0, 0, 0, 0, 0, 0]] -// ``` -func Pad(scope *Scope, input tf.Output, paddings tf.Output) (output tf.Output) { +// @compatibility(numpy) +// Equivalent to np.fft.ifft +// @end_compatibility +func IFFT(scope *Scope, input tf.Output) (output tf.Output) { if scope.Err() != nil { return } opspec := tf.OpSpec{ - Type: "Pad", + Type: "IFFT", Input: []tf.Input{ - input, paddings, + input, }, } op := scope.AddOperation(opspec) return op.Output(0) } -// Checks whether a resource handle-based variable has been initialized. +// 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: -// resource: the input resource handle. +// start: First entry in the range. +// stop: Last entry in the range. +// num: Number of values to generate. // -// Returns a scalar boolean which is true if the variable has been -// initialized. -func VarIsInitializedOp(scope *Scope, resource tf.Output) (is_initialized tf.Output) { +// 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: "VarIsInitializedOp", + Type: "LinSpace", Input: []tf.Input{ - resource, + start, stop, num, }, } op := scope.AddOperation(opspec) return op.Output(0) } -// StatelessRandomUniformAttr is an optional argument to StatelessRandomUniform. -type StatelessRandomUniformAttr func(optionalAttr) +// DestroyResourceOpAttr is an optional argument to DestroyResourceOp. +type DestroyResourceOpAttr func(optionalAttr) -// StatelessRandomUniformDtype sets the optional dtype attribute to value. +// DestroyResourceOpIgnoreLookupError sets the optional ignore_lookup_error attribute to value. // -// value: The type of the output. -// If not specified, defaults to DT_FLOAT -func StatelessRandomUniformDtype(value tf.DataType) StatelessRandomUniformAttr { +// 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["dtype"] = value + m["ignore_lookup_error"] = value } } -// Outputs deterministic pseudorandom random values from a uniform distribution. -// -// The generated values follow a uniform distribution in the range `[0, 1)`. The -// lower bound 0 is included in the range, while the upper bound 1 is excluded. +// Deletes the resource specified by the handle. // -// The outputs are a deterministic function of `shape` and `seed`. +// All subsequent operations using the resource will result in a NotFound +// error status. // // Arguments: -// shape: The shape of the output tensor. -// seed: 2 seeds (shape [2]). +// resource: handle to the resource to delete. // -// Returns Random values with specified shape. -func StatelessRandomUniform(scope *Scope, shape tf.Output, seed tf.Output, optional ...StatelessRandomUniformAttr) (output tf.Output) { +// Returns the created operation. +func DestroyResourceOp(scope *Scope, resource tf.Output, optional ...DestroyResourceOpAttr) (o *tf.Operation) { if scope.Err() != nil { return } @@ -9237,79 +9891,140 @@ func StatelessRandomUniform(scope *Scope, shape tf.Output, seed tf.Output, optio a(attrs) } opspec := tf.OpSpec{ - Type: "StatelessRandomUniform", + Type: "DestroyResourceOp", Input: []tf.Input{ - shape, seed, + resource, }, Attrs: attrs, } - op := scope.AddOperation(opspec) - return op.Output(0) + return scope.AddOperation(opspec) } -// Makes its input available to the next iteration. +// LRNAttr is an optional argument to LRN. +type LRNAttr func(optionalAttr) + +// LRNDepthRadius sets the optional depth_radius attribute to value. // -// Arguments: -// data: The tensor to be made available to the next iteration. +// value: 0-D. Half-width of the 1-D normalization window. +// If not specified, defaults to 5 +func LRNDepthRadius(value int64) LRNAttr { + return func(m optionalAttr) { + m["depth_radius"] = value + } +} + +// LRNBias sets the optional bias attribute to value. // -// Returns The same tensor as `data`. -func NextIteration(scope *Scope, data tf.Output) (output tf.Output) { +// value: An offset (usually positive to avoid dividing by 0). +// If not specified, defaults to 1 +func LRNBias(value float32) LRNAttr { + return func(m optionalAttr) { + m["bias"] = value + } +} + +// LRNAlpha sets the optional alpha attribute to value. +// +// value: A scale factor, usually positive. +// If not specified, defaults to 1 +func LRNAlpha(value float32) LRNAttr { + return func(m optionalAttr) { + m["alpha"] = value + } +} + +// LRNBeta sets the optional beta attribute to value. +// +// value: An exponent. +// If not specified, defaults to 0.5 +func LRNBeta(value float32) LRNAttr { + return func(m optionalAttr) { + m["beta"] = value + } +} + +// Local Response Normalization. +// +// The 4-D `input` tensor is treated as a 3-D array of 1-D vectors (along the last +// dimension), and each vector is normalized independently. Within a given vector, +// each component is divided by the weighted, squared sum of inputs within +// `depth_radius`. In detail, +// +// sqr_sum[a, b, c, d] = +// sum(input[a, b, c, d - depth_radius : d + depth_radius + 1] ** 2) +// output = input / (bias + alpha * sqr_sum) ** beta +// +// For details, see [Krizhevsky et al., ImageNet classification with deep +// convolutional neural networks (NIPS 2012)](http://papers.nips.cc/paper/4824-imagenet-classification-with-deep-convolutional-neural-networks). +// +// Arguments: +// input: 4-D. +func LRN(scope *Scope, input tf.Output, optional ...LRNAttr) (output tf.Output) { if scope.Err() != nil { return } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } opspec := tf.OpSpec{ - Type: "NextIteration", + Type: "LRN", Input: []tf.Input{ - data, + input, }, + Attrs: attrs, } op := scope.AddOperation(opspec) return op.Output(0) } -// Output a fact about factorials. -func Fact(scope *Scope) (fact tf.Output) { +// Creates a dataset that zips together `input_datasets`. +func ZipDataset(scope *Scope, input_datasets []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: "Fact", + Type: "ZipDataset", + Input: []tf.Input{ + tf.OutputList(input_datasets), + }, + Attrs: attrs, } op := scope.AddOperation(opspec) return op.Output(0) } -// AngleAttr is an optional argument to Angle. -type AngleAttr func(optionalAttr) +// ResourceSparseApplyAdagradAttr is an optional argument to ResourceSparseApplyAdagrad. +type ResourceSparseApplyAdagradAttr func(optionalAttr) -// AngleTout sets the optional Tout attribute to value. -// If not specified, defaults to DT_FLOAT -func AngleTout(value tf.DataType) AngleAttr { +// ResourceSparseApplyAdagradUseLocking 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 ResourceSparseApplyAdagradUseLocking(value bool) ResourceSparseApplyAdagradAttr { return func(m optionalAttr) { - m["Tout"] = value + m["use_locking"] = value } } -// Returns the argument of a complex number. -// -// Given a tensor `input` of complex numbers, this operation returns a tensor of -// type `float` that is the argument of each element in `input`. All elements in -// `input` must be complex numbers of the form \\(a + bj\\), where *a* -// is the real part and *b* is the imaginary part. -// -// The argument returned by this operation is of the form \\(atan2(b, a)\\). +// Update relevant entries in '*var' and '*accum' according to the adagrad scheme. // -// For example: +// That is for rows we have grad for, we update var and accum as follows: +// accum += grad * grad +// var -= lr * grad * (1 / sqrt(accum)) // -// ``` -// # tensor 'input' is [-2.25 + 4.75j, 3.25 + 5.75j] -// tf.angle(input) ==> [2.0132, 1.056] -// ``` +// Arguments: +// var_: Should be from a Variable(). +// accum: Should be from a Variable(). +// lr: Learning rate. Must be a scalar. +// grad: The gradient. +// indices: A vector of indices into the first dimension of var and accum. // -// @compatibility(numpy) -// Equivalent to np.angle. -// @end_compatibility -func Angle(scope *Scope, input tf.Output, optional ...AngleAttr) (output tf.Output) { +// Returns the created operation. +func ResourceSparseApplyAdagrad(scope *Scope, var_ tf.Output, accum tf.Output, lr tf.Output, grad tf.Output, indices tf.Output, optional ...ResourceSparseApplyAdagradAttr) (o *tf.Operation) { if scope.Err() != nil { return } @@ -9318,74 +10033,351 @@ func Angle(scope *Scope, input tf.Output, optional ...AngleAttr) (output tf.Outp a(attrs) } opspec := tf.OpSpec{ - Type: "Angle", + Type: "ResourceSparseApplyAdagrad", Input: []tf.Input{ - input, + var_, accum, lr, grad, indices, }, Attrs: attrs, } - op := scope.AddOperation(opspec) - return op.Output(0) + return scope.AddOperation(opspec) } -// VarHandleOpAttr is an optional argument to VarHandleOp. -type VarHandleOpAttr func(optionalAttr) - -// VarHandleOpContainer sets the optional container attribute to value. +// 2D real-valued fast Fourier transform. // -// value: the container this variable is placed in. -// If not specified, defaults to "" -func VarHandleOpContainer(value string) VarHandleOpAttr { - return func(m optionalAttr) { - m["container"] = value +// Computes the 2-dimensional discrete Fourier transform of a real-valued signal +// over the inner-most 2 dimensions of `input`. +// +// Since the DFT of a real signal is Hermitian-symmetric, `RFFT2D` only returns the +// `fft_length / 2 + 1` unique components of the FFT for the inner-most dimension +// of `output`: the zero-frequency term, followed by the `fft_length / 2` +// positive-frequency terms. +// +// Along each axis `RFFT2D` is computed on, if `fft_length` is smaller than the +// corresponding dimension of `input`, the dimension is cropped. If it is larger, +// the dimension is padded with zeros. +// +// Arguments: +// input: A float32 tensor. +// fft_length: An int32 tensor of shape [2]. The FFT length for each dimension. +// +// Returns A complex64 tensor of the same rank as `input`. The inner-most 2 +// dimensions of `input` are replaced with their 2D Fourier transform. The +// inner-most dimension contains `fft_length / 2 + 1` unique frequency +// components. +// +// @compatibility(numpy) +// Equivalent to np.fft.rfft2 +// @end_compatibility +func RFFT2D(scope *Scope, input tf.Output, fft_length tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "RFFT2D", + Input: []tf.Input{ + input, fft_length, + }, } + op := scope.AddOperation(opspec) + return op.Output(0) } -// VarHandleOpSharedName sets the optional shared_name attribute to value. +// ResizeAreaAttr is an optional argument to ResizeArea. +type ResizeAreaAttr func(optionalAttr) + +// ResizeAreaAlignCorners sets the optional align_corners attribute to value. // -// value: the name by which this variable is referred to. -// If not specified, defaults to "" -func VarHandleOpSharedName(value string) VarHandleOpAttr { +// value: If true, the centers of the 4 corner pixels of the input and output tensors are +// aligned, preserving the values at the corner pixels. Defaults to false. +// If not specified, defaults to false +func ResizeAreaAlignCorners(value bool) ResizeAreaAttr { return func(m optionalAttr) { - m["shared_name"] = value + m["align_corners"] = value } } -// Creates a handle to a Variable resource. +// Resize `images` to `size` using area interpolation. +// +// Input images can be of different types but output images are always float. +// +// The range of pixel values for the output image might be slightly different +// from the range for the input image because of limited numerical precision. +// To guarantee an output range, for example `[0.0, 1.0]`, apply +// `tf.clip_by_value` to the output. +// +// Each output pixel is computed by first transforming the pixel's footprint into +// the input tensor and then averaging the pixels that intersect the footprint. An +// input pixel's contribution to the average is weighted by the fraction of its +// area that intersects the footprint. This is the same as OpenCV's INTER_AREA. // // Arguments: -// dtype: the type of this variable. Must agree with the dtypes -// of all ops using this variable. -// shape: The (possibly partially specified) shape of this variable. -func VarHandleOp(scope *Scope, dtype tf.DataType, shape tf.Shape, optional ...VarHandleOpAttr) (resource tf.Output) { +// images: 4-D with shape `[batch, height, width, channels]`. +// size: = A 1-D int32 Tensor of 2 elements: `new_height, new_width`. The +// new size for the images. +// +// Returns 4-D with shape +// `[batch, new_height, new_width, channels]`. +func ResizeArea(scope *Scope, images tf.Output, size tf.Output, optional ...ResizeAreaAttr) (resized_images tf.Output) { if scope.Err() != nil { return } - attrs := map[string]interface{}{"dtype": dtype, "shape": shape} + attrs := map[string]interface{}{} for _, a := range optional { a(attrs) } opspec := tf.OpSpec{ - Type: "VarHandleOp", - + Type: "ResizeArea", + Input: []tf.Input{ + images, size, + }, Attrs: attrs, } op := scope.AddOperation(opspec) return op.Output(0) } -// Elementwise computes the bitwise XOR of `x` and `y`. +// Pads a tensor with zeros. // -// The result will have those bits set, that are different in `x` and `y`. The -// computation is performed on the underlying representations of `x` and `y`. -func BitwiseXor(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) { +// This operation pads a `input` with zeros according to the `paddings` you +// specify. `paddings` is an integer tensor with shape `[Dn, 2]`, where n is the +// rank of `input`. For each dimension D of `input`, `paddings[D, 0]` indicates +// how many zeros to add before the contents of `input` in that dimension, and +// `paddings[D, 1]` indicates how many zeros to add after the contents of `input` +// in that dimension. +// +// The padded size of each dimension D of the output is: +// +// `paddings(D, 0) + input.dim_size(D) + paddings(D, 1)` +// +// For example: +// +// ``` +// # 't' is [[1, 1], [2, 2]] +// # 'paddings' is [[1, 1], [2, 2]] +// # rank of 't' is 2 +// pad(t, paddings) ==> [[0, 0, 0, 0, 0, 0] +// [0, 0, 1, 1, 0, 0] +// [0, 0, 2, 2, 0, 0] +// [0, 0, 0, 0, 0, 0]] +// ``` +func Pad(scope *Scope, input tf.Output, paddings tf.Output) (output tf.Output) { if scope.Err() != nil { return } opspec := tf.OpSpec{ - Type: "BitwiseXor", + Type: "Pad", Input: []tf.Input{ - x, y, + input, paddings, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Checks whether a resource handle-based variable has been initialized. +// +// Arguments: +// resource: the input resource handle. +// +// Returns a scalar boolean which is true if the variable has been +// initialized. +func VarIsInitializedOp(scope *Scope, resource tf.Output) (is_initialized tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "VarIsInitializedOp", + Input: []tf.Input{ + resource, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// StatelessRandomUniformAttr is an optional argument to StatelessRandomUniform. +type StatelessRandomUniformAttr func(optionalAttr) + +// StatelessRandomUniformDtype sets the optional dtype attribute to value. +// +// value: The type of the output. +// If not specified, defaults to DT_FLOAT +func StatelessRandomUniformDtype(value tf.DataType) StatelessRandomUniformAttr { + return func(m optionalAttr) { + m["dtype"] = value + } +} + +// Outputs deterministic pseudorandom random values from a uniform distribution. +// +// The generated values follow a uniform distribution in the range `[0, 1)`. The +// lower bound 0 is included in the range, while the upper bound 1 is excluded. +// +// The outputs are a deterministic function of `shape` and `seed`. +// +// Arguments: +// shape: The shape of the output tensor. +// seed: 2 seeds (shape [2]). +// +// Returns Random values with specified shape. +func StatelessRandomUniform(scope *Scope, shape tf.Output, seed tf.Output, optional ...StatelessRandomUniformAttr) (output tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "StatelessRandomUniform", + Input: []tf.Input{ + shape, seed, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Makes its input available to the next iteration. +// +// Arguments: +// data: The tensor to be made available to the next iteration. +// +// Returns The same tensor as `data`. +func NextIteration(scope *Scope, data tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "NextIteration", + Input: []tf.Input{ + data, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Output a fact about factorials. +func Fact(scope *Scope) (fact tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "Fact", + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// AngleAttr is an optional argument to Angle. +type AngleAttr func(optionalAttr) + +// AngleTout sets the optional Tout attribute to value. +// If not specified, defaults to DT_FLOAT +func AngleTout(value tf.DataType) AngleAttr { + return func(m optionalAttr) { + m["Tout"] = value + } +} + +// Returns the argument of a complex number. +// +// Given a tensor `input` of complex numbers, this operation returns a tensor of +// type `float` that is the argument of each element in `input`. All elements in +// `input` must be complex numbers of the form \\(a + bj\\), where *a* +// is the real part and *b* is the imaginary part. +// +// The argument returned by this operation is of the form \\(atan2(b, a)\\). +// +// For example: +// +// ``` +// # tensor 'input' is [-2.25 + 4.75j, 3.25 + 5.75j] +// tf.angle(input) ==> [2.0132, 1.056] +// ``` +// +// @compatibility(numpy) +// Equivalent to np.angle. +// @end_compatibility +func Angle(scope *Scope, input tf.Output, optional ...AngleAttr) (output tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "Angle", + Input: []tf.Input{ + input, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// VarHandleOpAttr is an optional argument to VarHandleOp. +type VarHandleOpAttr func(optionalAttr) + +// VarHandleOpContainer sets the optional container attribute to value. +// +// value: the container this variable is placed in. +// If not specified, defaults to "" +func VarHandleOpContainer(value string) VarHandleOpAttr { + return func(m optionalAttr) { + m["container"] = value + } +} + +// VarHandleOpSharedName sets the optional shared_name attribute to value. +// +// value: the name by which this variable is referred to. +// If not specified, defaults to "" +func VarHandleOpSharedName(value string) VarHandleOpAttr { + return func(m optionalAttr) { + m["shared_name"] = value + } +} + +// Creates a handle to a Variable resource. +// +// Arguments: +// dtype: the type of this variable. Must agree with the dtypes +// of all ops using this variable. +// shape: The (possibly partially specified) shape of this variable. +func VarHandleOp(scope *Scope, dtype tf.DataType, shape tf.Shape, optional ...VarHandleOpAttr) (resource tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"dtype": dtype, "shape": shape} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "VarHandleOp", + + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Elementwise computes the bitwise XOR of `x` and `y`. +// +// The result will have those bits set, that are different in `x` and `y`. The +// computation is performed on the underlying representations of `x` and `y`. +func BitwiseXor(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "BitwiseXor", + Input: []tf.Input{ + x, y, }, } op := scope.AddOperation(opspec) @@ -10117,6 +11109,50 @@ func TensorArrayGatherV3(scope *Scope, handle tf.Output, indices tf.Output, flow return op.Output(0) } +// Mutually reduces multiple tensors of identical type and shape. +func CollectiveReduce(scope *Scope, input tf.Output, group_size int64, group_key int64, instance_key int64, merge_op string, final_op string, subdiv_offsets []int64) (data tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"group_size": group_size, "group_key": group_key, "instance_key": instance_key, "merge_op": merge_op, "final_op": final_op, "subdiv_offsets": subdiv_offsets} + opspec := tf.OpSpec{ + Type: "CollectiveReduce", + Input: []tf.Input{ + input, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// This op consumes a lock created by `MutexLock`. +// +// This op exists to consume a tensor created by `MutexLock` (other than +// direct control dependencies). It should be the only that consumes the tensor, +// and will raise an error if it is not. Its only purpose is to keep the +// mutex lock tensor alive until it is consumed by this op. +// +// **NOTE**: This operation must run on the same device as its input. This may +// be enforced via the `colocate_with` mechanism. +// +// Arguments: +// mutex_lock: A tensor returned by `MutexLock`. +// +// Returns the created operation. +func ConsumeMutexLock(scope *Scope, mutex_lock tf.Output) (o *tf.Operation) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "ConsumeMutexLock", + Input: []tf.Input{ + mutex_lock, + }, + } + return scope.AddOperation(opspec) +} + // Returns x / y element-wise for integer types. // // Truncation designates that negative numbers will round fractional quantities @@ -10217,7 +11253,7 @@ func SkipDataset(scope *Scope, input_dataset tf.Output, count tf.Output, output_ // Computes the maximum along segments of a tensor. // -// Read @{$math_ops#segmentation$the section on segmentation} for an explanation of +// Read @{$math_ops#Segmentation$the section on segmentation} for an explanation of // segments. // // Computes a tensor such that @@ -10266,6 +11302,21 @@ func Tanh(scope *Scope, x tf.Output) (y tf.Output) { return op.Output(0) } +// Receives a tensor value broadcast from another device. +func CollectiveBcastRecv(scope *Scope, T tf.DataType, group_size int64, group_key int64, instance_key int64, shape tf.Shape) (data tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"T": T, "group_size": group_size, "group_key": group_key, "instance_key": instance_key, "shape": shape} + opspec := tf.OpSpec{ + Type: "CollectiveBcastRecv", + + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Decode web-safe base64-encoded strings. // // Input may or may not have padding at the end. See EncodeBase64 for padding. @@ -10462,324 +11513,143 @@ func SparseReduceSum(scope *Scope, input_indices tf.Output, input_values tf.Outp return op.Output(0) } -// 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) * -// y + truncate_mod(x, y) = x`. -// -// *NOTE*: `TruncateMod` supports broadcasting. More about broadcasting -// [here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html) -func TruncateMod(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "TruncateMod", - Input: []tf.Input{ - x, y, - }, +// VariableShapeAttr is an optional argument to VariableShape. +type VariableShapeAttr func(optionalAttr) + +// VariableShapeOutType sets the optional out_type attribute to value. +// If not specified, defaults to DT_INT32 +func VariableShapeOutType(value tf.DataType) VariableShapeAttr { + return func(m optionalAttr) { + m["out_type"] = value } - op := scope.AddOperation(opspec) - return op.Output(0) } -// Inverse 2D real-valued fast Fourier transform. -// -// Computes the inverse 2-dimensional discrete Fourier transform of a real-valued -// signal over the inner-most 2 dimensions of `input`. -// -// The inner-most 2 dimensions of `input` are assumed to be the result of `RFFT2D`: -// The inner-most dimension contains the `fft_length / 2 + 1` unique components of -// the DFT of a real-valued signal. If `fft_length` is not provided, it is computed -// from the size of the inner-most 2 dimensions of `input`. If the FFT length used -// to compute `input` is odd, it should be provided since it cannot be inferred -// properly. -// -// Along each axis `IRFFT2D` is computed on, if `fft_length` (or -// `fft_length / 2 + 1` for the inner-most dimension) is smaller than the -// corresponding dimension of `input`, the dimension is cropped. If it is larger, -// the dimension is padded with zeros. +// Returns the shape of the variable pointed to by `resource`. // -// Arguments: -// input: A complex64 tensor. -// fft_length: An int32 tensor of shape [2]. The FFT length for each dimension. +// This operation returns a 1-D integer tensor representing the shape of `input`. // -// Returns A float32 tensor of the same rank as `input`. The inner-most 2 -// dimensions of `input` are replaced with the `fft_length` samples of their -// inverse 2D Fourier transform. +// For example: // -// @compatibility(numpy) -// Equivalent to np.fft.irfft2 -// @end_compatibility -func IRFFT2D(scope *Scope, input tf.Output, fft_length tf.Output) (output tf.Output) { +// ``` +// # 't' is [[[1, 1, 1], [2, 2, 2]], [[3, 3, 3], [4, 4, 4]]] +// shape(t) ==> [2, 2, 3] +// ``` +func VariableShape(scope *Scope, input tf.Output, optional ...VariableShapeAttr) (output tf.Output) { if scope.Err() != nil { return } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } opspec := tf.OpSpec{ - Type: "IRFFT2D", + Type: "VariableShape", Input: []tf.Input{ - input, fft_length, + input, }, + Attrs: attrs, } op := scope.AddOperation(opspec) return op.Output(0) } -// DecodeJpegAttr is an optional argument to DecodeJpeg. -type DecodeJpegAttr func(optionalAttr) - -// DecodeJpegChannels sets the optional channels attribute to value. -// -// value: Number of color channels for the decoded image. -// If not specified, defaults to 0 -func DecodeJpegChannels(value int64) DecodeJpegAttr { - return func(m optionalAttr) { - m["channels"] = value - } -} - -// DecodeJpegRatio sets the optional ratio attribute to value. -// -// value: Downscaling ratio. -// If not specified, defaults to 1 -func DecodeJpegRatio(value int64) DecodeJpegAttr { - return func(m optionalAttr) { - m["ratio"] = value - } -} +// SparseToSparseSetOperationAttr is an optional argument to SparseToSparseSetOperation. +type SparseToSparseSetOperationAttr func(optionalAttr) -// DecodeJpegFancyUpscaling sets the optional fancy_upscaling attribute to value. -// -// value: If true use a slower but nicer upscaling of the -// chroma planes (yuv420/422 only). +// SparseToSparseSetOperationValidateIndices sets the optional validate_indices attribute to value. // If not specified, defaults to true -func DecodeJpegFancyUpscaling(value bool) DecodeJpegAttr { - return func(m optionalAttr) { - m["fancy_upscaling"] = value - } -} - -// DecodeJpegTryRecoverTruncated sets the optional try_recover_truncated attribute to value. -// -// value: If true try to recover an image from truncated input. -// If not specified, defaults to false -func DecodeJpegTryRecoverTruncated(value bool) DecodeJpegAttr { - return func(m optionalAttr) { - m["try_recover_truncated"] = value - } -} - -// DecodeJpegAcceptableFraction sets the optional acceptable_fraction attribute to value. -// -// value: The minimum required fraction of lines before a truncated -// input is accepted. -// If not specified, defaults to 1 -func DecodeJpegAcceptableFraction(value float32) DecodeJpegAttr { - return func(m optionalAttr) { - m["acceptable_fraction"] = value - } -} - -// DecodeJpegDctMethod sets the optional dct_method attribute to value. -// -// value: string specifying a hint about the algorithm used for -// decompression. Defaults to "" which maps to a system-specific -// default. Currently valid values are ["INTEGER_FAST", -// "INTEGER_ACCURATE"]. The hint may be ignored (e.g., the internal -// jpeg library changes to a version that does not have that specific -// option.) -// If not specified, defaults to "" -func DecodeJpegDctMethod(value string) DecodeJpegAttr { +func SparseToSparseSetOperationValidateIndices(value bool) SparseToSparseSetOperationAttr { return func(m optionalAttr) { - m["dct_method"] = value + m["validate_indices"] = value } } -// Decode a JPEG-encoded image to a uint8 tensor. -// -// The attr `channels` indicates the desired number of color channels for the -// decoded image. +// Applies set operation along last dimension of 2 `SparseTensor` inputs. // -// Accepted values are: +// See SetOperationOp::SetOperationFromContext for values of `set_operation`. // -// * 0: Use the number of channels in the JPEG-encoded image. -// * 1: output a grayscale image. -// * 3: output an RGB image. +// If `validate_indices` is `True`, `SparseToSparseSetOperation` validates the +// order and range of `set1` and `set2` indices. // -// If needed, the JPEG-encoded image is transformed to match the requested number -// of color channels. +// Input `set1` is a `SparseTensor` represented by `set1_indices`, `set1_values`, +// and `set1_shape`. For `set1` ranked `n`, 1st `n-1` dimensions must be the same +// as `set2`. Dimension `n` contains values in a set, duplicates are allowed but +// ignored. // -// The attr `ratio` allows downscaling the image by an integer factor during -// decoding. Allowed values are: 1, 2, 4, and 8. This is much faster than -// downscaling the image later. +// Input `set2` is a `SparseTensor` represented by `set2_indices`, `set2_values`, +// and `set2_shape`. For `set2` ranked `n`, 1st `n-1` dimensions must be the same +// as `set1`. Dimension `n` contains values in a set, duplicates are allowed but +// ignored. // +// If `validate_indices` is `True`, this op validates the order and range of `set1` +// and `set2` indices. // -// This op also supports decoding PNGs and non-animated GIFs since the interface is -// the same, though it is cleaner to use `tf.image.decode_image`. +// Output `result` is a `SparseTensor` represented by `result_indices`, +// `result_values`, and `result_shape`. For `set1` and `set2` ranked `n`, this +// has rank `n` and the same 1st `n-1` dimensions as `set1` and `set2`. The `nth` +// dimension contains the result of `set_operation` applied to the corresponding +// `[0...n-1]` dimension of `set`. // // Arguments: -// contents: 0-D. The JPEG-encoded image. +// set1_indices: 2D `Tensor`, indices of a `SparseTensor`. Must be in row-major +// order. +// set1_values: 1D `Tensor`, values of a `SparseTensor`. Must be in row-major +// order. +// set1_shape: 1D `Tensor`, shape of a `SparseTensor`. `set1_shape[0...n-1]` must +// be the same as `set2_shape[0...n-1]`, `set1_shape[n]` is the +// max set size across `0...n-1` dimensions. +// set2_indices: 2D `Tensor`, indices of a `SparseTensor`. Must be in row-major +// order. +// set2_values: 1D `Tensor`, values of a `SparseTensor`. Must be in row-major +// order. +// set2_shape: 1D `Tensor`, shape of a `SparseTensor`. `set2_shape[0...n-1]` must +// be the same as `set1_shape[0...n-1]`, `set2_shape[n]` is the +// max set size across `0...n-1` dimensions. // -// Returns 3-D with shape `[height, width, channels]`.. -func DecodeJpeg(scope *Scope, contents tf.Output, optional ...DecodeJpegAttr) (image tf.Output) { +// +// Returns 2D indices of a `SparseTensor`.1D values of a `SparseTensor`.1D `Tensor` shape of a `SparseTensor`. `result_shape[0...n-1]` is +// the same as the 1st `n-1` dimensions of `set1` and `set2`, `result_shape[n]` +// is the max result set size across all `0...n-1` dimensions. +func SparseToSparseSetOperation(scope *Scope, set1_indices tf.Output, set1_values tf.Output, set1_shape tf.Output, set2_indices tf.Output, set2_values tf.Output, set2_shape tf.Output, set_operation string, optional ...SparseToSparseSetOperationAttr) (result_indices tf.Output, result_values tf.Output, result_shape tf.Output) { if scope.Err() != nil { return } - attrs := map[string]interface{}{} + attrs := map[string]interface{}{"set_operation": set_operation} for _, a := range optional { a(attrs) } opspec := tf.OpSpec{ - Type: "DecodeJpeg", + Type: "SparseToSparseSetOperation", Input: []tf.Input{ - contents, + set1_indices, set1_values, set1_shape, set2_indices, set2_values, set2_shape, }, Attrs: attrs, } op := scope.AddOperation(opspec) - return op.Output(0) + return op.Output(0), op.Output(1), op.Output(2) } -// Transforms a vector of brain.Example protos (as strings) into typed tensors. +// Computes softmax cross entropy cost and gradients to backpropagate. +// +// Unlike `SoftmaxCrossEntropyWithLogits`, this operation does not accept +// a matrix of label probabilities, but rather a single label per row +// of features. This label is considered to have probability 1.0 for the +// given row. +// +// Inputs are the logits, not probabilities. // // Arguments: -// serialized: A vector containing a batch of binary serialized Example protos. -// names: A vector containing the names of the serialized protos. -// May contain, for example, table key (descriptive) names for the -// corresponding serialized protos. These are purely useful for debugging -// purposes, and the presence of values here has no effect on the output. -// May also be an empty vector if no names are available. -// If non-empty, this vector must be the same length as "serialized". -// sparse_keys: A list of Nsparse string Tensors (scalars). -// The keys expected in the Examples' features associated with sparse values. -// dense_keys: A list of Ndense string Tensors (scalars). -// The keys expected in the Examples' features associated with dense values. -// dense_defaults: A list of Ndense Tensors (some may be empty). -// dense_defaults[j] provides default values -// when the example's feature_map lacks dense_key[j]. If an empty Tensor is -// provided for dense_defaults[j], then the Feature dense_keys[j] is required. -// The input type is inferred from dense_defaults[j], even when it's empty. -// If dense_defaults[j] is not empty, and dense_shapes[j] is fully defined, -// then the shape of dense_defaults[j] must match that of dense_shapes[j]. -// If dense_shapes[j] has an undefined major dimension (variable strides dense -// feature), dense_defaults[j] must contain a single element: -// the padding element. -// sparse_types: A list of Nsparse types; the data types of data in each Feature -// given in sparse_keys. -// Currently the ParseExample supports DT_FLOAT (FloatList), -// DT_INT64 (Int64List), and DT_STRING (BytesList). -// dense_shapes: A list of Ndense shapes; the shapes of data in each Feature -// given in dense_keys. -// The number of elements in the Feature corresponding to dense_key[j] -// must always equal dense_shapes[j].NumEntries(). -// If dense_shapes[j] == (D0, D1, ..., DN) then the shape of output -// Tensor dense_values[j] will be (|serialized|, D0, D1, ..., DN): -// The dense outputs are just the inputs row-stacked by batch. -// This works for dense_shapes[j] = (-1, D1, ..., DN). In this case -// the shape of the output Tensor dense_values[j] will be -// (|serialized|, M, D1, .., DN), where M is the maximum number of blocks -// of elements of length D1 * .... * DN, across all minibatch entries -// in the input. Any minibatch entry with less than M blocks of elements of -// length D1 * ... * DN will be padded with the corresponding default_value -// scalar element along the second dimension. -func ParseExample(scope *Scope, serialized tf.Output, names tf.Output, sparse_keys []tf.Output, dense_keys []tf.Output, dense_defaults []tf.Output, sparse_types []tf.DataType, dense_shapes []tf.Shape) (sparse_indices []tf.Output, sparse_values []tf.Output, sparse_shapes []tf.Output, dense_values []tf.Output) { +// features: batch_size x num_classes matrix +// labels: batch_size vector with values in [0, num_classes). +// This is the label for the given minibatch entry. +// +// Returns Per example loss (batch_size vector).backpropagated gradients (batch_size x num_classes matrix). +func SparseSoftmaxCrossEntropyWithLogits(scope *Scope, features tf.Output, labels tf.Output) (loss tf.Output, backprop tf.Output) { if scope.Err() != nil { return } - attrs := map[string]interface{}{"sparse_types": sparse_types, "dense_shapes": dense_shapes} opspec := tf.OpSpec{ - Type: "ParseExample", - Input: []tf.Input{ - serialized, names, tf.OutputList(sparse_keys), tf.OutputList(dense_keys), tf.OutputList(dense_defaults), - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - if scope.Err() != nil { - return - } - var idx int - var err error - if sparse_indices, idx, err = makeOutputList(op, idx, "sparse_indices"); err != nil { - scope.UpdateErr("ParseExample", err) - return - } - if sparse_values, idx, err = makeOutputList(op, idx, "sparse_values"); err != nil { - scope.UpdateErr("ParseExample", err) - return - } - if sparse_shapes, idx, err = makeOutputList(op, idx, "sparse_shapes"); err != nil { - scope.UpdateErr("ParseExample", err) - return - } - if dense_values, idx, err = makeOutputList(op, idx, "dense_values"); err != nil { - scope.UpdateErr("ParseExample", err) - return - } - return sparse_indices, sparse_values, sparse_shapes, dense_values -} - -// VariableShapeAttr is an optional argument to VariableShape. -type VariableShapeAttr func(optionalAttr) - -// VariableShapeOutType sets the optional out_type attribute to value. -// If not specified, defaults to DT_INT32 -func VariableShapeOutType(value tf.DataType) VariableShapeAttr { - return func(m optionalAttr) { - m["out_type"] = value - } -} - -// Returns the shape of the variable pointed to by `resource`. -// -// This operation returns a 1-D integer tensor representing the shape of `input`. -// -// For example: -// -// ``` -// # 't' is [[[1, 1, 1], [2, 2, 2]], [[3, 3, 3], [4, 4, 4]]] -// shape(t) ==> [2, 2, 3] -// ``` -func VariableShape(scope *Scope, input tf.Output, optional ...VariableShapeAttr) (output tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "VariableShape", - Input: []tf.Input{ - input, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// Computes softmax cross entropy cost and gradients to backpropagate. -// -// Unlike `SoftmaxCrossEntropyWithLogits`, this operation does not accept -// a matrix of label probabilities, but rather a single label per row -// of features. This label is considered to have probability 1.0 for the -// given row. -// -// Inputs are the logits, not probabilities. -// -// Arguments: -// features: batch_size x num_classes matrix -// labels: batch_size vector with values in [0, num_classes). -// This is the label for the given minibatch entry. -// -// Returns Per example loss (batch_size vector).backpropagated gradients (batch_size x num_classes matrix). -func SparseSoftmaxCrossEntropyWithLogits(scope *Scope, features tf.Output, labels tf.Output) (loss tf.Output, backprop tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "SparseSoftmaxCrossEntropyWithLogits", + Type: "SparseSoftmaxCrossEntropyWithLogits", Input: []tf.Input{ features, labels, }, @@ -11279,6 +12149,130 @@ func TensorArrayV3(scope *Scope, size tf.Output, dtype tf.DataType, optional ... return op.Output(0), op.Output(1) } +// Runs multiple additive regression ensemble predictors on input instances and +// +// computes the logits. It is designed to be used during prediction. +// It traverses all the trees and calculates the final score for each instance. +// +// Arguments: +// +// bucketized_features: A list of rank 1 Tensors containing bucket id for each +// feature. +// logits_dimension: scalar, dimension of the logits, to be used for partial logits +// shape. +// +// Returns Output rank 2 Tensor containing logits for each example. +func BoostedTreesPredict(scope *Scope, tree_ensemble_handle tf.Output, bucketized_features []tf.Output, logits_dimension int64) (logits tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"logits_dimension": logits_dimension} + opspec := tf.OpSpec{ + Type: "BoostedTreesPredict", + Input: []tf.Input{ + tree_ensemble_handle, tf.OutputList(bucketized_features), + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// MatrixSolveLsAttr is an optional argument to MatrixSolveLs. +type MatrixSolveLsAttr func(optionalAttr) + +// MatrixSolveLsFast sets the optional fast attribute to value. +// If not specified, defaults to true +func MatrixSolveLsFast(value bool) MatrixSolveLsAttr { + return func(m optionalAttr) { + m["fast"] = value + } +} + +// Solves one or more linear least-squares problems. +// +// `matrix` is a tensor of shape `[..., M, N]` whose inner-most 2 dimensions +// form real or complex matrices of size `[M, N]`. `Rhs` is a tensor of the same +// type as `matrix` and shape `[..., M, K]`. +// The output is a tensor shape `[..., N, K]` where each output matrix solves +// each of the equations +// `matrix[..., :, :]` * `output[..., :, :]` = `rhs[..., :, :]` +// in the least squares sense. +// +// We use the following notation for (complex) matrix and right-hand sides +// in the batch: +// +// `matrix`=\\(A \in \mathbb{C}^{m \times n}\\), +// `rhs`=\\(B \in \mathbb{C}^{m \times k}\\), +// `output`=\\(X \in \mathbb{C}^{n \times k}\\), +// `l2_regularizer`=\\(\lambda \in \mathbb{R}\\). +// +// If `fast` is `True`, then the solution is computed by solving the normal +// equations using Cholesky decomposition. Specifically, if \\(m \ge n\\) then +// \\(X = (A^H A + \lambda I)^{-1} A^H B\\), which solves the least-squares +// problem \\(X = \mathrm{argmin}_{Z \in \Re^{n \times k} } ||A Z - B||_F^2 + \lambda ||Z||_F^2\\). +// If \\(m \lt n\\) then `output` is computed as +// \\(X = A^H (A A^H + \lambda I)^{-1} B\\), which (for \\(\lambda = 0\\)) is the +// minimum-norm solution to the under-determined linear system, i.e. +// \\(X = \mathrm{argmin}_{Z \in \mathbb{C}^{n \times k} } ||Z||_F^2 \\), +// subject to \\(A Z = B\\). Notice that the fast path is only numerically stable +// when \\(A\\) is numerically full rank and has a condition number +// \\(\mathrm{cond}(A) \lt \frac{1}{\sqrt{\epsilon_{mach} } }\\) or \\(\lambda\\) is +// sufficiently large. +// +// If `fast` is `False` an algorithm based on the numerically robust complete +// orthogonal decomposition is used. This computes the minimum-norm +// least-squares solution, even when \\(A\\) is rank deficient. This path is +// typically 6-7 times slower than the fast path. If `fast` is `False` then +// `l2_regularizer` is ignored. +// +// Arguments: +// matrix: Shape is `[..., M, N]`. +// rhs: Shape is `[..., M, K]`. +// l2_regularizer: Scalar tensor. +// +// @compatibility(numpy) +// Equivalent to np.linalg.lstsq +// @end_compatibility +// +// Returns Shape is `[..., N, K]`. +func MatrixSolveLs(scope *Scope, matrix tf.Output, rhs tf.Output, l2_regularizer tf.Output, optional ...MatrixSolveLsAttr) (output tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "MatrixSolveLs", + Input: []tf.Input{ + matrix, rhs, l2_regularizer, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Elementwise computes the bitwise OR of `x` and `y`. +// +// The result will have those bits set, that are set in `x`, `y` or both. The +// computation is performed on the underlying representations of `x` and `y`. +func BitwiseOr(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "BitwiseOr", + Input: []tf.Input{ + x, y, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // MaxPool3DAttr is an optional argument to MaxPool3D. type MaxPool3DAttr func(optionalAttr) @@ -11326,6 +12320,17 @@ func MaxPool3D(scope *Scope, input tf.Output, ksize []int64, strides []int64, pa return op.Output(0) } +// Conv3DBackpropInputAttr is an optional argument to Conv3DBackpropInput. +type Conv3DBackpropInputAttr func(optionalAttr) + +// Conv3DBackpropInputDilations sets the optional dilations attribute to value. +// If not specified, defaults to +func Conv3DBackpropInputDilations(value []int64) Conv3DBackpropInputAttr { + return func(m optionalAttr) { + m["dilations"] = value + } +} + // Computes the gradients of 3-D convolution with respect to the input. // // DEPRECATED at GraphDef version 10: Use Conv3DBackpropInputV2 @@ -11339,11 +12344,14 @@ func MaxPool3D(scope *Scope, input tf.Output, ksize []int64, strides []int64, pa // strides: 1-D tensor of length 5. The stride of the sliding window for each // dimension of `input`. Must have `strides[0] = strides[4] = 1`. // padding: The type of padding algorithm to use. -func Conv3DBackpropInput(scope *Scope, input tf.Output, filter tf.Output, out_backprop tf.Output, strides []int64, padding string) (output tf.Output) { +func Conv3DBackpropInput(scope *Scope, input tf.Output, filter tf.Output, out_backprop tf.Output, strides []int64, padding string, optional ...Conv3DBackpropInputAttr) (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: "Conv3DBackpropInput", Input: []tf.Input{ @@ -11471,6 +12479,47 @@ func MutableHashTableOfTensorsV2(scope *Scope, key_dtype tf.DataType, value_dtyp return op.Output(0) } +// Subtracts sparse updates from the variable referenced by `resource`. +// +// This operation computes +// +// # Scalar indices +// ref[indices, ...] -= updates[...] +// +// # Vector indices (for each i) +// ref[indices[i], ...] -= updates[i, ...] +// +// # High rank indices (for each i, ..., j) +// ref[indices[i, ..., j], ...] -= updates[i, ..., j, ...] +// +// Duplicate entries are handled correctly: if multiple `indices` reference +// the same location, their contributions add. +// +// Requires `updates.shape = indices.shape + ref.shape[1:]` or `updates.shape = []`. +// +//
+// +//
+// +// Arguments: +// resource: Should be from a `Variable` node. +// indices: A tensor of indices into the first dimension of `ref`. +// updates: A tensor of updated values to add to `ref`. +// +// Returns the created operation. +func ResourceScatterSub(scope *Scope, resource tf.Output, indices tf.Output, updates tf.Output) (o *tf.Operation) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "ResourceScatterSub", + Input: []tf.Input{ + resource, indices, updates, + }, + } + return scope.AddOperation(opspec) +} + // Inverse 2D fast Fourier transform. // // Computes the inverse 2-dimensional discrete Fourier transform over the @@ -12069,9 +13118,8 @@ type FusedResizeAndPadConv2DAttr func(optionalAttr) // FusedResizeAndPadConv2DResizeAlignCorners sets the optional resize_align_corners attribute to value. // -// value: If true, rescale input by (new_height - 1) / (height - 1), -// which exactly aligns the 4 corners of images and resized images. If false, rescale -// by new_height / height. Treat similarly the width dimension. +// value: If true, the centers of the 4 corner pixels of the input and output tensors are +// aligned, preserving the values at the corner pixels. Defaults to false. // If not specified, defaults to false func FusedResizeAndPadConv2DResizeAlignCorners(value bool) FusedResizeAndPadConv2DAttr { return func(m optionalAttr) { @@ -12123,32 +13171,110 @@ func FusedResizeAndPadConv2D(scope *Scope, input tf.Output, size tf.Output, padd return op.Output(0) } -// Inverse 3D fast Fourier transform. +// Returns a list of tensors with the same shapes and contents as the input // -// Computes the inverse 3-dimensional discrete Fourier transform over the -// inner-most 3 dimensions of `input`. +// tensors. // -// Arguments: -// input: A complex64 tensor. +// This op can be used to override the gradient for complicated functions. For +// example, suppose y = f(x) and we wish to apply a custom function g for backprop +// such that dx = g(dy). In Python, // -// Returns A complex64 tensor of the same shape as `input`. The inner-most 3 -// dimensions of `input` are replaced with their inverse 3D Fourier transform. +// ```python +// with tf.get_default_graph().gradient_override_map( +// {'IdentityN': 'OverrideGradientWithG'}): +// y, _ = identity_n([f(x), x]) // -// @compatibility(numpy) -// Equivalent to np.fft.ifftn with 3 dimensions. -// @end_compatibility -func IFFT3D(scope *Scope, input tf.Output) (output tf.Output) { +// @tf.RegisterGradient('OverrideGradientWithG') +// def ApplyG(op, dy, _): +// return [None, g(dy)] # Do not backprop to f(x). +// ``` +func IdentityN(scope *Scope, input []tf.Output) (output []tf.Output) { if scope.Err() != nil { return } opspec := tf.OpSpec{ - Type: "IFFT3D", + Type: "IdentityN", Input: []tf.Input{ - input, + tf.OutputList(input), }, } op := scope.AddOperation(opspec) - return op.Output(0) + if scope.Err() != nil { + return + } + var idx int + var err error + if output, idx, err = makeOutputList(op, idx, "output"); err != nil { + scope.UpdateErr("IdentityN", err) + return + } + return output +} + +// ResourceApplyCenteredRMSPropAttr is an optional argument to ResourceApplyCenteredRMSProp. +type ResourceApplyCenteredRMSPropAttr func(optionalAttr) + +// ResourceApplyCenteredRMSPropUseLocking sets the optional use_locking attribute to value. +// +// value: If `True`, updating of the var, mg, ms, and mom tensors is +// protected by a lock; otherwise the behavior is undefined, but may exhibit less +// contention. +// If not specified, defaults to false +func ResourceApplyCenteredRMSPropUseLocking(value bool) ResourceApplyCenteredRMSPropAttr { + return func(m optionalAttr) { + m["use_locking"] = value + } +} + +// Update '*var' according to the centered RMSProp algorithm. +// +// The centered RMSProp algorithm uses an estimate of the centered second moment +// (i.e., the variance) for normalization, as opposed to regular RMSProp, which +// uses the (uncentered) second moment. This often helps with training, but is +// slightly more expensive in terms of computation and memory. +// +// Note that in dense implementation of this algorithm, mg, ms, and mom will +// update even if the grad is zero, but in this sparse implementation, mg, ms, +// and mom will not update in iterations during which the grad is zero. +// +// mean_square = decay * mean_square + (1-decay) * gradient ** 2 +// mean_grad = decay * mean_grad + (1-decay) * gradient +// +// Delta = learning_rate * gradient / sqrt(mean_square + epsilon - mean_grad ** 2) +// +// mg <- rho * mg_{t-1} + (1-rho) * grad +// ms <- rho * ms_{t-1} + (1-rho) * grad * grad +// mom <- momentum * mom_{t-1} + lr * grad / sqrt(ms - mg * mg + epsilon) +// var <- var - mom +// +// Arguments: +// var_: Should be from a Variable(). +// mg: Should be from a Variable(). +// ms: Should be from a Variable(). +// mom: Should be from a Variable(). +// lr: Scaling factor. Must be a scalar. +// rho: Decay rate. Must be a scalar. +// +// epsilon: Ridge term. Must be a scalar. +// grad: The gradient. +// +// Returns the created operation. +func ResourceApplyCenteredRMSProp(scope *Scope, var_ tf.Output, mg tf.Output, ms tf.Output, mom tf.Output, lr tf.Output, rho tf.Output, momentum tf.Output, epsilon tf.Output, grad tf.Output, optional ...ResourceApplyCenteredRMSPropAttr) (o *tf.Operation) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "ResourceApplyCenteredRMSProp", + Input: []tf.Input{ + var_, mg, ms, mom, lr, rho, momentum, epsilon, grad, + }, + Attrs: attrs, + } + return scope.AddOperation(opspec) } // Adds `bias` to `value`. @@ -12825,6 +13951,117 @@ func ResourceApplyAddSign(scope *Scope, var_ tf.Output, m tf.Output, lr tf.Outpu return scope.AddOperation(opspec) } +// CudnnRNNBackpropAttr is an optional argument to CudnnRNNBackprop. +type CudnnRNNBackpropAttr func(optionalAttr) + +// CudnnRNNBackpropRnnMode sets the optional rnn_mode attribute to value. +// If not specified, defaults to "lstm" +func CudnnRNNBackpropRnnMode(value string) CudnnRNNBackpropAttr { + return func(m optionalAttr) { + m["rnn_mode"] = value + } +} + +// CudnnRNNBackpropInputMode sets the optional input_mode attribute to value. +// If not specified, defaults to "linear_input" +func CudnnRNNBackpropInputMode(value string) CudnnRNNBackpropAttr { + return func(m optionalAttr) { + m["input_mode"] = value + } +} + +// CudnnRNNBackpropDirection sets the optional direction attribute to value. +// If not specified, defaults to "unidirectional" +func CudnnRNNBackpropDirection(value string) CudnnRNNBackpropAttr { + return func(m optionalAttr) { + m["direction"] = value + } +} + +// CudnnRNNBackpropDropout sets the optional dropout attribute to value. +// If not specified, defaults to 0 +func CudnnRNNBackpropDropout(value float32) CudnnRNNBackpropAttr { + return func(m optionalAttr) { + m["dropout"] = value + } +} + +// CudnnRNNBackpropSeed sets the optional seed attribute to value. +// If not specified, defaults to 0 +func CudnnRNNBackpropSeed(value int64) CudnnRNNBackpropAttr { + return func(m optionalAttr) { + m["seed"] = value + } +} + +// CudnnRNNBackpropSeed2 sets the optional seed2 attribute to value. +// If not specified, defaults to 0 +func CudnnRNNBackpropSeed2(value int64) CudnnRNNBackpropAttr { + return func(m optionalAttr) { + m["seed2"] = value + } +} + +// Backprop step of CudnnRNN. +// +// Compute the backprop of both data and weights in a RNN. +// +// rnn_mode: Indicates the type of the RNN model. +// input_mode: Indicate whether there is a linear projection between the input and +// The actual computation before the first layer. 'skip_input' is only allowed +// when input_size == num_units; 'auto_select' implies 'skip_input' when +// input_size == num_units; otherwise, it implies 'linear_input'. +// direction: Indicates whether a bidirectional model will be used. +// dir = (direction == bidirectional) ? 2 : 1 +// dropout: dropout probability. When set to 0., dropout is disabled. +// seed: the 1st part of a seed to initialize dropout. +// seed2: the 2nd part of a seed to initialize dropout. +// input: a 3-D tensor with the shape of [seq_length, batch_size, input_size]. +// input_h: a 3-D tensor with the shape of [num_layer * dir, batch_size, +// num_units]. +// input_c: For LSTM, a 3-D tensor with the shape of +// [num_layer * dir, batch, num_units]. For other models, it is ignored. +// params: a 1-D tensor that contains the weights and biases in an opaque layout. +// The size must be created through CudnnRNNParamsSize, and initialized +// separately. Note that they might not be compatible across different +// generations. So it is a good idea to save and restore +// output: a 3-D tensor with the shape of [seq_length, batch_size, +// dir * num_units]. +// output_h: the same shape has input_h. +// output_c: the same shape as input_c for LSTM. An empty tensor for other models. +// output_backprop: A 3-D tensor with the same shape as output in the forward pass. +// output_h_backprop: A 3-D tensor with the same shape as output_h in the forward +// pass. +// output_c_backprop: A 3-D tensor with the same shape as output_c in the forward +// pass. +// reserve_space: The same reserve_space produced in for forward operation. +// input_backprop: The backprop to input in the forward pass. Has the same shape +// as input. +// input_h_backprop: The backprop to input_h in the forward pass. Has the same +// shape as input_h. +// input_c_backprop: The backprop to input_c in the forward pass. Has the same +// shape as input_c. +// params_backprop: The backprop to the params buffer in the forward pass. Has the +// same shape as params. +func CudnnRNNBackprop(scope *Scope, input tf.Output, input_h tf.Output, input_c tf.Output, params tf.Output, output tf.Output, output_h tf.Output, output_c tf.Output, output_backprop tf.Output, output_h_backprop tf.Output, output_c_backprop tf.Output, reserve_space tf.Output, optional ...CudnnRNNBackpropAttr) (input_backprop tf.Output, input_h_backprop tf.Output, input_c_backprop tf.Output, params_backprop tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "CudnnRNNBackprop", + Input: []tf.Input{ + input, input_h, input_c, params, output, output_h, output_c, output_backprop, output_h_backprop, output_c_backprop, reserve_space, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0), op.Output(1), op.Output(2), op.Output(3) +} + // FractionalMaxPoolGradAttr is an optional argument to FractionalMaxPoolGrad. type FractionalMaxPoolGradAttr func(optionalAttr) @@ -12923,6 +14160,107 @@ func ResourceApplyAdagradDA(scope *Scope, var_ tf.Output, gradient_accumulator t return scope.AddOperation(opspec) } +// CudnnRNNCanonicalToParamsAttr is an optional argument to CudnnRNNCanonicalToParams. +type CudnnRNNCanonicalToParamsAttr func(optionalAttr) + +// CudnnRNNCanonicalToParamsRnnMode sets the optional rnn_mode attribute to value. +// If not specified, defaults to "lstm" +func CudnnRNNCanonicalToParamsRnnMode(value string) CudnnRNNCanonicalToParamsAttr { + return func(m optionalAttr) { + m["rnn_mode"] = value + } +} + +// CudnnRNNCanonicalToParamsInputMode sets the optional input_mode attribute to value. +// If not specified, defaults to "linear_input" +func CudnnRNNCanonicalToParamsInputMode(value string) CudnnRNNCanonicalToParamsAttr { + return func(m optionalAttr) { + m["input_mode"] = value + } +} + +// CudnnRNNCanonicalToParamsDirection sets the optional direction attribute to value. +// If not specified, defaults to "unidirectional" +func CudnnRNNCanonicalToParamsDirection(value string) CudnnRNNCanonicalToParamsAttr { + return func(m optionalAttr) { + m["direction"] = value + } +} + +// CudnnRNNCanonicalToParamsDropout sets the optional dropout attribute to value. +// If not specified, defaults to 0 +func CudnnRNNCanonicalToParamsDropout(value float32) CudnnRNNCanonicalToParamsAttr { + return func(m optionalAttr) { + m["dropout"] = value + } +} + +// CudnnRNNCanonicalToParamsSeed sets the optional seed attribute to value. +// If not specified, defaults to 0 +func CudnnRNNCanonicalToParamsSeed(value int64) CudnnRNNCanonicalToParamsAttr { + return func(m optionalAttr) { + m["seed"] = value + } +} + +// CudnnRNNCanonicalToParamsSeed2 sets the optional seed2 attribute to value. +// If not specified, defaults to 0 +func CudnnRNNCanonicalToParamsSeed2(value int64) CudnnRNNCanonicalToParamsAttr { + return func(m optionalAttr) { + m["seed2"] = value + } +} + +// Converts CudnnRNN params from canonical form to usable form. +// +// Writes a set of weights into the opaque params buffer so they can be used in +// upcoming training or inferences. +// +// Note that the params buffer may not be compatible across different GPUs. So any +// save and restoration should be converted to and from the canonical weights and +// biases. +// +// num_layers: Specifies the number of layers in the RNN model. +// num_units: Specifies the size of the hidden state. +// input_size: Specifies the size of the input state. +// weights: the canonical form of weights that can be used for saving +// and restoration. They are more likely to be compatible across different +// generations. +// biases: the canonical form of biases that can be used for saving +// and restoration. They are more likely to be compatible across different +// generations. +// num_params: number of parameter sets for all layers. +// Each layer may contain multiple parameter sets, with each set consisting of +// a weight matrix and a bias vector. +// rnn_mode: Indicates the type of the RNN model. +// input_mode: Indicate whether there is a linear projection between the input and +// The actual computation before the first layer. 'skip_input' is only allowed +// when input_size == num_units; 'auto_select' implies 'skip_input' when +// input_size == num_units; otherwise, it implies 'linear_input'. +// direction: Indicates whether a bidirectional model will be used. +// dir = (direction == bidirectional) ? 2 : 1 +// dropout: dropout probability. When set to 0., dropout is disabled. +// seed: the 1st part of a seed to initialize dropout. +// seed2: the 2nd part of a seed to initialize dropout. +func CudnnRNNCanonicalToParams(scope *Scope, num_layers tf.Output, num_units tf.Output, input_size tf.Output, weights []tf.Output, biases []tf.Output, optional ...CudnnRNNCanonicalToParamsAttr) (params tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "CudnnRNNCanonicalToParams", + Input: []tf.Input{ + num_layers, num_units, input_size, tf.OutputList(weights), tf.OutputList(biases), + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // SparseReduceMaxSparseAttr is an optional argument to SparseReduceMaxSparse. type SparseReduceMaxSparseAttr func(optionalAttr) @@ -13117,6 +14455,47 @@ func ReadFile(scope *Scope, filename tf.Output) (contents tf.Output) { return op.Output(0) } +// Multiplies sparse updates into the variable referenced by `resource`. +// +// This operation computes +// +// # Scalar indices +// ref[indices, ...] *= updates[...] +// +// # Vector indices (for each i) +// ref[indices[i], ...] *= updates[i, ...] +// +// # High rank indices (for each i, ..., j) +// ref[indices[i, ..., j], ...] *= updates[i, ..., j, ...] +// +// Duplicate entries are handled correctly: if multiple `indices` reference +// the same location, their contributions multiply. +// +// Requires `updates.shape = indices.shape + ref.shape[1:]` or `updates.shape = []`. +// +//
+// +//
+// +// Arguments: +// resource: Should be from a `Variable` node. +// indices: A tensor of indices into the first dimension of `ref`. +// updates: A tensor of updated values to add to `ref`. +// +// Returns the created operation. +func ResourceScatterMul(scope *Scope, resource tf.Output, indices tf.Output, updates tf.Output) (o *tf.Operation) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "ResourceScatterMul", + Input: []tf.Input{ + resource, indices, updates, + }, + } + return scope.AddOperation(opspec) +} + // Computes sigmoid of `x` element-wise. // // Specifically, `y = 1 / (1 + exp(-x))`. @@ -13134,6 +14513,30 @@ func Sigmoid(scope *Scope, x tf.Output) (y tf.Output) { return op.Output(0) } +// Updates specified rows with values in `v`. +// +// Computes `x[i, :] = v; return x`. +// +// Arguments: +// x: A tensor of type `T`. +// i: A vector. Indices into the left-most dimension of `x`. +// v: A `Tensor` of type T. Same dimension sizes as x except the first dimension, which must be the same as i's size. +// +// Returns A `Tensor` of type T. An alias of `x`. The content of `y` is undefined if there are duplicates in `i`. +func InplaceUpdate(scope *Scope, x tf.Output, i tf.Output, v tf.Output) (y tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "InplaceUpdate", + Input: []tf.Input{ + x, i, v, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // FusedBatchNormAttr is an optional argument to FusedBatchNorm. type FusedBatchNormAttr func(optionalAttr) @@ -13344,22 +14747,59 @@ func FractionalAvgPoolGrad(scope *Scope, orig_input_tensor_shape tf.Output, out_ return op.Output(0) } -// Concatenates tensors along one dimension. -// -// Arguments: -// concat_dim: 0-D. The dimension along which to concatenate. Must be in the -// range [0, rank(values)). -// values: The `N` Tensors to concatenate. Their ranks and types must match, -// and their sizes must match in all dimensions except `concat_dim`. -// -// 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 Concat(scope *Scope, concat_dim tf.Output, values []tf.Output) (output tf.Output) { - if scope.Err() != nil { - return +// BoostedTreesEnsembleResourceHandleOpAttr is an optional argument to BoostedTreesEnsembleResourceHandleOp. +type BoostedTreesEnsembleResourceHandleOpAttr func(optionalAttr) + +// BoostedTreesEnsembleResourceHandleOpContainer sets the optional container attribute to value. +// If not specified, defaults to "" +func BoostedTreesEnsembleResourceHandleOpContainer(value string) BoostedTreesEnsembleResourceHandleOpAttr { + return func(m optionalAttr) { + m["container"] = value } - opspec := tf.OpSpec{ +} + +// BoostedTreesEnsembleResourceHandleOpSharedName sets the optional shared_name attribute to value. +// If not specified, defaults to "" +func BoostedTreesEnsembleResourceHandleOpSharedName(value string) BoostedTreesEnsembleResourceHandleOpAttr { + return func(m optionalAttr) { + m["shared_name"] = value + } +} + +// Creates a handle to a BoostedTreesEnsembleResource +func BoostedTreesEnsembleResourceHandleOp(scope *Scope, optional ...BoostedTreesEnsembleResourceHandleOpAttr) (resource tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "BoostedTreesEnsembleResourceHandleOp", + + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Concatenates tensors along one dimension. +// +// Arguments: +// concat_dim: 0-D. The dimension along which to concatenate. Must be in the +// range [0, rank(values)). +// values: The `N` Tensors to concatenate. Their ranks and types must match, +// and their sizes must match in all dimensions except `concat_dim`. +// +// 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 Concat(scope *Scope, concat_dim tf.Output, values []tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ Type: "Concat", Input: []tf.Input{ concat_dim, tf.OutputList(values), @@ -13477,6 +14917,30 @@ func MaxPoolGradGrad(scope *Scope, orig_input tf.Output, orig_output tf.Output, return op.Output(0) } +// Returns the last element of the input list as well as a list with all but that element. +// +// Fails if the list is empty. +// +// input_handle: the input list +// tensor: the withdrawn last element of the list +// element_dtype: the type of elements in the list +// element_shape: the shape of the output tensor +func TensorListPopBack(scope *Scope, input_handle tf.Output, element_dtype tf.DataType) (output_handle tf.Output, tensor tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"element_dtype": element_dtype} + opspec := tf.OpSpec{ + Type: "TensorListPopBack", + Input: []tf.Input{ + input_handle, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0), op.Output(1) +} + // Returns element-wise integer closest to x. // // If the result is midway between two representable values, @@ -13574,160 +15038,6 @@ func OrderedMapUnstageNoKey(scope *Scope, indices tf.Output, dtypes []tf.DataTyp return key, values } -// MaxPool3DGradGradAttr is an optional argument to MaxPool3DGradGrad. -type MaxPool3DGradGradAttr func(optionalAttr) - -// MaxPool3DGradGradDataFormat sets the optional data_format attribute to value. -// -// value: The data format of the input and output data. With the -// default format "NDHWC", the data is stored in the order of: -// [batch, in_depth, in_height, in_width, in_channels]. -// Alternatively, the format could be "NCDHW", the data storage order is: -// [batch, in_channels, in_depth, in_height, in_width]. -// If not specified, defaults to "NDHWC" -func MaxPool3DGradGradDataFormat(value string) MaxPool3DGradGradAttr { - 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: Output backprop of shape `[batch, depth, rows, cols, channels]`. -// ksize: 1-D tensor of length 5. The size of the window for each dimension of -// the input tensor. Must have `ksize[0] = ksize[4] = 1`. -// strides: 1-D tensor of length 5. The stride of the sliding window for each -// dimension of `input`. Must have `strides[0] = strides[4] = 1`. -// padding: The type of padding algorithm to use. -// -// Returns Gradients of gradients w.r.t. the input to `max_pool`. -func MaxPool3DGradGrad(scope *Scope, orig_input tf.Output, orig_output tf.Output, grad tf.Output, ksize []int64, strides []int64, padding string, optional ...MaxPool3DGradGradAttr) (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: "MaxPool3DGradGrad", - Input: []tf.Input{ - orig_input, orig_output, grad, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// Conv3DBackpropFilterV2Attr is an optional argument to Conv3DBackpropFilterV2. -type Conv3DBackpropFilterV2Attr func(optionalAttr) - -// Conv3DBackpropFilterV2DataFormat sets the optional data_format attribute to value. -// -// value: The data format of the input and output data. With the -// default format "NDHWC", the data is stored in the order of: -// [batch, in_depth, in_height, in_width, in_channels]. -// Alternatively, the format could be "NCDHW", the data storage order is: -// [batch, in_channels, in_depth, in_height, in_width]. -// If not specified, defaults to "NDHWC" -func Conv3DBackpropFilterV2DataFormat(value string) Conv3DBackpropFilterV2Attr { - return func(m optionalAttr) { - m["data_format"] = value - } -} - -// Conv3DBackpropFilterV2Dilations sets the optional dilations attribute to value. -// -// value: 1-D tensor of length 5. The dilation factor for each dimension of -// `input`. If set to k > 1, there will be k-1 skipped cells between each -// filter element on that dimension. The dimension order is determined by the -// value of `data_format`, see above for details. Dilations in the batch and -// depth dimensions must be 1. -// If not specified, defaults to -func Conv3DBackpropFilterV2Dilations(value []int64) Conv3DBackpropFilterV2Attr { - return func(m optionalAttr) { - m["dilations"] = value - } -} - -// Computes the gradients of 3-D convolution with respect to the filter. -// -// Arguments: -// input: Shape `[batch, depth, rows, cols, in_channels]`. -// filter_sizes: An integer vector representing the tensor shape of `filter`, -// where `filter` is a 5-D -// `[filter_depth, filter_height, filter_width, in_channels, out_channels]` -// tensor. -// out_backprop: Backprop signal of shape `[batch, out_depth, out_rows, out_cols, -// out_channels]`. -// strides: 1-D tensor of length 5. The stride of the sliding window for each -// dimension of `input`. Must have `strides[0] = strides[4] = 1`. -// padding: The type of padding algorithm to use. -func Conv3DBackpropFilterV2(scope *Scope, input tf.Output, filter_sizes tf.Output, out_backprop tf.Output, strides []int64, padding string, optional ...Conv3DBackpropFilterV2Attr) (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: "Conv3DBackpropFilterV2", - Input: []tf.Input{ - input, filter_sizes, out_backprop, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// Execute a sub graph on a remote processor. -// -// The graph specifications(such as graph itself, input tensors and output names) -// are stored as a serialized protocol buffer of RemoteFusedGraphExecuteInfo -// as serialized_remote_fused_graph_execute_info. -// The specifications will be passed to a dedicated registered -// remote fused graph executor. The executor will send the graph specifications -// to a remote processor and execute that graph. The execution results -// will be passed to consumer nodes as outputs of this node. -// -// Arguments: -// inputs: Arbitrary number of tensors with arbitrary data types -// -// serialized_remote_fused_graph_execute_info: Serialized protocol buffer -// of RemoteFusedGraphExecuteInfo which contains graph specifications. -// -// Returns Arbitrary number of tensors with arbitrary data types -func RemoteFusedGraphExecute(scope *Scope, inputs []tf.Output, Toutputs []tf.DataType, serialized_remote_fused_graph_execute_info string) (outputs []tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{"Toutputs": Toutputs, "serialized_remote_fused_graph_execute_info": serialized_remote_fused_graph_execute_info} - opspec := tf.OpSpec{ - Type: "RemoteFusedGraphExecute", - Input: []tf.Input{ - tf.OutputList(inputs), - }, - 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("RemoteFusedGraphExecute", err) - return - } - return outputs -} - // SerializeManySparseAttr is an optional argument to SerializeManySparse. type SerializeManySparseAttr func(optionalAttr) @@ -13882,6 +15192,14 @@ func DecodeCSVNaValue(value string) DecodeCSVAttr { } } +// DecodeCSVSelectCols sets the optional select_cols attribute to value. +// If not specified, defaults to <> +func DecodeCSVSelectCols(value []int64) DecodeCSVAttr { + return func(m optionalAttr) { + m["select_cols"] = value + } +} + // Convert CSV records to tensors. Each column maps to one tensor. // // RFC 4180 format is expected for the CSV records. @@ -14199,41 +15517,239 @@ 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 +// Returns element-wise remainder of division. This emulates C semantics in that // -// 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 +// the result here is consistent with a truncating divide. E.g. `truncate(x / y) * +// y + truncate_mod(x, y) = x`. // -// REQUIRES: value >= 0 -func StageSizeMemoryLimit(value int64) StageSizeAttr { - return func(m optionalAttr) { - m["memory_limit"] = value +// *NOTE*: `TruncateMod` supports broadcasting. More about broadcasting +// [here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html) +func TruncateMod(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) { + if scope.Err() != nil { + return } -} - -// 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 + opspec := tf.OpSpec{ + Type: "TruncateMod", + Input: []tf.Input{ + x, y, + }, } + op := scope.AddOperation(opspec) + return op.Output(0) } -// StageSizeSharedName sets the optional shared_name attribute to value. -// If not specified, defaults to "" -func StageSizeSharedName(value string) StageSizeAttr { - return func(m optionalAttr) { +// Inverse 2D real-valued fast Fourier transform. +// +// Computes the inverse 2-dimensional discrete Fourier transform of a real-valued +// signal over the inner-most 2 dimensions of `input`. +// +// The inner-most 2 dimensions of `input` are assumed to be the result of `RFFT2D`: +// The inner-most dimension contains the `fft_length / 2 + 1` unique components of +// the DFT of a real-valued signal. If `fft_length` is not provided, it is computed +// from the size of the inner-most 2 dimensions of `input`. If the FFT length used +// to compute `input` is odd, it should be provided since it cannot be inferred +// properly. +// +// Along each axis `IRFFT2D` is computed on, if `fft_length` (or +// `fft_length / 2 + 1` for the inner-most dimension) is smaller than the +// corresponding dimension of `input`, the dimension is cropped. If it is larger, +// the dimension is padded with zeros. +// +// Arguments: +// input: A complex64 tensor. +// fft_length: An int32 tensor of shape [2]. The FFT length for each dimension. +// +// Returns A float32 tensor of the same rank as `input`. The inner-most 2 +// dimensions of `input` are replaced with the `fft_length` samples of their +// inverse 2D Fourier transform. +// +// @compatibility(numpy) +// Equivalent to np.fft.irfft2 +// @end_compatibility +func IRFFT2D(scope *Scope, input tf.Output, fft_length tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "IRFFT2D", + Input: []tf.Input{ + input, fft_length, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// DecodeJpegAttr is an optional argument to DecodeJpeg. +type DecodeJpegAttr func(optionalAttr) + +// DecodeJpegChannels sets the optional channels attribute to value. +// +// value: Number of color channels for the decoded image. +// If not specified, defaults to 0 +func DecodeJpegChannels(value int64) DecodeJpegAttr { + return func(m optionalAttr) { + m["channels"] = value + } +} + +// DecodeJpegRatio sets the optional ratio attribute to value. +// +// value: Downscaling ratio. +// If not specified, defaults to 1 +func DecodeJpegRatio(value int64) DecodeJpegAttr { + return func(m optionalAttr) { + m["ratio"] = value + } +} + +// DecodeJpegFancyUpscaling sets the optional fancy_upscaling attribute to value. +// +// value: If true use a slower but nicer upscaling of the +// chroma planes (yuv420/422 only). +// If not specified, defaults to true +func DecodeJpegFancyUpscaling(value bool) DecodeJpegAttr { + return func(m optionalAttr) { + m["fancy_upscaling"] = value + } +} + +// DecodeJpegTryRecoverTruncated sets the optional try_recover_truncated attribute to value. +// +// value: If true try to recover an image from truncated input. +// If not specified, defaults to false +func DecodeJpegTryRecoverTruncated(value bool) DecodeJpegAttr { + return func(m optionalAttr) { + m["try_recover_truncated"] = value + } +} + +// DecodeJpegAcceptableFraction sets the optional acceptable_fraction attribute to value. +// +// value: The minimum required fraction of lines before a truncated +// input is accepted. +// If not specified, defaults to 1 +func DecodeJpegAcceptableFraction(value float32) DecodeJpegAttr { + return func(m optionalAttr) { + m["acceptable_fraction"] = value + } +} + +// DecodeJpegDctMethod sets the optional dct_method attribute to value. +// +// value: string specifying a hint about the algorithm used for +// decompression. Defaults to "" which maps to a system-specific +// default. Currently valid values are ["INTEGER_FAST", +// "INTEGER_ACCURATE"]. The hint may be ignored (e.g., the internal +// jpeg library changes to a version that does not have that specific +// option.) +// If not specified, defaults to "" +func DecodeJpegDctMethod(value string) DecodeJpegAttr { + return func(m optionalAttr) { + m["dct_method"] = value + } +} + +// Decode a JPEG-encoded image to a uint8 tensor. +// +// The attr `channels` indicates the desired number of color channels for the +// decoded image. +// +// Accepted values are: +// +// * 0: Use the number of channels in the JPEG-encoded image. +// * 1: output a grayscale image. +// * 3: output an RGB image. +// +// If needed, the JPEG-encoded image is transformed to match the requested number +// of color channels. +// +// The attr `ratio` allows downscaling the image by an integer factor during +// decoding. Allowed values are: 1, 2, 4, and 8. This is much faster than +// downscaling the image later. +// +// +// This op also supports decoding PNGs and non-animated GIFs since the interface is +// the same, though it is cleaner to use `tf.image.decode_image`. +// +// Arguments: +// contents: 0-D. The JPEG-encoded image. +// +// Returns 3-D with shape `[height, width, channels]`.. +func DecodeJpeg(scope *Scope, contents tf.Output, optional ...DecodeJpegAttr) (image tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "DecodeJpeg", + Input: []tf.Input{ + contents, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Serializes the tree ensemble to a proto. +// +// Arguments: +// tree_ensemble_handle: Handle to the tree ensemble. +// +// Returns Stamp token of the tree ensemble resource.Serialized proto of the ensemble. +func BoostedTreesSerializeEnsemble(scope *Scope, tree_ensemble_handle tf.Output) (stamp_token tf.Output, tree_ensemble_serialized tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "BoostedTreesSerializeEnsemble", + Input: []tf.Input{ + tree_ensemble_handle, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0), op.Output(1) +} + +// 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 } } @@ -14340,7 +15856,7 @@ func ReaderNumRecordsProducedV2(scope *Scope, reader_handle tf.Output) (records_ // Computes the sum along segments of a tensor. // -// Read @{$math_ops#segmentation$the section on segmentation} for an explanation of +// Read @{$math_ops#Segmentation$the section on segmentation} for an explanation of // segments. // // Computes a tensor such that @@ -14396,6 +15912,99 @@ func TextLineDataset(scope *Scope, filenames tf.Output, compression_type tf.Outp return op.Output(0) } +// CudnnRNNParamsSizeAttr is an optional argument to CudnnRNNParamsSize. +type CudnnRNNParamsSizeAttr func(optionalAttr) + +// CudnnRNNParamsSizeRnnMode sets the optional rnn_mode attribute to value. +// If not specified, defaults to "lstm" +func CudnnRNNParamsSizeRnnMode(value string) CudnnRNNParamsSizeAttr { + return func(m optionalAttr) { + m["rnn_mode"] = value + } +} + +// CudnnRNNParamsSizeInputMode sets the optional input_mode attribute to value. +// If not specified, defaults to "linear_input" +func CudnnRNNParamsSizeInputMode(value string) CudnnRNNParamsSizeAttr { + return func(m optionalAttr) { + m["input_mode"] = value + } +} + +// CudnnRNNParamsSizeDirection sets the optional direction attribute to value. +// If not specified, defaults to "unidirectional" +func CudnnRNNParamsSizeDirection(value string) CudnnRNNParamsSizeAttr { + return func(m optionalAttr) { + m["direction"] = value + } +} + +// CudnnRNNParamsSizeDropout sets the optional dropout attribute to value. +// If not specified, defaults to 0 +func CudnnRNNParamsSizeDropout(value float32) CudnnRNNParamsSizeAttr { + return func(m optionalAttr) { + m["dropout"] = value + } +} + +// CudnnRNNParamsSizeSeed sets the optional seed attribute to value. +// If not specified, defaults to 0 +func CudnnRNNParamsSizeSeed(value int64) CudnnRNNParamsSizeAttr { + return func(m optionalAttr) { + m["seed"] = value + } +} + +// CudnnRNNParamsSizeSeed2 sets the optional seed2 attribute to value. +// If not specified, defaults to 0 +func CudnnRNNParamsSizeSeed2(value int64) CudnnRNNParamsSizeAttr { + return func(m optionalAttr) { + m["seed2"] = value + } +} + +// Computes size of weights that can be used by a Cudnn RNN model. +// +// Return the params size that can be used by the Cudnn RNN model. Subsequent +// weight allocation and initialization should use this size. +// +// num_layers: Specifies the number of layers in the RNN model. +// num_units: Specifies the size of the hidden state. +// input_size: Specifies the size of the input state. +// rnn_mode: Indicates the type of the RNN model. +// input_mode: Indicate whether there is a linear projection between the input and +// The actual computation before the first layer. 'skip_input' is only allowed +// when input_size == num_units; 'auto_select' implies 'skip_input' when +// input_size == num_units; otherwise, it implies 'linear_input'. +// direction: Indicates whether a bidirectional model will be used. +// dir = (direction == bidirectional) ? 2 : 1 +// dropout: dropout probability. When set to 0., dropout is disabled. +// seed: the 1st part of a seed to initialize dropout. +// seed2: the 2nd part of a seed to initialize dropout. +// params_size: The size of the params buffer that should be allocated and +// initialized for this RNN model. Note that this params buffer may not be +// compatible across GPUs. Please use CudnnRNNParamsWeights and +// CudnnRNNParamsBiases to save and restore them in a way that is compatible +// across different runs. +func CudnnRNNParamsSize(scope *Scope, num_layers tf.Output, num_units tf.Output, input_size tf.Output, T tf.DataType, S tf.DataType, optional ...CudnnRNNParamsSizeAttr) (params_size tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"T": T, "S": S} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "CudnnRNNParamsSize", + Input: []tf.Input{ + num_layers, num_units, input_size, + }, + 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 @@ -14420,7 +16029,31 @@ func SparseSegmentMeanGrad(scope *Scope, grad tf.Output, indices tf.Output, segm return op.Output(0) } -// Returns the truth value of (x >= y) element-wise. +// Returns the set of files matching one or more glob patterns. +// +// Note that this routine only supports wildcard characters in the +// basename portion of the pattern, not in the directory portion. +// Note also that the order of filenames returned can be non-deterministic. +// +// Arguments: +// pattern: Shell wildcard pattern(s). Scalar or vector of type string. +// +// Returns A vector of matching filenames. +func MatchingFiles(scope *Scope, pattern tf.Output) (filenames tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "MatchingFiles", + Input: []tf.Input{ + pattern, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Returns the truth value of (x >= y) element-wise. // // *NOTE*: `GreaterEqual` supports broadcasting. More about broadcasting // [here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html) @@ -14821,6 +16454,46 @@ func QuantizedConv2D(scope *Scope, input tf.Output, filter tf.Output, min_input return op.Output(0), op.Output(1), op.Output(2) } +// StatelessMultinomialAttr is an optional argument to StatelessMultinomial. +type StatelessMultinomialAttr func(optionalAttr) + +// StatelessMultinomialOutputDtype sets the optional output_dtype attribute to value. +// If not specified, defaults to DT_INT64 +func StatelessMultinomialOutputDtype(value tf.DataType) StatelessMultinomialAttr { + return func(m optionalAttr) { + m["output_dtype"] = value + } +} + +// Draws samples from a multinomial distribution. +// +// Arguments: +// logits: 2-D Tensor with shape `[batch_size, num_classes]`. Each slice `[i, :]` +// represents the unnormalized log probabilities for all classes. +// num_samples: 0-D. Number of independent samples to draw for each row slice. +// seed: 2 seeds (shape [2]). +// +// Returns 2-D Tensor with shape `[batch_size, num_samples]`. Each slice `[i, :]` +// contains the drawn class labels with range `[0, num_classes)`. +func StatelessMultinomial(scope *Scope, logits tf.Output, num_samples tf.Output, seed tf.Output, optional ...StatelessMultinomialAttr) (output tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "StatelessMultinomial", + Input: []tf.Input{ + logits, num_samples, seed, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // ResourceGatherAttr is an optional argument to ResourceGather. type ResourceGatherAttr func(optionalAttr) @@ -15165,366 +16838,21 @@ func AllCandidateSamplerSeed2(value int64) AllCandidateSamplerAttr { // the ID of a sampled candidate.A batch_size * num_true matrix, representing // the number of times each candidate is expected to occur in a batch // of sampled candidates. If unique=true, then this is a probability.A vector of length num_sampled, for each sampled -// candidate representing the number of times the candidate is expected -// to occur in a batch of sampled candidates. If unique=true, then this is a -// probability. -func AllCandidateSampler(scope *Scope, true_classes tf.Output, num_true int64, num_sampled int64, unique bool, optional ...AllCandidateSamplerAttr) (sampled_candidates tf.Output, true_expected_count tf.Output, sampled_expected_count tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{"num_true": num_true, "num_sampled": num_sampled, "unique": unique} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "AllCandidateSampler", - Input: []tf.Input{ - true_classes, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0), op.Output(1), op.Output(2) -} - -// Saves the input tensors to disk. -// -// The size of `tensor_names` must match the number of tensors in `data`. `data[i]` -// is written to `filename` with name `tensor_names[i]`. -// -// See also `SaveSlices`. -// -// Arguments: -// filename: Must have a single element. The name of the file to which we write -// the tensor. -// tensor_names: Shape `[N]`. The names of the tensors to be saved. -// data: `N` tensors to save. -// -// Returns the created operation. -func Save(scope *Scope, filename tf.Output, tensor_names tf.Output, data []tf.Output) (o *tf.Operation) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "Save", - Input: []tf.Input{ - filename, tensor_names, tf.OutputList(data), - }, - } - return scope.AddOperation(opspec) -} - -// Returns element-wise remainder of division. When `x < 0` xor `y < 0` is -// -// true, this follows Python semantics in that the result here is consistent -// with a flooring divide. E.g. `floor(x / y) * y + mod(x, y) = x`. -// -// *NOTE*: `FloorMod` supports broadcasting. More about broadcasting -// [here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html) -func FloorMod(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "FloorMod", - Input: []tf.Input{ - x, y, - }, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// SparseTensorDenseMatMulAttr is an optional argument to SparseTensorDenseMatMul. -type SparseTensorDenseMatMulAttr func(optionalAttr) - -// SparseTensorDenseMatMulAdjointA sets the optional adjoint_a attribute to value. -// -// value: Use the adjoint of A in the matrix multiply. If A is complex, this -// is transpose(conj(A)). Otherwise it's transpose(A). -// If not specified, defaults to false -func SparseTensorDenseMatMulAdjointA(value bool) SparseTensorDenseMatMulAttr { - return func(m optionalAttr) { - m["adjoint_a"] = value - } -} - -// SparseTensorDenseMatMulAdjointB sets the optional adjoint_b attribute to value. -// -// value: Use the adjoint of B in the matrix multiply. If B is complex, this -// is transpose(conj(B)). Otherwise it's transpose(B). -// If not specified, defaults to false -func SparseTensorDenseMatMulAdjointB(value bool) SparseTensorDenseMatMulAttr { - return func(m optionalAttr) { - m["adjoint_b"] = value - } -} - -// Multiply SparseTensor (of rank 2) "A" by dense matrix "B". -// -// No validity checking is performed on the indices of A. However, the following -// input format is recommended for optimal behavior: -// -// if adjoint_a == false: -// A should be sorted in lexicographically increasing order. Use SparseReorder -// if you're not sure. -// if adjoint_a == true: -// A should be sorted in order of increasing dimension 1 (i.e., "column major" -// order instead of "row major" order). -// -// Arguments: -// a_indices: 2-D. The `indices` of the `SparseTensor`, size `[nnz, 2]` Matrix. -// a_values: 1-D. The `values` of the `SparseTensor`, size `[nnz]` Vector. -// a_shape: 1-D. The `shape` of the `SparseTensor`, size `[2]` Vector. -// b: 2-D. A dense Matrix. -func SparseTensorDenseMatMul(scope *Scope, a_indices tf.Output, a_values tf.Output, a_shape tf.Output, b tf.Output, optional ...SparseTensorDenseMatMulAttr) (product tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "SparseTensorDenseMatMul", - Input: []tf.Input{ - a_indices, a_values, a_shape, b, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// Deserialize and concatenate `SparseTensors` from a serialized minibatch. -// -// The input `serialized_sparse` must be a string matrix of shape `[N x 3]` where -// `N` is the minibatch size and the rows correspond to packed outputs of -// `SerializeSparse`. The ranks of the original `SparseTensor` objects -// must all match. When the final `SparseTensor` is created, it has rank one -// higher than the ranks of the incoming `SparseTensor` objects -// (they have been concatenated along a new row dimension). -// -// The output `SparseTensor` object's shape values for all dimensions but the -// first are the max across the input `SparseTensor` objects' shape values -// for the corresponding dimensions. Its first shape value is `N`, the minibatch -// size. -// -// The input `SparseTensor` objects' indices are assumed ordered in -// standard lexicographic order. If this is not the case, after this -// step run `SparseReorder` to restore index ordering. -// -// For example, if the serialized input is a `[2 x 3]` matrix representing two -// original `SparseTensor` objects: -// -// index = [ 0] -// [10] -// [20] -// values = [1, 2, 3] -// shape = [50] -// -// and -// -// index = [ 2] -// [10] -// values = [4, 5] -// shape = [30] -// -// then the final deserialized `SparseTensor` will be: -// -// index = [0 0] -// [0 10] -// [0 20] -// [1 2] -// [1 10] -// values = [1, 2, 3, 4, 5] -// shape = [2 50] -// -// Arguments: -// serialized_sparse: 2-D, The `N` serialized `SparseTensor` objects. -// Must have 3 columns. -// dtype: The `dtype` of the serialized `SparseTensor` objects. -func DeserializeManySparse(scope *Scope, serialized_sparse tf.Output, dtype tf.DataType) (sparse_indices tf.Output, sparse_values tf.Output, sparse_shape tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{"dtype": dtype} - opspec := tf.OpSpec{ - Type: "DeserializeManySparse", - Input: []tf.Input{ - serialized_sparse, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0), op.Output(1), op.Output(2) -} - -// StringJoinAttr is an optional argument to StringJoin. -type StringJoinAttr func(optionalAttr) - -// StringJoinSeparator sets the optional separator attribute to value. -// -// value: string, an optional join separator. -// If not specified, defaults to "" -func StringJoinSeparator(value string) StringJoinAttr { - return func(m optionalAttr) { - m["separator"] = value - } -} - -// Joins the strings in the given list of string tensors into one tensor; -// -// with the given separator (default is an empty separator). -// -// Arguments: -// inputs: A list of string tensors. The tensors must all have the same shape, -// or be scalars. Scalars may be mixed in; these will be broadcast to the shape -// of non-scalar inputs. -func StringJoin(scope *Scope, inputs []tf.Output, optional ...StringJoinAttr) (output tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "StringJoin", - Input: []tf.Input{ - tf.OutputList(inputs), - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// Returns immutable tensor from memory region. -// -// The current implementation memmaps the tensor from a file. -// -// Arguments: -// dtype: Type of the returned tensor. -// shape: Shape of the returned tensor. -// memory_region_name: Name of readonly memory region used by the tensor, see -// NewReadOnlyMemoryRegionFromFile in tensorflow::Env. -func ImmutableConst(scope *Scope, dtype tf.DataType, shape tf.Shape, memory_region_name string) (tensor tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{"dtype": dtype, "shape": shape, "memory_region_name": memory_region_name} - opspec := tf.OpSpec{ - Type: "ImmutableConst", - - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// Inverse real-valued fast Fourier transform. -// -// Computes the inverse 1-dimensional discrete Fourier transform of a real-valued -// signal over the inner-most dimension of `input`. -// -// The inner-most dimension of `input` is assumed to be the result of `RFFT`: the -// `fft_length / 2 + 1` unique components of the DFT of a real-valued signal. If -// `fft_length` is not provided, it is computed from the size of the inner-most -// dimension of `input` (`fft_length = 2 * (inner - 1)`). If the FFT length used to -// compute `input` is odd, it should be provided since it cannot be inferred -// properly. -// -// Along the axis `IRFFT` is computed on, if `fft_length / 2 + 1` is smaller -// than the corresponding dimension of `input`, the dimension is cropped. If it is -// larger, the dimension is padded with zeros. -// -// Arguments: -// input: A complex64 tensor. -// fft_length: An int32 tensor of shape [1]. The FFT length. -// -// Returns A float32 tensor of the same rank as `input`. The inner-most -// dimension of `input` is replaced with the `fft_length` samples of its inverse -// 1D Fourier transform. -// -// @compatibility(numpy) -// Equivalent to np.fft.irfft -// @end_compatibility -func IRFFT(scope *Scope, input tf.Output, fft_length tf.Output) (output tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "IRFFT", - Input: []tf.Input{ - input, fft_length, - }, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// Concatenates a list of `SparseTensor` along the specified dimension. -// -// Concatenation is with respect to the dense versions of these sparse tensors. -// It is assumed that each input is a `SparseTensor` whose elements are ordered -// along increasing dimension number. -// -// All inputs' shapes must match, except for the concat dimension. The -// `indices`, `values`, and `shapes` lists must have the same length. -// -// The output shape is identical to the inputs', except along the concat -// dimension, where it is the sum of the inputs' sizes along that dimension. -// -// The output elements will be resorted to preserve the sort order along -// increasing dimension number. -// -// This op runs in `O(M log M)` time, where `M` is the total number of non-empty -// values across all inputs. This is due to the need for an internal sort in -// order to concatenate efficiently across an arbitrary dimension. -// -// For example, if `concat_dim = 1` and the inputs are -// -// sp_inputs[0]: shape = [2, 3] -// [0, 2]: "a" -// [1, 0]: "b" -// [1, 1]: "c" -// -// sp_inputs[1]: shape = [2, 4] -// [0, 1]: "d" -// [0, 2]: "e" -// -// then the output will be -// -// shape = [2, 7] -// [0, 2]: "a" -// [0, 4]: "d" -// [0, 5]: "e" -// [1, 0]: "b" -// [1, 1]: "c" -// -// Graphically this is equivalent to doing -// -// [ a] concat [ d e ] = [ a d e ] -// [b c ] [ ] [b c ] -// -// Arguments: -// indices: 2-D. Indices of each input `SparseTensor`. -// values: 1-D. Non-empty values of each `SparseTensor`. -// shapes: 1-D. Shapes of each `SparseTensor`. -// concat_dim: Dimension to concatenate along. Must be in range [-rank, rank), -// where rank is the number of dimensions in each input `SparseTensor`. -// -// Returns 2-D. Indices of the concatenated `SparseTensor`.1-D. Non-empty values of the concatenated `SparseTensor`.1-D. Shape of the concatenated `SparseTensor`. -func SparseConcat(scope *Scope, indices []tf.Output, values []tf.Output, shapes []tf.Output, concat_dim int64) (output_indices tf.Output, output_values tf.Output, output_shape tf.Output) { +// candidate representing the number of times the candidate is expected +// to occur in a batch of sampled candidates. If unique=true, then this is a +// probability. +func AllCandidateSampler(scope *Scope, true_classes tf.Output, num_true int64, num_sampled int64, unique bool, optional ...AllCandidateSamplerAttr) (sampled_candidates tf.Output, true_expected_count tf.Output, sampled_expected_count tf.Output) { if scope.Err() != nil { return } - attrs := map[string]interface{}{"concat_dim": concat_dim} + attrs := map[string]interface{}{"num_true": num_true, "num_sampled": num_sampled, "unique": unique} + for _, a := range optional { + a(attrs) + } opspec := tf.OpSpec{ - Type: "SparseConcat", + Type: "AllCandidateSampler", Input: []tf.Input{ - tf.OutputList(indices), tf.OutputList(values), tf.OutputList(shapes), + true_classes, }, Attrs: attrs, } @@ -15532,358 +16860,415 @@ func SparseConcat(scope *Scope, indices []tf.Output, values []tf.Output, shapes return op.Output(0), op.Output(1), op.Output(2) } -// Generates sparse cross from a list of sparse and dense tensors. -// -// The op takes two lists, one of 2D `SparseTensor` and one of 2D `Tensor`, each -// representing features of one feature column. It outputs a 2D `SparseTensor` with -// the batchwise crosses of these features. -// -// For example, if the inputs are -// -// inputs[0]: SparseTensor with shape = [2, 2] -// [0, 0]: "a" -// [1, 0]: "b" -// [1, 1]: "c" -// -// inputs[1]: SparseTensor with shape = [2, 1] -// [0, 0]: "d" -// [1, 0]: "e" -// -// inputs[2]: Tensor [["f"], ["g"]] -// -// then the output will be -// -// shape = [2, 2] -// [0, 0]: "a_X_d_X_f" -// [1, 0]: "b_X_e_X_g" -// [1, 1]: "c_X_e_X_g" +// Saves the input tensors to disk. // -// if hashed_output=true then the output will be +// The size of `tensor_names` must match the number of tensors in `data`. `data[i]` +// is written to `filename` with name `tensor_names[i]`. // -// shape = [2, 2] -// [0, 0]: FingerprintCat64( -// Fingerprint64("f"), FingerprintCat64( -// Fingerprint64("d"), Fingerprint64("a"))) -// [1, 0]: FingerprintCat64( -// Fingerprint64("g"), FingerprintCat64( -// Fingerprint64("e"), Fingerprint64("b"))) -// [1, 1]: FingerprintCat64( -// Fingerprint64("g"), FingerprintCat64( -// Fingerprint64("e"), Fingerprint64("c"))) +// See also `SaveSlices`. // // Arguments: -// indices: 2-D. Indices of each input `SparseTensor`. -// values: 1-D. values of each `SparseTensor`. -// shapes: 1-D. Shapes of each `SparseTensor`. -// dense_inputs: 2-D. Columns represented by dense `Tensor`. -// hashed_output: If true, returns the hash of the cross instead of the string. -// This will allow us avoiding string manipulations. -// num_buckets: It is used if hashed_output is true. -// output = hashed_value%num_buckets if num_buckets > 0 else hashed_value. -// hash_key: Specify the hash_key that will be used by the `FingerprintCat64` -// function to combine the crosses fingerprints. +// filename: Must have a single element. The name of the file to which we write +// the tensor. +// tensor_names: Shape `[N]`. The names of the tensors to be saved. +// data: `N` tensors to save. // +// Returns the created operation. +func Save(scope *Scope, filename tf.Output, tensor_names tf.Output, data []tf.Output) (o *tf.Operation) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "Save", + Input: []tf.Input{ + filename, tensor_names, tf.OutputList(data), + }, + } + return scope.AddOperation(opspec) +} + +// Returns element-wise remainder of division. When `x < 0` xor `y < 0` is // +// true, this follows Python semantics in that the result here is consistent +// with a flooring divide. E.g. `floor(x / y) * y + mod(x, y) = x`. // -// Returns 2-D. Indices of the concatenated `SparseTensor`.1-D. Non-empty values of the concatenated or hashed -// `SparseTensor`.1-D. Shape of the concatenated `SparseTensor`. -func SparseCross(scope *Scope, indices []tf.Output, values []tf.Output, shapes []tf.Output, dense_inputs []tf.Output, hashed_output bool, num_buckets int64, hash_key int64, out_type tf.DataType, internal_type tf.DataType) (output_indices tf.Output, output_values tf.Output, output_shape tf.Output) { +// *NOTE*: `FloorMod` supports broadcasting. More about broadcasting +// [here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html) +func FloorMod(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) { if scope.Err() != nil { return } - attrs := map[string]interface{}{"hashed_output": hashed_output, "num_buckets": num_buckets, "hash_key": hash_key, "out_type": out_type, "internal_type": internal_type} opspec := tf.OpSpec{ - Type: "SparseCross", + Type: "FloorMod", Input: []tf.Input{ - tf.OutputList(indices), tf.OutputList(values), tf.OutputList(shapes), tf.OutputList(dense_inputs), + x, y, }, - Attrs: attrs, } op := scope.AddOperation(opspec) - return op.Output(0), op.Output(1), op.Output(2) + return op.Output(0) } -// Concatenates quantized tensors along one dimension. +// SparseTensorDenseMatMulAttr is an optional argument to SparseTensorDenseMatMul. +type SparseTensorDenseMatMulAttr func(optionalAttr) + +// SparseTensorDenseMatMulAdjointA sets the optional adjoint_a attribute to value. // -// Arguments: -// concat_dim: 0-D. The dimension along which to concatenate. Must be in the -// range [0, rank(values)). -// values: The `N` Tensors to concatenate. Their ranks and types must match, -// and their sizes must match in all dimensions except `concat_dim`. -// input_mins: The minimum scalar values for each of the input tensors. -// input_maxes: The maximum scalar values for each of the input tensors. +// value: Use the adjoint of A in the matrix multiply. If A is complex, this +// is transpose(conj(A)). Otherwise it's transpose(A). +// If not specified, defaults to false +func SparseTensorDenseMatMulAdjointA(value bool) SparseTensorDenseMatMulAttr { + return func(m optionalAttr) { + m["adjoint_a"] = value + } +} + +// SparseTensorDenseMatMulAdjointB sets the optional adjoint_b attribute to value. // -// 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.The float value that the minimum quantized output value represents.The float value that the maximum quantized output value represents. -func QuantizedConcat(scope *Scope, concat_dim tf.Output, values []tf.Output, input_mins []tf.Output, input_maxes []tf.Output) (output tf.Output, output_min tf.Output, output_max tf.Output) { +// value: Use the adjoint of B in the matrix multiply. If B is complex, this +// is transpose(conj(B)). Otherwise it's transpose(B). +// If not specified, defaults to false +func SparseTensorDenseMatMulAdjointB(value bool) SparseTensorDenseMatMulAttr { + return func(m optionalAttr) { + m["adjoint_b"] = value + } +} + +// Multiply SparseTensor (of rank 2) "A" by dense matrix "B". +// +// No validity checking is performed on the indices of A. However, the following +// input format is recommended for optimal behavior: +// +// if adjoint_a == false: +// A should be sorted in lexicographically increasing order. Use SparseReorder +// if you're not sure. +// if adjoint_a == true: +// A should be sorted in order of increasing dimension 1 (i.e., "column major" +// order instead of "row major" order). +// +// Arguments: +// a_indices: 2-D. The `indices` of the `SparseTensor`, size `[nnz, 2]` Matrix. +// a_values: 1-D. The `values` of the `SparseTensor`, size `[nnz]` Vector. +// a_shape: 1-D. The `shape` of the `SparseTensor`, size `[2]` Vector. +// b: 2-D. A dense Matrix. +func SparseTensorDenseMatMul(scope *Scope, a_indices tf.Output, a_values tf.Output, a_shape tf.Output, b tf.Output, optional ...SparseTensorDenseMatMulAttr) (product tf.Output) { if scope.Err() != nil { return } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } opspec := tf.OpSpec{ - Type: "QuantizedConcat", + Type: "SparseTensorDenseMatMul", Input: []tf.Input{ - concat_dim, tf.OutputList(values), tf.OutputList(input_mins), tf.OutputList(input_maxes), + a_indices, a_values, a_shape, b, }, + Attrs: attrs, } op := scope.AddOperation(opspec) - return op.Output(0), op.Output(1), op.Output(2) + return op.Output(0) } -// Slice a `SparseTensor` based on the `start` and `size`. +// Deserialize and concatenate `SparseTensors` from a serialized minibatch. // -// For example, if the input is +// The input `serialized_sparse` must be a string matrix of shape `[N x 3]` where +// `N` is the minibatch size and the rows correspond to packed outputs of +// `SerializeSparse`. The ranks of the original `SparseTensor` objects +// must all match. When the final `SparseTensor` is created, it has rank one +// higher than the ranks of the incoming `SparseTensor` objects +// (they have been concatenated along a new row dimension). // -// input_tensor = shape = [2, 7] -// [ a d e ] -// [b c ] +// The output `SparseTensor` object's shape values for all dimensions but the +// first are the max across the input `SparseTensor` objects' shape values +// for the corresponding dimensions. Its first shape value is `N`, the minibatch +// size. // -// Graphically the output tensors are: +// The input `SparseTensor` objects' indices are assumed ordered in +// standard lexicographic order. If this is not the case, after this +// step run `SparseReorder` to restore index ordering. // -// sparse_slice([0, 0], [2, 4]) = shape = [2, 4] -// [ a ] -// [b c ] +// For example, if the serialized input is a `[2 x 3]` matrix representing two +// original `SparseTensor` objects: // -// sparse_slice([0, 4], [2, 3]) = shape = [2, 3] -// [ d e ] -// [ ] +// index = [ 0] +// [10] +// [20] +// values = [1, 2, 3] +// shape = [50] // -// Arguments: -// indices: 2-D tensor represents the indices of the sparse tensor. -// values: 1-D tensor represents the values of the sparse tensor. -// shape: 1-D. tensor represents the shape of the sparse tensor. -// start: 1-D. tensor represents the start of the slice. -// size: 1-D. tensor represents the size of the slice. -// output indices: A list of 1-D tensors represents the indices of the output -// sparse tensors. +// and // -// Returns A list of 1-D tensors represents the values of the output sparse -// tensors.A list of 1-D tensors represents the shape of the output sparse -// tensors. -func SparseSlice(scope *Scope, indices tf.Output, values tf.Output, shape tf.Output, start tf.Output, size tf.Output) (output_indices tf.Output, output_values tf.Output, output_shape tf.Output) { +// index = [ 2] +// [10] +// values = [4, 5] +// shape = [30] +// +// then the final deserialized `SparseTensor` will be: +// +// index = [0 0] +// [0 10] +// [0 20] +// [1 2] +// [1 10] +// values = [1, 2, 3, 4, 5] +// shape = [2 50] +// +// Arguments: +// serialized_sparse: 2-D, The `N` serialized `SparseTensor` objects. +// Must have 3 columns. +// dtype: The `dtype` of the serialized `SparseTensor` objects. +func DeserializeManySparse(scope *Scope, serialized_sparse tf.Output, dtype tf.DataType) (sparse_indices tf.Output, sparse_values tf.Output, sparse_shape tf.Output) { if scope.Err() != nil { return } + attrs := map[string]interface{}{"dtype": dtype} opspec := tf.OpSpec{ - Type: "SparseSlice", + Type: "DeserializeManySparse", Input: []tf.Input{ - indices, values, shape, start, size, + serialized_sparse, }, + Attrs: attrs, } op := scope.AddOperation(opspec) return op.Output(0), op.Output(1), op.Output(2) } -// Adds up a `SparseTensor` and a dense `Tensor`, producing a dense `Tensor`. +// StringJoinAttr is an optional argument to StringJoin. +type StringJoinAttr func(optionalAttr) + +// StringJoinSeparator sets the optional separator attribute to value. // -// This Op does not require `a_indices` be sorted in standard lexicographic order. +// value: string, an optional join separator. +// If not specified, defaults to "" +func StringJoinSeparator(value string) StringJoinAttr { + return func(m optionalAttr) { + m["separator"] = value + } +} + +// Joins the strings in the given list of string tensors into one tensor; +// +// with the given separator (default is an empty separator). // // Arguments: -// a_indices: 2-D. The `indices` of the `SparseTensor`, with shape `[nnz, ndims]`. -// a_values: 1-D. The `values` of the `SparseTensor`, with shape `[nnz]`. -// a_shape: 1-D. The `shape` of the `SparseTensor`, with shape `[ndims]`. -// b: `ndims`-D Tensor. With shape `a_shape`. -func SparseTensorDenseAdd(scope *Scope, a_indices tf.Output, a_values tf.Output, a_shape tf.Output, b tf.Output) (output tf.Output) { +// inputs: A list of string tensors. The tensors must all have the same shape, +// or be scalars. Scalars may be mixed in; these will be broadcast to the shape +// of non-scalar inputs. +func StringJoin(scope *Scope, inputs []tf.Output, optional ...StringJoinAttr) (output tf.Output) { if scope.Err() != nil { return } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } opspec := tf.OpSpec{ - Type: "SparseTensorDenseAdd", + Type: "StringJoin", Input: []tf.Input{ - a_indices, a_values, a_shape, b, + tf.OutputList(inputs), }, + Attrs: attrs, } op := scope.AddOperation(opspec) return op.Output(0) } -// Returns the set of files matching one or more glob patterns. +// Returns immutable tensor from memory region. // -// Note that this routine only supports wildcard characters in the -// basename portion of the pattern, not in the directory portion. +// The current implementation memmaps the tensor from a file. // // Arguments: -// pattern: Shell wildcard pattern(s). Scalar or vector of type string. -// -// Returns A vector of matching filenames. -func MatchingFiles(scope *Scope, pattern tf.Output) (filenames tf.Output) { +// dtype: Type of the returned tensor. +// shape: Shape of the returned tensor. +// memory_region_name: Name of readonly memory region used by the tensor, see +// NewReadOnlyMemoryRegionFromFile in tensorflow::Env. +func ImmutableConst(scope *Scope, dtype tf.DataType, shape tf.Shape, memory_region_name string) (tensor tf.Output) { if scope.Err() != nil { return } + attrs := map[string]interface{}{"dtype": dtype, "shape": shape, "memory_region_name": memory_region_name} opspec := tf.OpSpec{ - Type: "MatchingFiles", - Input: []tf.Input{ - pattern, - }, + Type: "ImmutableConst", + + Attrs: attrs, } op := scope.AddOperation(opspec) return op.Output(0) } -// MatrixSolveLsAttr is an optional argument to MatrixSolveLs. -type MatrixSolveLsAttr func(optionalAttr) - -// MatrixSolveLsFast sets the optional fast attribute to value. -// If not specified, defaults to true -func MatrixSolveLsFast(value bool) MatrixSolveLsAttr { - return func(m optionalAttr) { - m["fast"] = value - } -} - -// Solves one or more linear least-squares problems. -// -// `matrix` is a tensor of shape `[..., M, N]` whose inner-most 2 dimensions -// form real or complex matrices of size `[M, N]`. `Rhs` is a tensor of the same -// type as `matrix` and shape `[..., M, K]`. -// The output is a tensor shape `[..., N, K]` where each output matrix solves -// each of the equations -// `matrix[..., :, :]` * `output[..., :, :]` = `rhs[..., :, :]` -// in the least squares sense. -// -// We use the following notation for (complex) matrix and right-hand sides -// in the batch: +// Inverse real-valued fast Fourier transform. // -// `matrix`=\\(A \in \mathbb{C}^{m \times n}\\), -// `rhs`=\\(B \in \mathbb{C}^{m \times k}\\), -// `output`=\\(X \in \mathbb{C}^{n \times k}\\), -// `l2_regularizer`=\\(\lambda \in \mathbb{R}\\). +// Computes the inverse 1-dimensional discrete Fourier transform of a real-valued +// signal over the inner-most dimension of `input`. // -// If `fast` is `True`, then the solution is computed by solving the normal -// equations using Cholesky decomposition. Specifically, if \\(m \ge n\\) then -// \\(X = (A^H A + \lambda I)^{-1} A^H B\\), which solves the least-squares -// problem \\(X = \mathrm{argmin}_{Z \in \Re^{n \times k} } ||A Z - B||_F^2 + -// \lambda ||Z||_F^2\\). If \\(m \lt n\\) then `output` is computed as -// \\(X = A^H (A A^H + \lambda I)^{-1} B\\), which (for \\(\lambda = 0\\)) is the -// minimum-norm solution to the under-determined linear system, i.e. -// \\(X = \mathrm{argmin}_{Z \in \mathbb{C}^{n \times k} } ||Z||_F^2 \\), -// subject to \\(A Z = B\\). Notice that the fast path is only numerically stable -// when \\(A\\) is numerically full rank and has a condition number -// \\(\mathrm{cond}(A) \lt \frac{1}{\sqrt{\epsilon_{mach} } }\\) or\\(\lambda\\) is -// sufficiently large. +// The inner-most dimension of `input` is assumed to be the result of `RFFT`: the +// `fft_length / 2 + 1` unique components of the DFT of a real-valued signal. If +// `fft_length` is not provided, it is computed from the size of the inner-most +// dimension of `input` (`fft_length = 2 * (inner - 1)`). If the FFT length used to +// compute `input` is odd, it should be provided since it cannot be inferred +// properly. // -// If `fast` is `False` an algorithm based on the numerically robust complete -// orthogonal decomposition is used. This computes the minimum-norm -// least-squares solution, even when \\(A\\) is rank deficient. This path is -// typically 6-7 times slower than the fast path. If `fast` is `False` then -// `l2_regularizer` is ignored. +// Along the axis `IRFFT` is computed on, if `fft_length / 2 + 1` is smaller +// than the corresponding dimension of `input`, the dimension is cropped. If it is +// larger, the dimension is padded with zeros. // // Arguments: -// matrix: Shape is `[..., M, N]`. -// rhs: Shape is `[..., M, K]`. -// l2_regularizer: Scalar tensor. +// input: A complex64 tensor. +// fft_length: An int32 tensor of shape [1]. The FFT length. +// +// Returns A float32 tensor of the same rank as `input`. The inner-most +// dimension of `input` is replaced with the `fft_length` samples of its inverse +// 1D Fourier transform. // // @compatibility(numpy) -// Equivalent to np.linalg.lstsq +// Equivalent to np.fft.irfft // @end_compatibility -// -// Returns Shape is `[..., N, K]`. -func MatrixSolveLs(scope *Scope, matrix tf.Output, rhs tf.Output, l2_regularizer tf.Output, optional ...MatrixSolveLsAttr) (output tf.Output) { +func IRFFT(scope *Scope, input tf.Output, fft_length tf.Output) (output tf.Output) { if scope.Err() != nil { return } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } opspec := tf.OpSpec{ - Type: "MatrixSolveLs", + Type: "IRFFT", Input: []tf.Input{ - matrix, rhs, l2_regularizer, + input, fft_length, }, - Attrs: attrs, } op := scope.AddOperation(opspec) return op.Output(0) } -// Elementwise computes the bitwise OR of `x` and `y`. +// Concatenates a list of `SparseTensor` along the specified dimension. // -// The result will have those bits set, that are set in `x`, `y` or both. The -// computation is performed on the underlying representations of `x` and `y`. -func BitwiseOr(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) { +// Concatenation is with respect to the dense versions of these sparse tensors. +// It is assumed that each input is a `SparseTensor` whose elements are ordered +// along increasing dimension number. +// +// All inputs' shapes must match, except for the concat dimension. The +// `indices`, `values`, and `shapes` lists must have the same length. +// +// The output shape is identical to the inputs', except along the concat +// dimension, where it is the sum of the inputs' sizes along that dimension. +// +// The output elements will be resorted to preserve the sort order along +// increasing dimension number. +// +// This op runs in `O(M log M)` time, where `M` is the total number of non-empty +// values across all inputs. This is due to the need for an internal sort in +// order to concatenate efficiently across an arbitrary dimension. +// +// For example, if `concat_dim = 1` and the inputs are +// +// sp_inputs[0]: shape = [2, 3] +// [0, 2]: "a" +// [1, 0]: "b" +// [1, 1]: "c" +// +// sp_inputs[1]: shape = [2, 4] +// [0, 1]: "d" +// [0, 2]: "e" +// +// then the output will be +// +// shape = [2, 7] +// [0, 2]: "a" +// [0, 4]: "d" +// [0, 5]: "e" +// [1, 0]: "b" +// [1, 1]: "c" +// +// Graphically this is equivalent to doing +// +// [ a] concat [ d e ] = [ a d e ] +// [b c ] [ ] [b c ] +// +// Arguments: +// indices: 2-D. Indices of each input `SparseTensor`. +// values: 1-D. Non-empty values of each `SparseTensor`. +// shapes: 1-D. Shapes of each `SparseTensor`. +// concat_dim: Dimension to concatenate along. Must be in range [-rank, rank), +// where rank is the number of dimensions in each input `SparseTensor`. +// +// Returns 2-D. Indices of the concatenated `SparseTensor`.1-D. Non-empty values of the concatenated `SparseTensor`.1-D. Shape of the concatenated `SparseTensor`. +func SparseConcat(scope *Scope, indices []tf.Output, values []tf.Output, shapes []tf.Output, concat_dim int64) (output_indices tf.Output, output_values tf.Output, output_shape tf.Output) { if scope.Err() != nil { return } + attrs := map[string]interface{}{"concat_dim": concat_dim} opspec := tf.OpSpec{ - Type: "BitwiseOr", + Type: "SparseConcat", Input: []tf.Input{ - x, y, + tf.OutputList(indices), tf.OutputList(values), tf.OutputList(shapes), }, + Attrs: attrs, } op := scope.AddOperation(opspec) - return op.Output(0) -} - -// SparseToSparseSetOperationAttr is an optional argument to SparseToSparseSetOperation. -type SparseToSparseSetOperationAttr func(optionalAttr) - -// SparseToSparseSetOperationValidateIndices sets the optional validate_indices attribute to value. -// If not specified, defaults to true -func SparseToSparseSetOperationValidateIndices(value bool) SparseToSparseSetOperationAttr { - return func(m optionalAttr) { - m["validate_indices"] = value - } + return op.Output(0), op.Output(1), op.Output(2) } -// Applies set operation along last dimension of 2 `SparseTensor` inputs. +// Generates sparse cross from a list of sparse and dense tensors. // -// See SetOperationOp::SetOperationFromContext for values of `set_operation`. +// The op takes two lists, one of 2D `SparseTensor` and one of 2D `Tensor`, each +// representing features of one feature column. It outputs a 2D `SparseTensor` with +// the batchwise crosses of these features. // -// If `validate_indices` is `True`, `SparseToSparseSetOperation` validates the -// order and range of `set1` and `set2` indices. +// For example, if the inputs are // -// Input `set1` is a `SparseTensor` represented by `set1_indices`, `set1_values`, -// and `set1_shape`. For `set1` ranked `n`, 1st `n-1` dimensions must be the same -// as `set2`. Dimension `n` contains values in a set, duplicates are allowed but -// ignored. +// inputs[0]: SparseTensor with shape = [2, 2] +// [0, 0]: "a" +// [1, 0]: "b" +// [1, 1]: "c" // -// Input `set2` is a `SparseTensor` represented by `set2_indices`, `set2_values`, -// and `set2_shape`. For `set2` ranked `n`, 1st `n-1` dimensions must be the same -// as `set1`. Dimension `n` contains values in a set, duplicates are allowed but -// ignored. +// inputs[1]: SparseTensor with shape = [2, 1] +// [0, 0]: "d" +// [1, 0]: "e" // -// If `validate_indices` is `True`, this op validates the order and range of `set1` -// and `set2` indices. +// inputs[2]: Tensor [["f"], ["g"]] // -// Output `result` is a `SparseTensor` represented by `result_indices`, -// `result_values`, and `result_shape`. For `set1` and `set2` ranked `n`, this -// has rank `n` and the same 1st `n-1` dimensions as `set1` and `set2`. The `nth` -// dimension contains the result of `set_operation` applied to the corresponding -// `[0...n-1]` dimension of `set`. +// then the output will be +// +// shape = [2, 2] +// [0, 0]: "a_X_d_X_f" +// [1, 0]: "b_X_e_X_g" +// [1, 1]: "c_X_e_X_g" +// +// if hashed_output=true then the output will be +// +// shape = [2, 2] +// [0, 0]: FingerprintCat64( +// Fingerprint64("f"), FingerprintCat64( +// Fingerprint64("d"), Fingerprint64("a"))) +// [1, 0]: FingerprintCat64( +// Fingerprint64("g"), FingerprintCat64( +// Fingerprint64("e"), Fingerprint64("b"))) +// [1, 1]: FingerprintCat64( +// Fingerprint64("g"), FingerprintCat64( +// Fingerprint64("e"), Fingerprint64("c"))) // // Arguments: -// set1_indices: 2D `Tensor`, indices of a `SparseTensor`. Must be in row-major -// order. -// set1_values: 1D `Tensor`, values of a `SparseTensor`. Must be in row-major -// order. -// set1_shape: 1D `Tensor`, shape of a `SparseTensor`. `set1_shape[0...n-1]` must -// be the same as `set2_shape[0...n-1]`, `set1_shape[n]` is the -// max set size across `0...n-1` dimensions. -// set2_indices: 2D `Tensor`, indices of a `SparseTensor`. Must be in row-major -// order. -// set2_values: 1D `Tensor`, values of a `SparseTensor`. Must be in row-major -// order. -// set2_shape: 1D `Tensor`, shape of a `SparseTensor`. `set2_shape[0...n-1]` must -// be the same as `set1_shape[0...n-1]`, `set2_shape[n]` is the -// max set size across `0...n-1` dimensions. +// indices: 2-D. Indices of each input `SparseTensor`. +// values: 1-D. values of each `SparseTensor`. +// shapes: 1-D. Shapes of each `SparseTensor`. +// dense_inputs: 2-D. Columns represented by dense `Tensor`. +// hashed_output: If true, returns the hash of the cross instead of the string. +// This will allow us avoiding string manipulations. +// num_buckets: It is used if hashed_output is true. +// output = hashed_value%num_buckets if num_buckets > 0 else hashed_value. +// hash_key: Specify the hash_key that will be used by the `FingerprintCat64` +// function to combine the crosses fingerprints. // // -// Returns 2D indices of a `SparseTensor`.1D values of a `SparseTensor`.1D `Tensor` shape of a `SparseTensor`. `result_shape[0...n-1]` is -// the same as the 1st `n-1` dimensions of `set1` and `set2`, `result_shape[n]` -// is the max result set size across all `0...n-1` dimensions. -func SparseToSparseSetOperation(scope *Scope, set1_indices tf.Output, set1_values tf.Output, set1_shape tf.Output, set2_indices tf.Output, set2_values tf.Output, set2_shape tf.Output, set_operation string, optional ...SparseToSparseSetOperationAttr) (result_indices tf.Output, result_values tf.Output, result_shape tf.Output) { +// +// Returns 2-D. Indices of the concatenated `SparseTensor`.1-D. Non-empty values of the concatenated or hashed +// `SparseTensor`.1-D. Shape of the concatenated `SparseTensor`. +func SparseCross(scope *Scope, indices []tf.Output, values []tf.Output, shapes []tf.Output, dense_inputs []tf.Output, hashed_output bool, num_buckets int64, hash_key int64, out_type tf.DataType, internal_type tf.DataType) (output_indices tf.Output, output_values tf.Output, output_shape tf.Output) { if scope.Err() != nil { return } - attrs := map[string]interface{}{"set_operation": set_operation} - for _, a := range optional { - a(attrs) - } + attrs := map[string]interface{}{"hashed_output": hashed_output, "num_buckets": num_buckets, "hash_key": hash_key, "out_type": out_type, "internal_type": internal_type} opspec := tf.OpSpec{ - Type: "SparseToSparseSetOperation", + Type: "SparseCross", Input: []tf.Input{ - set1_indices, set1_values, set1_shape, set2_indices, set2_values, set2_shape, + tf.OutputList(indices), tf.OutputList(values), tf.OutputList(shapes), tf.OutputList(dense_inputs), }, Attrs: attrs, } @@ -15891,71 +17276,75 @@ func SparseToSparseSetOperation(scope *Scope, set1_indices tf.Output, set1_value return op.Output(0), op.Output(1), op.Output(2) } -// Computes numerical negative value element-wise. +// Concatenates quantized tensors along one dimension. // -// I.e., \\(y = -x\\). -func Neg(scope *Scope, x tf.Output) (y tf.Output) { +// Arguments: +// concat_dim: 0-D. The dimension along which to concatenate. Must be in the +// range [0, rank(values)). +// values: The `N` Tensors to concatenate. Their ranks and types must match, +// and their sizes must match in all dimensions except `concat_dim`. +// input_mins: The minimum scalar values for each of the input tensors. +// input_maxes: The maximum scalar values for each of the input tensors. +// +// 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.The float value that the minimum quantized output value represents.The float value that the maximum quantized output value represents. +func QuantizedConcat(scope *Scope, concat_dim tf.Output, values []tf.Output, input_mins []tf.Output, input_maxes []tf.Output) (output tf.Output, output_min tf.Output, output_max tf.Output) { if scope.Err() != nil { return } opspec := tf.OpSpec{ - Type: "Neg", + Type: "QuantizedConcat", Input: []tf.Input{ - x, + concat_dim, tf.OutputList(values), tf.OutputList(input_mins), tf.OutputList(input_maxes), }, } op := scope.AddOperation(opspec) - 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 - } + return op.Output(0), op.Output(1), op.Output(2) } -// Fake-quantize the 'inputs' tensor of type float via global float scalars `min` +// Slice a `SparseTensor` based on the `start` and `size`. // -// and `max` to 'outputs' tensor of same shape as `inputs`. +// For example, if the input is // -// `[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. +// input_tensor = shape = [2, 7] +// [ a d e ] +// [b c ] // -// 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) { +// Graphically the output tensors are: +// +// sparse_slice([0, 0], [2, 4]) = shape = [2, 4] +// [ a ] +// [b c ] +// +// sparse_slice([0, 4], [2, 3]) = shape = [2, 3] +// [ d e ] +// [ ] +// +// Arguments: +// indices: 2-D tensor represents the indices of the sparse tensor. +// values: 1-D tensor represents the values of the sparse tensor. +// shape: 1-D. tensor represents the shape of the sparse tensor. +// start: 1-D. tensor represents the start of the slice. +// size: 1-D. tensor represents the size of the slice. +// output indices: A list of 1-D tensors represents the indices of the output +// sparse tensors. +// +// Returns A list of 1-D tensors represents the values of the output sparse +// tensors.A list of 1-D tensors represents the shape of the output sparse +// tensors. +func SparseSlice(scope *Scope, indices tf.Output, values tf.Output, shape tf.Output, start tf.Output, size tf.Output) (output_indices tf.Output, output_values tf.Output, output_shape tf.Output) { if scope.Err() != nil { return } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } opspec := tf.OpSpec{ - Type: "FakeQuantWithMinMaxVars", + Type: "SparseSlice", Input: []tf.Input{ - inputs, min, max, + indices, values, shape, start, size, }, - Attrs: attrs, } op := scope.AddOperation(opspec) - return op.Output(0) + return op.Output(0), op.Output(1), op.Output(2) } // Returns the element-wise min of two SparseTensors. @@ -16494,11 +17883,8 @@ func DenseToSparseSetOperation(scope *Scope, set1 tf.Output, set2_indices tf.Out // Subtracts a value from the current value of a variable. // -// Any ReadVariableOp which depends directly or indirectly on this assign is -// guaranteed to see the incremented value or a subsequent newer one. -// -// Outputs the incremented value, which can be used to totally order the -// increments to this variable. +// Any ReadVariableOp with a control dependency on this op is guaranteed to +// see the decremented value or a subsequent newer one. // // Arguments: // resource: handle to the resource in which to store the variable. @@ -16583,9 +17969,8 @@ type QuantizedResizeBilinearAttr func(optionalAttr) // QuantizedResizeBilinearAlignCorners sets the optional align_corners attribute to value. // -// value: If true, rescale input by (new_height - 1) / (height - 1), which -// exactly aligns the 4 corners of images and resized images. If false, rescale -// by new_height / height. Treat similarly the width dimension. +// value: If true, the centers of the 4 corner pixels of the input and output tensors are +// aligned, preserving the values at the corner pixels. Defaults to false. // If not specified, defaults to false func QuantizedResizeBilinearAlignCorners(value bool) QuantizedResizeBilinearAttr { return func(m optionalAttr) { @@ -16627,7 +18012,7 @@ func QuantizedResizeBilinear(scope *Scope, images tf.Output, size tf.Output, min // Computes the minimum along segments of a tensor. // -// Read @{$math_ops#segmentation$the section on segmentation} for an explanation of +// Read @{$math_ops#Segmentation$the section on segmentation} for an explanation of // segments. // // Computes a tensor such that @@ -16666,8 +18051,8 @@ type SdcaOptimizerAttr func(optionalAttr) // SdcaOptimizerAdaptative sets the optional adaptative attribute to value. // -// value: Whether to use Adapative SDCA for the inner loop. -// If not specified, defaults to false +// value: Whether to use Adaptive SDCA for the inner loop. +// If not specified, defaults to true func SdcaOptimizerAdaptative(value bool) SdcaOptimizerAttr { return func(m optionalAttr) { m["adaptative"] = value @@ -17237,7 +18622,7 @@ func EagerPyFunc(scope *Scope, input []tf.Output, token string, Tout []tf.DataTy // Duplicate entries are handled correctly: if multiple `indices` reference // the same location, their contributions add. // -// Requires `updates.shape = indices.shape + ref.shape[1:]`. +// Requires `updates.shape = indices.shape + ref.shape[1:]` or `updates.shape = []`. // //
// @@ -17515,6 +18900,43 @@ func RandomGamma(scope *Scope, shape tf.Output, alpha tf.Output, optional ...Ran return op.Output(0) } +// Computes the product 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 product of all +// entries belonging to a segment such that: +// +// \\(output_i = \prod_j data_j\\) where the product is over `j` such +// that `segment_ids[j] == i`. +// +// If there is no entry for a given segment ID `i`, it outputs 1. +// +// Arguments: +// +// segment_ids: A 1-D tensor whose rank is equal to the rank of `data`'s +// first dimension. +// +// +// Returns Has same shape as data, except for dimension 0 which +// has size `num_segments`. +func UnsortedSegmentProd(scope *Scope, data tf.Output, segment_ids tf.Output, num_segments tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "UnsortedSegmentProd", + Input: []tf.Input{ + data, segment_ids, num_segments, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // RandomUniformIntAttr is an optional argument to RandomUniformInt. type RandomUniformIntAttr func(optionalAttr) @@ -18010,6 +19432,39 @@ func MatrixBandPart(scope *Scope, input tf.Output, num_lower tf.Output, num_uppe return op.Output(0) } +// Counts the number of occurrences of each value in an integer array. +// +// Outputs a vector with length `size` and the same dtype as `weights`. If +// `weights` are empty, then index `i` stores the number of times the value `i` is +// counted in `arr`. If `weights` are non-empty, then index `i` stores the sum of +// the value in `weights` at each index where the corresponding value in `arr` is +// `i`. +// +// Values in `arr` outside of the range [0, size) are ignored. +// +// Arguments: +// arr: int32 `Tensor`. +// size: non-negative int32 scalar `Tensor`. +// weights: is an int32, int64, float32, or float64 `Tensor` with the same +// shape as `arr`, or a length-0 `Tensor`, in which case it acts as all weights +// equal to 1. +// +// Returns 1D `Tensor` with length equal to `size`. The counts or summed weights for +// each value in the range [0, size). +func Bincount(scope *Scope, arr tf.Output, size tf.Output, weights tf.Output) (bins tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "Bincount", + Input: []tf.Input{ + arr, size, weights, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // CumsumAttr is an optional argument to Cumsum. type CumsumAttr func(optionalAttr) @@ -18508,6 +19963,57 @@ func RequantizationRange(scope *Scope, input tf.Output, input_min tf.Output, inp return op.Output(0), op.Output(1) } +// Rolls the elements of a tensor along an axis. +// +// The elements are shifted positively (towards larger indices) by the offset of +// `shift` along the dimension of `axis`. Negative `shift` values will shift +// elements in the opposite direction. Elements that roll passed the last position +// will wrap around to the first and vice versa. Multiple shifts along multiple +// axes may be specified. +// +// For example: +// +// ``` +// # 't' is [0, 1, 2, 3, 4] +// roll(t, shift=2, axis=0) ==> [3, 4, 0, 1, 2] +// +// # shifting along multiple dimensions +// # 't' is [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]] +// roll(t, shift=[1, -2], axis=[0, 1]) ==> [[7, 8, 9, 5, 6], [2, 3, 4, 0, 1]] +// +// # shifting along the same axis multiple times +// # 't' is [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]] +// roll(t, shift=[2, -3], axis=[1, 1]) ==> [[1, 2, 3, 4, 0], [6, 7, 8, 9, 5]] +// ``` +// +// Arguments: +// +// shift: Dimension must be 0-D or 1-D. `shift[i]` specifies the number of places by which +// elements are shifted positively (towards larger indices) along the dimension +// specified by `axis[i]`. Negative shifts will roll the elements in the opposite +// direction. +// axis: Dimension must be 0-D or 1-D. `axis[i]` specifies the dimension that the shift +// `shift[i]` should occur. If the same axis is referenced more than once, the +// total shift for that axis will be the sum of all the shifts that belong to that +// axis. +// +// Returns Has the same shape and size as the input. The elements are shifted +// positively (towards larger indices) by the offsets of `shift` along the +// dimensions of `axis`. +func Roll(scope *Scope, input tf.Output, shift tf.Output, axis tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "Roll", + Input: []tf.Input{ + input, shift, axis, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // MapPeekAttr is an optional argument to MapPeek. type MapPeekAttr func(optionalAttr) @@ -18646,6 +20152,68 @@ func Bucketize(scope *Scope, input tf.Output, boundaries []float32) (output tf.O return op.Output(0) } +// Calculates gains for each feature and returns the best possible split information for the feature. +// +// The split information is the best threshold (bucket id), gains and left/right node contributions per node for each feature. +// +// It is possible that not all nodes can be split on each feature. Hence, the list of possible nodes can differ between the features. Therefore, we return `node_ids_list` for each feature, containing the list of nodes that this feature can be used to split. +// +// In this manner, the output is the best split per features and per node, so that it needs to be combined later to produce the best split for each node (among all possible features). +// +// The length of output lists are all of the same length, `num_features`. +// The output shapes are compatible in a way that the first dimension of all tensors of all lists are the same and equal to the number of possible split nodes for each feature. +// +// Arguments: +// node_id_range: A Rank 1 tensor (shape=[2]) to specify the range [first, last) of node ids to process within `stats_summary_list`. The nodes are iterated between the two nodes specified by the tensor, as like `for node_id in range(node_id_range[0], node_id_range[1])` (Note that the last index node_id_range[1] is exclusive). +// stats_summary_list: A list of Rank 3 tensor (#shape=[max_splits, bucket, 2]) for accumulated stats summary (gradient/hessian) per node per buckets for each feature. The first dimension of the tensor is the maximum number of splits, and thus not all elements of it will be used, but only the indexes specified by node_ids will be used. +// l1: l1 regularization factor on leaf weights, per instance based. +// l2: l2 regularization factor on leaf weights, per instance based. +// tree_complexity: adjustment to the gain, per leaf based. +// min_node_weight: mininum avg of hessians in a node before required for the node to be considered for splitting. +// max_splits: the number of nodes that can be split in the whole tree. Used as a dimension of output tensors. +// +// Returns An output list of Rank 1 tensors indicating possible split node ids for each feature. The length of the list is num_features, but each tensor has different size as each feature provides different possible nodes. See above for details like shapes and sizes.An output list of Rank 1 tensors indicating the best gains for each feature to split for certain nodes. See above for details like shapes and sizes.An output list of Rank 1 tensors indicating the bucket id to compare with (as a threshold) for split in each node. See above for details like shapes and sizes.A list of Rank 2 tensors indicating the contribution of the left nodes when branching from parent nodes (given by the tensor element in the output node_ids_list) to the left direction by the given threshold for each feature. This value will be used to make the left node value by adding to the parent node value. Second dimension size is 1 for 1-dimensional logits, but would be larger for multi-class problems. See above for details like shapes and sizes.A list of Rank 2 tensors, with the same shape/conditions as left_node_contribs_list, but just that the value is for the right node. +func BoostedTreesCalculateBestGainsPerFeature(scope *Scope, node_id_range tf.Output, stats_summary_list []tf.Output, l1 tf.Output, l2 tf.Output, tree_complexity tf.Output, min_node_weight tf.Output, max_splits int64) (node_ids_list []tf.Output, gains_list []tf.Output, thresholds_list []tf.Output, left_node_contribs_list []tf.Output, right_node_contribs_list []tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"max_splits": max_splits} + opspec := tf.OpSpec{ + Type: "BoostedTreesCalculateBestGainsPerFeature", + Input: []tf.Input{ + node_id_range, tf.OutputList(stats_summary_list), l1, l2, tree_complexity, min_node_weight, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + if scope.Err() != nil { + return + } + var idx int + var err error + if node_ids_list, idx, err = makeOutputList(op, idx, "node_ids_list"); err != nil { + scope.UpdateErr("BoostedTreesCalculateBestGainsPerFeature", err) + return + } + if gains_list, idx, err = makeOutputList(op, idx, "gains_list"); err != nil { + scope.UpdateErr("BoostedTreesCalculateBestGainsPerFeature", err) + return + } + if thresholds_list, idx, err = makeOutputList(op, idx, "thresholds_list"); err != nil { + scope.UpdateErr("BoostedTreesCalculateBestGainsPerFeature", err) + return + } + if left_node_contribs_list, idx, err = makeOutputList(op, idx, "left_node_contribs_list"); err != nil { + scope.UpdateErr("BoostedTreesCalculateBestGainsPerFeature", err) + return + } + if right_node_contribs_list, idx, err = makeOutputList(op, idx, "right_node_contribs_list"); err != nil { + scope.UpdateErr("BoostedTreesCalculateBestGainsPerFeature", err) + return + } + return node_ids_list, gains_list, thresholds_list, left_node_contribs_list, right_node_contribs_list +} + // EncodePngAttr is an optional argument to EncodePng. type EncodePngAttr func(optionalAttr) @@ -18760,9 +20328,8 @@ type ResizeBilinearGradAttr func(optionalAttr) // ResizeBilinearGradAlignCorners sets the optional align_corners attribute to value. // -// value: If true, rescale grads by (orig_height - 1) / (height - 1), which -// exactly aligns the 4 corners of grads and original_image. If false, rescale by -// orig_height / height. Treat similarly the width dimension. +// value: If true, the centers of the 4 corner pixels of the input and grad tensors are +// aligned. Defaults to false. // If not specified, defaults to false func ResizeBilinearGradAlignCorners(value bool) ResizeBilinearGradAttr { return func(m optionalAttr) { @@ -19159,16 +20726,44 @@ func Dequantize(scope *Scope, input tf.Output, min_range tf.Output, max_range tf // Flips all bits elementwise. // -// The result will have exactly those bits set, that are not set in `x`. The -// computation is performed on the underlying representation of x. -func Invert(scope *Scope, x tf.Output) (y tf.Output) { +// The result will have exactly those bits set, that are not set in `x`. The +// computation is performed on the underlying representation of x. +func Invert(scope *Scope, x tf.Output) (y tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "Invert", + Input: []tf.Input{ + x, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Inverse 3D fast Fourier transform. +// +// Computes the inverse 3-dimensional discrete Fourier transform over the +// inner-most 3 dimensions of `input`. +// +// Arguments: +// input: A complex64 tensor. +// +// Returns A complex64 tensor of the same shape as `input`. The inner-most 3 +// dimensions of `input` are replaced with their inverse 3D Fourier transform. +// +// @compatibility(numpy) +// Equivalent to np.fft.ifftn with 3 dimensions. +// @end_compatibility +func IFFT3D(scope *Scope, input tf.Output) (output tf.Output) { if scope.Err() != nil { return } opspec := tf.OpSpec{ - Type: "Invert", + Type: "IFFT3D", Input: []tf.Input{ - x, + input, }, } op := scope.AddOperation(opspec) @@ -19397,6 +20992,47 @@ func TanhGrad(scope *Scope, y tf.Output, dy tf.Output) (z tf.Output) { return op.Output(0) } +// Reduces sparse updates into the variable referenced by `resource` using the `max` operation. +// +// This operation computes +// +// # Scalar indices +// ref[indices, ...] = max(ref[indices, ...], updates[...]) +// +// # Vector indices (for each i) +// ref[indices[i], ...] = max(ref[indices[i], ...], updates[i, ...]) +// +// # High rank indices (for each i, ..., j) +// ref[indices[i, ..., j], ...] = max(ref[indices[i, ..., j], ...], updates[i, ..., j, ...]) +// +// Duplicate entries are handled correctly: if multiple `indices` reference +// the same location, their contributions are combined. +// +// Requires `updates.shape = indices.shape + ref.shape[1:]` or `updates.shape = []`. +// +//
+// +//
+// +// Arguments: +// resource: Should be from a `Variable` node. +// indices: A tensor of indices into the first dimension of `ref`. +// updates: A tensor of updated values to add to `ref`. +// +// Returns the created operation. +func ResourceScatterMax(scope *Scope, resource tf.Output, indices tf.Output, updates tf.Output) (o *tf.Operation) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "ResourceScatterMax", + Input: []tf.Input{ + resource, indices, updates, + }, + } + return scope.AddOperation(opspec) +} + // Outputs a `Summary` protocol buffer with scalar values. // // The input `tags` and `values` must have the same shape. The generated summary @@ -19523,7 +21159,7 @@ func ImageSummaryBadColor(value tf.Tensor) ImageSummaryAttr { // generated sequentially as '*tag*/image/0', '*tag*/image/1', etc. // // The `bad_color` argument is the color to use in the generated images for -// non-finite input values. It is a `unit8` 1-D tensor of length `channels`. +// non-finite input values. It is a `uint8` 1-D tensor of length `channels`. // Each element must be in the range `[0, 255]` (It represents the value of a // pixel in the output image). Non-finite values in the input tensor are // replaced by this tensor in the output image. The default value is the color @@ -19781,6 +21417,88 @@ func AddSparseToTensorsMap(scope *Scope, sparse_indices tf.Output, sparse_values return op.Output(0) } +// Returns a list list which has the passed-in `Tensor` as last element and the other elements of the given list in `input_handle`. +// +// tensor: The tensor to put on the list. +// input_handle: The old list. +// output_handle: A list with the elements of the old list followed by tensor. +// element_dtype: the type of elements in the list. +// element_shape: a shape compatible with that of elements in the list. +func TensorListPushBack(scope *Scope, input_handle tf.Output, tensor tf.Output) (output_handle tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "TensorListPushBack", + Input: []tf.Input{ + input_handle, tensor, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Returns the number of tensors in the input tensor list. +// +// input_handle: the input list +// length: the number of tensors in the list +func TensorListLength(scope *Scope, input_handle tf.Output) (length tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "TensorListLength", + Input: []tf.Input{ + input_handle, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// The shape of the elements of the given list, as a tensor. +// +// input_handle: the list +// element_shape: the shape of elements of the list +func TensorListElementShape(scope *Scope, input_handle tf.Output, shape_type tf.DataType) (element_shape tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"shape_type": shape_type} + opspec := tf.OpSpec{ + Type: "TensorListElementShape", + Input: []tf.Input{ + input_handle, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Returns the item in the list with the given index. +// +// input_handle: the list +// index: the position in the list from which an element will be retrieved +// item: the element at that position +// +// +func TensorListGetItem(scope *Scope, input_handle tf.Output, index tf.Output, element_dtype tf.DataType) (item tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"element_dtype": element_dtype} + opspec := tf.OpSpec{ + Type: "TensorListGetItem", + Input: []tf.Input{ + input_handle, index, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Computes the matrix exponential of one or more square matrices: // // exp(A) = \sum_{n=0}^\infty A^n/n! @@ -19816,6 +21534,46 @@ func MatrixExponential(scope *Scope, input tf.Output) (output tf.Output) { return op.Output(0) } +// Computes the matrix logarithm of one or more square matrices: +// +// +// log(exp(A)) = A +// +// This op is only defined for complex matrices. If A is positive-definite and +// real, then casting to a complex matrix, taking the logarithm and casting back +// to a real matrix will give the correct result. +// +// This function computes the matrix logarithm using the Schur-Parlett algorithm. +// Details of the algorithm can be found in Section 11.6.2 of: +// Nicholas J. Higham, Functions of Matrices: Theory and Computation, SIAM 2008. +// ISBN 978-0-898716-46-7. +// +// The input is a tensor of shape `[..., M, M]` whose inner-most 2 dimensions +// form square matrices. The output is a tensor of the same shape as the input +// containing the exponential for all input submatrices `[..., :, :]`. +// +// Arguments: +// input: Shape is `[..., M, M]`. +// +// Returns Shape is `[..., M, M]`. +// +// @compatibility(scipy) +// Equivalent to scipy.linalg.logm +// @end_compatibility +func MatrixLogarithm(scope *Scope, input tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "MatrixLogarithm", + Input: []tf.Input{ + input, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // QueueDequeueUpToV2Attr is an optional argument to QueueDequeueUpToV2. type QueueDequeueUpToV2Attr func(optionalAttr) @@ -19995,7 +21753,8 @@ func All(scope *Scope, input tf.Output, axis tf.Output, optional ...AllAttr) (ou // SelfAdjointEig. // // The result is a [..., M+1, M] matrix with [..., 0,:] containing the -// eigenvalues, and subsequent [...,1:, :] containing the eigenvectors. +// eigenvalues, and subsequent [...,1:, :] containing the eigenvectors. The eigenvalues +// are sorted in non-decreasing order. // // Arguments: // input: Shape is `[..., M, M]`. @@ -20053,7 +21812,8 @@ func SelfAdjointEigV2ComputeV(value bool) SelfAdjointEigV2Attr { // Computes the eigen decomposition of one or more square self-adjoint matrices. // // Computes the eigenvalues and (optionally) eigenvectors of each inner matrix in -// `input` such that `input[..., :, :] = v[..., :, :] * diag(e[..., :])`. +// `input` such that `input[..., :, :] = v[..., :, :] * diag(e[..., :])`. The eigenvalues +// are sorted in non-decreasing order. // // ```python // # a is a tensor. @@ -20236,7 +21996,7 @@ func QueueEnqueueManyV2(scope *Scope, handle tf.Output, components []tf.Output, // Computes the product along segments of a tensor. // -// Read @{$math_ops#segmentation$the section on segmentation} for an explanation of +// Read @{$math_ops#Segmentation$the section on segmentation} for an explanation of // segments. // // Computes a tensor such that @@ -21110,25 +22870,6 @@ func ReaderReadUpToV2(scope *Scope, reader_handle tf.Output, queue_handle tf.Out return op.Output(0), op.Output(1) } -// Restore a Reader to its initial clean state. -// -// Arguments: -// reader_handle: Handle to a Reader. -// -// Returns the created operation. -func ReaderResetV2(scope *Scope, reader_handle tf.Output) (o *tf.Operation) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "ReaderResetV2", - Input: []tf.Input{ - reader_handle, - }, - } - return scope.AddOperation(opspec) -} - // ResourceApplyAdamAttr is an optional argument to ResourceApplyAdam. type ResourceApplyAdamAttr func(optionalAttr) @@ -21218,9 +22959,8 @@ type ResizeBicubicGradAttr func(optionalAttr) // ResizeBicubicGradAlignCorners sets the optional align_corners attribute to value. // -// value: If true, rescale grads by (orig_height - 1) / (height - 1), which -// exactly aligns the 4 corners of grads and original_image. If false, rescale by -// orig_height / height. Treat similarly the width dimension. +// value: If true, the centers of the 4 corner pixels of the input and grad tensors are +// aligned. Defaults to false. // If not specified, defaults to false func ResizeBicubicGradAlignCorners(value bool) ResizeBicubicGradAttr { return func(m optionalAttr) { @@ -21262,9 +23002,8 @@ type ResizeNearestNeighborAttr func(optionalAttr) // ResizeNearestNeighborAlignCorners sets the optional align_corners attribute to value. // -// value: If true, rescale input by (new_height - 1) / (height - 1), which -// exactly aligns the 4 corners of images and resized images. If false, rescale -// by new_height / height. Treat similarly the width dimension. +// value: If true, the centers of the 4 corner pixels of the input and output tensors are +// aligned, preserving the values at the corner pixels. Defaults to false. // If not specified, defaults to false func ResizeNearestNeighborAlignCorners(value bool) ResizeNearestNeighborAttr { return func(m optionalAttr) { @@ -21305,9 +23044,8 @@ type ResizeNearestNeighborGradAttr func(optionalAttr) // ResizeNearestNeighborGradAlignCorners sets the optional align_corners attribute to value. // -// value: If true, rescale grads by (orig_height - 1) / (height - 1), which -// exactly aligns the 4 corners of grads and original_image. If false, rescale by -// orig_height / height. Treat similarly the width dimension. +// value: If true, the centers of the 4 corner pixels of the input and grad tensors are +// aligned. Defaults to false. // If not specified, defaults to false func ResizeNearestNeighborGradAlignCorners(value bool) ResizeNearestNeighborGradAttr { return func(m optionalAttr) { @@ -21556,112 +23294,6 @@ func DecodeGif(scope *Scope, contents tf.Output) (image tf.Output) { return op.Output(0) } -// ResourceApplyCenteredRMSPropAttr is an optional argument to ResourceApplyCenteredRMSProp. -type ResourceApplyCenteredRMSPropAttr func(optionalAttr) - -// ResourceApplyCenteredRMSPropUseLocking sets the optional use_locking attribute to value. -// -// value: If `True`, updating of the var, mg, ms, and mom tensors is -// protected by a lock; otherwise the behavior is undefined, but may exhibit less -// contention. -// If not specified, defaults to false -func ResourceApplyCenteredRMSPropUseLocking(value bool) ResourceApplyCenteredRMSPropAttr { - return func(m optionalAttr) { - m["use_locking"] = value - } -} - -// Update '*var' according to the centered RMSProp algorithm. -// -// The centered RMSProp algorithm uses an estimate of the centered second moment -// (i.e., the variance) for normalization, as opposed to regular RMSProp, which -// uses the (uncentered) second moment. This often helps with training, but is -// slightly more expensive in terms of computation and memory. -// -// Note that in dense implementation of this algorithm, mg, ms, and mom will -// update even if the grad is zero, but in this sparse implementation, mg, ms, -// and mom will not update in iterations during which the grad is zero. -// -// mean_square = decay * mean_square + (1-decay) * gradient ** 2 -// mean_grad = decay * mean_grad + (1-decay) * gradient -// -// Delta = learning_rate * gradient / sqrt(mean_square + epsilon - mean_grad ** 2) -// -// mg <- rho * mg_{t-1} + (1-rho) * grad -// ms <- rho * ms_{t-1} + (1-rho) * grad * grad -// mom <- momentum * mom_{t-1} + lr * grad / sqrt(ms - mg * mg + epsilon) -// var <- var - mom -// -// Arguments: -// var_: Should be from a Variable(). -// mg: Should be from a Variable(). -// ms: Should be from a Variable(). -// mom: Should be from a Variable(). -// lr: Scaling factor. Must be a scalar. -// rho: Decay rate. Must be a scalar. -// -// epsilon: Ridge term. Must be a scalar. -// grad: The gradient. -// -// Returns the created operation. -func ResourceApplyCenteredRMSProp(scope *Scope, var_ tf.Output, mg tf.Output, ms tf.Output, mom tf.Output, lr tf.Output, rho tf.Output, momentum tf.Output, epsilon tf.Output, grad tf.Output, optional ...ResourceApplyCenteredRMSPropAttr) (o *tf.Operation) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "ResourceApplyCenteredRMSProp", - Input: []tf.Input{ - var_, mg, ms, mom, lr, rho, momentum, epsilon, grad, - }, - Attrs: attrs, - } - return scope.AddOperation(opspec) -} - -// Returns a list of tensors with the same shapes and contents as the input -// -// tensors. -// -// This op can be used to override the gradient for complicated functions. For -// example, suppose y = f(x) and we wish to apply a custom function g for backprop -// such that dx = g(dy). In Python, -// -// ```python -// with tf.get_default_graph().gradient_override_map( -// {'IdentityN': 'OverrideGradientWithG'}): -// y, _ = identity_n([f(x), x]) -// -// @tf.RegisterGradient('OverrideGradientWithG') -// def ApplyG(op, dy, _): -// return [None, g(dy)] # Do not backprop to f(x). -// ``` -func IdentityN(scope *Scope, input []tf.Output) (output []tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "IdentityN", - Input: []tf.Input{ - tf.OutputList(input), - }, - } - op := scope.AddOperation(opspec) - if scope.Err() != nil { - return - } - var idx int - var err error - if output, idx, err = makeOutputList(op, idx, "output"); err != nil { - scope.UpdateErr("IdentityN", err) - return - } - return output -} - // Computes the gradient of the sigmoid of `x` wrt its input. // // Specifically, `grad = dy * y * (1 - y)`, where `y = sigmoid(x)`, and @@ -21706,6 +23338,58 @@ func HSVToRGB(scope *Scope, images tf.Output) (output tf.Output) { return op.Output(0) } +// Retrieves the tree ensemble resource stamp token, number of trees and growing statistics. +// +// Arguments: +// tree_ensemble_handle: Handle to the tree ensemble. +// +// Returns Stamp token of the tree ensemble resource.The number of trees in the tree ensemble resource.The number of trees that were finished successfully.The number of layers we attempted to build (but not necessarily succeeded).Rank size 2 tensor that contains start and end ids of the nodes in the latest +// layer. +func BoostedTreesGetEnsembleStates(scope *Scope, tree_ensemble_handle tf.Output) (stamp_token tf.Output, num_trees tf.Output, num_finalized_trees tf.Output, num_attempted_layers tf.Output, last_layer_nodes_range tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "BoostedTreesGetEnsembleStates", + Input: []tf.Input{ + tree_ensemble_handle, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0), op.Output(1), op.Output(2), op.Output(3), op.Output(4) +} + +// Gets the next output from the given iterator. +// +// This operation is a synchronous version IteratorGetNext. It should only be used +// in situations where the iterator does not block the calling thread, or where +// the calling thread is not a member of the thread pool used to execute parallel +// operations (e.g. in eager mode). +func IteratorGetNextSync(scope *Scope, iterator tf.Output, output_types []tf.DataType, output_shapes []tf.Shape) (components []tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"output_types": output_types, "output_shapes": output_shapes} + opspec := tf.OpSpec{ + Type: "IteratorGetNextSync", + Input: []tf.Input{ + iterator, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + if scope.Err() != nil { + return + } + var idx int + var err error + if components, idx, err = makeOutputList(op, idx, "components"); err != nil { + scope.UpdateErr("IteratorGetNextSync", err) + return + } + return components +} + // SampleDistortedBoundingBoxV2Attr is an optional argument to SampleDistortedBoundingBoxV2. type SampleDistortedBoundingBoxV2Attr func(optionalAttr) @@ -22254,6 +23938,83 @@ func NonMaxSuppressionV2(scope *Scope, boxes tf.Output, scores tf.Output, max_ou return op.Output(0) } +// EncodeProtoAttr is an optional argument to EncodeProto. +type EncodeProtoAttr func(optionalAttr) + +// EncodeProtoDescriptorSource sets the optional descriptor_source attribute to value. +// If not specified, defaults to "local://" +func EncodeProtoDescriptorSource(value string) EncodeProtoAttr { + return func(m optionalAttr) { + m["descriptor_source"] = value + } +} + +// The op serializes protobuf messages provided in the input tensors. +// +// The types of the tensors in `values` must match the schema for the +// fields specified in `field_names`. All the tensors in `values` must +// have a common shape prefix, *batch_shape*. +// +// The `sizes` tensor specifies repeat counts for each field. The repeat +// count (last dimension) of a each tensor in `values` must be greater +// than or equal to corresponding repeat count in `sizes`. +// +// A `message_type` name must be provided to give context for the field +// names. The actual message descriptor can be looked up either in the +// linked-in descriptor pool or a filename provided by the caller using +// the `descriptor_source` attribute. +// +// The `descriptor_source` attribute selects a source of protocol +// descriptors to consult when looking up `message_type`. This may be a +// filename containing a serialized `FileDescriptorSet` message, +// or the special value `local://`, in which case only descriptors linked +// into the code will be searched; the filename can be on any filesystem +// accessible to TensorFlow. +// +// You can build a `descriptor_source` file using the `--descriptor_set_out` +// and `--include_imports` options to the protocol compiler `protoc`. +// +// The `local://` database only covers descriptors linked into the +// code via C++ libraries, not Python imports. You can link in a proto descriptor +// by creating a cc_library target with alwayslink=1. +// +// There are a few special cases in the value mapping: +// +// Submessage and group fields must be pre-serialized as TensorFlow strings. +// +// TensorFlow lacks support for unsigned int64s, so they must be +// represented as `tf.int64` with the same twos-complement bit pattern +// (the obvious way). +// +// Unsigned int32 values can be represented exactly with `tf.int64`, or +// with sign wrapping if the input is of type `tf.int32`. +// +// Arguments: +// sizes: Tensor of int32 with shape `[batch_shape, len(field_names)]`. +// values: List of tensors containing values for the corresponding field. +// field_names: List of strings containing proto field names. +// message_type: Name of the proto message type to decode. +// +// Returns Tensor of serialized protos with shape `batch_shape`. +func EncodeProto(scope *Scope, sizes tf.Output, values []tf.Output, field_names []string, message_type string, optional ...EncodeProtoAttr) (bytes tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"field_names": field_names, "message_type": message_type} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "EncodeProto", + Input: []tf.Input{ + sizes, tf.OutputList(values), + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Creates a TensorArray for storing the gradients of values in the given handle. // // If the given TensorArray gradient already exists, returns a reference to it. @@ -22283,67 +24044,471 @@ func NonMaxSuppressionV2(scope *Scope, boxes tf.Output, scores tf.Output, max_ou // gradient nodes may accidentally flow through the same accumulator TensorArray. // This double counts and generally breaks the TensorArray gradient flow. // -// The solution is to identify which gradient call this particular -// TensorArray gradient is being called in. This is performed by identifying -// a unique string (e.g. "gradients", "gradients_1", ...) from the input -// gradient Tensor's name. This string is used as a suffix when creating -// the TensorArray gradient object here (the attribute `source`). +// The solution is to identify which gradient call this particular +// TensorArray gradient is being called in. This is performed by identifying +// a unique string (e.g. "gradients", "gradients_1", ...) from the input +// gradient Tensor's name. This string is used as a suffix when creating +// the TensorArray gradient object here (the attribute `source`). +// +// The attribute `source` is added as a suffix to the forward TensorArray's +// name when performing the creation / lookup, so that each separate gradient +// calculation gets its own TensorArray accumulator. +// +// Arguments: +// handle: The handle to the forward TensorArray. +// flow_in: A float scalar that enforces proper chaining of operations. +// source: The gradient source string, used to decide which gradient TensorArray +// to return. +func TensorArrayGradV3(scope *Scope, handle tf.Output, flow_in tf.Output, source string) (grad_handle tf.Output, flow_out tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"source": source} + opspec := tf.OpSpec{ + Type: "TensorArrayGradV3", + Input: []tf.Input{ + handle, flow_in, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0), op.Output(1) +} + +// DecodeProtoV2Attr is an optional argument to DecodeProtoV2. +type DecodeProtoV2Attr func(optionalAttr) + +// DecodeProtoV2DescriptorSource sets the optional descriptor_source attribute to value. +// +// value: Either the special value `local://` or a path to a file containing +// a serialized `FileDescriptorSet`. +// If not specified, defaults to "local://" +func DecodeProtoV2DescriptorSource(value string) DecodeProtoV2Attr { + return func(m optionalAttr) { + m["descriptor_source"] = value + } +} + +// DecodeProtoV2MessageFormat sets the optional message_format attribute to value. +// +// value: Either `binary` or `text`. +// If not specified, defaults to "binary" +func DecodeProtoV2MessageFormat(value string) DecodeProtoV2Attr { + return func(m optionalAttr) { + m["message_format"] = value + } +} + +// DecodeProtoV2Sanitize sets the optional sanitize attribute to value. +// +// value: Whether to sanitize the result or not. +// If not specified, defaults to false +func DecodeProtoV2Sanitize(value bool) DecodeProtoV2Attr { + return func(m optionalAttr) { + m["sanitize"] = value + } +} + +// The op extracts fields from a serialized protocol buffers message into tensors. +// +// The `decode_proto` op extracts fields from a serialized protocol buffers +// message into tensors. The fields in `field_names` are decoded and converted +// to the corresponding `output_types` if possible. +// +// A `message_type` name must be provided to give context for the field +// names. The actual message descriptor can be looked up either in the +// linked-in descriptor pool or a filename provided by the caller using +// the `descriptor_source` attribute. +// +// Each output tensor is a dense tensor. This means that it is padded to +// hold the largest number of repeated elements seen in the input +// minibatch. (The shape is also padded by one to prevent zero-sized +// dimensions). The actual repeat counts for each example in the +// minibatch can be found in the `sizes` output. In many cases the output +// of `decode_proto` is fed immediately into tf.squeeze if missing values +// are not a concern. When using tf.squeeze, always pass the squeeze +// dimension explicitly to avoid surprises. +// +// For the most part, the mapping between Proto field types and +// TensorFlow dtypes is straightforward. However, there are a few +// special cases: +// +// - A proto field that contains a submessage or group can only be converted +// to `DT_STRING` (the serialized submessage). This is to reduce the +// complexity of the API. The resulting string can be used as input +// to another instance of the decode_proto op. +// +// - TensorFlow lacks support for unsigned integers. The ops represent uint64 +// types as a `DT_INT64` with the same twos-complement bit pattern +// (the obvious way). Unsigned int32 values can be represented exactly by +// specifying type `DT_INT64`, or using twos-complement if the caller +// specifies `DT_INT32` in the `output_types` attribute. +// +// The `descriptor_source` attribute selects a source of protocol +// descriptors to consult when looking up `message_type`. This may be a +// filename containing a serialized `FileDescriptorSet` message, +// or the special value `local://`, in which case only descriptors linked +// into the code will be searched; the filename can be on any filesystem +// accessible to TensorFlow. +// +// You can build a `descriptor_source` file using the `--descriptor_set_out` +// and `--include_imports` options to the protocol compiler `protoc`. +// +// The `local://` database only covers descriptors linked into the +// code via C++ libraries, not Python imports. You can link in a proto descriptor +// by creating a cc_library target with alwayslink=1. +// +// Both binary and text proto serializations are supported, and can be +// chosen using the `format` attribute. +// +// Arguments: +// bytes: Tensor of serialized protos with shape `batch_shape`. +// message_type: Name of the proto message type to decode. +// field_names: List of strings containing proto field names. +// output_types: List of TF types to use for the respective field in field_names. +// +// Returns Tensor of int32 with shape `[batch_shape, len(field_names)]`. +// Each entry is the number of values found for the corresponding field. +// Optional fields may have 0 or 1 values.List of tensors containing values for the corresponding field. +// `values[i]` has datatype `output_types[i]` +// and shape `[batch_shape, max(sizes[...,i])]`. +func DecodeProtoV2(scope *Scope, bytes tf.Output, message_type string, field_names []string, output_types []tf.DataType, optional ...DecodeProtoV2Attr) (sizes tf.Output, values []tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"message_type": message_type, "field_names": field_names, "output_types": output_types} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "DecodeProtoV2", + Input: []tf.Input{ + bytes, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + if scope.Err() != nil { + return + } + var idx int + var err error + sizes = op.Output(idx) + if values, idx, err = makeOutputList(op, idx, "values"); err != nil { + scope.UpdateErr("DecodeProtoV2", err) + return + } + return sizes, values +} + +// Creates a dataset that splits a SparseTensor into elements row-wise. +func SparseTensorSliceDataset(scope *Scope, indices tf.Output, values tf.Output, dense_shape tf.Output) (handle tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "SparseTensorSliceDataset", + Input: []tf.Input{ + indices, values, dense_shape, + }, + } + 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) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "RealDiv", + Input: []tf.Input{ + x, y, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Adds v into specified rows of x. +// +// Computes y = x; y[i, :] += v; return y. +// +// Arguments: +// x: A `Tensor` of type T. +// i: A vector. Indices into the left-most dimension of `x`. +// v: A `Tensor` of type T. Same dimension sizes as x except the first dimension, which must be the same as i's size. +// +// Returns A `Tensor` of type T. An alias of `x`. The content of `y` is undefined if there are duplicates in `i`. +func InplaceAdd(scope *Scope, x tf.Output, i tf.Output, v tf.Output) (y tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "InplaceAdd", + Input: []tf.Input{ + x, i, v, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Restore a Reader to its initial clean state. +// +// Arguments: +// reader_handle: Handle to a Reader. +// +// Returns the created operation. +func ReaderResetV2(scope *Scope, reader_handle tf.Output) (o *tf.Operation) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "ReaderResetV2", + Input: []tf.Input{ + reader_handle, + }, + } + return scope.AddOperation(opspec) +} + +// RpcAttr is an optional argument to Rpc. +type RpcAttr func(optionalAttr) + +// RpcProtocol sets the optional protocol attribute to value. +// +// value: RPC protocol to use. Empty string means use the default protocol. +// Options include 'grpc'. +// If not specified, defaults to "" +func RpcProtocol(value string) RpcAttr { + return func(m optionalAttr) { + m["protocol"] = value + } +} + +// RpcFailFast sets the optional fail_fast attribute to value. +// +// value: `boolean`. If `true` (default), then failures to connect +// (i.e., the server does not immediately respond) cause an RPC failure. +// If not specified, defaults to true +func RpcFailFast(value bool) RpcAttr { + return func(m optionalAttr) { + m["fail_fast"] = value + } +} + +// RpcTimeoutInMs sets the optional timeout_in_ms attribute to value. +// +// value: `int`. If `0` (default), then the kernel will run the RPC +// request and only time out if the RPC deadline passes or the session times out. +// If this value is greater than `0`, then the op will raise an exception if +// the RPC takes longer than `timeout_in_ms`. +// If not specified, defaults to 0 +func RpcTimeoutInMs(value int64) RpcAttr { + return func(m optionalAttr) { + m["timeout_in_ms"] = value + } +} + +// Perform batches of RPC requests. +// +// This op asynchronously performs either a single RPC request, or a batch +// of requests. RPC requests are defined by three main parameters: +// +// - `address` (the host+port or BNS address of the request) +// - `method` (the RPC method name for the request) +// - `request` (the serialized proto string, or vector of strings, +// of the RPC request argument). +// +// For example, if you have an RPC service running on port localhost:2345, +// and its interface is configured with the following proto declaration: +// +// ``` +// service MyService { +// rpc MyMethod(MyRequestProto) returns (MyResponseProto) { +// } +// }; +// ``` +// +// then call this op with arguments: +// +// ``` +// address = "localhost:2345" +// method = "MyService/MyMethod" +// ``` +// +// The `request` tensor is a string tensor representing serialized `MyRequestProto` +// strings; and the output string tensor `response` will have the same shape +// and contain (upon successful completion) corresponding serialized +// `MyResponseProto` strings. +// +// For example, to send a single, empty, `MyRequestProto`, call +// this op with `request = ""`. To send 5 **parallel** empty requests, +// call this op with `request = ["", "", "", "", ""]`. +// +// More generally, one can create a batch of `MyRequestProto` serialized protos +// from regular batched tensors using the `encode_proto` op, and convert +// the response `MyResponseProto` serialized protos to batched tensors +// using the `decode_proto` op. // -// The attribute `source` is added as a suffix to the forward TensorArray's -// name when performing the creation / lookup, so that each separate gradient -// calculation gets its own TensorArray accumulator. +// **NOTE** Working with serialized proto strings is faster than instantiating +// actual proto objects in memory, so no performance degradation is expected +// compared to writing custom kernels for this workflow. +// +// If the connection fails or the remote worker returns an error +// status, the op reraises this exception locally. +// +// See the `TryRpc` op if you prefer to handle RPC failures manually in the graph. // // Arguments: -// handle: The handle to the forward TensorArray. -// flow_in: A float scalar that enforces proper chaining of operations. -// source: The gradient source string, used to decide which gradient TensorArray -// to return. -func TensorArrayGradV3(scope *Scope, handle tf.Output, flow_in tf.Output, source string) (grad_handle tf.Output, flow_out tf.Output) { +// address: `0-D` or `1-D`. The address (i.e. host_name:port) of the RPC server. +// If this tensor has more than 1 element, then multiple parallel rpc requests +// are sent. This argument broadcasts with `method` and `request`. +// method: `0-D` or `1-D`. The method address on the RPC server. +// If this tensor has more than 1 element, then multiple parallel rpc requests +// are sent. This argument broadcasts with `address` and `request`. +// request: `0-D` or `1-D`. Serialized proto strings: the rpc request argument. +// If this tensor has more than 1 element, then multiple parallel rpc requests +// are sent. This argument broadcasts with `address` and `method`. +// +// Returns Same shape as `request`. Serialized proto strings: the rpc responses. +func Rpc(scope *Scope, address tf.Output, method tf.Output, request tf.Output, optional ...RpcAttr) (response tf.Output) { if scope.Err() != nil { return } - attrs := map[string]interface{}{"source": source} + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } opspec := tf.OpSpec{ - Type: "TensorArrayGradV3", + Type: "Rpc", Input: []tf.Input{ - handle, flow_in, + address, method, request, }, Attrs: attrs, } op := scope.AddOperation(opspec) - return op.Output(0), op.Output(1) + return op.Output(0) } -// Creates a dataset that splits a SparseTensor into elements row-wise. -func SparseTensorSliceDataset(scope *Scope, indices tf.Output, values tf.Output, dense_shape tf.Output) (handle tf.Output) { +// OrderedMapStageAttr is an optional argument to OrderedMapStage. +type OrderedMapStageAttr func(optionalAttr) + +// OrderedMapStageCapacity sets the optional capacity attribute to value. +// +// value: Maximum number of elements in the Staging Area. If > 0, inserts +// on the container will block when the capacity is reached. +// If not specified, defaults to 0 +// +// REQUIRES: value >= 0 +func OrderedMapStageCapacity(value int64) OrderedMapStageAttr { + return func(m optionalAttr) { + m["capacity"] = value + } +} + +// OrderedMapStageMemoryLimit sets the optional memory_limit attribute to value. +// If not specified, defaults to 0 +// +// REQUIRES: value >= 0 +func OrderedMapStageMemoryLimit(value int64) OrderedMapStageAttr { + return func(m optionalAttr) { + m["memory_limit"] = value + } +} + +// OrderedMapStageContainer sets the optional container attribute to value. +// +// value: If non-empty, this queue is placed in the given container. Otherwise, +// a default container is used. +// If not specified, defaults to "" +func OrderedMapStageContainer(value string) OrderedMapStageAttr { + return func(m optionalAttr) { + m["container"] = value + } +} + +// OrderedMapStageSharedName sets the optional shared_name attribute to value. +// +// value: It is necessary to match this name to the matching Unstage Op. +// If not specified, defaults to "" +func OrderedMapStageSharedName(value string) OrderedMapStageAttr { + return func(m optionalAttr) { + m["shared_name"] = value + } +} + +// Stage (key, values) in the underlying container which behaves like a ordered +// +// associative container. Elements are ordered by key. +// +// Arguments: +// key: int64 +// +// values: a list of tensors +// dtypes A list of data types that inserted values should adhere to. +// +// +// Returns the created operation. +func OrderedMapStage(scope *Scope, key tf.Output, indices tf.Output, values []tf.Output, dtypes []tf.DataType, optional ...OrderedMapStageAttr) (o *tf.Operation) { if scope.Err() != nil { return } + attrs := map[string]interface{}{"dtypes": dtypes} + for _, a := range optional { + a(attrs) + } opspec := tf.OpSpec{ - Type: "SparseTensorSliceDataset", + Type: "OrderedMapStage", Input: []tf.Input{ - indices, values, dense_shape, + key, indices, tf.OutputList(values), }, + Attrs: attrs, } - op := scope.AddOperation(opspec) - return op.Output(0) + return scope.AddOperation(opspec) } -// Returns x / y element-wise for real types. +// StackPushV2Attr is an optional argument to StackPushV2. +type StackPushV2Attr func(optionalAttr) + +// StackPushV2SwapMemory sets the optional swap_memory attribute to value. // -// If `x` and `y` are reals, this will return the floating-point division. +// value: Swap `elem` to CPU. Default to false. +// If not specified, defaults to false +func StackPushV2SwapMemory(value bool) StackPushV2Attr { + return func(m optionalAttr) { + m["swap_memory"] = value + } +} + +// Push an element onto the stack. // -// *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) { +// Arguments: +// handle: The handle to a stack. +// elem: The tensor to be pushed onto the stack. +// +// Returns The same tensor as the input 'elem'. +func StackPushV2(scope *Scope, handle tf.Output, elem tf.Output, optional ...StackPushV2Attr) (output tf.Output) { if scope.Err() != nil { return } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } opspec := tf.OpSpec{ - Type: "RealDiv", + Type: "StackPushV2", Input: []tf.Input{ - x, y, + handle, elem, }, + Attrs: attrs, } op := scope.AddOperation(opspec) return op.Output(0) @@ -22368,11 +24533,8 @@ func ConcatenateDataset(scope *Scope, input_dataset tf.Output, another_dataset t // Adds a value to the current value of a variable. // -// Any ReadVariableOp which depends directly or indirectly on this assign is -// guaranteed to see the incremented value or a subsequent newer one. -// -// Outputs the incremented value, which can be used to totally order the -// increments to this variable. +// Any ReadVariableOp with a control dependency on this op is guaranteed to +// see the incremented value or a subsequent newer one. // // Arguments: // resource: handle to the resource in which to store the variable. @@ -22409,6 +24571,63 @@ func LatencyStatsDataset(scope *Scope, input_dataset tf.Output, tag tf.Output, o return op.Output(0) } +// MapSizeAttr is an optional argument to MapSize. +type MapSizeAttr func(optionalAttr) + +// MapSizeCapacity sets the optional capacity attribute to value. +// If not specified, defaults to 0 +// +// REQUIRES: value >= 0 +func MapSizeCapacity(value int64) MapSizeAttr { + return func(m optionalAttr) { + m["capacity"] = value + } +} + +// MapSizeMemoryLimit sets the optional memory_limit attribute to value. +// If not specified, defaults to 0 +// +// REQUIRES: value >= 0 +func MapSizeMemoryLimit(value int64) MapSizeAttr { + return func(m optionalAttr) { + m["memory_limit"] = value + } +} + +// MapSizeContainer sets the optional container attribute to value. +// If not specified, defaults to "" +func MapSizeContainer(value string) MapSizeAttr { + return func(m optionalAttr) { + m["container"] = value + } +} + +// MapSizeSharedName sets the optional shared_name attribute to value. +// If not specified, defaults to "" +func MapSizeSharedName(value string) MapSizeAttr { + return func(m optionalAttr) { + m["shared_name"] = value + } +} + +// Op returns the number of elements in the underlying container. +func MapSize(scope *Scope, dtypes []tf.DataType, optional ...MapSizeAttr) (size tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"dtypes": dtypes} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "MapSize", + + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Convert JSON-encoded Example records to binary protocol buffer strings. // // This op translates a tensor containing Example records, encoded using @@ -22605,77 +24824,13 @@ func TensorArrayGradV2(scope *Scope, handle tf.Output, flow_in tf.Output, source if scope.Err() != nil { return } - attrs := map[string]interface{}{"source": source} - opspec := tf.OpSpec{ - Type: "TensorArrayGradV2", - Input: []tf.Input{ - handle, flow_in, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// ResourceSparseApplyAdadeltaAttr is an optional argument to ResourceSparseApplyAdadelta. -type ResourceSparseApplyAdadeltaAttr func(optionalAttr) - -// ResourceSparseApplyAdadeltaUseLocking 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 ResourceSparseApplyAdadeltaUseLocking(value bool) ResourceSparseApplyAdadeltaAttr { - return func(m optionalAttr) { - m["use_locking"] = value - } -} - -// var: Should be from a Variable(). -// -// Arguments: -// -// accum: Should be from a Variable(). -// accum_update: : Should be from a Variable(). -// lr: Learning rate. Must be a scalar. -// rho: Decay factor. Must be a scalar. -// epsilon: Constant factor. Must be a scalar. -// grad: The gradient. -// indices: A vector of indices into the first dimension of var and accum. -// -// Returns the created operation. -func ResourceSparseApplyAdadelta(scope *Scope, var_ tf.Output, accum tf.Output, accum_update tf.Output, lr tf.Output, rho tf.Output, epsilon tf.Output, grad tf.Output, indices tf.Output, optional ...ResourceSparseApplyAdadeltaAttr) (o *tf.Operation) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "ResourceSparseApplyAdadelta", - Input: []tf.Input{ - var_, accum, accum_update, lr, rho, epsilon, grad, indices, - }, - Attrs: attrs, - } - return scope.AddOperation(opspec) -} - -// Identity op for gradient debugging. -// -// This op is hidden from public in Python. It is used by TensorFlow Debugger to -// register gradient tensors for gradient debugging. -// This op operates on non-reference-type tensors. -func DebugGradientIdentity(scope *Scope, input tf.Output) (output tf.Output) { - if scope.Err() != nil { - return - } + attrs := map[string]interface{}{"source": source} opspec := tf.OpSpec{ - Type: "DebugGradientIdentity", + Type: "TensorArrayGradV2", Input: []tf.Input{ - input, + handle, flow_in, }, + Attrs: attrs, } op := scope.AddOperation(opspec) return op.Output(0) @@ -23099,6 +25254,35 @@ func MakeIterator(scope *Scope, dataset tf.Output, iterator tf.Output) (o *tf.Op return scope.AddOperation(opspec) } +// Makes the summary of accumulated stats for the batch. +// +// The summary stats contains gradients and hessians accumulated into the corresponding node and bucket for each example. +// +// Arguments: +// node_ids: int32 Rank 1 Tensor containing node ids, which each example falls into for the requested layer. +// gradients: float32; Rank 2 Tensor (shape=[#examples, 1]) for gradients. +// hessians: float32; Rank 2 Tensor (shape=[#examples, 1]) for hessians. +// bucketized_features_list: int32 list of Rank 1 Tensors, each containing the bucketized feature (for each feature column). +// max_splits: int; the maximum number of splits possible in the whole tree. +// num_buckets: int; equals to the maximum possible value of bucketized feature. +// +// Returns output Rank 4 Tensor (shape=[#features, #splits, #buckets, 2]) containing accumulated stats put into the corresponding node and bucket. The first index of 4th dimension refers to gradients, and the second to hessians. +func BoostedTreesMakeStatsSummary(scope *Scope, node_ids tf.Output, gradients tf.Output, hessians tf.Output, bucketized_features_list []tf.Output, max_splits int64, num_buckets int64) (stats_summary tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"max_splits": max_splits, "num_buckets": num_buckets} + opspec := tf.OpSpec{ + Type: "BoostedTreesMakeStatsSummary", + Input: []tf.Input{ + node_ids, gradients, hessians, tf.OutputList(bucketized_features_list), + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Adjust the contrast of one or more images. // // `images` is a tensor of at least 3 dimensions. The last 3 dimensions are @@ -23263,42 +25447,6 @@ func IteratorFromStringHandle(scope *Scope, string_handle tf.Output, optional .. return op.Output(0) } -// Computes arctangent of `y/x` element-wise, respecting signs of the arguments. -// -// This is the angle \( \theta \in [-\pi, \pi] \) such that -// \[ x = r \cos(\theta) \] -// and -// \[ y = r \sin(\theta) \] -// where \(r = \sqrt(x^2 + y^2) \). -func Atan2(scope *Scope, y tf.Output, x tf.Output) (z tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "Atan2", - Input: []tf.Input{ - y, x, - }, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - -// Return a tensor with the same shape and contents as the input tensor or value. -func Identity(scope *Scope, input tf.Output) (output tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "Identity", - Input: []tf.Input{ - input, - }, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - // Gather slices from `params` axis `axis` according to `indices`. // // `indices` must be an integer tensor of any dimension (usually 0-D or 1-D). @@ -23323,6 +25471,10 @@ func Identity(scope *Scope, input tf.Output) (output tf.Output) { // //
// +// Note that on CPU, if an out of bound index is found, an error is returned. +// On GPU, if an out of bound index is found, a 0 is stored in the +// corresponding output value. +// // Arguments: // params: The tensor from which to gather values. Must be at least rank // `axis + 1`. @@ -23819,6 +25971,28 @@ func InvGrad(scope *Scope, y tf.Output, dy tf.Output) (z tf.Output) { return op.Output(0) } +// List of the given size with empty elements. +// +// element_shape: the shape of the future elements of the list +// num_elements: the number of elements to reserve +// handle: the output list +// element_dtype: the desired type of elements in the list. +func TensorListReserve(scope *Scope, element_shape tf.Output, num_elements tf.Output, element_dtype tf.DataType) (handle tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"element_dtype": element_dtype} + opspec := tf.OpSpec{ + Type: "TensorListReserve", + Input: []tf.Input{ + element_shape, num_elements, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // PriorityQueueV2Attr is an optional argument to PriorityQueueV2. type PriorityQueueV2Attr func(optionalAttr) @@ -24358,122 +26532,6 @@ func StackV2(scope *Scope, max_size tf.Output, elem_type tf.DataType, optional . return op.Output(0) } -// OrderedMapStageAttr is an optional argument to OrderedMapStage. -type OrderedMapStageAttr func(optionalAttr) - -// OrderedMapStageCapacity sets the optional capacity attribute to value. -// -// value: Maximum number of elements in the Staging Area. If > 0, inserts -// on the container will block when the capacity is reached. -// If not specified, defaults to 0 -// -// REQUIRES: value >= 0 -func OrderedMapStageCapacity(value int64) OrderedMapStageAttr { - return func(m optionalAttr) { - m["capacity"] = value - } -} - -// OrderedMapStageMemoryLimit sets the optional memory_limit attribute to value. -// If not specified, defaults to 0 -// -// REQUIRES: value >= 0 -func OrderedMapStageMemoryLimit(value int64) OrderedMapStageAttr { - return func(m optionalAttr) { - m["memory_limit"] = value - } -} - -// OrderedMapStageContainer sets the optional container attribute to value. -// -// value: If non-empty, this queue is placed in the given container. Otherwise, -// a default container is used. -// If not specified, defaults to "" -func OrderedMapStageContainer(value string) OrderedMapStageAttr { - return func(m optionalAttr) { - m["container"] = value - } -} - -// OrderedMapStageSharedName sets the optional shared_name attribute to value. -// -// value: It is necessary to match this name to the matching Unstage Op. -// If not specified, defaults to "" -func OrderedMapStageSharedName(value string) OrderedMapStageAttr { - return func(m optionalAttr) { - m["shared_name"] = value - } -} - -// Stage (key, values) in the underlying container which behaves like a ordered -// -// associative container. Elements are ordered by key. -// -// Arguments: -// key: int64 -// -// values: a list of tensors -// dtypes A list of data types that inserted values should adhere to. -// -// -// Returns the created operation. -func OrderedMapStage(scope *Scope, key tf.Output, indices tf.Output, values []tf.Output, dtypes []tf.DataType, optional ...OrderedMapStageAttr) (o *tf.Operation) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{"dtypes": dtypes} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "OrderedMapStage", - Input: []tf.Input{ - key, indices, tf.OutputList(values), - }, - Attrs: attrs, - } - return scope.AddOperation(opspec) -} - -// StackPushV2Attr is an optional argument to StackPushV2. -type StackPushV2Attr func(optionalAttr) - -// StackPushV2SwapMemory sets the optional swap_memory attribute to value. -// -// value: Swap `elem` to CPU. Default to false. -// If not specified, defaults to false -func StackPushV2SwapMemory(value bool) StackPushV2Attr { - return func(m optionalAttr) { - m["swap_memory"] = value - } -} - -// Push an element onto the stack. -// -// Arguments: -// handle: The handle to a stack. -// elem: The tensor to be pushed onto the stack. -// -// Returns The same tensor as the input 'elem'. -func StackPushV2(scope *Scope, handle tf.Output, elem tf.Output, optional ...StackPushV2Attr) (output tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "StackPushV2", - Input: []tf.Input{ - handle, elem, - }, - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - // FusedBatchNormGradV2Attr is an optional argument to FusedBatchNormGradV2. type FusedBatchNormGradV2Attr func(optionalAttr) @@ -24559,25 +26617,135 @@ type DecodeCompressedAttr func(optionalAttr) // If not specified, defaults to "" func DecodeCompressedCompressionType(value string) DecodeCompressedAttr { return func(m optionalAttr) { - m["compression_type"] = value + m["compression_type"] = value + } +} + +// Decompress strings. +// +// This op decompresses each element of the `bytes` input `Tensor`, which +// is assumed to be compressed using the given `compression_type`. +// +// The `output` is a string `Tensor` of the same shape as `bytes`, +// each element containing the decompressed data from the corresponding +// element in `bytes`. +// +// Arguments: +// bytes: A Tensor of string which is compressed. +// +// Returns A Tensor with the same shape as input `bytes`, uncompressed +// from bytes. +func DecodeCompressed(scope *Scope, bytes tf.Output, optional ...DecodeCompressedAttr) (output tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "DecodeCompressed", + Input: []tf.Input{ + bytes, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// CudnnRNNAttr is an optional argument to CudnnRNN. +type CudnnRNNAttr func(optionalAttr) + +// CudnnRNNRnnMode sets the optional rnn_mode attribute to value. +// If not specified, defaults to "lstm" +func CudnnRNNRnnMode(value string) CudnnRNNAttr { + return func(m optionalAttr) { + m["rnn_mode"] = value + } +} + +// CudnnRNNInputMode sets the optional input_mode attribute to value. +// If not specified, defaults to "linear_input" +func CudnnRNNInputMode(value string) CudnnRNNAttr { + return func(m optionalAttr) { + m["input_mode"] = value + } +} + +// CudnnRNNDirection sets the optional direction attribute to value. +// If not specified, defaults to "unidirectional" +func CudnnRNNDirection(value string) CudnnRNNAttr { + return func(m optionalAttr) { + m["direction"] = value + } +} + +// CudnnRNNDropout sets the optional dropout attribute to value. +// If not specified, defaults to 0 +func CudnnRNNDropout(value float32) CudnnRNNAttr { + return func(m optionalAttr) { + m["dropout"] = value + } +} + +// CudnnRNNSeed sets the optional seed attribute to value. +// If not specified, defaults to 0 +func CudnnRNNSeed(value int64) CudnnRNNAttr { + return func(m optionalAttr) { + m["seed"] = value + } +} + +// CudnnRNNSeed2 sets the optional seed2 attribute to value. +// If not specified, defaults to 0 +func CudnnRNNSeed2(value int64) CudnnRNNAttr { + return func(m optionalAttr) { + m["seed2"] = value + } +} + +// CudnnRNNIsTraining sets the optional is_training attribute to value. +// If not specified, defaults to true +func CudnnRNNIsTraining(value bool) CudnnRNNAttr { + return func(m optionalAttr) { + m["is_training"] = value } } -// Decompress strings. -// -// This op decompresses each element of the `bytes` input `Tensor`, which -// is assumed to be compressed using the given `compression_type`. -// -// The `output` is a string `Tensor` of the same shape as `bytes`, -// each element containing the decompressed data from the corresponding -// element in `bytes`. -// -// Arguments: -// bytes: A Tensor of string which is compressed. -// -// Returns A Tensor with the same shape as input `bytes`, uncompressed -// from bytes. -func DecodeCompressed(scope *Scope, bytes tf.Output, optional ...DecodeCompressedAttr) (output tf.Output) { +// A RNN backed by cuDNN. +// +// Computes the RNN from the input and initial states, with respect to the params +// buffer. +// +// rnn_mode: Indicates the type of the RNN model. +// input_mode: Indicate whether there is a linear projection between the input and +// The actual computation before the first layer. 'skip_input' is only allowed +// when input_size == num_units; 'auto_select' implies 'skip_input' when +// input_size == num_units; otherwise, it implies 'linear_input'. +// direction: Indicates whether a bidirectional model will be used. +// dir = (direction == bidirectional) ? 2 : 1 +// dropout: dropout probability. When set to 0., dropout is disabled. +// seed: the 1st part of a seed to initialize dropout. +// seed2: the 2nd part of a seed to initialize dropout. +// input: a 3-D tensor with the shape of [seq_length, batch_size, input_size]. +// input_h: a 3-D tensor with the shape of [num_layer * dir, batch_size, +// num_units]. +// input_c: For LSTM, a 3-D tensor with the shape of +// [num_layer * dir, batch, num_units]. For other models, it is ignored. +// params: a 1-D tensor that contains the weights and biases in an opaque layout. +// The size must be created through CudnnRNNParamsSize, and initialized +// separately. Note that they might not be compatible across different +// generations. So it is a good idea to save and restore +// output: a 3-D tensor with the shape of [seq_length, batch_size, +// dir * num_units]. +// output_h: the same shape has input_h. +// output_c: the same shape as input_c for LSTM. An empty tensor for other models. +// is_training: Indicates whether this operation is used for inferenece or +// training. +// reserve_space: an opaque tensor that can be used in backprop calculation. It +// is only produced if is_training is false. +func CudnnRNN(scope *Scope, input tf.Output, input_h tf.Output, input_c tf.Output, params tf.Output, optional ...CudnnRNNAttr) (output tf.Output, output_h tf.Output, output_c tf.Output, reserve_space tf.Output) { if scope.Err() != nil { return } @@ -24586,14 +26754,14 @@ func DecodeCompressed(scope *Scope, bytes tf.Output, optional ...DecodeCompresse a(attrs) } opspec := tf.OpSpec{ - Type: "DecodeCompressed", + Type: "CudnnRNN", Input: []tf.Input{ - bytes, + input, input_h, input_c, params, }, Attrs: attrs, } op := scope.AddOperation(opspec) - return op.Output(0) + return op.Output(0), op.Output(1), op.Output(2), op.Output(3) } // Compare values of `input` to `threshold` and pack resulting bits into a `uint8`. @@ -24689,6 +26857,47 @@ func TensorArrayScatterV3(scope *Scope, handle tf.Output, indices tf.Output, val return op.Output(0) } +// EmptyAttr is an optional argument to Empty. +type EmptyAttr func(optionalAttr) + +// EmptyInit sets the optional init attribute to value. +// +// value: If True, initialize the returned tensor with the default value of dtype. Otherwise, the implementation is free not to initializethe tensor's content. +// If not specified, defaults to false +func EmptyInit(value bool) EmptyAttr { + return func(m optionalAttr) { + m["init"] = value + } +} + +// Creates a tensor with the given shape. +// +// This operation creates a tensor of `shape` and `dtype`. +// +// Arguments: +// shape: 1-D. Represents the shape of the output tensor. +// +// +// Returns A `Tensor` of type `T`. +func Empty(scope *Scope, shape tf.Output, dtype tf.DataType, optional ...EmptyAttr) (output tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"dtype": dtype} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "Empty", + Input: []tf.Input{ + shape, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // TensorArrayConcatV3Attr is an optional argument to TensorArrayConcatV3. type TensorArrayConcatV3Attr func(optionalAttr) @@ -24806,6 +27015,27 @@ func ParameterizedTruncatedNormal(scope *Scope, shape tf.Output, means tf.Output return op.Output(0) } +// Sets the index-th position of the list to contain the given tensor. +// +// input_handle: the list +// index: the position in the list to which the tensor will be assigned +// item: the element to be assigned to that position +// output_handle: the new list, with the element in the proper position +// +func TensorListSetItem(scope *Scope, input_handle tf.Output, index tf.Output, item tf.Output) (output_handle tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "TensorListSetItem", + Input: []tf.Input{ + input_handle, index, item, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Returns a diagonal tensor with a given diagonal values. // // Given a `diagonal`, this operation returns a tensor with the `diagonal` and @@ -25350,6 +27580,27 @@ func TensorArrayScatterV2(scope *Scope, handle tf.Output, indices tf.Output, val return op.Output(0) } +// Creates a tree ensemble model and returns a handle to it. +// +// Arguments: +// tree_ensemble_handle: Handle to the tree ensemble resource to be created. +// stamp_token: Token to use as the initial value of the resource stamp. +// tree_ensemble_serialized: Serialized proto of the tree ensemble. +// +// Returns the created operation. +func BoostedTreesCreateEnsemble(scope *Scope, tree_ensemble_handle tf.Output, stamp_token tf.Output, tree_ensemble_serialized tf.Output) (o *tf.Operation) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "BoostedTreesCreateEnsemble", + Input: []tf.Input{ + tree_ensemble_handle, stamp_token, tree_ensemble_serialized, + }, + } + return scope.AddOperation(opspec) +} + // Applies sparse addition to `input` using individual values or slices // // from `updates` according to indices `indices`. The updates are non-aliasing: @@ -25946,63 +28197,6 @@ func MapUnstage(scope *Scope, key tf.Output, indices tf.Output, dtypes []tf.Data return values } -// MapSizeAttr is an optional argument to MapSize. -type MapSizeAttr func(optionalAttr) - -// MapSizeCapacity sets the optional capacity attribute to value. -// If not specified, defaults to 0 -// -// REQUIRES: value >= 0 -func MapSizeCapacity(value int64) MapSizeAttr { - return func(m optionalAttr) { - m["capacity"] = value - } -} - -// MapSizeMemoryLimit sets the optional memory_limit attribute to value. -// If not specified, defaults to 0 -// -// REQUIRES: value >= 0 -func MapSizeMemoryLimit(value int64) MapSizeAttr { - return func(m optionalAttr) { - m["memory_limit"] = value - } -} - -// MapSizeContainer sets the optional container attribute to value. -// If not specified, defaults to "" -func MapSizeContainer(value string) MapSizeAttr { - return func(m optionalAttr) { - m["container"] = value - } -} - -// MapSizeSharedName sets the optional shared_name attribute to value. -// If not specified, defaults to "" -func MapSizeSharedName(value string) MapSizeAttr { - return func(m optionalAttr) { - m["shared_name"] = value - } -} - -// Op returns the number of elements in the underlying container. -func MapSize(scope *Scope, dtypes []tf.DataType, optional ...MapSizeAttr) (size tf.Output) { - if scope.Err() != nil { - return - } - attrs := map[string]interface{}{"dtypes": dtypes} - for _, a := range optional { - a(attrs) - } - opspec := tf.OpSpec{ - Type: "MapSize", - - Attrs: attrs, - } - op := scope.AddOperation(opspec) - return op.Output(0) -} - // MapIncompleteSizeAttr is an optional argument to MapIncompleteSize. type MapIncompleteSizeAttr func(optionalAttr) @@ -26230,6 +28424,120 @@ func ShapeN(scope *Scope, input []tf.Output, optional ...ShapeNAttr) (output []t return output } +// CudnnRNNParamsToCanonicalAttr is an optional argument to CudnnRNNParamsToCanonical. +type CudnnRNNParamsToCanonicalAttr func(optionalAttr) + +// CudnnRNNParamsToCanonicalRnnMode sets the optional rnn_mode attribute to value. +// If not specified, defaults to "lstm" +func CudnnRNNParamsToCanonicalRnnMode(value string) CudnnRNNParamsToCanonicalAttr { + return func(m optionalAttr) { + m["rnn_mode"] = value + } +} + +// CudnnRNNParamsToCanonicalInputMode sets the optional input_mode attribute to value. +// If not specified, defaults to "linear_input" +func CudnnRNNParamsToCanonicalInputMode(value string) CudnnRNNParamsToCanonicalAttr { + return func(m optionalAttr) { + m["input_mode"] = value + } +} + +// CudnnRNNParamsToCanonicalDirection sets the optional direction attribute to value. +// If not specified, defaults to "unidirectional" +func CudnnRNNParamsToCanonicalDirection(value string) CudnnRNNParamsToCanonicalAttr { + return func(m optionalAttr) { + m["direction"] = value + } +} + +// CudnnRNNParamsToCanonicalDropout sets the optional dropout attribute to value. +// If not specified, defaults to 0 +func CudnnRNNParamsToCanonicalDropout(value float32) CudnnRNNParamsToCanonicalAttr { + return func(m optionalAttr) { + m["dropout"] = value + } +} + +// CudnnRNNParamsToCanonicalSeed sets the optional seed attribute to value. +// If not specified, defaults to 0 +func CudnnRNNParamsToCanonicalSeed(value int64) CudnnRNNParamsToCanonicalAttr { + return func(m optionalAttr) { + m["seed"] = value + } +} + +// CudnnRNNParamsToCanonicalSeed2 sets the optional seed2 attribute to value. +// If not specified, defaults to 0 +func CudnnRNNParamsToCanonicalSeed2(value int64) CudnnRNNParamsToCanonicalAttr { + return func(m optionalAttr) { + m["seed2"] = value + } +} + +// Retrieves CudnnRNN params in canonical form. +// +// Retrieves a set of weights from the opaque params buffer that can be saved and +// restored in a way compatible with future runs. +// +// Note that the params buffer may not be compatible across different GPUs. So any +// save and restoration should be converted to and from the canonical weights and +// biases. +// +// num_layers: Specifies the number of layers in the RNN model. +// num_units: Specifies the size of the hidden state. +// input_size: Specifies the size of the input state. +// num_params: number of parameter sets for all layers. +// Each layer may contain multiple parameter sets, with each set consisting of +// a weight matrix and a bias vector. +// weights: the canonical form of weights that can be used for saving +// and restoration. They are more likely to be compatible across different +// generations. +// biases: the canonical form of biases that can be used for saving +// and restoration. They are more likely to be compatible across different +// generations. +// rnn_mode: Indicates the type of the RNN model. +// input_mode: Indicate whether there is a linear projection between the input and +// The actual computation before the first layer. 'skip_input' is only allowed +// when input_size == num_units; 'auto_select' implies 'skip_input' when +// input_size == num_units; otherwise, it implies 'linear_input'. +// direction: Indicates whether a bidirectional model will be used. +// dir = (direction == bidirectional) ? 2 : 1 +// dropout: dropout probability. When set to 0., dropout is disabled. +// seed: the 1st part of a seed to initialize dropout. +// seed2: the 2nd part of a seed to initialize dropout. +func CudnnRNNParamsToCanonical(scope *Scope, num_layers tf.Output, num_units tf.Output, input_size tf.Output, params tf.Output, num_params int64, optional ...CudnnRNNParamsToCanonicalAttr) (weights []tf.Output, biases []tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"num_params": num_params} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "CudnnRNNParamsToCanonical", + Input: []tf.Input{ + num_layers, num_units, input_size, params, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + if scope.Err() != nil { + return + } + var idx int + var err error + if weights, idx, err = makeOutputList(op, idx, "weights"); err != nil { + scope.UpdateErr("CudnnRNNParamsToCanonical", err) + return + } + if biases, idx, err = makeOutputList(op, idx, "biases"); err != nil { + scope.UpdateErr("CudnnRNNParamsToCanonical", err) + return + } + return weights, biases +} + // UniformCandidateSamplerAttr is an optional argument to UniformCandidateSampler. type UniformCandidateSamplerAttr func(optionalAttr) @@ -26424,52 +28732,174 @@ func CTCGreedyDecoder(scope *Scope, inputs tf.Output, sequence_length tf.Output, }, Attrs: attrs, } - op := scope.AddOperation(opspec) - return op.Output(0), op.Output(1), op.Output(2), op.Output(3) + op := scope.AddOperation(opspec) + return op.Output(0), op.Output(1), op.Output(2), op.Output(3) +} + +// Forwards `data` to the output port determined by `pred`. +// +// If `pred` is true, the `data` input is forwarded to `output_true`. Otherwise, +// the data goes to `output_false`. +// +// See also `RefSwitch` and `Merge`. +// +// Arguments: +// data: The tensor to be forwarded to the appropriate output. +// pred: A scalar that specifies which output port will receive data. +// +// Returns If `pred` is false, data will be forwarded to this output.If `pred` is true, data will be forwarded to this output. +func Switch(scope *Scope, data tf.Output, pred tf.Output) (output_false tf.Output, output_true tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "Switch", + Input: []tf.Input{ + data, pred, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0), op.Output(1) +} + +// Add all input tensors element wise. +// +// Arguments: +// inputs: Must all be the same size and shape. +func AddN(scope *Scope, inputs []tf.Output) (sum tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "AddN", + Input: []tf.Input{ + tf.OutputList(inputs), + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// TryRpcAttr is an optional argument to TryRpc. +type TryRpcAttr func(optionalAttr) + +// TryRpcProtocol sets the optional protocol attribute to value. +// +// value: RPC protocol to use. Empty string means use the default protocol. +// Options include 'grpc'. +// If not specified, defaults to "" +func TryRpcProtocol(value string) TryRpcAttr { + return func(m optionalAttr) { + m["protocol"] = value + } +} + +// TryRpcFailFast sets the optional fail_fast attribute to value. +// +// value: `boolean`. If `true` (default), then failures to connect +// (i.e., the server does not immediately respond) cause an RPC failure. +// If not specified, defaults to true +func TryRpcFailFast(value bool) TryRpcAttr { + return func(m optionalAttr) { + m["fail_fast"] = value + } +} + +// TryRpcTimeoutInMs sets the optional timeout_in_ms attribute to value. +// +// value: `int`. If `0` (default), then the kernel will run the RPC +// request and only time out if the RPC deadline passes or the session times out. +// If this value is greater than `0`, then the op will raise an exception if +// the RPC takes longer than `timeout_in_ms`. +// If not specified, defaults to 0 +func TryRpcTimeoutInMs(value int64) TryRpcAttr { + return func(m optionalAttr) { + m["timeout_in_ms"] = value + } } -// Forwards `data` to the output port determined by `pred`. +// Perform batches of RPC requests. // -// If `pred` is true, the `data` input is forwarded to `output_true`. Otherwise, -// the data goes to `output_false`. +// This op asynchronously performs either a single RPC request, or a batch +// of requests. RPC requests are defined by three main parameters: // -// See also `RefSwitch` and `Merge`. +// - `address` (the host+port or BNS address of the request) +// - `method` (the method name for the request) +// - `request` (the serialized proto string, or vector of strings, +// of the RPC request argument). // -// Arguments: -// data: The tensor to be forwarded to the appropriate output. -// pred: A scalar that specifies which output port will receive data. +// For example, if you have an RPC service running on port localhost:2345, +// and its interface is configured with the following proto declaration: // -// Returns If `pred` is false, data will be forwarded to this output.If `pred` is true, data will be forwarded to this output. -func Switch(scope *Scope, data tf.Output, pred tf.Output) (output_false tf.Output, output_true tf.Output) { - if scope.Err() != nil { - return - } - opspec := tf.OpSpec{ - Type: "Switch", - Input: []tf.Input{ - data, pred, - }, - } - op := scope.AddOperation(opspec) - return op.Output(0), op.Output(1) -} - -// Add all input tensors element wise. +// ``` +// service MyService { +// rpc MyMethod(MyRequestProto) returns (MyResponseProto) { +// } +// }; +// ``` +// +// then call this op with arguments: +// +// ``` +// address = "localhost:2345" +// method = "MyService/MyMethod" +// ``` +// +// The `request` tensor is a string tensor representing serialized `MyRequestProto` +// strings; and the output string tensor `response` will have the same shape +// and contain (upon successful completion) corresponding serialized +// `MyResponseProto` strings. +// +// For example, to send a single, empty, `MyRequestProto`, call +// this op with `request = ""`. To send 5 **parallel** empty requests, +// call this op with `request = ["", "", "", "", ""]`. +// +// More generally, one can create a batch of `MyRequestProto` serialized protos +// from regular batched tensors using the `encode_proto` op, and convert +// the response `MyResponseProto` serialized protos to batched tensors +// using the `decode_proto` op. +// +// **NOTE** Working with serialized proto strings is faster than instantiating +// actual proto objects in memory, so no performance degradation is expected +// compared to writing custom kernels for this workflow. +// +// Unlike the standard `Rpc` op, if the connection fails or the remote worker +// returns an error status, this op does **not** reraise the exception. +// Instead, the `status_code` and `status_message` entry for the corresponding RPC +// call is set with the error returned from the RPC call. The `response` tensor +// will contain valid response values for those minibatch entries whose RPCs did +// not fail; the rest of the entries will have empty strings. // // Arguments: -// inputs: Must all be the same size and shape. -func AddN(scope *Scope, inputs []tf.Output) (sum tf.Output) { +// address: `0-D` or `1-D`. The address (i.e. host_name:port) of the RPC server. +// If this tensor has more than 1 element, then multiple parallel rpc requests +// are sent. This argument broadcasts with `method` and `request`. +// method: `0-D` or `1-D`. The method address on the RPC server. +// If this tensor has more than 1 element, then multiple parallel rpc requests +// are sent. This argument broadcasts with `address` and `request`. +// request: `0-D` or `1-D`. Serialized proto strings: the rpc request argument. +// If this tensor has more than 1 element, then multiple parallel rpc requests +// are sent. This argument broadcasts with `address` and `method`. +// +// Returns Same shape as `request`. Serialized proto strings: the rpc responses.Same shape as `request`. Values correspond to tensorflow Status enum codes.Same shape as `request`. Values correspond to Status messages +// returned from the RPC calls. +func TryRpc(scope *Scope, address tf.Output, method tf.Output, request tf.Output, optional ...TryRpcAttr) (response tf.Output, status_code tf.Output, status_message tf.Output) { if scope.Err() != nil { return } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } opspec := tf.OpSpec{ - Type: "AddN", + Type: "TryRpc", Input: []tf.Input{ - tf.OutputList(inputs), + address, method, request, }, + Attrs: attrs, } op := scope.AddOperation(opspec) - return op.Output(0) + return op.Output(0), op.Output(1), op.Output(2) } // EnterAttr is an optional argument to Enter. @@ -26907,6 +29337,64 @@ func ParseSingleExample(scope *Scope, serialized tf.Output, dense_defaults []tf. return sparse_indices, sparse_values, sparse_shapes, dense_values } +// Deserializes a serialized tree ensemble config and replaces current tree +// +// ensemble. +// +// Arguments: +// tree_ensemble_handle: Handle to the tree ensemble. +// stamp_token: Token to use as the new value of the resource stamp. +// tree_ensemble_serialized: Serialized proto of the ensemble. +// +// Returns the created operation. +func BoostedTreesDeserializeEnsemble(scope *Scope, tree_ensemble_handle tf.Output, stamp_token tf.Output, tree_ensemble_serialized tf.Output) (o *tf.Operation) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "BoostedTreesDeserializeEnsemble", + Input: []tf.Input{ + tree_ensemble_handle, stamp_token, tree_ensemble_serialized, + }, + } + return scope.AddOperation(opspec) +} + +// Runs multiple additive regression ensemble predictors on input instances and +// +// computes the update to cached logits. It is designed to be used during training. +// It traverses the trees starting from cached tree id and cached node id and +// calculates the updates to be pushed to the cache. +// +// Arguments: +// +// cached_tree_ids: Rank 1 Tensor containing cached tree ids which is the starting +// tree of prediction. +// cached_node_ids: Rank 1 Tensor containing cached node id which is the starting +// node of prediction. +// bucketized_features: A list of rank 1 Tensors containing bucket id for each +// feature. +// logits_dimension: scalar, dimension of the logits, to be used for partial logits +// shape. +// +// Returns Rank 2 Tensor containing logits update (with respect to cached +// values stored) for each example.Rank 1 Tensor containing new tree ids for each example.Rank 1 Tensor containing new node ids in the new tree_ids. +func BoostedTreesTrainingPredict(scope *Scope, tree_ensemble_handle tf.Output, cached_tree_ids tf.Output, cached_node_ids tf.Output, bucketized_features []tf.Output, logits_dimension int64) (partial_logits tf.Output, tree_ids tf.Output, node_ids tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"logits_dimension": logits_dimension} + opspec := tf.OpSpec{ + Type: "BoostedTreesTrainingPredict", + Input: []tf.Input{ + tree_ensemble_handle, cached_tree_ids, cached_node_ids, tf.OutputList(bucketized_features), + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0), op.Output(1), op.Output(2) +} + // Elementwise computes the bitwise AND of `x` and `y`. // // The result will have those bits set, that are set in both `x` and `y`. The @@ -26943,6 +29431,44 @@ func LeftShift(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) { return op.Output(0) } +// TensorListStackAttr is an optional argument to TensorListStack. +type TensorListStackAttr func(optionalAttr) + +// TensorListStackNumElements sets the optional num_elements attribute to value. +// If not specified, defaults to -1 +func TensorListStackNumElements(value int64) TensorListStackAttr { + return func(m optionalAttr) { + m["num_elements"] = value + } +} + +// Stacks all tensors in the list. +// +// Requires that all tensors have the same shape. +// +// input_handle: the input list +// tensor: the gathered result +// num_elements: optional. If not -1, the number of elements in the list. +// +func TensorListStack(scope *Scope, input_handle tf.Output, element_dtype tf.DataType, optional ...TensorListStackAttr) (tensor tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"element_dtype": element_dtype} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "TensorListStack", + Input: []tf.Input{ + input_handle, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Elementwise computes the bitwise right-shift of `x` and `y`. // // Performs a logical shift for unsigned integer types, and an arithmetic shift @@ -26992,6 +29518,175 @@ func AdjustHue(scope *Scope, images tf.Output, delta tf.Output) (output tf.Outpu return op.Output(0) } +// BatchAttr is an optional argument to Batch. +type BatchAttr func(optionalAttr) + +// BatchMaxEnqueuedBatches sets the optional max_enqueued_batches attribute to value. +// If not specified, defaults to 10 +func BatchMaxEnqueuedBatches(value int64) BatchAttr { + return func(m optionalAttr) { + m["max_enqueued_batches"] = value + } +} + +// BatchAllowedBatchSizes sets the optional allowed_batch_sizes attribute to value. +// If not specified, defaults to <> +func BatchAllowedBatchSizes(value []int64) BatchAttr { + return func(m optionalAttr) { + m["allowed_batch_sizes"] = value + } +} + +// BatchContainer sets the optional container attribute to value. +// If not specified, defaults to "" +func BatchContainer(value string) BatchAttr { + return func(m optionalAttr) { + m["container"] = value + } +} + +// BatchSharedName sets the optional shared_name attribute to value. +// If not specified, defaults to "" +func BatchSharedName(value string) BatchAttr { + return func(m optionalAttr) { + m["shared_name"] = value + } +} + +// BatchBatchingQueue sets the optional batching_queue attribute to value. +// If not specified, defaults to "" +func BatchBatchingQueue(value string) BatchAttr { + return func(m optionalAttr) { + m["batching_queue"] = value + } +} + +// Batches all input tensors nondeterministically. +// +// When many instances of this Op are being run concurrently with the same +// container/shared_name in the same device, some will output zero-shaped Tensors +// and others will output Tensors of size up to max_batch_size. +// +// All Tensors in in_tensors are batched together (so, for example, labels and +// features should be batched with a single instance of this operation. +// +// Each invocation of batch emits an `id` scalar which will be used to identify +// this particular invocation when doing unbatch or its gradient. +// +// Each op which emits a non-empty batch will also emit a non-empty batch_index +// Tensor, which, is a [K, 3] matrix where each row contains the invocation's id, +// start, and length of elements of each set of Tensors present in batched_tensors. +// +// Batched tensors are concatenated along the first dimension, and all tensors in +// in_tensors must have the first dimension of the same size. +// +// in_tensors: The tensors to be batched. +// num_batch_threads: Number of scheduling threads for processing batches of work. +// Determines the number of batches processed in parallel. +// max_batch_size: Batch sizes will never be bigger than this. +// batch_timeout_micros: Maximum number of microseconds to wait before outputting +// an incomplete batch. +// allowed_batch_sizes: Optional list of allowed batch sizes. If left empty, does +// nothing. Otherwise, supplies a list of batch sizes, causing the op to pad +// batches up to one of those sizes. The entries must increase monotonically, and +// the final entry must equal max_batch_size. +// grad_timeout_micros: The timeout to use for the gradient. See Unbatch. +// batched_tensors: Either empty tensors or a batch of concatenated Tensors. +// batch_index: If out_tensors is non-empty, has information to invert it. +// container: Controls the scope of sharing of this batch. +// id: always contains a scalar with a unique ID for this invocation of Batch. +// shared_name: Concurrently running instances of batch in the same device with the +// same container and shared_name will batch their elements together. If left +// empty, the op name will be used as the shared name. +// T: the types of tensors to be batched. +func Batch(scope *Scope, in_tensors []tf.Output, num_batch_threads int64, max_batch_size int64, batch_timeout_micros int64, grad_timeout_micros int64, optional ...BatchAttr) (batched_tensors []tf.Output, batch_index tf.Output, id tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"num_batch_threads": num_batch_threads, "max_batch_size": max_batch_size, "batch_timeout_micros": batch_timeout_micros, "grad_timeout_micros": grad_timeout_micros} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "Batch", + Input: []tf.Input{ + tf.OutputList(in_tensors), + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + if scope.Err() != nil { + return + } + var idx int + var err error + if batched_tensors, idx, err = makeOutputList(op, idx, "batched_tensors"); err != nil { + scope.UpdateErr("Batch", err) + return + } + batch_index = op.Output(idx) + id = op.Output(idx) + return batched_tensors, batch_index, id +} + +// UnbatchAttr is an optional argument to Unbatch. +type UnbatchAttr func(optionalAttr) + +// UnbatchContainer sets the optional container attribute to value. +// If not specified, defaults to "" +func UnbatchContainer(value string) UnbatchAttr { + return func(m optionalAttr) { + m["container"] = value + } +} + +// UnbatchSharedName sets the optional shared_name attribute to value. +// If not specified, defaults to "" +func UnbatchSharedName(value string) UnbatchAttr { + return func(m optionalAttr) { + m["shared_name"] = value + } +} + +// Reverses the operation of Batch for a single output Tensor. +// +// An instance of Unbatch either receives an empty batched_tensor, in which case it +// asynchronously waits until the values become available from a concurrently +// running instance of Unbatch with the same container and shared_name, or receives +// a non-empty batched_tensor in which case it finalizes all other concurrently +// running instances and outputs its own element from the batch. +// +// batched_tensor: The possibly transformed output of Batch. The size of the first +// dimension should remain unchanged by the transformations for the operation to +// work. +// batch_index: The matching batch_index obtained from Batch. +// id: The id scalar emitted by Batch. +// unbatched_tensor: The Tensor corresponding to this execution. +// timeout_micros: Maximum amount of time (in microseconds) to wait to receive the +// batched input tensor associated with a given invocation of the op. +// container: Container to control resource sharing. +// shared_name: Instances of Unbatch with the same container and shared_name are +// assumed to possibly belong to the same batch. If left empty, the op name will +// be used as the shared name. +func Unbatch(scope *Scope, batched_tensor tf.Output, batch_index tf.Output, id tf.Output, timeout_micros int64, optional ...UnbatchAttr) (unbatched_tensor tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{"timeout_micros": timeout_micros} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "Unbatch", + Input: []tf.Input{ + batched_tensor, batch_index, id, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // AvgPool3DGradAttr is an optional argument to AvgPool3DGrad. type AvgPool3DGradAttr func(optionalAttr) @@ -27204,6 +29899,60 @@ func ParseSingleSequenceExample(scope *Scope, serialized tf.Output, feature_list return context_sparse_indices, context_sparse_values, context_sparse_shapes, context_dense_values, feature_list_sparse_indices, feature_list_sparse_values, feature_list_sparse_shapes, feature_list_dense_values } +// UnbatchGradAttr is an optional argument to UnbatchGrad. +type UnbatchGradAttr func(optionalAttr) + +// UnbatchGradContainer sets the optional container attribute to value. +// If not specified, defaults to "" +func UnbatchGradContainer(value string) UnbatchGradAttr { + return func(m optionalAttr) { + m["container"] = value + } +} + +// UnbatchGradSharedName sets the optional shared_name attribute to value. +// If not specified, defaults to "" +func UnbatchGradSharedName(value string) UnbatchGradAttr { + return func(m optionalAttr) { + m["shared_name"] = value + } +} + +// Gradient of Unbatch. +// +// Acts like Batch but using the given batch_index index of batching things as they +// become available. This ensures that the gradients are propagated back in the +// same session which did the forward pass. +// +// original_input: The input to the Unbatch operation this is the gradient of. +// batch_index: The batch_index given to the Unbatch operation this is the gradient +// of. +// grad: The downstream gradient. +// id: The id scalar emitted by Batch. +// batched_grad: The return value, either an empty tensor or the batched gradient. +// container: Container to control resource sharing. +// shared_name: Instances of UnbatchGrad with the same container and shared_name +// are assumed to possibly belong to the same batch. If left empty, the op name +// will be used as the shared name. +func UnbatchGrad(scope *Scope, original_input tf.Output, batch_index tf.Output, grad tf.Output, id tf.Output, optional ...UnbatchGradAttr) (batched_grad tf.Output) { + if scope.Err() != nil { + return + } + attrs := map[string]interface{}{} + for _, a := range optional { + a(attrs) + } + opspec := tf.OpSpec{ + Type: "UnbatchGrad", + Input: []tf.Input{ + original_input, batch_index, grad, id, + }, + Attrs: attrs, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // DecodeWavAttr is an optional argument to DecodeWav. type DecodeWavAttr func(optionalAttr) @@ -27309,6 +30058,60 @@ func ParallelConcat(scope *Scope, values []tf.Output, shape tf.Shape) (output tf return op.Output(0) } +// Subtracts `v` into specified rows of `x`. +// +// Computes y = x; y[i, :] -= v; return y. +// +// Arguments: +// x: A `Tensor` of type T. +// i: A vector. Indices into the left-most dimension of `x`. +// v: A `Tensor` of type T. Same dimension sizes as x except the first dimension, which must be the same as i's size. +// +// Returns A `Tensor` of type T. An alias of `x`. The content of `y` is undefined if there are duplicates in `i`. +func InplaceSub(scope *Scope, x tf.Output, i tf.Output, v tf.Output) (y tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "InplaceSub", + Input: []tf.Input{ + x, i, v, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + +// Converts a flat index or array of flat indices into a tuple of +// +// coordinate arrays. +// +// @compatibility(numpy) +// Equivalent to np.unravel_index +// @end_compatibility +// +// Arguments: +// indices: An 0-D or 1-D `int` Tensor whose elements are indices into the +// flattened version of an array of dimensions dims. +// dims: An 1-D `int` Tensor. The shape of the array to use for unraveling +// indices. +// +// Returns An 2-D (or 1-D if indices is 0-D) tensor where each row has the +// same shape as the indices array. +func UnravelIndex(scope *Scope, indices tf.Output, dims tf.Output) (output tf.Output) { + if scope.Err() != nil { + return + } + opspec := tf.OpSpec{ + Type: "UnravelIndex", + Input: []tf.Input{ + indices, dims, + }, + } + op := scope.AddOperation(opspec) + return op.Output(0) +} + // Compute the lower regularized incomplete Gamma function `Q(a, x)`. // // The lower regularized incomplete Gamma function is defined as: @@ -27791,6 +30594,10 @@ func EditDistance(scope *Scope, hypothesis_indices tf.Output, hypothesis_values // // indices.shape[:-1] + params.shape[indices.shape[-1]:] // +// Note that on CPU, if an out of bound index is found, an error is returned. +// On GPU, if an out of bound index is found, a 0 is stored in the +// corresponding output value. +// // Some examples below. // // Simple indexing into a matrix: diff --git a/tensorflow/java/maven/libtensorflow/pom.xml b/tensorflow/java/maven/libtensorflow/pom.xml index c99d04869a714c95e78db6f14caab515a175cb38..9c1601753bdf5076ccb14423d339b760ae784f8d 100644 --- a/tensorflow/java/maven/libtensorflow/pom.xml +++ b/tensorflow/java/maven/libtensorflow/pom.xml @@ -6,7 +6,7 @@ org.tensorflow parentpom - 1.7.0 + 1.8.0-rc0 ../ libtensorflow diff --git a/tensorflow/java/maven/libtensorflow_jni/pom.xml b/tensorflow/java/maven/libtensorflow_jni/pom.xml index 4561c2c8ade2826f779ff20c2ae1702fc97fa797..3d013e12b0d2abdd64beb50b58b562dd2265f3f3 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.7.0 + 1.8.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 82a2b8e7694d15b7de921c1005ece30236d755ee..40e44af1f537c0b2e905b164d880fcc9461373d4 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.7.0 + 1.8.0-rc0 ../ libtensorflow_jni_gpu diff --git a/tensorflow/java/maven/pom.xml b/tensorflow/java/maven/pom.xml index 4c1ec0cc8032009e7b206537dd15f1fedece4855..82bfd0c73aece8f4c265d9e1344ff8edb4764b32 100644 --- a/tensorflow/java/maven/pom.xml +++ b/tensorflow/java/maven/pom.xml @@ -6,7 +6,7 @@ 4.0.0 org.tensorflow parentpom - 1.7.0 + 1.8.0-rc0 pom https://www.tensorflow.org diff --git a/tensorflow/java/maven/proto/pom.xml b/tensorflow/java/maven/proto/pom.xml index fcd8236bad315319cb1b8b57cf6ec5eb5f277705..0a2775a500cb90500a408f12c4ed042bff0efa02 100644 --- a/tensorflow/java/maven/proto/pom.xml +++ b/tensorflow/java/maven/proto/pom.xml @@ -6,7 +6,7 @@ org.tensorflow parentpom - 1.7.0 + 1.8.0-rc0 ../ proto diff --git a/tensorflow/java/maven/tensorflow/pom.xml b/tensorflow/java/maven/tensorflow/pom.xml index 241581713ad9b129d8df4a297e9e4a3e712117fc..61961432a7e4a478a0e80c1e3eea0db12fae6868 100644 --- a/tensorflow/java/maven/tensorflow/pom.xml +++ b/tensorflow/java/maven/tensorflow/pom.xml @@ -6,7 +6,7 @@ org.tensorflow parentpom - 1.7.0 + 1.8.0-rc0 ../ tensorflow diff --git a/tensorflow/java/src/gen/cc/source_writer.h b/tensorflow/java/src/gen/cc/source_writer.h index cb0e9270d3173fc34d712fcf4ee0af87fb54d321..f011acd30aae3963c7f6ac995838c401f54fde54 100644 --- a/tensorflow/java/src/gen/cc/source_writer.h +++ b/tensorflow/java/src/gen/cc/source_writer.h @@ -32,6 +32,7 @@ namespace java { class SourceWriter { public: SourceWriter(); + virtual ~SourceWriter(); // Indents following lines with white spaces. @@ -60,7 +61,7 @@ class SourceWriter { // The data might potentially contain newline characters, therefore it will // be scanned to ensure that each line is indented and prefixed properly, // making it a bit slower than Append(). - SourceWriter& Write(const StringPiece& text); + SourceWriter& Write(const StringPiece& str); // Writes a source code snippet read from a file. // diff --git a/tensorflow/java/src/main/java/org/tensorflow/examples/LabelImage.java b/tensorflow/java/src/main/java/org/tensorflow/examples/LabelImage.java index 489e95c3102557d7a75d83789c46106aa5aa3ed4..3948991c84d35009217f7c05844551fdcc49fb22 100644 --- a/tensorflow/java/src/main/java/org/tensorflow/examples/LabelImage.java +++ b/tensorflow/java/src/main/java/org/tensorflow/examples/LabelImage.java @@ -101,6 +101,7 @@ public class LabelImage { b.constant("mean", mean)), b.constant("scale", scale)); try (Session s = new Session(g)) { + // Generally, there may be multiple output tensors, all of them must be closed to prevent resource leaks. return s.runner().fetch(output.op().name()).run().get(0).expect(Float.class); } } @@ -110,6 +111,7 @@ public class LabelImage { try (Graph g = new Graph()) { g.importGraphDef(graphDef); try (Session s = new Session(g); + // Generally, there may be multiple output tensors, all of them must be closed to prevent resource leaks. Tensor result = s.runner().feed("input", image).fetch("output").run().get(0).expect(Float.class)) { final long[] rshape = result.shape(); diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD index 87b5c596ac8008fd7183a045c8d3362875f13a54..bb32f4bbe0fb8379d4d752d0e9e1a4c1ff07495e 100644 --- a/tensorflow/python/BUILD +++ b/tensorflow/python/BUILD @@ -4,16 +4,14 @@ # Public targets: # ":platform" - Low-level and platform-specific Python code. -package( - default_visibility = [ - "//engedu/ml/tf_from_scratch:__pkg__", - "//tensorflow:internal", - "//tensorflow/contrib/lite/toco/python:__pkg__", - "//tensorflow_models:__subpackages__", - # TODO(aselle): to pass open source test. - "//bazel_pip/tensorflow/contrib/lite/toco/python:__pkg__", - ], -) +package(default_visibility = [ + "//engedu/ml/tf_from_scratch:__pkg__", + "//tensorflow:internal", + "//tensorflow/contrib/lite/toco/python:__pkg__", + "//tensorflow_models:__subpackages__", + # TODO(aselle): to pass open source test. + "//bazel_pip/tensorflow/contrib/lite/toco/python:__pkg__", +]) licenses(["notice"]) # Apache 2.0 @@ -1794,6 +1792,16 @@ py_library( ], ) +py_library( + name = "cudnn_rnn_grad", + srcs = ["ops/cudnn_rnn_grad.py"], + srcs_version = "PY2AND3", + deps = [ + ":framework_for_generated_wrappers", + "//tensorflow/python:cudnn_rnn_ops_gen", + ], +) + py_library( name = "data_flow_grad", srcs = ["ops/data_flow_grad.py"], @@ -1937,7 +1945,8 @@ py_library( ":array_ops", ":constant_op", ":dtypes", - ":linalg_ops", + ":linalg_ops_gen", + ":linalg_ops_impl", ":math_ops", ":nn_ops", ":random_ops", @@ -1988,7 +1997,22 @@ py_library( ":array_ops", ":dtypes", ":framework_ops", + ":functional_ops", ":linalg_ops_gen", + ":linalg_ops_impl", + ":math_ops", + "//third_party/py/numpy", + ], +) + +py_library( + name = "linalg_ops_impl", + srcs = ["ops/linalg_ops_impl.py"], + srcs_version = "PY2AND3", + deps = [ + ":array_ops", + ":dtypes", + ":framework_ops", ":math_ops", "//third_party/py/numpy", ], @@ -2274,7 +2298,6 @@ py_library( ":clip_ops", ":framework_for_generated_wrappers", ":init_ops", - ":layers_base", ":math_ops", ":nn_ops", ":partitioned_variables", @@ -2468,6 +2491,7 @@ py_library( ":clip_ops", ":confusion_matrix", ":control_flow_ops", + ":cudnn_rnn_grad", ":data_flow_grad", ":data_flow_ops", ":framework_for_generated_wrappers", @@ -2551,6 +2575,30 @@ py_library( ], ) +py_library( + name = "summary_ops_v2", + srcs = ["ops/summary_ops_v2.py"], + srcs_version = "PY2AND3", + visibility = ["//tensorflow:internal"], + deps = [ + ":array_ops", + ":constant_op", + ":control_flow_ops", + ":dtypes", + ":framework_ops", + ":math_ops", + ":resource_variable_ops", + ":smart_cond", + ":summary_op_util", + ":summary_ops_gen", + ":training_util", + ":util", + "//tensorflow/core:protos_all_py", + "//tensorflow/python/eager:context", + "@six_archive//:six", + ], +) + py_library( name = "template", srcs = ["ops/template.py"], @@ -2770,7 +2818,7 @@ cuda_py_test( cuda_py_test( name = "image_ops_test", - size = "small", + size = "medium", srcs = ["ops/image_ops_test.py"], additional_deps = [ ":array_ops", @@ -2913,11 +2961,15 @@ py_library( name = "training", srcs = glob( ["training/**/*.py"], - exclude = ["**/*test*"], + exclude = [ + "**/*test*", + "training/training_util.py", # See :training_util + ], ), srcs_version = "PY2AND3", deps = [ ":array_ops", + ":array_ops_gen", ":checkpoint_ops_gen", ":client", ":control_flow_ops", @@ -2929,6 +2981,7 @@ py_library( ":framework_ops", ":gradients", ":init_ops", + ":distribute", ":io_ops", ":io_ops_gen", ":layers_base", @@ -2947,14 +3000,18 @@ py_library( ":string_ops", ":summary", ":training_ops_gen", + ":training_util", ":util", ":variable_scope", ":variables", + "//third_party/py/numpy", + "@six_archive//:six", + "//tensorflow/core:protos_all_py", "//tensorflow/python/eager:backprop", "//tensorflow/python/eager:context", "//tensorflow/python/ops/losses", - "//third_party/py/numpy", - "@six_archive//:six", + # `layers` dependency only exists due to the use of a small utility. + "//tensorflow/python/keras:layers", ], ) @@ -2982,6 +3039,69 @@ py_test( ], ) +py_library( + name = "device_util", + srcs = ["training/device_util.py"], + srcs_version = "PY2AND3", + deps = [ + ":device", + ":framework_ops", + "//tensorflow/python/eager:context", + ], +) + +py_library( + name = "distribute", + srcs = ["training/distribute.py"], + srcs_version = "PY2AND3", + deps = [ + ":array_ops", + ":control_flow_ops", + ":device_util", + ":framework_ops", + ":platform", + ":resource_variable_ops", + ":state_ops", + ":util", + ":variable_scope", + "//tensorflow/python/data", + "//tensorflow/python/ops/losses", + ], +) + +py_test( + name = "checkpointable_utils_test", + srcs = ["training/checkpointable_utils_test.py"], + srcs_version = "PY2AND3", + tags = [ + "no_windows", # TODO: needs investigation on Windows + "notsan", # b/74395663 + ], + deps = [ + ":checkpointable", + ":constant_op", + ":control_flow_ops", + ":dtypes", + ":framework_ops", + ":framework_test_lib", + ":init_ops", + ":resource_variable_ops", + ":session", + ":state_ops", + ":template", + ":training", + ":training_util", + ":variable_scope", + "//tensorflow/python/eager:backprop", + "//tensorflow/python/eager:context", + "//tensorflow/python/eager:function", + "//tensorflow/python/eager:test", + "//tensorflow/python/keras:engine", + "//tensorflow/python/keras:layers", + "@six_archive//:six", + ], +) + py_test( name = "distribute_test", size = "small", @@ -2989,7 +3109,7 @@ py_test( srcs_version = "PY2AND3", deps = [ ":client_testlib", - ":training", + ":distribute", ":variable_scope", ], ) @@ -3388,6 +3508,7 @@ tf_py_wrap_cc( "//tensorflow/core/profiler/internal:print_model_analysis", "//tensorflow/tools/graph_transforms:transform_graph_lib", "//tensorflow/python/eager:pywrap_tfe_lib", + "//tensorflow/python/eager:python_eager_op_gen", "//util/python:python_headers", ] + (tf_additional_lib_deps() + tf_additional_plugin_deps() + @@ -4194,6 +4315,25 @@ py_test( ], ) +py_library( + name = "training_util", + srcs = ["training/training_util.py"], + srcs_version = "PY2AND3", + deps = [ + ":dtypes", + ":framework", + ":framework_ops", + ":init_ops", + ":platform", + ":resource_variable_ops", + ":state_ops", + ":util", + ":variable_scope", + ":variables", + "//tensorflow/python/eager:context", + ], +) + py_test( name = "training_util_test", size = "small", @@ -4204,13 +4344,14 @@ py_test( ":framework", ":platform", ":training", + ":training_util", ":variables", ], ) tf_py_test( name = "input_test", - size = "small", + size = "medium", srcs = ["training/input_test.py"], additional_deps = [ ":array_ops", @@ -4233,6 +4374,7 @@ py_library( srcs_version = "PY2AND3", visibility = ["//visibility:public"], deps = [ + ":distribute", ":framework", ":framework_for_generated_wrappers", ":platform", @@ -4248,6 +4390,7 @@ py_library( srcs_version = "PY2AND3", visibility = ["//visibility:public"], deps = [ + ":client", ":constant_op", ":errors", ":framework", @@ -4260,6 +4403,7 @@ py_library( ":summary_op_util", ":summary_ops", ":summary_ops_gen", + ":summary_ops_v2", ":util", "//tensorflow/python/eager:context", "//third_party/py/numpy", @@ -4286,7 +4430,7 @@ py_tests( ":platform", ":platform_test", ":summary", - ":training", + ":summary_ops_v2", "//tensorflow/core:protos_all_py", ], ) @@ -4311,6 +4455,7 @@ py_library( ":variables", "//tensorflow/python/eager:context", "//tensorflow/python/estimator:util", + "//tensorflow/python/keras:engine", "//third_party/py/numpy", ], ) @@ -4347,6 +4492,7 @@ py_library( ":variables", "//tensorflow/python/eager:context", "//tensorflow/python/estimator:util", + "//tensorflow/python/keras:layers", "//third_party/py/numpy", "@six_archive//:six", ], diff --git a/tensorflow/python/__init__.py b/tensorflow/python/__init__.py index ab1d01a8351d63544b2c612ad228515d48975aca..13f8420a670fe64615037975139f3ee1f16820b6 100644 --- a/tensorflow/python/__init__.py +++ b/tensorflow/python/__init__.py @@ -149,6 +149,19 @@ from tensorflow.python.ops import tensor_array_ops from tensorflow.python.eager.context import executing_eagerly from tensorflow.python.framework.ops import enable_eager_execution +# Necessary for the symbols in this module to be taken into account by +# the namespace management system (API decorators). +from tensorflow.python.ops import rnn +from tensorflow.python.ops import rnn_cell + +# Required due to `rnn` and `rnn_cell` not being imported in `nn` directly +# (due to a circular dependency issue: rnn depends on layers). +nn.dynamic_rnn = rnn.dynamic_rnn +nn.static_rnn = rnn.static_rnn +nn.raw_rnn = rnn.raw_rnn +nn.bidirectional_dynamic_rnn = rnn.bidirectional_dynamic_rnn +nn.rnn_cell = rnn_cell + # Symbols whitelisted for export without documentation. # TODO(cwhipkey): review these and move to contrib, expose through # documentation, or remove. diff --git a/tensorflow/python/client/session.py b/tensorflow/python/client/session.py index 4c84d78f2e11922e4819e45aaee79374c8c5ec34..5507d011bb0746c84b868ca7efcc3e4f8d2e146a 100644 --- a/tensorflow/python/client/session.py +++ b/tensorflow/python/client/session.py @@ -1454,7 +1454,10 @@ class BaseSession(SessionInterface): self._session._session, self._handle, args, status, None) def __del__(self): - if self._handle is not None: + # NOTE(mrry): It is possible that `self._session.__del__()` could be + # called before this destructor, in which case `self._session._session` + # will be `None`. + if self._handle is not None and self._session._session is not None: with errors.raise_exception_on_not_ok_status() as status: if self._session._created_with_new_api: tf_session.TF_SessionReleaseCallable( 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 6442eb9ff554e61829796fb904342072d1846a32..f7d7d085c974fa217ed30708723cb1b887034ca0 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 @@ -69,6 +69,54 @@ class ListFilesDatasetOpTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(itr.get_next()) + def testSimpleDirectoryNotShuffled(self): + filenames = ['b', 'c', 'a'] + self._touchTempFiles(filenames) + + dataset = dataset_ops.Dataset.list_files( + path.join(self.tmp_dir, '*'), shuffle=False) + with self.test_session() as sess: + itr = dataset.make_one_shot_iterator() + next_element = itr.get_next() + + for filename in sorted(filenames): + self.assertEqual(compat.as_bytes(path.join(self.tmp_dir, filename)), + sess.run(next_element)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(itr.get_next()) + + def testFixedSeedResultsInRepeatableOrder(self): + filenames = ['a', 'b', 'c'] + self._touchTempFiles(filenames) + + dataset = dataset_ops.Dataset.list_files( + path.join(self.tmp_dir, '*'), shuffle=True, seed=37) + with self.test_session() as sess: + itr = dataset.make_initializable_iterator() + next_element = itr.get_next() + + full_filenames = [compat.as_bytes(path.join(self.tmp_dir, filename)) + for filename in filenames] + + all_produced_filenames = [] + for _ in range(3): + produced_filenames = [] + sess.run(itr.initializer) + try: + while True: + produced_filenames.append(sess.run(next_element)) + except errors.OutOfRangeError: + pass + all_produced_filenames.append(produced_filenames) + + # Each run should produce the same set of filenames, which may be + # different from the order of `full_filenames`. + self.assertItemsEqual(full_filenames, all_produced_filenames[0]) + # However, the different runs should produce filenames in the same order + # as each other. + self.assertEqual(all_produced_filenames[0], all_produced_filenames[1]) + self.assertEqual(all_produced_filenames[0], all_produced_filenames[2]) + def testEmptyDirectoryInitializer(self): filename_placeholder = array_ops.placeholder(dtypes.string, shape=[]) dataset = dataset_ops.Dataset.list_files(filename_placeholder) 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 0791c614fa88700fdf2d0d673e168fc9784731a5..1ad0b9de5e76e3edd66303ab4666108f43a27428 100644 --- a/tensorflow/python/data/kernel_tests/map_dataset_op_test.py +++ b/tensorflow/python/data/kernel_tests/map_dataset_op_test.py @@ -624,6 +624,20 @@ class MapDatasetTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) + def testConstantOutput(self): + iterator = ( + dataset_ops.Dataset.range(10).map(lambda x: [x, "hello", 10]) + .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, b"hello", 10), sess.run(get_next)) + with self.assertRaises(errors.OutOfRangeError): + sess.run(get_next) + class MapDatasetBenchmark(test.Benchmark): diff --git a/tensorflow/python/data/ops/dataset_ops.py b/tensorflow/python/data/ops/dataset_ops.py index c28de3d05443f40526eb589e155cb58e98acb14a..bd9686f692102df4ef64f0e81c33d2dec1b2222c 100644 --- a/tensorflow/python/data/ops/dataset_ops.py +++ b/tensorflow/python/data/ops/dataset_ops.py @@ -571,9 +571,13 @@ class Dataset(object): return PrefetchDataset(self, buffer_size) @staticmethod - def list_files(file_pattern, shuffle=None): + def list_files(file_pattern, shuffle=None, seed=None): """A dataset of all files matching a pattern. + NOTE: The default behavior of this method is to return filenames in + a non-deterministic random shuffled order. Pass a `seed` or `shuffle=False` + to get results in a deterministic order. + Example: If we had the following files on our filesystem: - /path/to/dir/a.txt @@ -584,20 +588,18 @@ 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`. + 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: Dataset: A `Dataset` of strings corresponding to file names. """ - # 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) @@ -607,7 +609,7 @@ class Dataset(object): # 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) + dataset = dataset.shuffle(buffer_size, seed=seed) return dataset def repeat(self, count=None): @@ -1155,10 +1157,12 @@ class _GeneratorDataset(Dataset): if isinstance(ret, list): ret = tuple(ret) - # Convert any `SparseTensorValue`s to `SparseTensor`s. + # Convert any `SparseTensorValue`s to `SparseTensor`s and all other + # values to tensors. 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) + if sparse_tensor_lib.is_sparse(t) else ops.convert_to_tensor(t) + for t in nest.flatten(ret) ]) self._state_classes = sparse.get_classes(ret) @@ -1167,11 +1171,9 @@ class _GeneratorDataset(Dataset): 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)) - ]) + # Serialize any sparse tensors. + ret = nest.pack_sequence_as( + ret, [t for t in nest.flatten(sparse.serialize_sparse_tensors(ret))]) return nest.flatten(ret) self._init_func = tf_init_func @@ -1214,10 +1216,12 @@ class _GeneratorDataset(Dataset): if isinstance(ret, list): ret = tuple(ret) - # Convert any `SparseTensorValue`s to `SparseTensor`s. + # Convert any `SparseTensorValue`s to `SparseTensor`s and all other + # values to tensors. 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) + if sparse_tensor_lib.is_sparse(t) else ops.convert_to_tensor(t) + for t in nest.flatten(ret) ]) self._output_classes = sparse.get_classes(ret) @@ -1226,11 +1230,9 @@ class _GeneratorDataset(Dataset): 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)) - ]) + # Serialize any sparse tensors. + ret = nest.pack_sequence_as( + ret, [t for t in nest.flatten(sparse.serialize_sparse_tensors(ret))]) return nest.flatten(ret) self._next_func = tf_next_func @@ -1816,10 +1818,12 @@ class MapDataset(Dataset): if isinstance(ret, list): ret = tuple(ret) - # Convert any `SparseTensorValue`s to `SparseTensor`s. + # Convert any `SparseTensorValue`s to `SparseTensor`s and all other + # values to tensors. 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) + if sparse_tensor_lib.is_sparse(t) else ops.convert_to_tensor(t) + for t in nest.flatten(ret) ]) self._output_classes = sparse.get_classes(ret) @@ -1828,11 +1832,9 @@ class MapDataset(Dataset): 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)) - ]) + # Serialize any sparse tensors. + ret = nest.pack_sequence_as( + ret, [t for t in nest.flatten(sparse.serialize_sparse_tensors(ret))]) return nest.flatten(ret) self._map_func = tf_map_func diff --git a/tensorflow/python/debug/cli/readline_ui.py b/tensorflow/python/debug/cli/readline_ui.py index 151638789f7eb509a39a6f066e9ad9a0c92c1fe3..3296e45d07e9f385fdcda3f34804547b049dcf02 100644 --- a/tensorflow/python/debug/cli/readline_ui.py +++ b/tensorflow/python/debug/cli/readline_ui.py @@ -19,6 +19,8 @@ from __future__ import print_function import readline +import six + from tensorflow.python.debug.cli import base_ui from tensorflow.python.debug.cli import debugger_cli_common @@ -39,11 +41,7 @@ class ReadlineUI(base_ui.BaseUI): readline.set_completer(self._readline_complete) readline.parse_and_bind("tab: complete") - # For Python 2-3 compatibility. - try: - self._input = raw_input - except NameError: - self._input = input + self._input = six.moves.input def _readline_complete(self, text, state): context, prefix, except_last_word = self._analyze_tab_complete_input(text) diff --git a/tensorflow/python/debug/wrappers/grpc_wrapper.py b/tensorflow/python/debug/wrappers/grpc_wrapper.py index fb9494f57636e46e54ef230cf4803dbb6ccad0c7..1f9c8fa5a96b4d6826fae0870608e0e737c7cd88 100644 --- a/tensorflow/python/debug/wrappers/grpc_wrapper.py +++ b/tensorflow/python/debug/wrappers/grpc_wrapper.py @@ -21,6 +21,8 @@ import signal import sys import traceback +import six + # Google-internal import(s). from tensorflow.python.debug.lib import common from tensorflow.python.debug.wrappers import framework @@ -140,14 +142,9 @@ class GrpcDebugWrapperSession(framework.NonInteractiveDebugWrapperSession): 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() + response = six.moves.input( + "\nSIGINT received. Quit program? (Y/n): ").strip() if response in ("", "Y", "y"): sys.exit(0) elif response in ("N", "n"): diff --git a/tensorflow/python/debug/wrappers/hooks.py b/tensorflow/python/debug/wrappers/hooks.py index 6705cd31e291d2eab7aa8179e9b2b829f8970c18..5e4604fda4d7249a1244f12a533e1cb09e16782f 100644 --- a/tensorflow/python/debug/wrappers/hooks.py +++ b/tensorflow/python/debug/wrappers/hooks.py @@ -31,15 +31,18 @@ from tensorflow.python.training import session_run_hook class LocalCLIDebugHook(session_run_hook.SessionRunHook): """Command-line-interface debugger hook. - Can be used as a monitor/hook for `tf.train.MonitoredSession`s and - `tf.contrib.learn`'s `Estimator`s and `Experiment`s. + Can be used as a hook for `tf.train.MonitoredSession`s and + `tf.estimator.Estimator`s. Provides a substitute for + `tfdbg.LocalCLIDebugWrapperSession` in cases where the session is not directly + available. """ def __init__(self, ui_type="curses", dump_root=None, thread_name_filter=None): """Create a local debugger command-line interface (CLI) hook. Args: - ui_type: (str) user-interface type. + ui_type: (`str`) requested user-interface type. Currently supported: + (curses | readline). dump_root: (`str`) optional path to the dump root directory. Must be a directory that does not exist or an empty directory. If the directory does not exist, it will be created by the debugger core during debug @@ -153,8 +156,8 @@ class LocalCLIDebugHook(session_run_hook.SessionRunHook): class DumpingDebugHook(session_run_hook.SessionRunHook): """A debugger hook that dumps debug data to filesystem. - Can be used as a monitor/hook for `tf.train.MonitoredSession`s and - `tf.contrib.learn`'s `Estimator`s and `Experiment`s. + Can be used as a hook for `tf.train.MonitoredSession`s and + `tf.estimator.Estimator`s. """ def __init__(self, @@ -229,8 +232,8 @@ class GrpcDebugHook(session_run_hook.SessionRunHook): When the arguments of debug_utils.watch_graph changes, strongly consider changing arguments here too so that features are available to tflearn users. - Can be used as a monitor/hook for `tf.train.MonitoredSession`s and - `tf.contrib.learn`'s `Estimator`s and `Experiment`s. + Can be used as a hook for `tf.train.MonitoredSession`s and + `tf.estimator.Estimator`s. """ def __init__(self, diff --git a/tensorflow/python/eager/BUILD b/tensorflow/python/eager/BUILD index 8c0d3feeceab1bf29c1dabc668176a6ef7806421..b3268c9047e264b8264ae37b404b51be6a88962f 100644 --- a/tensorflow/python/eager/BUILD +++ b/tensorflow/python/eager/BUILD @@ -142,6 +142,8 @@ cuda_py_test( ":tape", ":test", "//tensorflow/python:clip_ops", + "//tensorflow/python:init_ops", + "//tensorflow/python:layers", "//tensorflow/python:math_ops", "//tensorflow/python:resource_variable_ops", ], diff --git a/tensorflow/python/eager/function.py b/tensorflow/python/eager/function.py index 5168ad3b18f623588b7804f597fa3b816de147f3..0f1170bb42094f4d6565db15f1203f7d08ba89c7 100644 --- a/tensorflow/python/eager/function.py +++ b/tensorflow/python/eager/function.py @@ -38,6 +38,7 @@ from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import gradients_impl +from tensorflow.python.ops import resource_variable_ops from tensorflow.python.util import compat from tensorflow.python.util import nest from tensorflow.python.util import tf_decorator @@ -69,9 +70,22 @@ def capture_value(tensor_map, value, dtype, name): captured_value = graph_placeholder( dtype=dtype or value.dtype, shape=value.shape, name=name) if captured_value.dtype == dtypes_module.resource: - handle_data = value._handle_data # pylint: disable=protected-access - captured_value._handle_data = handle_data # pylint: disable=protected-access + if ops._USE_C_SHAPES: # pylint: disable=protected-access + if isinstance(value, ops.EagerTensor): + handle_data = value._handle_data # pylint: disable=protected-access + else: + handle_data = resource_variable_ops.get_resource_handle_data(value) + else: + handle_data = value._handle_data # pylint: disable=protected-access if handle_data is not None and handle_data.is_set: + # pylint: disable=protected-access + if ops._USE_C_SHAPES: + pywrap_tensorflow.TFE_SetResourceHandleShapeAndType( + captured_value.graph._c_graph, captured_value._as_tf_output(), + handle_data.SerializeToString()) + else: + captured_value._handle_data = handle_data + # pylint: enable=protected-access # Ensure that shapes and dtypes are propagated. shapes, types = zip(*[(pair.shape, pair.dtype) for pair in handle_data.shape_and_type]) diff --git a/tensorflow/python/eager/function_test.py b/tensorflow/python/eager/function_test.py index 9af197981bde309160781fa5821152962e5383bb..1828c987f4343015e042359e5ba2f715ee63ad83 100644 --- a/tensorflow/python/eager/function_test.py +++ b/tensorflow/python/eager/function_test.py @@ -29,9 +29,12 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import function as tf_function from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape +from tensorflow.python.framework import test_util +from tensorflow.python.layers import convolutional 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 init_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import resource_variable_ops from tensorflow.python.ops import variable_scope @@ -39,6 +42,7 @@ from tensorflow.python.ops import variables from tensorflow.python.training import gradient_descent +@test_util.with_c_shapes class FunctionTest(test.TestCase): def testBasic(self): @@ -104,6 +108,7 @@ class FunctionTest(test.TestCase): matmul = function.defun(math_ops.matmul) pair = collections.namedtuple('pair', ['a', 'b']) + def a_times_b(inputs): return matmul(inputs.a['a'], inputs.b['b']) @@ -312,6 +317,7 @@ class FunctionTest(test.TestCase): x = variable_scope.get_variable( 'v', initializer=constant_op.constant(1.0)) return x * constant_op.constant(2.0) + with self.assertRaisesRegexp(ValueError, 'No trainable variables were accessed'): backprop.implicit_val_and_grad(f)() @@ -581,6 +587,7 @@ class FunctionTest(test.TestCase): with ops.name_scope('foo'): v = resource_variable_ops.ResourceVariable(0.0, name='bar') self.assertEqual(v.name, 'foo/bar:0') + create_variable() def testVariableNamesRespectNameScopesWithDefunInGraph(self): @@ -590,10 +597,27 @@ class FunctionTest(test.TestCase): with ops.name_scope('foo'): v = resource_variable_ops.ResourceVariable([1.0, 2.0], name='bar') self.assertEqual(v.name, 'foo/bar:0') + with ops.get_default_graph().as_default(): create_variable() + def testLayerInDefun(self): + conv = convolutional.Conv2D( + filters=1, + kernel_size=2, + kernel_initializer=init_ops.ones_initializer(), + bias_initializer=init_ops.zeros_initializer()) + + @function.defun + def model(x): + return conv(x) + + x = array_ops.ones([1, 2, 2, 1]) + y = model(x) + self.assertAllEqual([[[[4.0]]]], y.numpy()) + +@test_util.with_c_shapes class AutomaticControlDependenciesTest(test.TestCase): def testBasic(self): diff --git a/tensorflow/python/eager/pywrap_tensor.cc b/tensorflow/python/eager/pywrap_tensor.cc index 519814b979e00dd7c9df41eacbe1edc02c9d88e8..b5b4e394e33bd3adcf8e90fa4c35f87fbbb5a155 100644 --- a/tensorflow/python/eager/pywrap_tensor.cc +++ b/tensorflow/python/eager/pywrap_tensor.cc @@ -60,42 +60,6 @@ TFE_TensorHandle* NumpyToTensorHandle(PyObject* obj) { } } -// Casts data referred to by `handle` from type `src_type_enum` to type -// `dst_type_enum`. -TFE_TensorHandle* EagerCast(TFE_Context* ctx, TFE_TensorHandle* handle, - TF_DataType src_type_enum, - TF_DataType dst_type_enum, TF_Status* out_status) { - if (ctx == nullptr) return nullptr; - const char* op_name = "Cast"; - const char* device_name = "/job:localhost/replica:0/task:0/device:CPU:0"; - TFE_Op* op = TFE_NewOp(ctx, op_name, out_status); -#define RETURN_ERROR \ - { \ - TFE_DeleteOp(op); \ - return nullptr; \ - } - if (TF_GetCode(out_status) != TF_OK) RETURN_ERROR - TFE_OpSetDevice(op, device_name, out_status); - if (TF_GetCode(out_status) != TF_OK) RETURN_ERROR - TFE_OpAddInput(op, handle, out_status); - if (TF_GetCode(out_status) != TF_OK) RETURN_ERROR - TFE_OpSetAttrType(op, "SrcT", src_type_enum); - TFE_OpSetAttrType(op, "DstT", dst_type_enum); - TFE_TensorHandle* output = nullptr; - int num_outputs = 1; - TFE_Execute(op, &output, &num_outputs, out_status); - if (TF_GetCode(out_status) != TF_OK || num_outputs != 1 || - output == nullptr) { - if (output != nullptr) { - TFE_DeleteTensorHandle(output); - } - RETURN_ERROR - } - TFE_DeleteOp(op); - return output; -#undef RETURN_ERROR -} - TFE_TensorHandle* CopyToDevice(TFE_TensorHandle* handle, PyObject* ctx, PyObject* dev) { const char* device = ""; @@ -161,6 +125,100 @@ PyObject* PyIntFromDataType(TF_DataType l) { } // namespace +namespace tensorflow { +// Casts data referred to by `handle` from type `src_type_enum` to type +// `dst_type_enum`. +TFE_TensorHandle* EagerCast(TFE_Context* ctx, TFE_TensorHandle* handle, + TF_DataType src_type_enum, + TF_DataType dst_type_enum, TF_Status* out_status) { + if (ctx == nullptr) return nullptr; + const char* op_name = "Cast"; + const char* device_name = "/job:localhost/replica:0/task:0/device:CPU:0"; + TFE_Op* op = TFE_NewOp(ctx, op_name, out_status); +#define RETURN_ERROR \ + { \ + TFE_DeleteOp(op); \ + return nullptr; \ + } + if (TF_GetCode(out_status) != TF_OK) RETURN_ERROR + TFE_OpSetDevice(op, device_name, out_status); + if (TF_GetCode(out_status) != TF_OK) RETURN_ERROR + TFE_OpAddInput(op, handle, out_status); + if (TF_GetCode(out_status) != TF_OK) RETURN_ERROR + TFE_OpSetAttrType(op, "SrcT", src_type_enum); + TFE_OpSetAttrType(op, "DstT", dst_type_enum); + TFE_TensorHandle* output = nullptr; + int num_outputs = 1; + TFE_Execute(op, &output, &num_outputs, out_status); + if (TF_GetCode(out_status) != TF_OK || num_outputs != 1 || + output == nullptr) { + if (output != nullptr) { + TFE_DeleteTensorHandle(output); + } + RETURN_ERROR + } + TFE_DeleteOp(op); + return output; +#undef RETURN_ERROR +} + +TFE_TensorHandle* ConvertToEagerTensor(PyObject* value, PyObject* dtype) { + int desired_dtype = -1; + if (dtype != Py_None) { + if (!PyIntToDataType(dtype, &desired_dtype)) { + PyErr_SetString(PyExc_TypeError, + tensorflow::strings::StrCat( + "Expecting a DataType value for dtype. Got ", + Py_TYPE(dtype)->tp_name) + .c_str()); + return nullptr; + } + } + if (PyArray_Check(value)) { + int desired_np_dtype = -1; + if (desired_dtype >= 0) { + if (!tensorflow::TF_DataType_to_PyArray_TYPE( + static_cast(desired_dtype), &desired_np_dtype) + .ok()) { + PyErr_SetString(PyExc_TypeError, + tensorflow::strings::StrCat( + "Invalid dtype argument value ", desired_dtype) + .c_str()); + return nullptr; + } + } + PyArrayObject* array = reinterpret_cast(value); + int current_np_dtype = PyArray_TYPE(array); + auto safe_value = tensorflow::make_safe(static_cast(nullptr)); + if ((desired_np_dtype >= 0 && desired_np_dtype != current_np_dtype) || + !PyArray_ISCARRAY(array)) { + int new_dtype = + desired_np_dtype >= 0 ? desired_np_dtype : current_np_dtype; + safe_value = tensorflow::make_safe( + PyArray_FromAny(value, PyArray_DescrFromType(new_dtype), 0, 0, + NPY_ARRAY_CARRAY | NPY_ARRAY_FORCECAST, nullptr)); + if (PyErr_Occurred()) return nullptr; + if (safe_value == nullptr) { + PyErr_SetString(PyExc_ValueError, "Error while casting a numpy value"); + return nullptr; + } + value = safe_value.get(); + } + return NumpyToTensorHandle(value); + } else { + tensorflow::Tensor t; + // TODO(josh11b): Have PySeqToTensor set python errors instead of + // returning Status. + auto cppstatus = tensorflow::PySeqToTensor(value, dtype, &t); + if (!cppstatus.ok()) { + PyErr_SetString(PyExc_ValueError, cppstatus.error_message().c_str()); + return nullptr; + } + return TFE_NewTensorHandle(t); + } +} +} // namespace tensorflow + extern "C" { static const int kMaxEagerTensorParentSize = 64; @@ -230,61 +288,16 @@ int EagerTensor_init(EagerTensor* self, PyObject* args, PyObject* kwds) { return -1; } } - tensorflow::Safe_TFE_TensorHandlePtr handle = - tensorflow::make_safe(static_cast(nullptr)); PyErr_Clear(); - if (PyArray_Check(value)) { - int desired_np_dtype = -1; - if (desired_dtype >= 0) { - if (!tensorflow::TF_DataType_to_PyArray_TYPE( - static_cast(desired_dtype), &desired_np_dtype) - .ok()) { - PyErr_SetString(PyExc_TypeError, - tensorflow::strings::StrCat( - "Invalid dtype argument value ", desired_dtype) - .c_str()); - return -1; - } - } - PyArrayObject* array = reinterpret_cast(value); - int current_np_dtype = PyArray_TYPE(array); - auto safe_value = tensorflow::make_safe(static_cast(nullptr)); - if ((desired_np_dtype >= 0 && desired_np_dtype != current_np_dtype) || - !PyArray_ISCARRAY(array)) { - int new_dtype = - desired_np_dtype >= 0 ? desired_np_dtype : current_np_dtype; - safe_value = tensorflow::make_safe( - PyArray_FromAny(value, PyArray_DescrFromType(new_dtype), 0, 0, - NPY_ARRAY_CARRAY | NPY_ARRAY_FORCECAST, nullptr)); - if (PyErr_Occurred()) return -1; - if (safe_value == nullptr) { - PyErr_SetString(PyExc_ValueError, "Error while casting a numpy value"); - return -1; - } - value = safe_value.get(); - } - handle = tensorflow::make_safe(NumpyToTensorHandle(value)); - } else { - tensorflow::Tensor t; - // TODO(josh11b): Have PySeqToTensor set python errors instead of - // returning Status. - auto cppstatus = tensorflow::PySeqToTensor(value, dtype, &t); - if (!cppstatus.ok()) { - PyErr_SetString(PyExc_ValueError, cppstatus.error_message().c_str()); - return -1; - } - handle = tensorflow::make_safe(TFE_NewTensorHandle(t)); - } - if (PyErr_Occurred()) return -1; - if (handle == nullptr) { - PyErr_SetString(PyExc_ValueError, "Error while creating an EagerTensor"); - return -1; - } + tensorflow::Safe_TFE_TensorHandlePtr handle = + tensorflow::make_safe(static_cast( + tensorflow::ConvertToEagerTensor(value, dtype))); + if (handle == nullptr) return -1; TF_DataType handle_dtype = TFE_TensorHandleDataType(handle.get()); if (desired_dtype >= 0 && desired_dtype != handle_dtype) { - handle = tensorflow::make_safe( - EagerCast(GetContext(context), handle.get(), handle_dtype, - static_cast(desired_dtype), self->status)); + handle = tensorflow::make_safe(tensorflow::EagerCast( + GetContext(context), handle.get(), handle_dtype, + static_cast(desired_dtype), self->status)); if (TF_GetCode(self->status) != TF_OK) { PyErr_SetString(PyExc_ValueError, tensorflow::strings::StrCat( @@ -701,12 +714,12 @@ PyObject* TFE_Py_InitEagerTensor(PyObject* base_class) { return reinterpret_cast(EagerTensorType); } -PyObject* TFE_Py_TensorShapeSlice(PyObject* tensor_list, int slice_dim) { - if (!PyList_Check(tensor_list)) { +PyObject* TFE_Py_TensorShapeSlice(PyObject* tensors, int slice_dim) { + if (!PyList_Check(tensors) && !PyTuple_Check(tensors)) { PyErr_SetString(PyExc_TypeError, tensorflow::strings::StrCat( - "tensor_list argument must be a list. Got \"", - Py_TYPE(tensor_list)->tp_name, "\"") + "tensors argument must be a list or a tuple. Got \"", + Py_TYPE(tensors)->tp_name, "\"") .c_str()); return nullptr; } @@ -720,14 +733,14 @@ PyObject* TFE_Py_TensorShapeSlice(PyObject* tensor_list, int slice_dim) { return nullptr; } - Py_ssize_t num_tensors = PyList_Size(tensor_list); + Py_ssize_t num_tensors = PySequence_Fast_GET_SIZE(tensors); int64_t num_tensors_int = static_cast(num_tensors); 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); + PyObject* tensor_obj = PySequence_Fast_GET_ITEM(tensors, i); if (!EagerTensor_CheckExact(tensor_obj)) { PyErr_SetString(PyExc_TypeError, tensorflow::strings::StrCat( diff --git a/tensorflow/python/eager/pywrap_tensor.h b/tensorflow/python/eager/pywrap_tensor.h index aa1efdd1b81cca9df0088c4cecedfe52f258d2bc..63ab1ed84d5ba3f280904be8dd202912e42241d0 100644 --- a/tensorflow/python/eager/pywrap_tensor.h +++ b/tensorflow/python/eager/pywrap_tensor.h @@ -22,4 +22,14 @@ limitations under the License. bool EagerTensor_CheckExact(const PyObject* o); tensorflow::int64 EagerTensor_id(const PyObject* tensor); +namespace tensorflow { +TFE_TensorHandle* ConvertToEagerTensor(PyObject* value, PyObject* dtype); + +// TODO(nareshmodi): Move EagerCast and ReadVariableOp (which use the C API to +// execute TFE Ops) to a separate common library. +TFE_TensorHandle* EagerCast(TFE_Context* ctx, TFE_TensorHandle* handle, + TF_DataType src_type_enum, + TF_DataType dst_type_enum, TF_Status* out_status); +} + #endif // TENSORFLOW_PYTHON_EAGER_PYWRAP_TENSOR_H_ diff --git a/tensorflow/python/eager/pywrap_tfe.h b/tensorflow/python/eager/pywrap_tfe.h index 32d731d0f68910b8e41a57cb32ae60c3ea6742f7..691b613e48b217c595fe0f3249c493facf756d47 100644 --- a/tensorflow/python/eager/pywrap_tfe.h +++ b/tensorflow/python/eager/pywrap_tfe.h @@ -186,16 +186,16 @@ PyObject* TFE_Py_RecordGradient(PyObject* op_name, PyObject* inputs, // Returns the set of variables watched by the given tape. PyObject* TFE_Py_TapeWatchedVariables(PyObject* tape); -// Returns an EagerTensor of dimension [len(`tensor_list`)] containing -// the `slice_dim`'th dimension of each tensor in `tensor_list`. In other words, +// Returns an EagerTensor of dimension [len(`tensors`)] containing +// the `slice_dim`'th dimension of each tensor in `tensors`. In other words, // TFE_Py_TensorShapeSlice takes a slice of dimensions of tensors in -// `tensor_list`. For example, if `tensor_list` contains tensors of with shapes +// `tensors`. For example, if `tensors` contains tensors of with shapes // [1, 2, 3], [4, 5], [6, 7, 8, 9], TFE_Py_TensorShapeSlice called with // `slice_dim` equal to 1 will return [2, 5, 7]. // On error, returns nullptr and sets python exception. -// REQUIRES: `tensor_list` is a python list of EagerTensors +// REQUIRES: `tensors` is a python list/tuple of EagerTensors // REQUIRES: `slice_dim` is non-negative and smaller than the rank of all -// tensors in `tensor_list`. -PyObject* TFE_Py_TensorShapeSlice(PyObject* tensor_list, int slice_dim); +// tensors in `tensors`. +PyObject* TFE_Py_TensorShapeSlice(PyObject* tensors, int slice_dim); #endif // TENSORFLOW_PYTHON_EAGER_PYWRAP_TFE_H_ diff --git a/tensorflow/python/eager/pywrap_tfe_src.cc b/tensorflow/python/eager/pywrap_tfe_src.cc index d99bd0b0ffe5ef8042b3b869f82a9952f1985abc..2bfa1f052cfe6550b25ad3bae7fa5c67a4e45be5 100644 --- a/tensorflow/python/eager/pywrap_tfe_src.cc +++ b/tensorflow/python/eager/pywrap_tfe_src.cc @@ -38,6 +38,54 @@ using tensorflow::strings::Printf; namespace { +struct InputInfo { + InputInfo(int i, bool is_list) : i(i), is_list(is_list) {} + + int i; + bool is_list = false; +}; + +using AttrToInputsMap = + tensorflow::gtl::FlatMap>; + +tensorflow::mutex all_attr_to_input_maps_lock( + tensorflow::LINKER_INITIALIZED); +tensorflow::gtl::FlatMap* GetAllAttrToInputsMaps() { + static auto* all_attr_to_input_maps = + new tensorflow::gtl::FlatMap; + return all_attr_to_input_maps; +} + +AttrToInputsMap* GetAttrToInputsMap(const tensorflow::OpDef& op_def) { + tensorflow::mutex_lock l(all_attr_to_input_maps_lock); + auto* all_attr_to_input_maps = GetAllAttrToInputsMaps(); + + auto* output = + tensorflow::gtl::FindPtrOrNull(*all_attr_to_input_maps, op_def.name()); + if (output != nullptr) { + return output; + } + + std::unique_ptr m(new AttrToInputsMap); + + // Store a list of InputIndex -> List of corresponding inputs. + for (int i = 0; i < op_def.input_arg_size(); i++) { + if (!op_def.input_arg(i).type_attr().empty()) { + auto it = m->find(op_def.input_arg(i).type_attr()); + if (it == m->end()) { + it = m->insert({op_def.input_arg(i).type_attr(), {}}).first; + } + it->second.emplace_back(i, !op_def.input_arg(i).number_attr().empty()); + } + } + + auto* retval = m.get(); + (*all_attr_to_input_maps)[op_def.name()] = m.release(); + + return retval; +} + struct FastPathOpExecInfo { TFE_Context* ctx; const char* device_name; @@ -53,6 +101,14 @@ struct FastPathOpExecInfo { // The op type name of the main op being executed. PyObject* op_name; PyObject* callbacks; + + // All the args passed into the FastPathOpExecInfo. + PyObject* args; + + // DTypes can come from another input that has the same attr. So build that + // map. + const AttrToInputsMap* attr_to_inputs_map; + tensorflow::gtl::FlatMap cached_dtypes; }; #define PARSE_VALUE(fn_name, type, check_fn, parse_fn) \ @@ -76,12 +132,29 @@ PARSE_VALUE(ParseIntValue, int, PyLong_Check, PyLong_AsLong) 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 +#if PY_MAJOR_VERSION < 3 +bool ParseInt64Value(const string& key, PyObject* py_value, TF_Status* status, + int64_t* value) { + if (PyInt_Check(py_value)) { + *value = static_cast(PyInt_AsLong(py_value)); + return true; + } else if (PyLong_Check(py_value)) { + *value = static_cast(PyLong_AsLong(py_value)); + return true; + } + TF_SetStatus( + status, TF_INVALID_ARGUMENT, + tensorflow::strings::StrCat("Expecting int or long value for attr ", key, + ", got ", py_value->ob_type->tp_name) + .c_str()); + return false; +} +#endif + Py_ssize_t TensorShapeNumDims(PyObject* value) { const auto size = PySequence_Size(value); if (size == -1) { @@ -234,7 +307,7 @@ bool SetOpAttrList( 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 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) { @@ -296,7 +369,7 @@ void SetOpAttrListDefault( 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]); + 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(); @@ -349,7 +422,7 @@ void SetOpAttrListDefault( 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 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) { @@ -369,7 +442,7 @@ void SetOpAttrListDefault( } 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]); + 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); } @@ -1399,10 +1472,39 @@ PyObject* GetPythonObjectFromString(const char* s) { #endif } +PyObject* GetPythonObjectFromInt(int num) { +#if PY_MAJOR_VERSION >= 3 + return PyLong_FromLong(num); +#else + return PyInt_FromLong(num); +#endif +} + bool CheckResourceVariable(PyObject* item) { return PyObject_TypeCheck(item, resource_variable_type); } +bool IsNumberType(PyObject* item) { +#if PY_MAJOR_VERSION >= 3 + return PyFloat_Check(item) || PyLong_Check(item); +#else + return PyFloat_Check(item) || PyInt_Check(item) || PyLong_Check(item); +#endif +} + +bool CheckOneInput(PyObject* item) { + if (EagerTensor_CheckExact(item) || CheckResourceVariable(item) || + PyArray_Check(item) || IsNumberType(item)) { + return true; + } + + // Sequences are not properly handled. Sequences with purely python numeric + // types work, but sequences with mixes of EagerTensors and python numeric + // types don't work. + // TODO(nareshmodi): fix + return false; +} + bool CheckInputsOk(PyObject* seq, int start_index, const tensorflow::OpDef& op_def) { for (int i = 0; i < op_def.input_arg_size(); i++) { @@ -1419,8 +1521,7 @@ bool CheckInputsOk(PyObject* seq, int start_index, } for (Py_ssize_t j = 0; j < PySequence_Fast_GET_SIZE(item); j++) { PyObject* inner_item = PySequence_Fast_GET_ITEM(item, j); - if (!EagerTensor_CheckExact(inner_item) && - !CheckResourceVariable(inner_item)) { + if (!CheckOneInput(inner_item)) { VLOG(1) << "Falling back to slow path for Op \"" << op_def.name() << "\", Input \"" << op_def.input_arg(i).name() << "\", Index " @@ -1430,7 +1531,7 @@ bool CheckInputsOk(PyObject* seq, int start_index, return false; } } - } else if (!EagerTensor_CheckExact(item) && !CheckResourceVariable(item)) { + } else if (!CheckOneInput(item)) { VLOG(1) << "Falling back to slow path for Op \"" << op_def.name() << "\", Input \"" << op_def.input_arg(i).name() @@ -1443,6 +1544,52 @@ bool CheckInputsOk(PyObject* seq, int start_index, return true; } +PyObject* MaybeGetDType(PyObject* item) { + if (EagerTensor_CheckExact(item)) { + tensorflow::Safe_PyObjectPtr py_dtype( + PyObject_GetAttrString(item, "dtype")); + return PyObject_GetAttrString(py_dtype.get(), "_type_enum"); + } + + if (CheckResourceVariable(item)) { + tensorflow::Safe_PyObjectPtr py_dtype( + PyObject_GetAttrString(item, "_dtype")); + return PyObject_GetAttrString(py_dtype.get(), "_type_enum"); + } + + return nullptr; +} + +PyObject* MaybeGetDTypeForAttr(const string& attr, + FastPathOpExecInfo* op_exec_info) { + auto cached_it = op_exec_info->cached_dtypes.find(attr); + if (cached_it != op_exec_info->cached_dtypes.end()) { + return GetPythonObjectFromInt(cached_it->second); + } + + auto it = op_exec_info->attr_to_inputs_map->find(attr); + if (it == op_exec_info->attr_to_inputs_map->end()) { + // No other inputs - this should never happen. + Py_RETURN_NONE; + } + + for (const auto& input_info : it->second) { + PyObject* item = PyTuple_GET_ITEM( + op_exec_info->args, kFastPathExecuteInputStartIndex + input_info.i); + if (input_info.is_list) { + for (int i = 0; i < PySequence_Fast_GET_SIZE(item); i++) { + auto* dtype = MaybeGetDType(PySequence_Fast_GET_ITEM(item, i)); + if (dtype != nullptr) return dtype; + } + } else { + auto* dtype = MaybeGetDType(item); + if (dtype != nullptr) return dtype; + } + } + + Py_RETURN_NONE; +} + bool OpDoesntRequireOutput(const string& op_name) { static tensorflow::gtl::FlatSet* ops_that_dont_require_outputs = new tensorflow::gtl::FlatSet({ @@ -1668,23 +1815,80 @@ bool ReadVariableOp(const FastPathOpExecInfo& parent_op_exec_info, // 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)) { +// +// NOTE: dtype_hint_getter must *always* return a PyObject that can be +// decref'd. So if no hint is found, Py_RETURN_NONE (which correctly +// increfs Py_None). +bool ConvertToTensor( + const FastPathOpExecInfo& op_exec_info, PyObject* input, + tensorflow::Safe_PyObjectPtr* output_handle, + // This gets a hint for this particular input. + const std::function& dtype_hint_getter, + // This sets the dtype after conversion is complete. + const std::function& dtype_setter, + TF_Status* status) { + if (EagerTensor_CheckExact(input)) { + Py_INCREF(input); + output_handle->reset(input); + return true; + } else if (CheckResourceVariable(input)) { return ReadVariableOp(op_exec_info, input, output_handle, status); } - Py_INCREF(input); - output_handle->reset(input); + // The hint comes from a supposedly similarly typed tensor. + tensorflow::Safe_PyObjectPtr dtype_hint(dtype_hint_getter()); + if (PyErr_Occurred()) { + return false; + } + + tensorflow::Safe_TFE_TensorHandlePtr handle = + tensorflow::make_safe(static_cast( + tensorflow::ConvertToEagerTensor(input, dtype_hint.get()))); + if (handle == nullptr) { + status->status = tensorflow::errors::InvalidArgument( + "Unable to convert value to tensor"); + return false; + } + + int desired_dtype = -1; + if (dtype_hint.get() != Py_None) { + if (!ParseTypeValue("", dtype_hint.get(), status, &desired_dtype)) { + status->status = tensorflow::errors::InvalidArgument( + "Expecting a DataType value for dtype. Got ", + Py_TYPE(dtype_hint.get())->tp_name); + } + } + + TF_DataType handle_dtype = TFE_TensorHandleDataType(handle.get()); + if (desired_dtype >= 0 && desired_dtype != handle_dtype) { + handle = tensorflow::make_safe( + tensorflow::EagerCast(op_exec_info.ctx, handle.get(), handle_dtype, + static_cast(desired_dtype), status)); + if (!status->status.ok()) return false; + + handle_dtype = TFE_TensorHandleDataType(handle.get()); + } + + if (handle_dtype != TF_INT32) { + // Note that this is a shallow copy and will share the underlying buffer + // if copying to the same device. + handle = tensorflow::make_safe(TFE_TensorHandleCopyToDevice( + handle.get(), op_exec_info.ctx, op_exec_info.device_name, status)); + if (!status->status.ok()) return false; + } + + output_handle->reset(EagerTensorFromHandle(handle.release())); + + dtype_setter(handle_dtype); 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, +bool AddInputToOp(FastPathOpExecInfo* op_exec_info, PyObject* input, + const bool add_type_attr, + const tensorflow::OpDef::ArgDef& input_arg, std::vector* flattened_attrs, std::vector* flattened_inputs, TFE_Op* op, TF_Status* status) { @@ -1693,18 +1897,30 @@ bool AddInputToOp(const FastPathOpExecInfo& op_exec_info, PyObject* input, // out of scope in this function. tensorflow::Safe_PyObjectPtr py_eager_tensor = nullptr; - if (!ConvertToTensor(op_exec_info, input, &py_eager_tensor, status)) { + if (!ConvertToTensor( + *op_exec_info, input, &py_eager_tensor, + [&]() { + if (input_arg.type() != tensorflow::DataType::DT_INVALID) { + return GetPythonObjectFromInt(input_arg.type()); + } + return MaybeGetDTypeForAttr(input_arg.type_attr(), op_exec_info); + }, + [&](const TF_DataType dtype) { + op_exec_info->cached_dtypes[input_arg.type_attr()] = + static_cast(dtype); + }, + status)) { return false; } TFE_TensorHandle* input_handle = EagerTensor_Handle(py_eager_tensor.get()); - if (input_arg != nullptr && !input_arg->type_attr().empty()) { + if (add_type_attr && !input_arg.type_attr().empty()) { auto dtype = TFE_TensorHandleDataType(input_handle); - TFE_OpSetAttrType(op, input_arg->type_attr().data(), dtype); + TFE_OpSetAttrType(op, input_arg.type_attr().data(), dtype); if (flattened_attrs != nullptr) { flattened_attrs->emplace_back( - GetPythonObjectFromString(input_arg->type_attr().data())); + GetPythonObjectFromString(input_arg.type_attr().data())); flattened_attrs->emplace_back(PyLong_FromLong(dtype)); } } @@ -1844,6 +2060,7 @@ PyObject* TFE_Py_FastPathExecute_C(PyObject*, PyObject* args) { op_exec_info.ctx = reinterpret_cast( PyCapsule_GetPointer(PyTuple_GET_ITEM(args, 0), nullptr)); + op_exec_info.args = args; if (op_exec_info.ctx == nullptr) { // The context hasn't been initialized. It will be in the slow path. @@ -1892,6 +2109,8 @@ PyObject* TFE_Py_FastPathExecute_C(PyObject*, PyObject* args) { return nullptr; } + op_exec_info.attr_to_inputs_map = GetAttrToInputsMap(*op_def); + TF_Status* status = TF_NewStatus(); TFE_Op* op = TFE_NewOp(op_exec_info.ctx, op_def->name().c_str(), status); auto cleaner = tensorflow::gtl::MakeCleanup([status, op] { @@ -1986,17 +2205,16 @@ PyObject* TFE_Py_FastPathExecute_C(PyObject*, PyObject* args) { if (len > 0) { // First item adds the type attr. - if (!AddInputToOp(op_exec_info, PySequence_Fast_GET_ITEM(input, 0), - &input_arg, flattened_attrs.get(), + if (!AddInputToOp(&op_exec_info, PySequence_Fast_GET_ITEM(input, 0), + true, 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, PySequence_Fast_GET_ITEM(input, j), - nullptr /* input_arg */, - nullptr /* flattened_attrs */, + if (!AddInputToOp(&op_exec_info, PySequence_Fast_GET_ITEM(input, j), + false, input_arg, nullptr /* flattened_attrs */, flattened_inputs.get(), op, status)) { return nullptr; } @@ -2018,7 +2236,8 @@ PyObject* TFE_Py_FastPathExecute_C(PyObject*, PyObject* args) { PyObject* py_input = PySequence_Fast_GET_ITEM(input, j); tensorflow::Safe_PyObjectPtr py_eager_tensor; if (!ConvertToTensor(op_exec_info, py_input, &py_eager_tensor, - status)) { + []() { Py_RETURN_NONE; }, + [](const TF_DataType& dtype) {}, status)) { return nullptr; } @@ -2048,8 +2267,9 @@ PyObject* TFE_Py_FastPathExecute_C(PyObject*, PyObject* args) { 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)) { + if (!AddInputToOp(&op_exec_info, input, true, input_arg, + flattened_attrs.get(), flattened_inputs.get(), op, + status)) { return nullptr; } } diff --git a/tensorflow/python/eager/tensor_test.py b/tensorflow/python/eager/tensor_test.py index 0bd5a5dbafd5ea8da21d4fb8a7dcae9fe23dd3d2..b044b30231603b0265aa1ef0320e9f1cfb303724 100644 --- a/tensorflow/python/eager/tensor_test.py +++ b/tensorflow/python/eager/tensor_test.py @@ -278,14 +278,9 @@ class TFETensorUtilTest(test_util.TensorFlowTestCase): with self.assertRaisesRegexp( TypeError, - r"tensor_list argument must be a list. Got \"EagerTensor\""): + r"tensors argument must be a list or a tuple. Got \"EagerTensor\""): pywrap_tensorflow.TFE_Py_TensorShapeSlice(t1, -2) - with self.assertRaisesRegexp( - TypeError, - r"tensor_list argument must be a list. Got \"tuple\""): - pywrap_tensorflow.TFE_Py_TensorShapeSlice((t1,), -2) - def testNegativeSliceDim(self): t1 = _create_tensor([1, 2], dtype=dtypes.int32) diff --git a/tensorflow/python/estimator/BUILD b/tensorflow/python/estimator/BUILD index 5d8b19223f000862aa46ad3a60796ae68bdec2f9..c6bb9b9be7cb8049acb2d6c3fe0f50a720e71f2e 100644 --- a/tensorflow/python/estimator/BUILD +++ b/tensorflow/python/estimator/BUILD @@ -7,7 +7,6 @@ package( licenses(["notice"]) # Apache 2.0 load("//tensorflow:tensorflow.bzl", "py_test") -load("//tensorflow:tensorflow.bzl", "cuda_py_test") py_library( name = "estimator_py", @@ -25,7 +24,6 @@ py_library( ":linear", ":model_fn", ":parsing_utils", - ":replicate_model_fn", ":run_config", ":training", "//tensorflow/python:util", @@ -207,6 +205,7 @@ py_test( "no_pip", "noasan", # test flakily times out in asan mode. "notsan", # b/67510291 + "optonly", # flakily times out in fastbuild ], deps = [ ":baseline", @@ -251,6 +250,7 @@ py_library( "//tensorflow/python:array_ops", "//tensorflow/python:boosted_trees_ops", "//tensorflow/python:data_flow_ops", + "//tensorflow/python:distribute", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", "//tensorflow/python:lookup_ops", @@ -327,6 +327,7 @@ py_library( "//tensorflow/python:client_testlib", "//tensorflow/python:constant_op", "//tensorflow/python:control_flow_ops", + "//tensorflow/python:distribute", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", "//tensorflow/python:math_ops", @@ -383,6 +384,7 @@ py_library( ":model_fn", ":optimizers", "//tensorflow/python:control_flow_ops", + "//tensorflow/python:distribute", "//tensorflow/python:framework_ops", "//tensorflow/python:init_ops", "//tensorflow/python:layers", @@ -466,6 +468,7 @@ py_library( "//tensorflow/core:protos_all_py", "//tensorflow/python:client", "//tensorflow/python:control_flow_ops", + "//tensorflow/python:distribute", "//tensorflow/python:framework_for_generated_wrappers", "//tensorflow/python:metrics", "//tensorflow/python:platform", @@ -743,6 +746,7 @@ py_library( "//tensorflow/python:client_testlib", "//tensorflow/python:control_flow_ops", "//tensorflow/python:data_flow_ops", + "//tensorflow/python:distribute", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", "//tensorflow/python:math_ops", @@ -904,68 +908,3 @@ py_test( "//tensorflow/python:training", ], ) - -py_library( - 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/ops/losses", - "@six_archive//:six", - ], -) - -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: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", - ":replicate_model_fn", - ], - tags = [ - "multi_gpu", - "noasan", # flaky time outs - "notsan", # flaky - ], -) diff --git a/tensorflow/python/estimator/canned/head.py b/tensorflow/python/estimator/canned/head.py index 189b81aeea80a987b5d9d70c0c5360703dd5a3f5..efa4bdf5980a34001e10bb7a1125e4434215f0ee 100644 --- a/tensorflow/python/estimator/canned/head.py +++ b/tensorflow/python/estimator/canned/head.py @@ -263,9 +263,12 @@ def _check_dense_labels_match_logits_and_reshape( if (dim1 is not None) and (dim1 != expected_labels_dimension): raise ValueError( 'Mismatched label shape. ' - 'Classifier configured with n_classes=%s. Received %s. ' - 'Suggested Fix: check your n_classes argument to the estimator ' - 'and/or the shape of your label.' % + 'Expected labels dimension=%s. Received %s. ' + 'Suggested Fix:' + 'If your classifier expects one-hot encoding label,' + 'check your n_classes argument to the estimator' + 'and/or the shape of your label.' + 'Otherwise, check the shape of your label.' % (expected_labels_dimension, dim1)) expected_labels_shape = array_ops.concat( [logits_shape[:-1], [expected_labels_dimension]], axis=0) @@ -1039,7 +1042,7 @@ class _BinaryLogisticHeadWithSigmoidCrossEntropyLoss(_Head): vocabulary_list=tuple(self._label_vocabulary), name='class_id_lookup').lookup(labels) labels = math_ops.to_float(labels) - labels = _assert_range(labels, 2) + labels = _assert_range(labels, n_classes=2) if self._loss_fn: unweighted_loss = _call_loss_fn( loss_fn=self._loss_fn, labels=labels, logits=logits, @@ -1447,12 +1450,12 @@ class _RegressionHeadWithMeanSquaredErrorLoss(_Head): def _assert_range(labels, n_classes, message=None): with ops.name_scope(None, 'assert_range', (labels,)): - assert_less = check_ops.assert_less( + assert_less = check_ops.assert_less_equal( labels, - ops.convert_to_tensor(n_classes, dtype=labels.dtype), - message=message or 'Label IDs must < n_classes') + ops.convert_to_tensor(n_classes - 1, dtype=labels.dtype), + message=message or 'Labels must <= n_classes - 1') assert_greater = check_ops.assert_non_negative( - labels, message=message or 'Label IDs must >= 0') + labels, message=message or 'Labels must >= 0') with ops.control_dependencies((assert_less, assert_greater)): return array_ops.identity(labels) diff --git a/tensorflow/python/estimator/canned/head_test.py b/tensorflow/python/estimator/canned/head_test.py index fe6ee07529bc0314618a7cc85926dbb39660a352..7da3df01dc48d9ed3aef8a030f7a516db3a5abeb 100644 --- a/tensorflow/python/estimator/canned/head_test.py +++ b/tensorflow/python/estimator/canned/head_test.py @@ -255,14 +255,14 @@ class MultiClassHeadWithSoftmaxCrossEntropyLoss(test.TestCase): logits=logits_placeholder, labels=labels_placeholder)[0] with self.test_session(): - with self.assertRaisesOpError('Label IDs must < n_classes'): + with self.assertRaisesOpError('Labels must <= n_classes - 1'): training_loss.eval({ labels_placeholder: labels_2x1_with_large_id, logits_placeholder: logits_2x3 }) with self.test_session(): - with self.assertRaisesOpError('Label IDs must >= 0'): + with self.assertRaisesOpError('Labels must >= 0'): training_loss.eval({ labels_placeholder: labels_2x1_with_negative_id, logits_placeholder: logits_2x3 @@ -2090,6 +2090,24 @@ class BinaryLogisticHeadWithSigmoidCrossEntropyLossTest(test.TestCase): expected_regularization_loss), }, summary_str) + def test_float_labels_invalid_values(self): + head = head_lib._binary_logistic_head_with_sigmoid_cross_entropy_loss() + + logits = np.array([[0.5], [-0.3]], dtype=np.float32) + labels = np.array([[1.2], [0.4]], dtype=np.float32) + features = {'x': np.array([[42]], dtype=np.float32)} + training_loss = head.create_loss( + features=features, + mode=model_fn.ModeKeys.TRAIN, + logits=logits, + labels=labels)[0] + with self.assertRaisesRegexp( + errors.InvalidArgumentError, + r'Labels must <= n_classes - 1'): + with self.test_session(): + _initialize_variables(self, monitored_session.Scaffold()) + training_loss.eval() + def test_float_labels_train_create_loss(self): head = head_lib._binary_logistic_head_with_sigmoid_cross_entropy_loss() diff --git a/tensorflow/python/estimator/estimator.py b/tensorflow/python/estimator/estimator.py index 4d3eff71ad2167315614c41b70f1127d51b12de3..887e9a474ab0471de64c53f24520c746e95195e9 100644 --- a/tensorflow/python/estimator/estimator.py +++ b/tensorflow/python/estimator/estimator.py @@ -216,7 +216,8 @@ class Estimator(object): else: self._session_config = self._config.session_config - self._device_fn = _get_replica_device_setter(self._config) + self._device_fn = self._config.device_fn or \ + _get_replica_device_setter(self._config) if model_fn is None: raise ValueError('model_fn must be provided to Estimator.') @@ -637,7 +638,7 @@ class Estimator(object): # pylint: disable=protected-access local_init_op = ( estimator_spec.scaffold.local_init_op or - monitored_session.Scaffold._default_local_init_op()) + monitored_session.Scaffold.default_local_init_op()) # pylint: enable=protected-access # Perform the export @@ -688,24 +689,16 @@ class Estimator(object): def _get_features_and_labels_from_input_fn(self, input_fn, mode): """Extracts the `features` and labels from return values of `input_fn`.""" - result = self._call_input_fn(input_fn, mode) - # TODO(anjalisridhar): What about the default DistributionStrategy? Perhaps - # using any input is alright in that case. There is also a - # has_dataset_or_queue_runner function that we may want to extend and use. - if (self._distribution is not None and - not isinstance(result, dataset_ops.Dataset) and - mode == model_fn_lib.ModeKeys.TRAIN): - raise ValueError('input_fn() must return a tf.data.Dataset when using a ' - 'DistributionStrategy.') input_hooks = [] - if isinstance(result, dataset_ops.Dataset): - if self._distribution is not None and mode == model_fn_lib.ModeKeys.TRAIN: - # TODO(josh11b): This is currently using a one-shot iterator, we - # will update this to an initializeable iterator once the - # necessory support for creating an initializable iterator is - # available. - result = self._distribution.distribute_dataset(result).get_next() - else: + if self._distribution is not None and mode == model_fn_lib.ModeKeys.TRAIN: + result = self._distribution.distribute_dataset( + lambda: self._call_input_fn(input_fn, mode)) + iterator = result.make_initializable_iterator() + input_hooks.append(_DatasetInitializerHook(iterator)) + result = iterator.get_next() + else: + result = self._call_input_fn(input_fn, mode) + if isinstance(result, dataset_ops.Dataset): iterator = result.make_initializable_iterator() input_hooks.append(_DatasetInitializerHook(iterator)) result = iterator.get_next() @@ -723,7 +716,7 @@ class Estimator(object): batch_length = batch_length or value.shape[0] if value.shape[0] != batch_length: raise ValueError('Batch length of predictions should be same. %s has ' - 'different batch length then others.' % key) + 'different batch length than others.' % key) return batch_length def _extract_keys(self, predictions, predict_keys): @@ -1264,7 +1257,8 @@ def _dict_to_str(dictionary): A `str` representing the `dictionary`. """ return ', '.join('%s = %s' % (k, v) - for k, v in sorted(six.iteritems(dictionary))) + for k, v in sorted(six.iteritems(dictionary)) + if not isinstance(v, six.binary_type)) def _write_dict_to_summary(output_dir, diff --git a/tensorflow/python/estimator/export/export.py b/tensorflow/python/estimator/export/export.py index 9206a4964b3b7a6e3cc1e0f9e965a197be78c4ba..41c1f5a2e25cd6ba1e93744b5b82ecc34c4375d0 100644 --- a/tensorflow/python/estimator/export/export.py +++ b/tensorflow/python/estimator/export/export.py @@ -74,8 +74,20 @@ class ServingInputReceiver(collections.namedtuple( raise ValueError('feature keys must be strings: {}.'.format(name)) if not (isinstance(tensor, ops.Tensor) or isinstance(tensor, sparse_tensor.SparseTensor)): - raise ValueError( + value_error = ValueError( 'feature {} must be a Tensor or SparseTensor.'.format(name)) + # NOTE(ericmc): This if-else block is a specific carve-out for + # LabeledTensor, which has a `.tensor` attribute and which is + # convertible to tf.Tensor via ops.convert_to_tensor. + # Allowing all types convertible to tf.Tensor is considered by soergel@ + # to be too permissive. + if hasattr(tensor, 'tensor'): + try: + ops.convert_to_tensor(tensor) + except TypeError: + raise value_error + else: + raise value_error if receiver_tensors is None: raise ValueError('receiver_tensors must be defined.') diff --git a/tensorflow/python/estimator/export/export_test.py b/tensorflow/python/estimator/export/export_test.py index eb9688bc973666554b6057f5f546b9a2d18461d6..c203be7dacf80043079b35cd8e89fc5b69dcaaa0 100644 --- a/tensorflow/python/estimator/export/export_test.py +++ b/tensorflow/python/estimator/export/export_test.py @@ -39,6 +39,21 @@ from tensorflow.python.saved_model import signature_constants from tensorflow.python.saved_model import signature_def_utils +class LabeledTensorMock(object): + """Mock class emulating LabeledTensor.""" + + def __init__(self): + self.tensor = constant_op.constant([1]) + + +def _convert_labeled_tensor_mock_to_tensor(value, *args, **kwargs): + return ops.internal_convert_to_tensor(value.tensor, *args, **kwargs) + + +ops.register_tensor_conversion_function(LabeledTensorMock, + _convert_labeled_tensor_mock_to_tensor) + + class ExportTest(test_util.TensorFlowTestCase): def test_serving_input_receiver_constructor(self): @@ -135,6 +150,11 @@ class ExportTest(test_util.TensorFlowTestCase): with self.assertRaises(ValueError): _ = export.ServingInputReceiver(feature, receiver_tensor) + def test_feature_labeled_tensor(self): + feature = LabeledTensorMock() + receiver_tensor = array_ops.placeholder(dtypes.string) + _ = export.ServingInputReceiver(feature, receiver_tensor) + def test_receiver_wrong_type(self): feature = constant_op.constant(5) receiver_tensor = "not a tensor" diff --git a/tensorflow/python/estimator/replicate_model_fn.py b/tensorflow/python/estimator/replicate_model_fn.py deleted file mode 100644 index 144d89abf3444062927d9261301fe50f4a63b280..0000000000000000000000000000000000000000 --- a/tensorflow/python/estimator/replicate_model_fn.py +++ /dev/null @@ -1,824 +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 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 deleted file mode 100644 index ad1f9c02b92d7b1ce929494f4b6fbf636762a7fd..0000000000000000000000000000000000000000 --- a/tensorflow/python/estimator/replicate_model_fn_test.py +++ /dev/null @@ -1,1739 +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 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 f62c9cece6a4d370532ca3077d679a54f38918f1..8162b249f1f0be6c901f09ff21f9a67f7c5e492f 100644 --- a/tensorflow/python/estimator/run_config.py +++ b/tensorflow/python/estimator/run_config.py @@ -27,11 +27,13 @@ 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.estimator import util from tensorflow.python.util import compat_internal from tensorflow.python.util.tf_export import tf_export _USE_DEFAULT = object() +_VALID_DEVICE_FN_ARGS = set(['op']) # A list of the property names in RunConfig that the user is allowed to change. _DEFAULT_REPLACEABLE_LIST = [ @@ -44,7 +46,8 @@ _DEFAULT_REPLACEABLE_LIST = [ 'keep_checkpoint_max', 'keep_checkpoint_every_n_hours', 'log_step_count_steps', - 'train_distribute' + 'train_distribute', + 'device_fn' ] _SAVE_CKPT_ERR = ( @@ -279,6 +282,11 @@ def _validate_properties(run_config): _validate('tf_random_seed', lambda seed: isinstance(seed, six.integer_types), message='tf_random_seed must be integer.') + _validate('device_fn', lambda device_fn: six.callable(device_fn) and + set(util.fn_args(device_fn)) == _VALID_DEVICE_FN_ARGS, + message='device_fn must be callable with exactly' + ' one argument "op".') + class TaskType(object): MASTER = 'master' @@ -302,7 +310,8 @@ class RunConfig(object): keep_checkpoint_max=5, keep_checkpoint_every_n_hours=10000, log_step_count_steps=100, - train_distribute=None): + train_distribute=None, + device_fn=None): """Constructs a RunConfig. All distributed training related properties `cluster_spec`, `is_chief`, @@ -316,7 +325,7 @@ class RunConfig(object): a list of task addresses. `task` has two attributes: `type` and `index`, where `type` can be any of - the task types in `cluster`. ` When `TF_CONFIG` contains said information, + the task types in `cluster`. When `TF_CONFIG` contains said information, the following properties are set on this class: * `cluster_spec` is parsed from `TF_CONFIG['cluster']`. Defaults to {}. If @@ -430,6 +439,10 @@ class RunConfig(object): `tf.contrib.distribute.DistributionStrategy`. If specified, then Estimator will distribute the user's model during training, according to the policy specified by that strategy. + device_fn: A callable invoked for every `Operation` that takes the + `Operation` and returns the device string. If `None`, defaults to + the device function returned by `tf.train.replica_device_setter` + with round-robin strategy. Raises: ValueError: If both `save_checkpoints_steps` and `save_checkpoints_secs` @@ -466,7 +479,8 @@ class RunConfig(object): keep_checkpoint_max=keep_checkpoint_max, keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours, log_step_count_steps=log_step_count_steps, - train_distribute=train_distribute) + train_distribute=train_distribute, + device_fn=device_fn) self._init_distributed_setting_from_environment_var(tf_config) @@ -568,6 +582,16 @@ class RunConfig(object): def cluster_spec(self): return self._cluster_spec + @property + def device_fn(self): + """Returns the device_fn. + + If device_fn is not `None`, it overrides the default + device function used in `Estimator`. + Otherwise the default one is used. + """ + return self._device_fn + @property def evaluation_master(self): return self._evaluation_master @@ -697,7 +721,8 @@ class RunConfig(object): - `keep_checkpoint_max`, - `keep_checkpoint_every_n_hours`, - `log_step_count_steps`, - - `train_distribute`. + - `train_distribute`, + - `device_fn`. In addition, either `save_checkpoints_steps` or `save_checkpoints_secs` can be set (should not be both). diff --git a/tensorflow/python/estimator/run_config_test.py b/tensorflow/python/estimator/run_config_test.py index a3eef4c53fd90a1ce69f3067d0b5c15909f43cec..c8b12605e1aaad11e114e4ace63697b93f3b2b92 100644 --- a/tensorflow/python/estimator/run_config_test.py +++ b/tensorflow/python/estimator/run_config_test.py @@ -42,6 +42,7 @@ _SESSION_CONFIG_ERR = 'session_config must be instance of ConfigProto' _KEEP_CKPT_MAX_ERR = 'keep_checkpoint_max should be >= 0' _KEEP_CKPT_HOURS_ERR = 'keep_checkpoint_every_n_hours should be > 0' _TF_RANDOM_SEED_ERR = 'tf_random_seed must be integer' +_DEVICE_FN_ERR = 'device_fn must be callable with exactly one argument "op".' _ONE_CHIEF_ERR = 'The "cluster" in TF_CONFIG must have only one "chief" node.' _ONE_MASTER_ERR = 'The "cluster" in TF_CONFIG must have only one "master" node.' _INVALID_TASK_TYPE_FOR_EVAL_MASTER = ( @@ -83,6 +84,7 @@ class RunConfigTest(test.TestCase): self.assertEqual(5, config.keep_checkpoint_max) self.assertEqual(10000, config.keep_checkpoint_every_n_hours) self.assertIsNone(config.service) + self.assertIsNone(config.device_fn) def test_model_dir(self): empty_config = run_config_lib.RunConfig() @@ -93,6 +95,7 @@ class RunConfigTest(test.TestCase): def test_replace_with_allowed_properties(self): session_config = config_pb2.ConfigProto(allow_soft_placement=True) + device_fn = lambda op: "/cpu:0" config = run_config_lib.RunConfig().replace( tf_random_seed=11, @@ -100,13 +103,15 @@ class RunConfigTest(test.TestCase): save_checkpoints_secs=14, session_config=session_config, keep_checkpoint_max=16, - keep_checkpoint_every_n_hours=17) + keep_checkpoint_every_n_hours=17, + device_fn=device_fn) self.assertEqual(11, config.tf_random_seed) self.assertEqual(12, config.save_summary_steps) self.assertEqual(14, config.save_checkpoints_secs) self.assertEqual(session_config, config.session_config) self.assertEqual(16, config.keep_checkpoint_max) self.assertEqual(17, config.keep_checkpoint_every_n_hours) + self.assertEqual(device_fn, config.device_fn) def test_replace_none_value(self): config = run_config_lib.RunConfig().replace( @@ -117,7 +122,8 @@ class RunConfigTest(test.TestCase): save_checkpoints_steps=None, session_config=None, keep_checkpoint_max=None, - keep_checkpoint_every_n_hours=None) + keep_checkpoint_every_n_hours=None, + device_fn=None) self.assertIsNone(config.tf_random_seed) self.assertIsNone(config.model_dir) self.assertIsNone(config.save_summary_steps) @@ -126,6 +132,7 @@ class RunConfigTest(test.TestCase): self.assertIsNone(config.session_config) self.assertIsNone(config.keep_checkpoint_max) self.assertIsNone(config.keep_checkpoint_every_n_hours) + self.assertIsNone(config.device_fn) def test_replace_with_disallowallowed_properties(self): config = run_config_lib.RunConfig() @@ -166,9 +173,12 @@ class RunConfigTest(test.TestCase): config.replace(keep_checkpoint_every_n_hours=0) with self.assertRaisesRegexp(ValueError, _TF_RANDOM_SEED_ERR): config.replace(tf_random_seed=1.0) + with self.assertRaisesRegexp(ValueError, _DEVICE_FN_ERR): + config.replace(device_fn=lambda x, y: 0) def test_init_with_allowed_properties(self): session_config = config_pb2.ConfigProto(allow_soft_placement=True) + device_fn = lambda op: "/cpu:0" config = run_config_lib.RunConfig( tf_random_seed=11, @@ -176,13 +186,15 @@ class RunConfigTest(test.TestCase): save_checkpoints_secs=14, session_config=session_config, keep_checkpoint_max=16, - keep_checkpoint_every_n_hours=17) + keep_checkpoint_every_n_hours=17, + device_fn=device_fn) self.assertEqual(11, config.tf_random_seed) self.assertEqual(12, config.save_summary_steps) self.assertEqual(14, config.save_checkpoints_secs) self.assertEqual(session_config, config.session_config) self.assertEqual(16, config.keep_checkpoint_max) self.assertEqual(17, config.keep_checkpoint_every_n_hours) + self.assertEqual(device_fn, config.device_fn) def test_init_none_value(self): config = run_config_lib.RunConfig( @@ -193,7 +205,8 @@ class RunConfigTest(test.TestCase): save_checkpoints_steps=None, session_config=None, keep_checkpoint_max=None, - keep_checkpoint_every_n_hours=None) + keep_checkpoint_every_n_hours=None, + device_fn=None) self.assertIsNone(config.tf_random_seed) self.assertIsNone(config.model_dir) self.assertIsNone(config.save_summary_steps) @@ -202,6 +215,7 @@ class RunConfigTest(test.TestCase): self.assertIsNone(config.session_config) self.assertIsNone(config.keep_checkpoint_max) self.assertIsNone(config.keep_checkpoint_every_n_hours) + self.assertIsNone(config.device_fn) def test_init_invalid_values(self): with self.assertRaisesRegexp(ValueError, _MODEL_DIR_ERR): @@ -220,6 +234,8 @@ class RunConfigTest(test.TestCase): run_config_lib.RunConfig(keep_checkpoint_every_n_hours=0) with self.assertRaisesRegexp(ValueError, _TF_RANDOM_SEED_ERR): run_config_lib.RunConfig(tf_random_seed=1.0) + with self.assertRaisesRegexp(ValueError, _DEVICE_FN_ERR): + run_config_lib.RunConfig(device_fn=lambda x: "/cpu:0") class RunConfigDistributedSettingTest(test.TestCase): diff --git a/tensorflow/python/estimator/training.py b/tensorflow/python/estimator/training.py index e38b765da52a7b6957a4fb8a02087c5d1fd5a781..9d271758f635869730c03bc2ac853b3493ec0cae 100644 --- a/tensorflow/python/estimator/training.py +++ b/tensorflow/python/estimator/training.py @@ -137,7 +137,7 @@ class TrainSpec( * 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 @@ -334,7 +334,8 @@ def train_and_evaluate(estimator, train_spec, eval_spec): can read and write). The only extra work to do is setting the environment variable `TF_CONFIG` properly for each worker correspondingly. - Also see: https://www.tensorflow.org/deploy/distributed + Also see + [Distributed TensorFlow](https://www.tensorflow.org/deploy/distributed). Setting environment variable depends on the platform. For example, on Linux, it can be done as follows (`$` is the shell prompt): diff --git a/tensorflow/python/feature_column/feature_column.py b/tensorflow/python/feature_column/feature_column.py index 3a315e5c2ea0d9607b5aa52715364d6bdf152e1c..c16c3cda4892b8017571c2b37736b85c80f3d8a4 100644 --- a/tensorflow/python/feature_column/feature_column.py +++ b/tensorflow/python/feature_column/feature_column.py @@ -135,6 +135,7 @@ import numpy as np import six +from tensorflow.python.eager import context from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib @@ -161,7 +162,6 @@ 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, @@ -408,58 +408,27 @@ def linear_model(features, ValueError: if an item in `feature_columns` is neither a `_DenseColumn` nor `_CategoricalColumn`. """ - feature_columns = _clean_feature_columns(feature_columns) - for column in feature_columns: - if not isinstance(column, (_DenseColumn, _CategoricalColumn)): - raise ValueError('Items of feature_columns must be either a _DenseColumn ' - 'or _CategoricalColumn. Given: {}'.format(column)) - weight_collections = list(weight_collections or []) - if ops.GraphKeys.GLOBAL_VARIABLES not in weight_collections: - weight_collections.append(ops.GraphKeys.GLOBAL_VARIABLES) - if ops.GraphKeys.MODEL_VARIABLES not in weight_collections: - weight_collections.append(ops.GraphKeys.MODEL_VARIABLES) - with variable_scope.variable_scope( - None, default_name='linear_model', values=features.values()): - weighted_sums = [] - ordered_columns = [] - builder = _LazyBuilder(features) - for column in sorted(feature_columns, key=lambda x: x.name): - with variable_scope.variable_scope( - None, default_name=column._var_scope_name): # pylint: disable=protected-access - ordered_columns.append(column) - weighted_sum = _create_weighted_sum( - column=column, - builder=builder, - units=units, - sparse_combiner=sparse_combiner, - weight_collections=weight_collections, - trainable=trainable) - weighted_sums.append(weighted_sum) - if cols_to_vars is not None: - # Retrieve the variables created. - cols_to_vars[column] = ops.get_collection( - ops.GraphKeys.GLOBAL_VARIABLES, - scope=variable_scope.get_variable_scope().name) - _verify_static_batch_size_equality(weighted_sums, ordered_columns) - predictions_no_bias = math_ops.add_n( - weighted_sums, name='weighted_sum_no_bias') - bias = variable_scope.get_variable( - 'bias_weights', - shape=[units], - initializer=init_ops.zeros_initializer(), - trainable=trainable, - collections=weight_collections) - predictions = nn_ops.bias_add( - predictions_no_bias, bias, name='weighted_sum') - if cols_to_vars is not None: - # Add the bias to cols_to_vars as well, converting the Variable or - # PartitionedVariable to a list of Variable's. - if (isinstance(bias, variables.Variable) or - resource_variable_ops.is_resource_variable(bias)): - cols_to_vars['bias'] = [bias] - else: # Must be a PartitionedVariable. - cols_to_vars['bias'] = list(bias) - return predictions + linear_model_layer = _LinearModel( + feature_columns=feature_columns, + units=units, + sparse_combiner=sparse_combiner, + weight_collections=weight_collections, + trainable=trainable, + name='linear_model') + retval = linear_model_layer(features) # pylint: disable=not-callable + if cols_to_vars is not None: + cols_to_vars.update(linear_model_layer.cols_to_vars()) + return retval + + +def _add_to_collections(var, weight_collections): + # TODO(rohanj): Explore adding a _get_variable_list method on `Variable` + # so that we don't have to do this check. + if isinstance(var, variables.PartitionedVariable): + for constituent_var in list(var): + ops.add_to_collections(weight_collections, constituent_var) + else: + ops.add_to_collections(weight_collections, var) class _FCLinearWrapper(base.Layer): @@ -482,12 +451,8 @@ class _FCLinearWrapper(base.Layer): self._units = units self._sparse_combiner = sparse_combiner self._weight_collections = weight_collections - self._state = {} def build(self, _): - self._state = self._feature_column._create_state( # pylint: disable=protected-access - self._weight_collections, self.add_variable) - if isinstance(self._feature_column, _CategoricalColumn): weight = self.add_variable( name='weights', @@ -501,7 +466,7 @@ class _FCLinearWrapper(base.Layer): shape=[num_elements, self._units], initializer=init_ops.zeros_initializer(), trainable=self.trainable) - ops.add_to_collections(self._weight_collections, weight) + _add_to_collections(weight, self._weight_collections) self._weight_var = weight self.built = True @@ -513,8 +478,7 @@ class _FCLinearWrapper(base.Layer): sparse_combiner=self._sparse_combiner, weight_collections=self._weight_collections, trainable=self.trainable, - weight_var=self._weight_var, - state=self._state) + weight_var=self._weight_var) return weighted_sum @@ -538,15 +502,29 @@ class _BiasLayer(base.Layer): shape=[self._units], initializer=init_ops.zeros_initializer(), trainable=self.trainable) - ops.add_to_collections(self._weight_collections, self._bias_variable) + _add_to_collections(self._bias_variable, self._weight_collections) self.built = True def call(self, _): return self._bias_variable +def _get_expanded_variable_list(variable): + if (isinstance(variable, variables.Variable) or + resource_variable_ops.is_resource_variable(variable)): + return [variable] # Single variable case. + else: # Must be a PartitionedVariable, so convert into a list. + return list(variable) + + +def _strip_leading_slashes(name): + return name.rsplit('/', 1)[-1] + + class _LinearModel(training.Model): """Creates a linear model using feature columns. + + See `linear_model` for details. """ def __init__(self, @@ -567,7 +545,10 @@ class _LinearModel(training.Model): for column in sorted(self._feature_columns, key=lambda x: x.name): with variable_scope.variable_scope( None, default_name=column._var_scope_name) as vs: # pylint: disable=protected-access - column_name = vs.name + # Having the fully expressed variable scope name ends up doubly + # expressing the outer scope (scope with which this method was called) + # in the name of the variable that would get created. + column_name = _strip_leading_slashes(vs.name) column_layer = _FCLinearWrapper(column, units, sparse_combiner, self._weight_collections, trainable, column_name, **kwargs) @@ -579,26 +560,45 @@ class _LinearModel(training.Model): weight_collections=self._weight_collections, name='bias_layer', **kwargs) + self._cols_to_vars = {} + + def cols_to_vars(self): + """Returns a dict mapping _FeatureColumns to variables. + + See `linear_model` for more information. + This is not populated till `call` is called i.e. layer is built. + """ + return self._cols_to_vars def call(self, features): - for column in self._feature_columns: - if not isinstance(column, (_DenseColumn, _CategoricalColumn)): - raise ValueError( - 'Items of feature_columns must be either a ' - '_DenseColumn or _CategoricalColumn. Given: {}'.format(column)) - weighted_sums = [] - ordered_columns = [] - builder = _LazyBuilder(features) - for layer in sorted(self._column_layers.values(), key=lambda x: x.name): - ordered_columns.append(layer._feature_column) # pylint: disable=protected-access - weighted_sum = layer(builder) - weighted_sums.append(weighted_sum) - - _verify_static_batch_size_equality(weighted_sums, ordered_columns) - predictions_no_bias = math_ops.add_n( - weighted_sums, name='weighted_sum_no_bias') - predictions = nn_ops.bias_add( - predictions_no_bias, self._bias_layer(builder), name='weighted_sum') # pylint: disable=not-callable + with variable_scope.variable_scope(self.name): + for column in self._feature_columns: + if not isinstance(column, (_DenseColumn, _CategoricalColumn)): + raise ValueError( + 'Items of feature_columns must be either a ' + '_DenseColumn or _CategoricalColumn. Given: {}'.format(column)) + weighted_sums = [] + ordered_columns = [] + builder = _LazyBuilder(features) + for layer in sorted(self._column_layers.values(), key=lambda x: x.name): + column = layer._feature_column # pylint: disable=protected-access + ordered_columns.append(column) + weighted_sum = layer(builder) + weighted_sums.append(weighted_sum) + self._cols_to_vars[column] = ops.get_collection( + ops.GraphKeys.GLOBAL_VARIABLES, scope=layer.scope_name) + + _verify_static_batch_size_equality(weighted_sums, ordered_columns) + predictions_no_bias = math_ops.add_n( + weighted_sums, name='weighted_sum_no_bias') + predictions = nn_ops.bias_add( + predictions_no_bias, + self._bias_layer( # pylint: disable=not-callable + builder, + scope=variable_scope.get_variable_scope()), # pylint: disable=not-callable + name='weighted_sum') + bias = self._bias_layer.variables[0] + self._cols_to_vars['bias'] = _get_expanded_variable_list(bias) return predictions def _add_layers(self, layers): @@ -805,11 +805,22 @@ def embedding_column( initializer = init_ops.truncated_normal_initializer( mean=0.0, stddev=1 / math.sqrt(dimension)) + embedding_shape = categorical_column._num_buckets, dimension # pylint: disable=protected-access + + def _creator(weight_collections, scope): + embedding_column_layer = _EmbeddingColumnLayer( + embedding_shape=embedding_shape, + initializer=initializer, + weight_collections=weight_collections, + trainable=trainable, + name='embedding_column_layer') + return embedding_column_layer(None, scope=scope) # pylint: disable=not-callable + return _EmbeddingColumn( categorical_column=categorical_column, dimension=dimension, combiner=combiner, - initializer=initializer, + layer_creator=_creator, ckpt_to_load_from=ckpt_to_load_from, tensor_name_in_ckpt=tensor_name_in_ckpt, max_norm=max_norm, @@ -932,6 +943,7 @@ def shared_embedding_columns( sorted_columns = sorted(categorical_columns, key=lambda x: x.name) c0 = sorted_columns[0] + num_buckets = c0._num_buckets # pylint: disable=protected-access if not isinstance(c0, _CategoricalColumn): raise ValueError( 'All categorical_columns must be subclasses of _CategoricalColumn. ' @@ -947,23 +959,45 @@ def shared_embedding_columns( 'the same type, or be weighted_categorical_column of the same type. ' 'Given column: {} of type: {} does not match given column: {} of ' 'type: {}'.format(c0, type(c0), c, type(c))) + if num_buckets != c._num_buckets: # pylint: disable=protected-access + raise ValueError( + 'To use shared_embedding_column, all categorical_columns must have ' + 'the same number of buckets. Given column: {} with buckets: {} does ' + 'not match column: {} with buckets: {}'.format( + c0, num_buckets, c, c._num_buckets)) # pylint: disable=protected-access if not shared_embedding_collection_name: shared_embedding_collection_name = '_'.join(c.name for c in sorted_columns) shared_embedding_collection_name += '_shared_embedding' + # Create the state (_SharedEmbeddingColumnLayer) here. + embedding_shape = num_buckets, dimension + + shared_embedding_column_layer = _EmbeddingColumnLayer( + embedding_shape=embedding_shape, + initializer=initializer, + weight_collections=[], + trainable=trainable, + name=shared_embedding_collection_name) + result = [] for column in categorical_columns: - result.append(_SharedEmbeddingColumn( - categorical_column=column, - dimension=dimension, - combiner=combiner, - initializer=initializer, - shared_embedding_collection_name=shared_embedding_collection_name, - ckpt_to_load_from=ckpt_to_load_from, - tensor_name_in_ckpt=tensor_name_in_ckpt, - max_norm=max_norm, - trainable=trainable)) + result.append( + _SharedEmbeddingColumn( + categorical_column=column, + initializer=initializer, + dimension=dimension, + combiner=combiner, + var_scope_name=shared_embedding_collection_name, + ckpt_to_load_from=ckpt_to_load_from, + tensor_name_in_ckpt=tensor_name_in_ckpt, + max_norm=max_norm, + trainable=trainable)) + + for single_result in result: + single_result._set_layer(shared_embedding_column_layer) # pylint: disable=protected-access + single_result._set_all_columns(result) # pylint: disable=protected-access + return result @@ -1720,6 +1754,57 @@ def crossed_column(keys, hash_bucket_size, hash_key=None): hash_key=hash_key) +# TODO(rohanj): Clearly define semantics of this layer. +class _EmbeddingColumnLayer(base.Layer): + """A layer that stores all the state required for a embedding column.""" + + def __init__(self, + embedding_shape, + initializer, + weight_collections=None, + trainable=True, + name=None, + **kwargs): + """Constructor. + + Args: + embedding_shape: Shape of the embedding variable used for lookup. + initializer: A variable initializer function to be used in embedding + variable initialization. If not specified, defaults to + `tf.truncated_normal_initializer` with mean `0.0` and standard deviation + `1/sqrt(dimension)`. + 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` (see `tf.Variable`). + name: Name of the layer + **kwargs: keyword named properties. + """ + super(_EmbeddingColumnLayer, self).__init__( + trainable=trainable, name=name, **kwargs) + self._embedding_shape = embedding_shape + self._initializer = initializer + self._weight_collections = weight_collections + + def build(self, _): + self._embedding_weight_var = self.add_variable( + name='embedding_weights', + shape=self._embedding_shape, + dtype=dtypes.float32, + initializer=self._initializer, + trainable=self.trainable) + # self.add_variable already appends to GLOBAL_VARIABLES collection. + if self._weight_collections and not context.executing_eagerly(): + for weight_collection in self._weight_collections: + if weight_collection != ops.GraphKeys.GLOBAL_VARIABLES: + _add_to_collections(self._embedding_weight_var, [weight_collection]) + self.built = True + + def call(self, _): + return self._embedding_weight_var + + class _FeatureColumn(object): """Represents a feature column abstraction. @@ -1793,18 +1878,13 @@ class _FeatureColumn(object): """ pass - def _create_state(self, weight_collections=None, creator=None): - """Returns an object that captures the state of the column. - - Args: - weight_collections: Collections to add the variable to - creator: Variable creator method called, if provided. + def _reset_config(self): + """Resets the configuration in the column. - Returns: - An object that encapsulates the state of the column. Can return None. + Some feature columns e.g. embedding or shared embedding columns might + have some state that is needed to be reset sometimes. Use this method + in that scenario. """ - del weight_collections, creator # Unused - return None class _DenseColumn(_FeatureColumn): @@ -1825,11 +1905,7 @@ class _DenseColumn(_FeatureColumn): pass @abc.abstractmethod - def _get_dense_tensor(self, - inputs, - weight_collections=None, - trainable=None, - state=None): + def _get_dense_tensor(self, inputs, weight_collections=None, trainable=None): """Returns a `Tensor`. The output of this function will be used by model-builder-functions. For @@ -1847,9 +1923,6 @@ class _DenseColumn(_FeatureColumn): will be created) are added. trainable: If `True` also add variables to the graph collection `GraphKeys.TRAINABLE_VARIABLES` (see @{tf.Variable}). - state: An object encapsulating the state of the column. Columns that - create state using the _create_state method would have that state - passed in to this method. Returns: `Tensor` of shape [batch_size] + `_variable_shape`. @@ -1863,8 +1936,7 @@ def _create_weighted_sum(column, sparse_combiner, weight_collections, trainable, - weight_var=None, - state=None): + weight_var=None): """Creates a weighted sum for a dense or sparse column for linear_model.""" if isinstance(column, _CategoricalColumn): return _create_categorical_column_weighted_sum( @@ -1882,8 +1954,7 @@ def _create_weighted_sum(column, units=units, weight_collections=weight_collections, trainable=trainable, - weight_var=weight_var, - state=state) + weight_var=weight_var) def _create_dense_column_weighted_sum(column, @@ -1891,20 +1962,12 @@ def _create_dense_column_weighted_sum(column, units, weight_collections, trainable, - weight_var=None, - state=None): + weight_var=None): """Create a weighted sum of a dense column for linear_model.""" - if state is not None: - tensor = column._get_dense_tensor( # pylint: disable=protected-access - builder, - weight_collections=weight_collections, - trainable=trainable, - state=state) - else: - tensor = column._get_dense_tensor( # pylint: disable=protected-access - builder, - weight_collections=weight_collections, - trainable=trainable) + tensor = column._get_dense_tensor( # pylint: disable=protected-access + builder, + weight_collections=weight_collections, + trainable=trainable) num_elements = column._variable_shape.num_elements() # pylint: disable=protected-access batch_size = array_ops.shape(tensor)[0] tensor = array_ops.reshape(tensor, shape=(batch_size, num_elements)) @@ -2367,10 +2430,10 @@ class _BucketizedColumn(_DenseColumn, _CategoricalColumn, class _EmbeddingColumn( _DenseColumn, _SequenceDenseColumn, - collections.namedtuple('_EmbeddingColumn', ( - 'categorical_column', 'dimension', 'combiner', 'initializer', - 'ckpt_to_load_from', 'tensor_name_in_ckpt', 'max_norm', 'trainable' - ))): + collections.namedtuple( + '_EmbeddingColumn', + ('categorical_column', 'dimension', 'combiner', 'layer_creator', + 'ckpt_to_load_from', 'tensor_name_in_ckpt', 'max_norm', 'trainable'))): """See `embedding_column`.""" @property @@ -2392,33 +2455,10 @@ class _EmbeddingColumn( self._shape = tensor_shape.vector(self.dimension) return self._shape - def _create_state(self, weight_collections=None, creator=None): - variables_map = {} - embedding_shape = (self.categorical_column._num_buckets, self.dimension) # pylint: disable=protected-access - if creator is not None: - embedding_weights = creator( - name='embedding_weights', - shape=embedding_shape, - dtype=dtypes.float32, - initializer=self.initializer, - trainable=self.trainable) - ops.add_to_collections(weight_collections, embedding_weights) - else: - embedding_weights = variable_scope.get_variable( - name='embedding_weights', - shape=embedding_shape, - dtype=dtypes.float32, - initializer=self.initializer, - trainable=self.trainable, - collections=weight_collections) - variables_map['embedding_weights'] = embedding_weights - return variables_map - def _get_dense_tensor_internal(self, inputs, weight_collections=None, - trainable=None, - state=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 @@ -2426,9 +2466,9 @@ class _EmbeddingColumn( sparse_ids = sparse_tensors.id_tensor sparse_weights = sparse_tensors.weight_tensor - if state is None: - state = self._create_state(weight_collections) - embedding_weights = state['embedding_weights'] + embedding_weights = self.layer_creator( + weight_collections=weight_collections, + scope=variable_scope.get_variable_scope()) if self.ckpt_to_load_from is not None: to_restore = embedding_weights @@ -2447,11 +2487,7 @@ class _EmbeddingColumn( name='%s_weights' % self.name, max_norm=self.max_norm) - def _get_dense_tensor(self, - inputs, - weight_collections=None, - trainable=None, - state=None): + def _get_dense_tensor(self, inputs, weight_collections=None, trainable=None): if isinstance(self.categorical_column, _SequenceCategoricalColumn): raise ValueError( 'In embedding_column: {}. ' @@ -2466,8 +2502,7 @@ class _EmbeddingColumn( return self._get_dense_tensor_internal( inputs=inputs, weight_collections=weight_collections, - trainable=trainable, - state=state) + trainable=trainable) def _get_sequence_dense_tensor( self, inputs, weight_collections=None, trainable=None): @@ -2491,13 +2526,20 @@ class _EmbeddingColumn( dense_tensor=dense_tensor, sequence_length=sequence_length) +def _get_graph_for_variable(var): + if isinstance(var, variables.PartitionedVariable): + return list(var)[0].graph + else: + return var.graph + + class _SharedEmbeddingColumn( _DenseColumn, - collections.namedtuple('_SharedEmbeddingColumn', ( - 'categorical_column', 'dimension', 'combiner', 'initializer', - 'shared_embedding_collection_name', 'ckpt_to_load_from', - 'tensor_name_in_ckpt', 'max_norm', 'trainable' - ))): + collections.namedtuple( + '_SharedEmbeddingColumn', + ('categorical_column', 'dimension', 'combiner', 'initializer', + 'var_scope_name', 'ckpt_to_load_from', 'tensor_name_in_ckpt', + 'max_norm', 'trainable'))): """See `embedding_column`.""" @property @@ -2508,7 +2550,7 @@ class _SharedEmbeddingColumn( @property def _var_scope_name(self): - return self.shared_embedding_collection_name + return self.var_scope_name @property def _parse_example_spec(self): @@ -2517,45 +2559,29 @@ class _SharedEmbeddingColumn( def _transform_feature(self, inputs): return inputs.get(self.categorical_column) + def _set_layer(self, layer): + self._layer = layer + + def _set_all_columns(self, all_columns): + self._all_columns = all_columns + + def _reset_config(self): + config = self._layer.get_config() + config['embedding_shape'] = ( + self.categorical_column._num_buckets, # pylint: disable=protected-access + self.dimension) + config['initializer'] = self.initializer + self._layer = self._layer.__class__.from_config(config) + for column in self._all_columns: + column._set_layer(self._layer) # pylint: disable=protected-access + @property def _variable_shape(self): if not hasattr(self, '_shape'): self._shape = tensor_shape.vector(self.dimension) return self._shape - def _create_state(self, weight_collections=None, creator=None): - variables_map = {} - shared_embedding_collection = ops.get_collection( - self.shared_embedding_collection_name) - if not shared_embedding_collection: - embedding_shape = (self.categorical_column._num_buckets, self.dimension) # pylint: disable=protected-access - if creator is not None: - embedding_weights = creator( - name='embedding_weights', - shape=embedding_shape, - dtype=dtypes.float32, - initializer=self.initializer, - trainable=self.trainable) - ops.add_to_collections(weight_collections, embedding_weights) - else: - embedding_weights = variable_scope.get_variable( - name='embedding_weights', - shape=embedding_shape, - dtype=dtypes.float32, - initializer=self.initializer, - trainable=self.trainable, - collections=weight_collections) - ops.add_to_collection(self.shared_embedding_collection_name, - embedding_weights) - variables_map['embedding_weights'] = embedding_weights - - return variables_map - - def _get_dense_tensor(self, - inputs, - weight_collections=None, - trainable=None, - state=None): + def _get_dense_tensor(self, inputs, weight_collections=None, trainable=None): # This method is called from a variable_scope with name _var_scope_name, # which is shared among all shared embeddings. Open a name_scope here, so # that the ops for different columns have distinct names. @@ -2566,38 +2592,17 @@ class _SharedEmbeddingColumn( sparse_ids = sparse_tensors.id_tensor sparse_weights = sparse_tensors.weight_tensor - embedding_shape = (self.categorical_column._num_buckets, self.dimension) # pylint: disable=protected-access - shared_embedding_collection = ops.get_collection( - self.shared_embedding_collection_name) - if shared_embedding_collection: - if len(shared_embedding_collection) > 1: - raise ValueError( - 'Collection {} can only contain one variable. ' - 'Suggested fix A: Choose a unique name for this collection. ' - 'Suggested fix B: Do not add any variables to this collection. ' - 'The feature_column library already adds a variable under the ' - 'hood.'.format(shared_embedding_collection)) - embedding_weights = shared_embedding_collection[0] - if embedding_weights.get_shape() != embedding_shape: - raise ValueError( - 'Shared embedding collection {} contains variable {} of ' - 'unexpected shape {}. Expected shape is {}. ' - 'Suggested fix A: Choose a unique name for this collection. ' - 'Suggested fix B: Do not add any variables to this collection. ' - 'The feature_column library already adds a variable under the ' - 'hood.'.format( - self.shared_embedding_collection_name, embedding_weights.name, - embedding_weights.get_shape(), embedding_shape)) - else: - embedding_weights = variable_scope.get_variable( - name='embedding_weights', - shape=embedding_shape, - dtype=dtypes.float32, - initializer=self.initializer, - trainable=self.trainable and trainable, - collections=weight_collections) - ops.add_to_collection( - self.shared_embedding_collection_name, embedding_weights) + embedding_weights = self._layer( + None, scope=variable_scope.get_variable_scope()) + # If we're in graph mode and this is called with a different graph, + # then we should reset. + if not context.executing_eagerly() and ( + ops.get_default_graph() != + _get_graph_for_variable(embedding_weights)): + self._reset_config() + embedding_weights = self._layer( + None, scope=variable_scope.get_variable_scope()) + if self.ckpt_to_load_from is not None: to_restore = embedding_weights if isinstance(to_restore, variables.PartitionedVariable): @@ -3147,6 +3152,9 @@ def _safe_embedding_lookup_sparse(embedding_weights, # Prune invalid ids and weights. sparse_ids, sparse_weights = _prune_invalid_ids(sparse_ids, sparse_weights) + if combiner != 'sum': + sparse_ids, sparse_weights = _prune_invalid_weights( + sparse_ids, sparse_weights) # Fill in dummy values for empty features, if necessary. sparse_ids, is_row_empty = sparse_ops.sparse_fill_empty_rows(sparse_ids, @@ -3195,13 +3203,23 @@ def _prune_invalid_ids(sparse_ids, sparse_weights): is_id_valid = math_ops.greater_equal(sparse_ids.values, 0) if sparse_weights is not None: is_id_valid = math_ops.logical_and( - is_id_valid, math_ops.greater(sparse_weights.values, 0)) + is_id_valid, + array_ops.ones_like(sparse_weights.values, dtype=dtypes.bool)) sparse_ids = sparse_ops.sparse_retain(sparse_ids, is_id_valid) if sparse_weights is not None: sparse_weights = sparse_ops.sparse_retain(sparse_weights, is_id_valid) return sparse_ids, sparse_weights +def _prune_invalid_weights(sparse_ids, sparse_weights): + """Prune invalid weights (< 0) from the input ids and weights.""" + if sparse_weights is not None: + is_weights_valid = math_ops.greater(sparse_weights.values, 0) + sparse_ids = sparse_ops.sparse_retain(sparse_ids, is_weights_valid) + sparse_weights = sparse_ops.sparse_retain(sparse_weights, is_weights_valid) + return sparse_ids, sparse_weights + + class _IndicatorColumn(_DenseColumn, _SequenceDenseColumn, collections.namedtuple('_IndicatorColumn', ['categorical_column'])): diff --git a/tensorflow/python/feature_column/feature_column_test.py b/tensorflow/python/feature_column/feature_column_test.py index 07588af37ee92eb2143d20eafa2874d794360fa4..d963dd9b551c0ebefbbf2677af75114abc59c084 100644 --- a/tensorflow/python/feature_column/feature_column_test.py +++ b/tensorflow/python/feature_column/feature_column_test.py @@ -345,7 +345,7 @@ class NumericColumnTest(test.TestCase): with ops.Graph().as_default(): features = {'price': [[1.], [5.]]} predictions = get_keras_linear_model_predictions(features, [price]) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() price_var = get_linear_model_column_var(price) with _initialized_session() as sess: self.assertAllClose([0.], bias.eval()) @@ -584,7 +584,7 @@ class BucketizedColumnTest(test.TestCase): features = {'price': [[-1.], [1.], [5.], [6.]]} predictions = get_keras_linear_model_predictions(features, [bucketized_price]) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() bucketized_price_var = get_linear_model_column_var(bucketized_price) with _initialized_session() as sess: self.assertAllClose([0.], bias.eval()) @@ -610,7 +610,7 @@ class BucketizedColumnTest(test.TestCase): features = {'price': [[-1., 1.], [5., 6.]]} predictions = get_keras_linear_model_predictions(features, [bucketized_price]) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() bucketized_price_var = get_linear_model_column_var(bucketized_price) with _initialized_session() as sess: self.assertAllClose([0.], bias.eval()) @@ -849,7 +849,7 @@ class HashedCategoricalColumnTest(test.TestCase): values=('marlo', 'skywalker', 'omar'), dense_shape=(2, 2)) }, (wire_column,)) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() wire_var = get_linear_model_column_var(wire_column) with _initialized_session(): self.assertAllClose((0.,), bias.eval()) @@ -1171,7 +1171,7 @@ class CrossedColumnTest(test.TestCase): values=['cA', 'cB', 'cC'], dense_shape=(2, 2)), }, (crossed,)) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() crossed_var = get_linear_model_column_var(crossed) with _initialized_session() as sess: self.assertAllClose((0.,), bias.eval()) @@ -1254,18 +1254,13 @@ def get_linear_model_column_var(column): 'linear_model/' + column.name)[0] -def get_keras_linear_model_bias(): - with variable_scope.variable_scope('linear_model', reuse=True): - with variable_scope.variable_scope('bias_layer', reuse=True): - return variable_scope.get_variable('bias_weights') - - def get_keras_linear_model_predictions(features, feature_columns, units=1, sparse_combiner='sum', weight_collections=None, - trainable=True): + trainable=True, + cols_to_vars=None): keras_linear_model = _LinearModel( feature_columns, units, @@ -1273,7 +1268,10 @@ def get_keras_linear_model_predictions(features, weight_collections, trainable, name='linear_model') - return keras_linear_model(features) # pylint: disable=not-callable + retval = keras_linear_model(features) # pylint: disable=not-callable + if cols_to_vars is not None: + cols_to_vars.update(keras_linear_model.cols_to_vars()) + return retval @test_util.with_c_api @@ -1511,6 +1509,28 @@ class LinearModelTest(test.TestCase): sess.run(bias.assign([5.])) self.assertAllClose([[1005.], [5010.]], predictions.eval()) + def test_sparse_combiner_with_negative_weights(self): + wire_cast = fc.categorical_column_with_hash_bucket('wire_cast', 4) + wire_cast_weights = fc.weighted_categorical_column(wire_cast, 'weights') + + with ops.Graph().as_default(): + wire_tensor = sparse_tensor.SparseTensor( + values=['omar', 'stringer', 'marlo'], # hashed to = [2, 0, 3] + indices=[[0, 0], [1, 0], [1, 1]], + dense_shape=[2, 2]) + features = { + 'wire_cast': wire_tensor, + 'weights': constant_op.constant([[1., 1., -1.0]]) + } + predictions = fc.linear_model( + features, [wire_cast_weights], sparse_combiner='sum') + bias = get_linear_model_bias() + wire_cast_var = get_linear_model_column_var(wire_cast) + with _initialized_session() as sess: + sess.run(wire_cast_var.assign([[10.], [100.], [1000.], [10000.]])) + sess.run(bias.assign([5.])) + self.assertAllClose([[1005.], [-9985.]], predictions.eval()) + def test_dense_multi_dimension_multi_output(self): price = fc.numeric_column('price', shape=2) with ops.Graph().as_default(): @@ -1955,7 +1975,7 @@ class _LinearModelTest(test.TestCase): with ops.Graph().as_default(): features = {'price': [[1.], [5.]]} predictions = get_keras_linear_model_predictions(features, [price]) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() price_var = get_linear_model_column_var(price) with _initialized_session() as sess: self.assertAllClose([0.], bias.eval()) @@ -1972,7 +1992,7 @@ class _LinearModelTest(test.TestCase): dense_shape=[2, 2]) features = {'wire_cast': wire_tensor} predictions = get_keras_linear_model_predictions(features, [wire_cast]) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() wire_cast_var = get_linear_model_column_var(wire_cast) with _initialized_session() as sess: self.assertAllClose([0.], bias.eval()) @@ -1992,7 +2012,7 @@ class _LinearModelTest(test.TestCase): features = {'wire_cast': wire_tensor, 'price': [[1.], [5.]]} predictions = get_keras_linear_model_predictions(features, [wire_cast, price]) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() wire_cast_var = get_linear_model_column_var(wire_cast) price_var = get_linear_model_column_var(price) with _initialized_session() as sess: @@ -2050,7 +2070,7 @@ class _LinearModelTest(test.TestCase): features = {dense_and_sparse_column.name: sp_tensor} predictions = get_keras_linear_model_predictions( features, [dense_and_sparse_column]) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() dense_and_sparse_column_var = get_linear_model_column_var( dense_and_sparse_column) with _initialized_session() as sess: @@ -2066,7 +2086,7 @@ class _LinearModelTest(test.TestCase): features = {'price': [[1.], [5.]]} predictions = get_keras_linear_model_predictions( features, [price], units=3) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() price_var = get_linear_model_column_var(price) with _initialized_session() as sess: self.assertAllClose(np.zeros((3,)), bias.eval()) @@ -2086,7 +2106,7 @@ class _LinearModelTest(test.TestCase): features = {'wire_cast': wire_tensor} predictions = get_keras_linear_model_predictions( features, [wire_cast], units=3) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() wire_cast_var = get_linear_model_column_var(wire_cast) with _initialized_session() as sess: self.assertAllClose(np.zeros((3,)), bias.eval()) @@ -2141,7 +2161,7 @@ class _LinearModelTest(test.TestCase): features = {'wire_cast': wire_tensor} predictions = get_keras_linear_model_predictions( features, [wire_cast], sparse_combiner='mean') - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() wire_cast_var = get_linear_model_column_var(wire_cast) with _initialized_session() as sess: sess.run(wire_cast_var.assign([[10.], [100.], [1000.], [10000.]])) @@ -2154,7 +2174,7 @@ class _LinearModelTest(test.TestCase): features = {'price': [[1., 2.], [5., 6.]]} predictions = get_keras_linear_model_predictions( features, [price], units=3) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() price_var = get_linear_model_column_var(price) with _initialized_session() as sess: self.assertAllClose(np.zeros((3,)), bias.eval()) @@ -2184,7 +2204,7 @@ class _LinearModelTest(test.TestCase): with ops.Graph().as_default(): features = {'price': [[[1., 2.]], [[5., 6.]]]} predictions = get_keras_linear_model_predictions(features, [price]) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() price_var = get_linear_model_column_var(price) with _initialized_session() as sess: self.assertAllClose([0.], bias.eval()) @@ -2200,7 +2220,7 @@ class _LinearModelTest(test.TestCase): features = {'price1': [[1., 2.], [5., 6.]], 'price2': [[3.], [4.]]} predictions = get_keras_linear_model_predictions(features, [price1, price2]) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() price1_var = get_linear_model_column_var(price1) price2_var = get_linear_model_column_var(price2) with _initialized_session() as sess: @@ -2213,6 +2233,45 @@ class _LinearModelTest(test.TestCase): sess.run(bias.assign([7.])) self.assertAllClose([[3217.], [4657.]], predictions.eval()) + def test_fills_cols_to_vars(self): + price1 = fc.numeric_column('price1', shape=2) + price2 = fc.numeric_column('price2') + with ops.Graph().as_default(): + features = {'price1': [[1., 2.], [5., 6.]], 'price2': [[3.], [4.]]} + cols_to_vars = {} + get_keras_linear_model_predictions( + features, [price1, price2], cols_to_vars=cols_to_vars) + bias = get_linear_model_bias() + price1_var = get_linear_model_column_var(price1) + price2_var = get_linear_model_column_var(price2) + self.assertAllEqual(cols_to_vars['bias'], [bias]) + self.assertAllEqual(cols_to_vars[price1], [price1_var]) + self.assertAllEqual(cols_to_vars[price2], [price2_var]) + + def test_fills_cols_to_vars_partitioned_variables(self): + price1 = fc.numeric_column('price1', shape=2) + price2 = fc.numeric_column('price2', shape=3) + with ops.Graph().as_default(): + features = { + 'price1': [[1., 2.], [6., 7.]], + 'price2': [[3., 4., 5.], [8., 9., 10.]] + } + cols_to_vars = {} + with variable_scope.variable_scope( + 'linear', + partitioner=partitioned_variables.fixed_size_partitioner(2, axis=0)): + get_keras_linear_model_predictions( + features, [price1, price2], cols_to_vars=cols_to_vars) + with _initialized_session(): + self.assertEqual([0.], cols_to_vars['bias'][0].eval()) + # Partitioning shards the [2, 1] price1 var into 2 [1, 1] Variables. + self.assertAllEqual([[0.]], cols_to_vars[price1][0].eval()) + self.assertAllEqual([[0.]], cols_to_vars[price1][1].eval()) + # Partitioning shards the [3, 1] price2 var into a [2, 1] Variable and + # a [1, 1] Variable. + self.assertAllEqual([[0.], [0.]], cols_to_vars[price2][0].eval()) + self.assertAllEqual([[0.]], cols_to_vars[price2][1].eval()) + def test_dense_collection(self): price = fc.numeric_column('price') with ops.Graph().as_default() as g: @@ -2220,7 +2279,7 @@ class _LinearModelTest(test.TestCase): get_keras_linear_model_predictions( features, [price], weight_collections=['my-vars']) my_vars = g.get_collection('my-vars') - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() price_var = get_linear_model_column_var(price) self.assertIn(bias, my_vars) self.assertIn(price_var, my_vars) @@ -2234,7 +2293,7 @@ class _LinearModelTest(test.TestCase): get_keras_linear_model_predictions( features, [wire_cast], weight_collections=['my-vars']) my_vars = g.get_collection('my-vars') - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() wire_cast_var = get_linear_model_column_var(wire_cast) self.assertIn(bias, my_vars) self.assertIn(wire_cast_var, my_vars) @@ -2244,7 +2303,7 @@ class _LinearModelTest(test.TestCase): with ops.Graph().as_default() as g: features = {'price': [[1.], [5.]]} get_keras_linear_model_predictions(features, [price]) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() price_var = get_linear_model_column_var(price) trainable_vars = g.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES) self.assertIn(bias, trainable_vars) @@ -2258,7 +2317,7 @@ class _LinearModelTest(test.TestCase): features = {'wire_cast': wire_tensor} get_keras_linear_model_predictions(features, [wire_cast]) trainable_vars = g.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() wire_cast_var = get_linear_model_column_var(wire_cast) self.assertIn(bias, trainable_vars) self.assertIn(wire_cast_var, trainable_vars) @@ -2405,7 +2464,7 @@ class _LinearModelTest(test.TestCase): coord = coordinator.Coordinator() threads = queue_runner_impl.start_queue_runners(sess, coord=coord) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() price_buckets_var = get_linear_model_column_var(price_buckets) body_style_var = get_linear_model_column_var(body_style) @@ -2448,7 +2507,7 @@ class _LinearModelTest(test.TestCase): net = get_keras_linear_model_predictions(features, [price_buckets, body_style]) with _initialized_session() as sess: - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() price_buckets_var = get_linear_model_column_var(price_buckets) body_style_var = get_linear_model_column_var(body_style) @@ -2487,7 +2546,7 @@ class _LinearModelTest(test.TestCase): net = get_keras_linear_model_predictions( features, [price_buckets, body_style, country]) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() price_buckets_var = get_linear_model_column_var(price_buckets) body_style_var = get_linear_model_column_var(body_style) with _initialized_session() as sess: @@ -2863,6 +2922,114 @@ class FunctionalInputLayerTest(test.TestCase): features['price2']: [[1.], [5.]], }) + def test_multiple_layers_with_same_embedding_column(self): + some_sparse_column = fc.categorical_column_with_hash_bucket( + 'sparse_feature', hash_bucket_size=5) + some_embedding_column = fc.embedding_column( + some_sparse_column, dimension=10) + + with ops.Graph().as_default(): + features = { + 'sparse_feature': [['a'], ['x']], + } + all_cols = [some_embedding_column] + fc.input_layer(features, all_cols) + fc.input_layer(features, all_cols) + # Make sure that 2 variables get created in this case. + self.assertEqual(2, len( + ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES))) + expected_var_names = [ + 'input_layer/sparse_feature_embedding/embedding_weights:0', + 'input_layer_1/sparse_feature_embedding/embedding_weights:0' + ] + self.assertItemsEqual( + expected_var_names, + [v.name for v in ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)]) + + def test_multiple_layers_with_same_shared_embedding_column(self): + categorical_column_a = fc.categorical_column_with_identity( + key='aaa', num_buckets=3) + categorical_column_b = fc.categorical_column_with_identity( + key='bbb', num_buckets=3) + embedding_dimension = 2 + embedding_column_b, embedding_column_a = fc.shared_embedding_columns( + [categorical_column_b, categorical_column_a], + dimension=embedding_dimension) + + with ops.Graph().as_default(): + features = { + 'aaa': + sparse_tensor.SparseTensor( + indices=((0, 0), (1, 0), (1, 1)), + values=(0, 1, 0), + dense_shape=(2, 2)), + 'bbb': + sparse_tensor.SparseTensor( + indices=((0, 0), (1, 0), (1, 1)), + values=(1, 2, 1), + dense_shape=(2, 2)), + } + all_cols = [embedding_column_a, embedding_column_b] + fc.input_layer(features, all_cols) + fc.input_layer(features, all_cols) + # Make sure that only 1 variable gets created in this case. + self.assertEqual(1, len( + ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES))) + self.assertItemsEqual( + ['input_layer/aaa_bbb_shared_embedding/embedding_weights:0'], + [v.name for v in ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)]) + + def test_multiple_layers_with_same_shared_embedding_column_diff_graphs(self): + categorical_column_a = fc.categorical_column_with_identity( + key='aaa', num_buckets=3) + categorical_column_b = fc.categorical_column_with_identity( + key='bbb', num_buckets=3) + embedding_dimension = 2 + embedding_column_b, embedding_column_a = fc.shared_embedding_columns( + [categorical_column_b, categorical_column_a], + dimension=embedding_dimension) + all_cols = [embedding_column_a, embedding_column_b] + + with ops.Graph().as_default(): + features = { + 'aaa': + sparse_tensor.SparseTensor( + indices=((0, 0), (1, 0), (1, 1)), + values=(0, 1, 0), + dense_shape=(2, 2)), + 'bbb': + sparse_tensor.SparseTensor( + indices=((0, 0), (1, 0), (1, 1)), + values=(1, 2, 1), + dense_shape=(2, 2)), + } + fc.input_layer(features, all_cols) + # Make sure that only 1 variable gets created in this case. + self.assertEqual(1, len( + ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES))) + + with ops.Graph().as_default(): + features1 = { + 'aaa': + sparse_tensor.SparseTensor( + indices=((0, 0), (1, 0), (1, 1)), + values=(0, 1, 0), + dense_shape=(2, 2)), + 'bbb': + sparse_tensor.SparseTensor( + indices=((0, 0), (1, 0), (1, 1)), + values=(1, 2, 1), + dense_shape=(2, 2)), + } + + fc.input_layer(features1, all_cols) + # Make sure that only 1 variable gets created in this case. + self.assertEqual(1, len( + ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES))) + self.assertItemsEqual( + ['input_layer/aaa_bbb_shared_embedding/embedding_weights:0'], + [v.name for v in ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)]) + def test_with_numpy_input_fn(self): embedding_values = ( (1., 2., 3., 4., 5.), # id 0 @@ -3558,7 +3725,7 @@ class VocabularyFileCategoricalColumnTest(test.TestCase): values=('marlo', 'skywalker', 'omar'), dense_shape=(2, 2)) }, (wire_column,)) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() wire_var = get_linear_model_column_var(wire_column) with _initialized_session(): self.assertAllClose((0.,), bias.eval()) @@ -3950,7 +4117,7 @@ class VocabularyListCategoricalColumnTest(test.TestCase): values=('marlo', 'skywalker', 'omar'), dense_shape=(2, 2)) }, (wire_column,)) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() wire_var = get_linear_model_column_var(wire_column) with _initialized_session(): self.assertAllClose((0.,), bias.eval()) @@ -4196,7 +4363,7 @@ class IdentityCategoricalColumnTest(test.TestCase): values=(0, 2, 1), dense_shape=(2, 2)) }, (column,)) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() weight_var = get_linear_model_column_var(column) with _initialized_session(): self.assertAllClose((0.,), bias.eval()) @@ -4482,7 +4649,6 @@ class EmbeddingColumnTest(test.TestCase): self.assertIs(categorical_column, embedding_column.categorical_column) self.assertEqual(embedding_dimension, embedding_column.dimension) self.assertEqual('mean', embedding_column.combiner) - self.assertIsNotNone(embedding_column.initializer) self.assertIsNone(embedding_column.ckpt_to_load_from) self.assertIsNone(embedding_column.tensor_name_in_ckpt) self.assertIsNone(embedding_column.max_norm) @@ -4507,7 +4673,6 @@ class EmbeddingColumnTest(test.TestCase): self.assertIs(categorical_column, embedding_column.categorical_column) self.assertEqual(embedding_dimension, embedding_column.dimension) self.assertEqual('my_combiner', embedding_column.combiner) - self.assertEqual('my_initializer', embedding_column.initializer()) self.assertEqual('my_ckpt', embedding_column.ckpt_to_load_from) self.assertEqual('my_ckpt_tensor', embedding_column.tensor_name_in_ckpt) self.assertEqual(42., embedding_column.max_norm) @@ -4538,7 +4703,6 @@ class EmbeddingColumnTest(test.TestCase): self.assertEqual(embedding_dimension, embedding_column.dimension) self.assertEqual('my_combiner', embedding_column.combiner) - self.assertEqual('my_initializer', embedding_column.initializer()) self.assertEqual('my_ckpt', embedding_column.ckpt_to_load_from) self.assertEqual('my_ckpt_tensor', embedding_column.tensor_name_in_ckpt) self.assertEqual(42., embedding_column.max_norm) @@ -4653,72 +4817,6 @@ class EmbeddingColumnTest(test.TestCase): self.assertAllEqual(embedding_values, global_vars[0].eval()) self.assertAllEqual(expected_lookups, embedding_lookup.eval()) - def test_get_dense_tensor_with_state(self): - # Inputs. - 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, 4), (3, 0)), - values=(2, 0, 1, 1), - dense_shape=(4, 5)) - - # Embedding variable. - 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 lookup result, using combiner='mean'. - expected_lookups = ( - # example 0, ids [2], embedding = [7, 11] - (7., 11.), - # example 1, ids [0, 1], embedding = mean([1, 2] + [3, 5]) = [2, 3.5] - (2., 3.5), - # example 2, ids [], embedding = [0, 0] - (0., 0.), - # example 3, ids [1], embedding = [3, 5] - (3., 5.), - ) - - # Build columns. - categorical_column = fc.categorical_column_with_identity( - key='aaa', num_buckets=vocabulary_size) - embedding_column = fc.embedding_column( - categorical_column, - dimension=embedding_dimension, - initializer=_initializer) - - # Create embedding_weights variable. - weight_collections = [ - ops.GraphKeys.GLOBAL_VARIABLES, ops.GraphKeys.MODEL_VARIABLES - ] - state = embedding_column._create_state(weight_collections) - - # Provide sparse input and get dense result. - embedding_lookup = embedding_column._get_dense_tensor( - _LazyBuilder({ - 'aaa': sparse_input - }), state=state) - - # Assert expected embedding variable and lookups. - global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) - self.assertItemsEqual( - ('embedding_weights:0',), tuple([v.name for v in global_vars])) - with _initialized_session(): - self.assertAllEqual(embedding_values, global_vars[0].eval()) - self.assertAllEqual(expected_lookups, embedding_lookup.eval()) - def test_get_dense_tensor_3d(self): # Inputs. vocabulary_size = 4 @@ -4773,8 +4871,8 @@ class EmbeddingColumnTest(test.TestCase): # Assert expected embedding variable and lookups. global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) - self.assertItemsEqual( - ('embedding_weights:0',), tuple([v.name for v in global_vars])) + self.assertItemsEqual(('embedding_weights:0',), + tuple([v.name for v in global_vars])) with _initialized_session(): self.assertAllEqual(embedding_values, global_vars[0].eval()) self.assertAllEqual(expected_lookups, embedding_lookup.eval()) @@ -4801,8 +4899,9 @@ class EmbeddingColumnTest(test.TestCase): }), weight_collections=('my_vars',)) # Assert expected embedding variable and lookups. - self.assertItemsEqual( - [], ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)) + global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) + self.assertItemsEqual(('embedding_weights:0',), + tuple([v.name for v in global_vars])) my_vars = ops.get_collection('my_vars') self.assertItemsEqual( ('embedding_weights:0',), tuple([v.name for v in my_vars])) @@ -5046,7 +5145,7 @@ class EmbeddingColumnTest(test.TestCase): categorical_column.name: sparse_input }, (embedding_column,)) expected_var_names = ( - 'linear_model/bias_layer/bias_weights:0', + 'linear_model/bias_weights:0', 'linear_model/aaa_embedding/weights:0', 'linear_model/aaa_embedding/embedding_weights:0', ) @@ -5058,7 +5157,7 @@ class EmbeddingColumnTest(test.TestCase): for v in ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES) } self.assertItemsEqual(expected_var_names, trainable_vars.keys()) - bias = trainable_vars['linear_model/bias_layer/bias_weights:0'] + bias = trainable_vars['linear_model/bias_weights:0'] embedding_weights = trainable_vars[ 'linear_model/aaa_embedding/embedding_weights:0'] linear_weights = trainable_vars['linear_model/aaa_embedding/weights:0'] @@ -5221,14 +5320,12 @@ class SharedEmbeddingColumnTest(test.TestCase): self.assertEqual(embedding_dimension, embedding_column_b.dimension) self.assertEqual('mean', embedding_column_a.combiner) self.assertEqual('mean', embedding_column_b.combiner) - self.assertIsNotNone(embedding_column_a.initializer) - self.assertIsNotNone(embedding_column_b.initializer) self.assertIsNone(embedding_column_a.ckpt_to_load_from) self.assertIsNone(embedding_column_b.ckpt_to_load_from) self.assertEqual('aaa_bbb_shared_embedding', - embedding_column_a.shared_embedding_collection_name) + embedding_column_a.var_scope_name) self.assertEqual('aaa_bbb_shared_embedding', - embedding_column_b.shared_embedding_collection_name) + embedding_column_b.var_scope_name) self.assertIsNone(embedding_column_a.tensor_name_in_ckpt) self.assertIsNone(embedding_column_b.tensor_name_in_ckpt) self.assertIsNone(embedding_column_a.max_norm) @@ -5274,12 +5371,10 @@ class SharedEmbeddingColumnTest(test.TestCase): self.assertEqual(embedding_dimension, embedding_column_b.dimension) self.assertEqual('my_combiner', embedding_column_a.combiner) self.assertEqual('my_combiner', embedding_column_b.combiner) - self.assertEqual('my_initializer', embedding_column_a.initializer()) - self.assertEqual('my_initializer', embedding_column_b.initializer()) self.assertEqual('shared_embedding_collection_name', - embedding_column_a.shared_embedding_collection_name) + embedding_column_a.var_scope_name) self.assertEqual('shared_embedding_collection_name', - embedding_column_b.shared_embedding_collection_name) + embedding_column_b.var_scope_name) self.assertEqual('my_ckpt', embedding_column_a.ckpt_to_load_from) self.assertEqual('my_ckpt', embedding_column_b.ckpt_to_load_from) self.assertEqual('my_ckpt_tensor', embedding_column_a.tensor_name_in_ckpt) @@ -5329,9 +5424,8 @@ class SharedEmbeddingColumnTest(test.TestCase): self.assertEqual(embedding_dimension, embedding_column_a.dimension) self.assertEqual('my_combiner', embedding_column_a.combiner) - self.assertEqual('my_initializer', embedding_column_a.initializer()) self.assertEqual('shared_embedding_collection_name', - embedding_column_a.shared_embedding_collection_name) + embedding_column_a.var_scope_name) self.assertEqual('my_ckpt', embedding_column_a.ckpt_to_load_from) self.assertEqual('my_ckpt_tensor', embedding_column_a.tensor_name_in_ckpt) self.assertEqual(42., embedding_column_a.max_norm) @@ -5515,80 +5609,6 @@ class SharedEmbeddingColumnTest(test.TestCase): self.assertAllEqual(expected_lookups_a, embedding_lookup_a.eval()) self.assertAllEqual(expected_lookups_b, embedding_lookup_b.eval()) - def test_get_dense_tensor_with_state(self): - # Inputs. - vocabulary_size = 3 - # -1 values are ignored. - input_a = np.array([ - [2, -1, -1], # example 0, ids [2] - [0, 1, -1] - ]) # example 1, ids [0, 1] - input_b = np.array([ - [0, -1, -1], # example 0, ids [0] - [-1, -1, -1] - ]) # example 1, ids [] - input_features = {'aaa': input_a, 'bbb': input_b} - - # Embedding variable. - 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 lookup result, using combiner='mean'. - expected_lookups_a = ( - # example 0: - (7., 11.), # ids [2], embedding = [7, 11] - # example 1: - (2., 3.5), # ids [0, 1], embedding = mean([1, 2] + [3, 5]) = [2, 3.5] - ) - expected_lookups_b = ( - # example 0: - (1., 2.), # ids [0], embedding = [1, 2] - # example 1: - (0., 0.), # ids [], embedding = [0, 0] - ) - - # Build columns. - categorical_column_a = fc.categorical_column_with_identity( - key='aaa', num_buckets=vocabulary_size) - categorical_column_b = fc.categorical_column_with_identity( - key='bbb', num_buckets=vocabulary_size) - embedding_column_a, embedding_column_b = fc.shared_embedding_columns( - [categorical_column_a, categorical_column_b], - dimension=embedding_dimension, - initializer=_initializer) - - # Create state. - weight_collections = [ - ops.GraphKeys.GLOBAL_VARIABLES, ops.GraphKeys.MODEL_VARIABLES - ] - state = embedding_column_a._create_state(weight_collections) - - # Provide sparse input and get dense result. - embedding_lookup_a = embedding_column_a._get_dense_tensor( - _LazyBuilder(input_features), state=state) - embedding_lookup_b = embedding_column_b._get_dense_tensor( - _LazyBuilder(input_features), state=state) - - # Assert expected embedding variable and lookups. - global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) - self.assertItemsEqual( - ('embedding_weights:0',), tuple([v.name for v in global_vars])) - embedding_var = global_vars[0] - with _initialized_session(): - self.assertAllEqual(embedding_values, embedding_var.eval()) - self.assertAllEqual(expected_lookups_a, embedding_lookup_a.eval()) - self.assertAllEqual(expected_lookups_b, embedding_lookup_b.eval()) - def test_get_dense_tensor_placeholder_inputs(self): # Inputs. vocabulary_size = 3 @@ -5774,7 +5794,7 @@ class SharedEmbeddingColumnTest(test.TestCase): # Linear weights do not follow the column name. But this is a rare use # case, and fixing it would add too much complexity to the code. expected_var_names = ( - 'linear_model/bias_layer/bias_weights:0', + 'linear_model/bias_weights:0', 'linear_model/aaa_bbb_shared_embedding/weights:0', 'linear_model/aaa_bbb_shared_embedding/embedding_weights:0', 'linear_model/aaa_bbb_shared_embedding_1/weights:0', @@ -5787,7 +5807,7 @@ class SharedEmbeddingColumnTest(test.TestCase): for v in ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES) } self.assertItemsEqual(expected_var_names, trainable_vars.keys()) - bias = trainable_vars['linear_model/bias_layer/bias_weights:0'] + bias = trainable_vars['linear_model/bias_weights:0'] embedding_weights = trainable_vars[ 'linear_model/aaa_bbb_shared_embedding/embedding_weights:0'] linear_weights_a = trainable_vars[ @@ -5890,10 +5910,7 @@ class SharedEmbeddingColumnTest(test.TestCase): tuple([v.name for v in trainable_vars])) else: self.assertItemsEqual([], tuple([v.name for v in trainable_vars])) - shared_embedding_vars = ops.get_collection('aaa_bbb_shared_embedding') - self.assertItemsEqual( - ['input_layer/aaa_bbb_shared_embedding/embedding_weights:0'], - tuple([v.name for v in shared_embedding_vars])) + shared_embedding_vars = global_vars with _initialized_session(): self.assertAllEqual(embedding_values, shared_embedding_vars[0].eval()) self.assertAllEqual(expected_lookups, input_layer.eval()) @@ -6125,7 +6142,7 @@ class WeightedCategoricalColumnTest(test.TestCase): values=(.5, 1., .1), dense_shape=(2, 2)) }, (column,)) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() weight_var = get_linear_model_column_var(column) with _initialized_session(): self.assertAllClose((0.,), bias.eval()) @@ -6164,14 +6181,16 @@ class WeightedCategoricalColumnTest(test.TestCase): key='ids', num_buckets=3), weight_feature_key='values') with ops.Graph().as_default(): - predictions = get_keras_linear_model_predictions({ - 'ids': - sparse_tensor.SparseTensorValue( - indices=((0, 0), (1, 0), (1, 1)), - values=(0, 2, 1), - dense_shape=(2, 2)), - 'values': ((.5,), (1.,)) - }, (column,)) + predictions = get_keras_linear_model_predictions( + { + 'ids': + sparse_tensor.SparseTensorValue( + indices=((0, 0), (1, 0), (1, 1)), + values=(0, 2, 1), + dense_shape=(2, 2)), + 'values': ((.5,), (1.,)) + }, (column,), + sparse_combiner='mean') with _initialized_session(): with self.assertRaisesRegexp(errors.OpError, 'Incompatible shapes'): predictions.eval() @@ -6190,7 +6209,7 @@ class WeightedCategoricalColumnTest(test.TestCase): dense_shape=(2, 2)), 'values': ((.5,), (1.,), (.1,)) }, (column,)) - bias = get_keras_linear_model_bias() + bias = get_linear_model_bias() weight_var = get_linear_model_column_var(column) with _initialized_session(): self.assertAllClose((0.,), bias.eval()) @@ -6255,13 +6274,16 @@ class WeightedCategoricalColumnTest(test.TestCase): key='ids', num_buckets=3), weight_feature_key='values') with ops.Graph().as_default(): - predictions = fc.linear_model({ - 'ids': sparse_tensor.SparseTensorValue( - indices=((0, 0), (1, 0), (1, 1)), - values=(0, 2, 1), - dense_shape=(2, 2)), - 'values': ((.5,), (1.,)) - }, (column,)) + predictions = fc.linear_model( + { + 'ids': + sparse_tensor.SparseTensorValue( + indices=((0, 0), (1, 0), (1, 1)), + values=(0, 2, 1), + dense_shape=(2, 2)), + 'values': ((.5,), (1.,)) + }, (column,), + sparse_combiner='mean') with _initialized_session(): with self.assertRaisesRegexp(errors.OpError, 'Incompatible shapes'): predictions.eval() diff --git a/tensorflow/python/framework/dtypes.py b/tensorflow/python/framework/dtypes.py index a31c424263677a880696adbae7e0b5edac4807d5..7f9ef53457ae060600067b946e686487f55adda1 100644 --- a/tensorflow/python/framework/dtypes.py +++ b/tensorflow/python/framework/dtypes.py @@ -297,6 +297,9 @@ class DType(object): def __hash__(self): return self._type_enum + def __reduce__(self): + return as_dtype, (self.name,) + @property def size(self): if (self._type_enum == types_pb2.DT_VARIANT or @@ -648,6 +651,11 @@ QUANTIZED_DTYPES = frozenset([ ]) tf_export("QUANTIZED_DTYPES").export_constant(__name__, "QUANTIZED_DTYPES") +_PYTHON_TO_TF = { + float: float32, + bool: bool, +} + @tf_export("as_dtype") def as_dtype(type_value): @@ -679,6 +687,11 @@ def as_dtype(type_value): except KeyError: pass + try: + return _PYTHON_TO_TF[type_value] + except KeyError: + pass + if isinstance(type_value, np.dtype): # The numpy dtype for strings is variable length. We can not compare # dtype with a single constant (np.string does not exist) to decide @@ -687,11 +700,13 @@ def as_dtype(type_value): if type_value.type == np.string_ or type_value.type == np.unicode_: return string - for key, val in _NP_TO_TF: - try: - if key == type_value: - return val - except TypeError as e: - raise TypeError("Cannot convert {} to a dtype. {}".format(type_value, e)) + if isinstance(type_value, (type, np.dtype)): + for key, val in _NP_TO_TF: + try: + if key == type_value: + return val + except TypeError as e: + raise TypeError("Cannot convert {} to a dtype. {}".format( + type_value, e)) raise TypeError("Cannot convert value %r to a TensorFlow DType." % type_value) diff --git a/tensorflow/python/framework/dtypes_test.py b/tensorflow/python/framework/dtypes_test.py index e49e2fda5d84da4f8f87fae73874351afe0a20f2..a873670e0461884d06cde1db4db2cf2db98fde3c 100644 --- a/tensorflow/python/framework/dtypes_test.py +++ b/tensorflow/python/framework/dtypes_test.py @@ -295,6 +295,20 @@ class TypesTest(test_util.TensorFlowTestCase): self.assertNotEqual(dtypes.int32, int) self.assertNotEqual(dtypes.float64, 2.1) + def testPythonTypesConversion(self): + self.assertIs(dtypes.float32, dtypes.as_dtype(float)) + self.assertIs(dtypes.bool, dtypes.as_dtype(bool)) + + def testReduce(self): + for enum in dtypes._TYPE_TO_STRING: + dtype = dtypes.DType(enum) + ctor, args = dtype.__reduce__() + self.assertEquals(ctor, dtypes.as_dtype) + self.assertEquals(args, (dtype.name,)) + reconstructed = ctor(*args) + self.assertEquals(reconstructed, dtype) + if __name__ == "__main__": googletest.main() + diff --git a/tensorflow/python/framework/graph_util_impl.py b/tensorflow/python/framework/graph_util_impl.py index 910364364c8be84b1a629dbdaae5e69443d07e75..394fac6c856197030f85aab5b11fa881eddf670d 100644 --- a/tensorflow/python/framework/graph_util_impl.py +++ b/tensorflow/python/framework/graph_util_impl.py @@ -285,7 +285,7 @@ def convert_variables_to_constants(sess, output_graph_def.node.extend([output_node]) output_graph_def.library.CopyFrom(inference_graph.library) - print("Converted %d variables to const ops." % how_many_converted) + logging.info("Converted %d variables to const ops.", how_many_converted) return output_graph_def diff --git a/tensorflow/python/framework/graph_util_test.py b/tensorflow/python/framework/graph_util_test.py index b618152b0256fd043dc7259960d867278ba55b0a..2dafb94ba7e4e779c8703096795f4e49139100ab 100644 --- a/tensorflow/python/framework/graph_util_test.py +++ b/tensorflow/python/framework/graph_util_test.py @@ -209,7 +209,7 @@ class DeviceFunctionsTest(test.TestCase): defun_node, 2.0, name="output_node") with session.Session() as sess: - init = variables.initialize_variables([variable_node]) + init = variables.variables_initializer([variable_node]) sess.run(init) output = sess.run(output_node) self.assertNear(4.0, output, 0.00001) diff --git a/tensorflow/python/framework/load_library.py b/tensorflow/python/framework/load_library.py index 535c6017f5fd0f8adf9ed091bd4477762e52b0e3..9a8477debb05fdf11d7d8c39adf9e5e55cc2caeb 100644 --- a/tensorflow/python/framework/load_library.py +++ b/tensorflow/python/framework/load_library.py @@ -58,7 +58,7 @@ def load_op_library(library_filename): op_list_str = py_tf.TF_GetOpList(lib_handle) op_list = op_def_pb2.OpList() op_list.ParseFromString(compat.as_bytes(op_list_str)) - wrappers = py_tf.GetPythonWrappers(op_list_str) + wrappers = py_tf.GetEagerPythonWrappers(op_list_str) # Delete the library handle to release any memory held in C # that are no longer needed. diff --git a/tensorflow/python/framework/meta_graph_test.py b/tensorflow/python/framework/meta_graph_test.py index 5d5fb037fc217849ea32102bf60796c47d565f3b..e5b157648e0830d400b11d5fe310f0fc1422e004 100644 --- a/tensorflow/python/framework/meta_graph_test.py +++ b/tensorflow/python/framework/meta_graph_test.py @@ -522,6 +522,31 @@ class ScopedMetaGraphTest(test.TestCase): actual_grad_value = sess.run(grad) self.assertEqual(expected_grad_value, actual_grad_value) + def testImportWhileLoopInWhileLoop(self): + # Create a simple while loop. + with ops.Graph().as_default(): + var = variables.Variable(0.0) + _, output = control_flow_ops.while_loop(lambda i, x: i < 5, + lambda i, x: (i + 1, x * 2.0), + [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() + + # Restore the MetaGraphDef in a while loop in a new graph. + with ops.Graph().as_default(): + + def body(i, _): + meta_graph.import_scoped_meta_graph(meta_graph_def) + return i + 1, ops.get_default_graph().get_tensor_by_name(output_name) + + _, x = control_flow_ops.while_loop(lambda i, x: i < 2, body, [0, 0.0], + name="") + with session.Session() as sess: + sess.run(variables.global_variables_initializer()) + sess.run(x) + def testScopedImportUnderNameScope(self): graph = ops.Graph() with graph.as_default(): diff --git a/tensorflow/python/framework/ops.py b/tensorflow/python/framework/ops.py index 662cda2a7d4822d92a7d10ac42012bc2675c5eac..8cd6820f6a5adce7f82ed2eb2e3eba2bca1e3f80 100644 --- a/tensorflow/python/framework/ops.py +++ b/tensorflow/python/framework/ops.py @@ -1385,6 +1385,22 @@ def register_tensor_conversion_function(base_type, if not callable(conversion_func): raise TypeError("conversion_func must be callable.") + # context._context is checked so that we don't inadvertently create it. + # This is because enable_eager_execution will fail when called from the main + # function if the context._context is already created, and the + # register_tensor_conversion_function calls happen when the module is + # imported. + if context._context is not None and context.executing_eagerly( + ) and isinstance(base_type, six.integer_types + ( + float, + np.ndarray, + )): + # TODO(nareshmodi): consider setting a context variable which disables the + # fastpath instead. + raise TypeError( + "Cannot register conversions for numpy arrays, python number types " + "when executing eagerly.") + try: funcs_at_priority = _tensor_conversion_func_registry[priority] except KeyError: diff --git a/tensorflow/python/framework/python_op_gen.cc b/tensorflow/python/framework/python_op_gen.cc index e5e3b821998718e7b87a95439a442df98ad7c997..ad6c36b4b1773ecbbb5be56f7eaf2d78af3d010c 100644 --- a/tensorflow/python/framework/python_op_gen.cc +++ b/tensorflow/python/framework/python_op_gen.cc @@ -98,7 +98,7 @@ bool IsOpWithUnderscorePrefix(const string& s) { // TODO(annarev): reduce usage of '*' imports and remove these from the // list. "fused_batch_norm", "histogram_fixed_width", "stack", - "batch_norm_with_global_normalization"}); + "batch_norm_with_global_normalization", "clip_by_value"}); return kUnderscoreOps->count(s) > 0; } diff --git a/tensorflow/python/framework/python_op_gen.i b/tensorflow/python/framework/python_op_gen.i index 26ec4e8e66b5d4e3be433c9e59f9b6034109d153..efcce2f20941790571199d358dfcfcc3c0283d08 100644 --- a/tensorflow/python/framework/python_op_gen.i +++ b/tensorflow/python/framework/python_op_gen.i @@ -16,10 +16,10 @@ limitations under the License. %include "tensorflow/python/platform/base.i" %{ -#include "tensorflow/python/framework/python_op_gen.h" +#include "tensorflow/python/eager/python_eager_op_gen.h" %} -// Input typemap for GetPythonWrappers. +// Input typemap for GetEagerPythonWrappers. // Accepts a python object of 'bytes' type, and converts it to // a const char* pointer and size_t length. The default typemap // going from python bytes to const char* tries to decode the @@ -37,5 +37,5 @@ limitations under the License. %ignoreall; -%unignore tensorflow::GetPythonWrappers; -%include "tensorflow/python/framework/python_op_gen.h" +%unignore tensorflow::GetEagerPythonWrappers; +%include "tensorflow/python/eager/python_eager_op_gen.h" diff --git a/tensorflow/python/framework/tensor_shape.py b/tensorflow/python/framework/tensor_shape.py index 26069d9d90e4a75cfe3988628f1407d6f327385b..0dd29460ed93aadf61ef1f1b2dbf1d7802ca4877 100644 --- a/tensorflow/python/framework/tensor_shape.py +++ b/tensorflow/python/framework/tensor_shape.py @@ -459,6 +459,9 @@ class Dimension(object): else: return self._value >= other.value + def __reduce__(self): + return Dimension, (self._value,) + def as_dimension(value): """Converts the given value to a Dimension. @@ -931,6 +934,9 @@ class TensorShape(object): return True return self._dims != other.dims + def __reduce__(self): + return TensorShape, (self._dims,) + def as_shape(shape): """Converts the given object to a TensorShape.""" diff --git a/tensorflow/python/framework/tensor_shape_test.py b/tensorflow/python/framework/tensor_shape_test.py index 4f239228332946d9a863be408f5967c282019852..9232d99a1f932b9e48cd7ddc125e353390064873 100644 --- a/tensorflow/python/framework/tensor_shape_test.py +++ b/tensorflow/python/framework/tensor_shape_test.py @@ -188,7 +188,7 @@ class DimensionTest(test_util.TensorFlowTestCase): def testUnsupportedType(self): with self.assertRaises(TypeError): tensor_shape.Dimension(dtypes.string) - + def testMod(self): four = tensor_shape.Dimension(4) nine = tensor_shape.Dimension(9) @@ -197,6 +197,14 @@ class DimensionTest(test_util.TensorFlowTestCase): self.assertEqual(nine % 4, 1) self.assertEqual(4 % nine, 4) + def testReduce(self): + dim = tensor_shape.Dimension(5) + ctor, args = dim.__reduce__() + self.assertEquals(ctor, tensor_shape.Dimension) + self.assertEquals(args, (5,)) + reconstructed = ctor(*args) + self.assertEquals(reconstructed, dim) + class ShapeTest(test_util.TensorFlowTestCase): @@ -422,5 +430,15 @@ class ShapeTest(test_util.TensorFlowTestCase): self.assertAllEqual([2, None, 4], tensor_shape.TensorShape( (2, None, 4)).as_list()) + def testReduce(self): + shape = tensor_shape.TensorShape([2, 3]) + ctor, args = shape.__reduce__() + self.assertEquals(ctor, tensor_shape.TensorShape) + self.assertEquals(args, ([tensor_shape.Dimension(2), + tensor_shape.Dimension(3)],)) + reconstructed = ctor(*args) + self.assertEquals(reconstructed, shape) + + if __name__ == "__main__": googletest.main() diff --git a/tensorflow/python/framework/test_util.py b/tensorflow/python/framework/test_util.py index bf00fa6439b82234e951598131d2d7ab579fb6c4..f954b9d6c73962712a6da818480981294a158b84 100644 --- a/tensorflow/python/framework/test_util.py +++ b/tensorflow/python/framework/test_util.py @@ -464,6 +464,30 @@ def with_c_api(cls): return cls +def with_c_shapes(cls): + """Adds methods that call original methods but with C API shapes enabled. + + Note this enables C shapes in new methods after running the test class's + setup method. + + Args: + cls: class to decorate + + Returns: + cls with new test methods added + """ + # If C shapes are 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 shapes by default + # without breaking these tests. + if ops._USE_C_SHAPES: + return cls + + for name, value in cls.__dict__.copy().items(): + if callable(value) and name.startswith("test"): + setattr(cls, name + "WithCShapes", enable_c_shapes(value)) + return cls + + def assert_no_new_pyobjects_executing_eagerly(f): """Decorator for asserting that no new Python objects persist after a test. @@ -615,45 +639,68 @@ def assert_no_garbage_created(f): def run_in_graph_and_eager_modes(__unused__=None, - graph=None, config=None, - use_gpu=False, - force_gpu=False, + use_gpu=True, reset_test=True, assert_no_eager_garbage=False): - """Runs the test in both graph and eager modes. + """Execute the decorated test with and without enabling eager execution. + + This function returns a decorator intended to be applied to test methods in + a @{tf.test.TestCase} class. Doing so will cause the contents of the test + method to be executed twice - once normally, and once with eager execution + enabled. This allows unittests to confirm the equivalence between eager + and graph execution (see @{tf.enable_eager_execution}). + + For example, consider the following unittest: + + ```python + class MyTests(tf.test.TestCase): + + @run_in_graph_and_eager_modes() + def test_foo(self): + x = tf.constant([1, 2]) + y = tf.constant([3, 4]) + z = tf.add(x, y) + self.assertAllEqual([4, 6], self.evaluate(z)) + + if __name__ == "__main__": + tf.test.main() + ``` + + This test validates that `tf.add()` has the same behavior when computed with + eager execution enabled as it does when constructing a TensorFlow graph and + executing the `z` tensor in a session. + Args: __unused__: Prevents sliently skipping tests. - graph: Optional graph to use during the returned session. config: An optional config_pb2.ConfigProto to use to configure the - session. - use_gpu: If True, attempt to run as many ops as possible on GPU. - force_gpu: If True, pin all ops to `/device:GPU:0`. - reset_test: If True, tearDown and SetUp the test case again. + session when executing graphs. + use_gpu: If True, attempt to run as many operations as possible on GPU. + reset_test: If True, tearDown and SetUp the test case between the two + executions of the test (once with and once without eager execution). assert_no_eager_garbage: If True, sets DEBUG_SAVEALL on the garbage collector and asserts that no extra garbage has been created when running - the test in eager mode. This will fail if there are reference cycles - (e.g. a = []; a.append(a)). Off by default because some tests may create - garbage for legitimate reasons (e.g. they define a class which inherits - from `object`), and because DEBUG_SAVEALL is sticky in some Python - interpreters (meaning that tests which rely on objects being collected - elsewhere in the unit test file will not work). Additionally, checks that - nothing still has a reference to Tensors that the test allocated. + the test with eager execution enabled. This will fail if there are + reference cycles (e.g. a = []; a.append(a)). Off by default because some + tests may create garbage for legitimate reasons (e.g. they define a class + which inherits from `object`), and because DEBUG_SAVEALL is sticky in some + Python interpreters (meaning that tests which rely on objects being + collected elsewhere in the unit test file will not work). Additionally, + checks that nothing still has a reference to Tensors that the test + allocated. Returns: - Returns a decorator that will run the decorated test function - using both a graph and using eager execution. + Returns a decorator that will run the decorated test method twice: + once by constructing and executing a graph in a session and once with + eager execution enabled. """ assert not __unused__, "Add () after run_in_graph_and_eager_modes." def decorator(f): - """Test method decorator.""" - def decorated(self, **kwargs): - """Decorated the test method.""" with context.graph_mode(): - with self.test_session(graph, config, use_gpu, force_gpu): + with self.test_session(use_gpu=use_gpu): f(self, **kwargs) if reset_test: @@ -663,27 +710,20 @@ def run_in_graph_and_eager_modes(__unused__=None, self._tempdir = None self.setUp() - def run_eager_mode(self, **kwargs): - if force_gpu: - gpu_name = gpu_device_name() - if not gpu_name: - gpu_name = "/device:GPU:0" - with context.device(gpu_name): - f(self) - elif use_gpu: - # TODO(xpan): Support softplacement and gpu by default when available. - f(self, **kwargs) - else: - with context.device("/device:CPU:0"): + def run_eagerly(self, **kwargs): + if not use_gpu: + with ops.device("/cpu:0"): f(self, **kwargs) + else: + f(self, **kwargs) if assert_no_eager_garbage: - run_eager_mode = assert_no_new_tensors( - assert_no_garbage_created(run_eager_mode)) + run_eagerly = assert_no_new_tensors( + assert_no_garbage_created(run_eagerly)) with context.eager_mode(): with ops.Graph().as_default(): - run_eager_mode(self, **kwargs) + run_eagerly(self, **kwargs) return decorated @@ -1364,7 +1404,9 @@ class TensorFlowTestCase(googletest.TestCase): " %s" % (a.shape, b.shape, msg)) same = (a == b) - if a.dtype == np.float32 or a.dtype == np.float64: + if (a.dtype in [ + np.float16, np.float32, np.float64, dtypes.bfloat16.as_numpy_dtype + ]): same = np.logical_or(same, np.logical_and(np.isnan(a), np.isnan(b))) if not np.all(same): # Prints more details than np.testing.assert_array_equal. diff --git a/tensorflow/python/grappler/layout_optimizer_test.py b/tensorflow/python/grappler/layout_optimizer_test.py index 5a84b16a23f567fba6d08aaefd3b816a76907735..e3dd4b0bdfbb28480c78b0add5947c79852fb44d 100644 --- a/tensorflow/python/grappler/layout_optimizer_test.py +++ b/tensorflow/python/grappler/layout_optimizer_test.py @@ -476,7 +476,7 @@ class LayoutOptimizerTest(test.TestCase): 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) + reduce_sum = math_ops.reduce_sum(conv, axis=[1, 2], keepdims=True) squeeze = array_ops.squeeze(reduce_sum, axis=[1, 2]) output = array_ops.identity(squeeze) @@ -506,7 +506,7 @@ class LayoutOptimizerTest(test.TestCase): 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) + reduce_sum = math_ops.reduce_sum(conv, axis=[0, 1, 2], keepdims=True) squeeze = array_ops.squeeze(reduce_sum, axis=[0, 1, 2]) output = array_ops.identity(squeeze) @@ -623,7 +623,7 @@ class LayoutOptimizerTest(test.TestCase): 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) + reduce_sum = math_ops.reduce_sum(conv, axis=[3], keepdims=True) output = array_ops.identity(reduce_sum) with session.Session(config=_get_config(False)) as sess: @@ -653,7 +653,7 @@ class LayoutOptimizerTest(test.TestCase): 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) + reduce_sum = math_ops.reduce_sum(conv, axis=[2], keepdims=True) output = array_ops.identity(reduce_sum) with session.Session(config=_get_config(False)) as sess: @@ -682,7 +682,7 @@ class LayoutOptimizerTest(test.TestCase): 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) + reduce_sum = math_ops.reduce_sum(conv, axis=[2, 3], keepdims=True) output = array_ops.identity(reduce_sum) with session.Session(config=_get_config(False)) as sess: diff --git a/tensorflow/python/keras/BUILD b/tensorflow/python/keras/BUILD index f6e1d0eec3c95410c9ee54bc505335c140b99c36..1c58553156e1fa42c1420afcb87d98a92630a65c 100755 --- a/tensorflow/python/keras/BUILD +++ b/tensorflow/python/keras/BUILD @@ -20,7 +20,6 @@ py_library( srcs = [ "__init__.py", "_impl/keras/__init__.py", - "_impl/keras/activations.py", "_impl/keras/applications/__init__.py", "_impl/keras/applications/densenet.py", "_impl/keras/applications/imagenet_utils.py", @@ -32,9 +31,6 @@ py_library( "_impl/keras/applications/vgg16.py", "_impl/keras/applications/vgg19.py", "_impl/keras/applications/xception.py", - "_impl/keras/backend.py", - "_impl/keras/callbacks.py", - "_impl/keras/constraints.py", "_impl/keras/datasets/__init__.py", "_impl/keras/datasets/boston_housing.py", "_impl/keras/datasets/cifar.py", @@ -44,49 +40,13 @@ py_library( "_impl/keras/datasets/imdb.py", "_impl/keras/datasets/mnist.py", "_impl/keras/datasets/reuters.py", - "_impl/keras/engine/__init__.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", - "_impl/keras/layers/advanced_activations.py", - "_impl/keras/layers/convolutional.py", - "_impl/keras/layers/convolutional_recurrent.py", - "_impl/keras/layers/core.py", - "_impl/keras/layers/embeddings.py", - "_impl/keras/layers/local.py", - "_impl/keras/layers/merge.py", - "_impl/keras/layers/noise.py", - "_impl/keras/layers/normalization.py", - "_impl/keras/layers/pooling.py", - "_impl/keras/layers/recurrent.py", - "_impl/keras/layers/serialization.py", - "_impl/keras/layers/wrappers.py", - "_impl/keras/losses.py", - "_impl/keras/metrics.py", - "_impl/keras/models.py", - "_impl/keras/optimizers.py", "_impl/keras/preprocessing/__init__.py", "_impl/keras/preprocessing/image.py", "_impl/keras/preprocessing/sequence.py", "_impl/keras/preprocessing/text.py", - "_impl/keras/regularizers.py", "_impl/keras/testing_utils.py", "_impl/keras/utils/__init__.py", - "_impl/keras/utils/conv_utils.py", - "_impl/keras/utils/data_utils.py", - "_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/vis_utils.py", @@ -136,7 +96,21 @@ py_library( ":empty_condition": [], "//conditions:default": [], }) + [ - "@six_archive//:six", + ":backend", + ":engine", + ":layers", + "//tensorflow/python/estimator", + "//tensorflow/python/estimator:model_fn", + "//tensorflow/python/saved_model", + "//tensorflow/python:training", + ], +) + +py_library( + name = "backend", + srcs = ["_impl/keras/backend.py"], + srcs_version = "PY2AND3", + deps = [ "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", "//tensorflow/python:check_ops", @@ -152,8 +126,6 @@ py_library( "//tensorflow/python:gradients", "//tensorflow/python:image_ops", "//tensorflow/python:init_ops", - "//tensorflow/python:layers", - "//tensorflow/python:layers_base", "//tensorflow/python:logging_ops", "//tensorflow/python:math_ops", "//tensorflow/python:metrics", @@ -168,13 +140,86 @@ py_library( "//tensorflow/python:tensor_array_grad", "//tensorflow/python:tensor_array_ops", "//tensorflow/python:tensor_shape", - "//tensorflow/python:training", "//tensorflow/python:util", "//tensorflow/python:variable_scope", "//tensorflow/python:variables", - "//tensorflow/python/estimator", - "//tensorflow/python/estimator:model_fn", - "//tensorflow/python/saved_model", + ], +) + +py_library( + name = "engine", + srcs = [ + "_impl/keras/activations.py", + "_impl/keras/callbacks.py", + "_impl/keras/constraints.py", + "_impl/keras/engine/__init__.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/initializers.py", + "_impl/keras/losses.py", + "_impl/keras/metrics.py", + "_impl/keras/models.py", + "_impl/keras/optimizers.py", + "_impl/keras/regularizers.py", + "_impl/keras/utils/data_utils.py", + "_impl/keras/utils/io_utils.py", + ], + srcs_version = "PY2AND3", + deps = [ + ":backend", + "//tensorflow/python/data", + "@six_archive//:six", + ], +) + +py_library( + name = "layers", + srcs = [ + "_impl/keras/layers/__init__.py", + "_impl/keras/layers/advanced_activations.py", + "_impl/keras/layers/convolutional.py", + "_impl/keras/layers/convolutional_recurrent.py", + "_impl/keras/layers/core.py", + "_impl/keras/layers/embeddings.py", + "_impl/keras/layers/local.py", + "_impl/keras/layers/merge.py", + "_impl/keras/layers/noise.py", + "_impl/keras/layers/normalization.py", + "_impl/keras/layers/pooling.py", + "_impl/keras/layers/recurrent.py", + "_impl/keras/layers/serialization.py", + "_impl/keras/layers/wrappers.py", + "_impl/keras/utils/conv_utils.py", + "_impl/keras/utils/generic_utils.py", + "_impl/keras/utils/layer_utils.py", + "_impl/keras/utils/tf_utils.py", + ], + srcs_version = "PY2AND3", + deps = [ + ":engine", + "//tensorflow/python:array_ops", + "//tensorflow/python:distribute", + "//tensorflow/python:dtypes", + "//tensorflow/python:embedding_ops", + "//tensorflow/python:framework_ops", + "//tensorflow/python:logging_ops", + "//tensorflow/python:math_ops", + "//tensorflow/python:nn", + "//tensorflow/python:random_ops", + "//tensorflow/python:sparse_ops", + "//tensorflow/python:sparse_tensor", + "//tensorflow/python:standard_ops", + "//tensorflow/python:tensor_array_ops", + "//tensorflow/python:tensor_shape", + "//tensorflow/python:util", ], ) @@ -605,7 +650,7 @@ py_test( py_test( name = "data_utils_test", - size = "medium", + size = "large", srcs = ["_impl/keras/utils/data_utils_test.py"], srcs_version = "PY2AND3", tags = [ @@ -792,7 +837,7 @@ py_test( py_test( name = "saving_test", - size = "small", + size = "medium", srcs = ["_impl/keras/engine/saving_test.py"], srcs_version = "PY2AND3", deps = [ diff --git a/tensorflow/python/keras/_impl/keras/activations.py b/tensorflow/python/keras/_impl/keras/activations.py index b518898ad8fe2fe7b859ec80714d610242d621dc..8def7ec49375c7ce23e8f2a24a4c3615d05ca9bb 100644 --- a/tensorflow/python/keras/_impl/keras/activations.py +++ b/tensorflow/python/keras/_impl/keras/activations.py @@ -22,10 +22,8 @@ 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.layers.base import Layer from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn -from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util.tf_export import tf_export @@ -136,12 +134,6 @@ def get(identifier): identifier = str(identifier) return deserialize(identifier) elif callable(identifier): - if isinstance(identifier, Layer): - logging.warning( - 'Do not pass a layer instance (such as {identifier}) as the ' - 'activation argument of another layer. Instead, advanced ' - 'activation layers should be used just like any other ' - 'layer in a model.'.format(identifier=identifier.__class__.__name__)) return identifier else: raise ValueError('Could not interpret ' diff --git a/tensorflow/python/keras/_impl/keras/backend.py b/tensorflow/python/keras/_impl/keras/backend.py index 3aac6a9065cfa6189db1a3d3b33648dc980161b6..449410fe082421193d178b768db2ad1eda183b36 100644 --- a/tensorflow/python/keras/_impl/keras/backend.py +++ b/tensorflow/python/keras/_impl/keras/backend.py @@ -24,6 +24,7 @@ from __future__ import print_function import collections import json import os +import weakref import numpy as np @@ -35,7 +36,6 @@ from tensorflow.python.framework import dtypes as dtypes_module from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor from tensorflow.python.framework import tensor_util -from tensorflow.python.layers import base as tf_base_layers from tensorflow.python.ops import array_ops from tensorflow.python.ops import clip_ops from tensorflow.python.ops import control_flow_ops @@ -55,7 +55,7 @@ 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 @@ -263,6 +263,12 @@ def set_image_data_format(data_format): _IMAGE_DATA_FORMAT = str(data_format) +# A global dictionary mapping graph objects to an index of counters used +# for various layer names in each graph. +# Allows to give unique autogenerated names to layers, in a graph-specific way. +PER_GRAPH_LAYER_NAME_UIDS = weakref.WeakKeyDictionary() + + @tf_export('keras.backend.get_uid') def get_uid(prefix=''): """Associates a string prefix with an integer counter in a TensorFlow graph. @@ -283,17 +289,16 @@ def get_uid(prefix=''): ``` """ graph = ops.get_default_graph() - if graph not in tf_base_layers.PER_GRAPH_LAYER_NAME_UIDS: - tf_base_layers.PER_GRAPH_LAYER_NAME_UIDS[graph] = collections.defaultdict( - int) - layer_name_uids = tf_base_layers.PER_GRAPH_LAYER_NAME_UIDS[graph] + if graph not in PER_GRAPH_LAYER_NAME_UIDS: + PER_GRAPH_LAYER_NAME_UIDS[graph] = collections.defaultdict(int) + layer_name_uids = PER_GRAPH_LAYER_NAME_UIDS[graph] layer_name_uids[prefix] += 1 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 + per_graph_layer_name_uids = PER_GRAPH_LAYER_NAME_UIDS keys = list(per_graph_layer_name_uids.keys()) for key in keys: del per_graph_layer_name_uids[key] @@ -1276,6 +1281,11 @@ def moving_average_update(x, value, momentum): Returns: An Operation to update the variable. """ + # `training` is higher-up than the Keras backend in the abstraction hierarchy. + # In particular, `training` depends on layers, and thus on Keras. + # moving_averages, being low-level ops, should not be part of the training + # module. + from tensorflow.python.training import moving_averages # pylint: disable=g-import-not-at-top return moving_averages.assign_moving_average( x, value, momentum, zero_debias=True) @@ -2750,8 +2760,7 @@ class Function(object): outputs: Output tensors to fetch. updates: Additional update ops to be run at function call. name: A name to help users identify what this function does. - session_kwargs: Arguments to `tf.Session.run()`: `fetches`, `feed_dict`, - `options`, `run_metadata` + session_kwargs: Arguments to `tf.Session.run()`: `fetches`, `feed_dict`. """ def __init__(self, inputs, outputs, updates=None, name=None, @@ -2785,19 +2794,76 @@ class Function(object): self.fetches = session_kwargs.pop('fetches', []) if not isinstance(self.fetches, list): self.fetches = [self.fetches] + # The main use case of `fetches` being passed to a model is the ability + # to run custom updates (since the outputs of fetches are never returned). + # This requires us to wrap fetches in `identity` ops. + self.fetches = [array_ops.identity(x) for x in self.fetches] self.session_kwargs = session_kwargs + if session_kwargs: + raise ValueError('Some keys in session_kwargs are not supported at this ' + 'time: %s', session_kwargs.keys()) + + self._callable_fn = None + self._feed_arrays = None + self._feed_symbols = None + self._symbol_vals = None + self._session = None + + def _make_callable(self, feed_arrays, feed_symbols, symbol_vals, session): + """Generates a callable that runs the graph. + + Arguments: + feed_arrays: List of input tensors to be fed Numpy arrays at runtime. + feed_symbols: List of input tensors to be fed symbolic tensors at runtime. + symbol_vals: List of symbolic tensors to be fed to `feed_symbols`. + session: Session to use to generate the callable. + + Returns: + Function that runs the graph according to the above options. + """ + # Prepare callable options. + callable_opts = config_pb2.CallableOptions() + # Handle external-data feed. + for x in feed_arrays: + callable_opts.feed.append(x.name) + if self.feed_dict: + for key in sorted(self.feed_dict.keys()): + callable_opts.feed.append(key.name) + # Handle symbolic feed. + for x, y in zip(feed_symbols, symbol_vals): + connection = callable_opts.tensor_connection.add() + if x.dtype != y.dtype: + y = math_ops.cast(y, dtype=x.dtype) + from_tensor = ops._as_graph_element(y) + if from_tensor is None: + from_tensor = y + connection.from_tensor = from_tensor.name # Data tensor + connection.to_tensor = x.name # Placeholder + # Handle fetches. + for x in self.outputs + self.fetches: + callable_opts.fetch.append(x.name) + # Handle updates. + callable_opts.target.append(self.updates_op.name) + # Create callable. + callable_fn = session._make_callable_from_options(callable_opts) + # Cache parameters corresponding to the generated callable, so that + # we can detect future mismatches and refresh the callable. + self._callable_fn = callable_fn + self._feed_arrays = feed_arrays + self._feed_symbols = feed_symbols + self._symbol_vals = symbol_vals + self._session = session + def __call__(self, inputs): if not isinstance(inputs, (list, tuple)): raise TypeError('`inputs` should be a list or tuple.') - if self.feed_dict: - feed_dict = self.feed_dict.copy() - else: - feed_dict = {} - session = get_session() - data_tensors_to_feed = [] + feed_arrays = [] + array_vals = [] + feed_symbols = [] + symbol_vals = [] for tensor, value in zip(self.inputs, inputs): if value is None: continue @@ -2806,23 +2872,31 @@ class Function(object): indices = np.concatenate((np.expand_dims(sparse_coo.row, 1), np.expand_dims(sparse_coo.col, 1)), 1) value = (indices, sparse_coo.data, sparse_coo.shape) - elif tensor_util.is_tensor(value): - data_tensors_to_feed.append((tensor, value)) + if tensor_util.is_tensor(value): + # Case: feeding symbolic tensor. + feed_symbols.append(tensor) + symbol_vals.append(value) else: - feed_dict[tensor] = value + # Case: feeding Numpy array. + feed_arrays.append(tensor) + # We need to do array conversion and type casting at this level, since + # `callable_fn` only supports exact matches. + array_vals.append(np.asarray(value, dtype=tensor.dtype.base_dtype.name)) + if self.feed_dict: + for key in sorted(self.feed_dict.keys()): + array_vals.append( + np.asarray(self.feed_dict[key], dtype=key.dtype.base_dtype.name)) - if data_tensors_to_feed: - # This is a *temporary* workaround (i.e. hack) to feed a symbolic tensor - # to `feed_dict`. It is very inefficient. It will be removed as soon - # as it becomes possible to pass symbolic tensors to `feed_dict`. - data_tensor_values = session.run([x[1] for x in data_tensors_to_feed]) - for i, v in enumerate(data_tensor_values): - feed_dict[data_tensors_to_feed[i][0]] = v + # Refresh callable if anything has changed. + if (self._callable_fn is None or + feed_arrays != self._feed_arrays or + symbol_vals != self._symbol_vals or + feed_symbols != self._feed_symbols or + session != self._session): + self._make_callable(feed_arrays, feed_symbols, symbol_vals, session) - fetches = self.outputs + [self.updates_op] + self.fetches - updated = session.run( - fetches=fetches, feed_dict=feed_dict, **self.session_kwargs) - return updated[:len(self.outputs)] + fetched = self._callable_fn(*array_vals) + return fetched[:len(self.outputs)] @tf_export('keras.backend.function') @@ -3374,7 +3448,7 @@ def categorical_crossentropy(target, output, from_logits=False): Returns: Output tensor. """ - # Note: nn.softmax_cross_entropy_with_logits + # Note: nn.softmax_cross_entropy_with_logits_v2 # expects logits, Keras expects probabilities. if not from_logits: # scale preds so that the class probas of each sample sum to 1 @@ -3438,7 +3512,7 @@ def binary_crossentropy(target, output, from_logits=False): Returns: A tensor. """ - # Note: nn.softmax_cross_entropy_with_logits + # Note: nn.sigmoid_cross_entropy_with_logits # expects logits, Keras expects probabilities. if not from_logits: # transform back to logits diff --git a/tensorflow/python/keras/_impl/keras/backend_test.py b/tensorflow/python/keras/_impl/keras/backend_test.py index fb4b2a0e1dc06c904d4b93038840dbf688d42ed4..de1ed467a2764a2b08269181bb8bcd615868d3b1 100644 --- a/tensorflow/python/keras/_impl/keras/backend_test.py +++ b/tensorflow/python/keras/_impl/keras/backend_test.py @@ -189,6 +189,39 @@ class BackendUtilsTest(test.TestCase): for y in ys: self.assertEqual(y.op.name[:12], 'StopGradient') + def test_function_tf_feed_symbols(self): + with self.test_session(): + # Test feeding a resource variable to `function`. + x1 = keras.backend.placeholder(shape=()) + x2 = keras.backend.placeholder(shape=()) + lr = keras.backend.learning_phase() # Include a placeholder_with_default. + + y1 = keras.backend.variable(10.) + y2 = 3 + + f = keras.backend.function( + inputs=[x1, x2, lr], + outputs=[x1 + 1, + keras.backend.in_train_phase(x2 + 2, x2 - 1)]) + outs = f([y1, y2, None]) # Use default learning_phase value. + self.assertEqual(outs, [11., 2.]) + outs = f([y1, y2, 1]) # Set learning phase value. + self.assertEqual(outs, [11., 5.]) + + # Test triggering a callable refresh by changing the input. + y3 = keras.backend.constant(20.) # Test with tensor + outs = f([y3, y2, None]) + self.assertEqual(outs, [21., 2.]) + + y4 = 4 # Test with non-symbol + outs = f([y4, y2, None]) + self.assertEqual(outs, [5., 2.]) + + # Test with a different dtype + y5 = keras.backend.constant(10., dtype='float64') + outs = f([y5, y2, None]) + self.assertEqual(outs, [11., 2.]) + def test_function_tf_fetches(self): # Additional operations can be passed to tf.Session().run() via its # `fetches` arguments. In contrast to `updates` argument of @@ -206,8 +239,9 @@ class BackendUtilsTest(test.TestCase): updates=[(x, x_placeholder + 1.)], fetches=[keras.backend.update(y, 5.)]) output = f([10., 20.]) - assert output == [30.] - assert keras.backend.get_session().run(fetches=[x, y]) == [11., 5.] + self.assertEqual(output, [30.]) + self.assertEqual( + keras.backend.get_session().run(fetches=[x, y]), [11., 5.]) def test_function_tf_feed_dict(self): # Additional substitutions can be passed to `tf.Session().run()` via its @@ -229,14 +263,16 @@ class BackendUtilsTest(test.TestCase): feed_dict=feed_dict, fetches=fetches) output = f([10.]) - assert output == [11.] - assert keras.backend.get_session().run(fetches=[x, y]) == [20., 30.] + self.assertEqual(output, [11.]) + self.assertEqual( + keras.backend.get_session().run(fetches=[x, y]), [20., 30.]) # updated value in feed_dict will be modified within the K.function() feed_dict[y_placeholder] = 4. output = f([20.]) - assert output == [21.] - assert keras.backend.get_session().run(fetches=[x, y]) == [30., 40.] + self.assertEqual(output, [21.]) + self.assertEqual( + keras.backend.get_session().run(fetches=[x, y]), [30., 40.]) class BackendVariableTest(test.TestCase): diff --git a/tensorflow/python/keras/_impl/keras/engine/base_layer.py b/tensorflow/python/keras/_impl/keras/engine/base_layer.py index 755607aafb9343f0c9f10c5f9394bedc2f8afd76..6c68d251275cad4a6f2d6a5bf959a42c041537f3 100644 --- a/tensorflow/python/keras/_impl/keras/engine/base_layer.py +++ b/tensorflow/python/keras/_impl/keras/engine/base_layer.py @@ -13,143 +13,145 @@ # limitations under the License. # ============================================================================== # pylint: disable=protected-access -"""Base layer code (`Layer`). -""" +"""Contains the base Layer class, from which all layers inherit.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function +import collections import inspect # Necessary supplement to tf_inspect to deal with variadic args. +import re +import numpy as np from six.moves import zip # pylint: disable=redefined-builtin 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.keras._impl.keras import backend as K +from tensorflow.python.keras._impl.keras import backend 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.ops import array_ops +from tensorflow.python.ops import init_ops +from tensorflow.python.ops import variable_scope as vs +from tensorflow.python.ops import variables as tf_variables +from tensorflow.python.training import checkpointable +from tensorflow.python.util import nest from tensorflow.python.util import tf_decorator from tensorflow.python.util import tf_inspect 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(checkpointable.CheckpointableBase): + """Base layer class. + This is the class from which all layers inherit. -@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) + A layer is a class implementing common neural networks operations, such + as convolution, batch norm, etc. These operations require managing weights, + losses, updates, and inter-layer connectivity. + + Users will just instantiate a layer and then treat it as a callable. + + We recommend that descendants of `Layer` implement the following methods: + * `__init__()`: Save configuration in member variables + * `build()`: Called once from `__call__`, when we know the shapes of inputs + and `dtype`. Should have the calls to `add_weight()`, and then + call the super's `build()` (which sets `self.built = True`, which is + nice in case the user wants to call `build()` manually before the + first `__call__`). + * `call()`: Called in `__call__` after making sure `build()` has been called + once. Should actually perform the logic of applying the layer to the + input tensors (which should be passed in as the first argument). + + Arguments: + trainable: Boolean, whether the layer's variables should be trainable. + name: String name of the layer. + dtype: Default dtype of the layer's weights (default of `None` means use the + type of the first input). + + Read-only properties: + name: The name of the layer (string). + dtype: Default dtype of the layer's weights (default of `None` means use the + type of the first input). + trainable_variables: List of trainable variables. + non_trainable_variables: List of non-trainable variables. + variables: List of all variables of this layer, trainable and + non-trainable. + updates: List of update ops of this layer. + losses: List of losses added by this layer. + trainable_weights: List of variables to be included in backprop. + non_trainable_weights: List of variables that should not be + included in backprop. + weights: The concatenation of the lists trainable_weights and + non_trainable_weights (in this order). + + Mutable properties: + trainable: Whether the layer should be trained (boolean). + input_spec: Optional (list of) `InputSpec` object(s) specifying the + constraints on inputs that can be accepted by the layer. """ - def __init__(self, **kwargs): + def __init__(self, trainable=True, name=None, dtype=None, **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', + 'activity_regularizer', } # 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')) + # Mutable properties + # Indicates whether the layer's weights are updated during training + # and whether the layer's updates are run during training + self.trainable = trainable + # A stateful layer is a layer whose updates are run during inference too, + # for instance stateful RNNs. + self.stateful = False + # Indicates whether `build` needs to be called upon layer call, to create + # the layer's weights. + self.built = False + # Provides information about which inputs are compatible with the layer. + self.input_spec = None + + self._init_set_name(name) + + activity_regularizer = kwargs.pop('activity_regularizer', None) + if activity_regularizer and context.executing_eagerly(): + raise ValueError( + ('Activity regularization is not supported when executing eagerly. ' + 'Got activity_regularizer=%s') % (activity_regularizer,)) + self._activity_regularizer = activity_regularizer + self._trainable_weights = [] + self._non_trainable_weights = [] + self._updates = [] + # When executing eagerly, _losses is a list of zero-argument lambdas which + # return tensors. When using graph execution, _losses is a list of ops. + self._losses = [] + self._dtype = None if dtype is None else dtypes.as_dtype(dtype).name + 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._uses_inputs_arg = True - # Add properties that are Keras-only for now. + # These lists will be filled via successive calls + # to self._add_inbound_node(). + self._inbound_nodes = [] + self._outbound_nodes = [] + self.supports_masking = False # Manage input shape information if passed. @@ -172,39 +174,418 @@ class Layer(tf_base_layers.Layer): else: self._initial_weights = None - def add_weight(self, - name, - shape, + def _init_set_name(self, name, zero_based=True): + if not name: + self._name = unique_layer_name( + to_snake_case(self.__class__.__name__), zero_based=zero_based) + else: + self._name = name + + @property + def dtype(self): + return self._dtype + + @property + def name(self): + return self._name + + @property + def activity_regularizer(self): + """Optional regularizer function for the output of this layer.""" + return self._activity_regularizer + + @activity_regularizer.setter + def activity_regularizer(self, regularizer): + """Optional regularizer function for the output of this layer.""" + self._activity_regularizer = regularizer + + @property + def trainable_weights(self): + return self._trainable_weights if self.trainable else [] + + @property + def non_trainable_weights(self): + if self.trainable: + return self._non_trainable_weights + else: + return self._trainable_weights + self._non_trainable_weights + + @property + def trainable_variables(self): + return self.trainable_weights + + @property + def non_trainable_variables(self): + return self.non_trainable_weights + + @property + def weights(self): + """Returns the list of all layer variables/weights. + + Returns: + A list of variables. + """ + return self.trainable_weights + self.non_trainable_weights + + @property + def variables(self): + """Returns the list of all layer variables/weights. + + Returns: + A list of variables. + """ + return self.weights + + @property + def updates(self): + if context.executing_eagerly(): + raise RuntimeError('Layer.updates not supported in Eager mode.') + if not self.trainable and not self.stateful: + return [] + return self._updates + + def add_update(self, updates, inputs=None): + """Add update op(s), potentially dependent on layer inputs. + + Weight updates (for instance, the updates of the moving mean and variance + in a BatchNormalization layer) may be dependent on the inputs passed + when calling a layer. Hence, when reusing the same layer on + different inputs `a` and `b`, some entries in `layer.updates` may be + dependent on `a` and some on `b`. This method automatically keeps track + of dependencies. + + The `get_updates_for` method allows to retrieve the updates relevant to a + specific set of inputs. + + This call is ignored when eager execution is enabled (in that case, variable + updates are run on the fly and thus do not need to be tracked for later + execution). + + Arguments: + updates: Update op, or list/tuple of update ops. + 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.executing_eagerly(): + return # Updates already applied when in eager mode. + + def process_update(x): + if isinstance(x, ops.Operation): + return x + elif hasattr(x, 'op'): + return x.op + else: + return ops.convert_to_tensor(x) + + updates = generic_utils.to_list(updates) + updates = [process_update(x) for x in updates] + self._updates += updates + if inputs is None: + for u in updates: + u._unconditional_update = True # pylint: disable=protected-access + else: + 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. + + Returns: + List of update ops of the layer that depend on `inputs`. + + Raises: + RuntimeError: If called 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 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 = 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): + """Losses which are associated with this `Layer`. + + Note that when executing eagerly, getting this property evaluates + regularizers. When using graph execution, variable regularization ops have + already been created and are simply returned here. + + Returns: + A list of tensors. + """ + 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. + return [regularizer() for regularizer in self._losses] + else: + return self._losses + + def add_loss(self, losses, inputs=None): + """Add loss tensor(s), potentially dependent on layer inputs. + + Some losses (for instance, activity regularization losses) may be dependent + on the inputs passed when calling a layer. Hence, when reusing the same + layer on different inputs `a` and `b`, some entries in `layer.losses` may + be dependent on `a` and some on `b`. This method automatically keeps track + of dependencies. + + The `get_losses_for` method allows to retrieve the losses relevant to a + specific set of inputs. + + Note that `add_loss` is not supported when executing eagerly. Instead, + variable regularizers may be added through `add_variable`. Activity + regularization is not supported directly (but such losses may be returned + from `Layer.call()`). + + Arguments: + losses: Loss tensor, or list/tuple of tensors. + 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.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 = generic_utils.to_list(losses) + self._losses += losses + if inputs is None: + for loss in losses: + loss._unconditional_loss = True # pylint: disable=protected-access + else: + for loss in losses: + loss._unconditional_loss = False # pylint: disable=protected-access + + def get_losses_for(self, inputs): + """Retrieves losses relevant to a specific set of inputs. + + Arguments: + inputs: Input tensor or list/tuple of input tensors. + + Returns: + List of loss tensors of the layer that depend on `inputs`. + + Raises: + RuntimeError: If called in Eager mode. + """ + if context.executing_eagerly(): + raise RuntimeError('Layer.get_losses_for not supported in Eager mode.') + + 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 = get_reachable_from_inputs(inputs, self.losses) + losses = [] + for loss in self.losses: + if loss in reachable: + losses.append(loss) + return losses + + def _name_scope(self): + return self.name + + def build(self, _): + """Creates the variables of the layer.""" + self.built = True + + def add_variable(self, *args, **kwargs): + """Alias for `add_weight`.""" + return self.add_weight(*args, **kwargs) + + def add_weight(self, name, shape, dtype=None, initializer=None, regularizer=None, trainable=True, - constraint=None): - """Adds a weight variable to the layer. + constraint=None, + partitioner=None, + use_resource=None, + getter=None): + """Adds a new variable to the layer, or gets an existing one; returns it. 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. + name: variable name. + shape: variable shape. + dtype: The type of the variable. Defaults to `self.dtype` or `float32`. + initializer: initializer instance (callable). + regularizer: regularizer instance (callable). + trainable: whether the variable should be part of the layer's + "trainable_variables" (e.g. variables, biases) + or "non_trainable_variables" (e.g. BatchNorm mean, stddev). + Note, if the current variable scope is marked as non-trainable + then this parameter is ignored and any added variables are also + marked as non-trainable. + constraint: constraint instance (callable). + partitioner: Partitioner to be passed to the `Checkpointable` API. + use_resource: Whether to use `ResourceVariable`. + getter: Variable getter argument to be passed to the `Checkpointable` API. Returns: - The created weight variable. + The created variable. Usually either a `Variable` or `ResourceVariable` + instance. If `partitioner` is not `None`, a `PartitionedVariable` + instance is returned. + + Raises: + RuntimeError: If called with partioned variable regularization and + eager execution is enabled. + ValueError: When giving unsupported dtype and no initializer. """ 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 + dtype = self.dtype or backend.floatx() + else: + dtype = dtypes.as_dtype(dtype) + initializer = initializers.get(initializer) + regularizer = regularizers.get(regularizer) + constraint = constraints.get(constraint) + + # Initialize variable when no initializer provided + if initializer is None: + # If dtype is DT_FLOAT, provide a uniform unit scaling initializer + if dtype.is_floating: + initializer = initializers.glorot_uniform() + # If dtype is DT_INT/DT_UINT, provide a default value `zero` + # If dtype is DT_BOOL, provide a default value `FALSE` + elif dtype.is_integer or dtype.is_unsigned or dtype.is_bool: + initializer = initializers.zeros() + # NOTES:Do we need to support for handling DT_STRING and DT_COMPLEX here? + else: + raise ValueError('An initializer for variable %s of type %s is required' + ' for layer %s' % (name, dtype.base_dtype, self.name)) + + variable = self._add_variable_with_custom_getter( + name=name, + shape=shape, + # TODO(allenl): a `make_variable` equivalent should be added as a + # `Checkpointable` method. + getter=getter or make_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, + use_resource=use_resource) + + if regularizer is not None: + # TODO(fchollet): in the future, this should be handled at the + # level of variable creation, and weight regularization losses + # should be variable attributes. + self._handle_weight_regularization(name, variable, regularizer) + + if trainable: + self._trainable_weights.append(variable) + else: + self._non_trainable_weights.append(variable) + return variable + + def _handle_weight_regularization(self, name, variable, regularizer): + # `init_graph` should point to the graph in which variable initialization + # will occur; it should be None if and only if initialization will take + # place in the eager context. + init_graph = None + if not context.executing_eagerly(): + default_graph = ops.get_default_graph() + if default_graph.building_function: + with ops.init_scope(): + # Retrieve the variables from the graph into which variables + # 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 not context.executing_eagerly(): + init_graph = ops.get_default_graph() + else: + # Initialization ops will not be lifted out of the default graph. + init_graph = default_graph + + if init_graph is not None: # pylint: disable=protected-access + # The variable was created and initialized in a graph. + if regularizer: + if isinstance(variable, tf_variables.PartitionedVariable): + for v in variable: + with ops.colocate_with(v.op): + with ops.name_scope(name + '/Regularizer'): + regularization = regularizer(v) + if regularization is not None: + self.add_loss(regularization) + else: + with ops.colocate_with(variable.op): + with ops.name_scope(name + '/Regularizer'): + regularization = regularizer(variable) + if regularization is not None: + self.add_loss(regularization) + elif regularizer: # initialization took place in an eager context + if isinstance(variable, tf_variables.PartitionedVariable): + raise RuntimeError( + 'Partitioned variable regularization is not yet ' + '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. + # This makes losses responsive to variable updates when executing + # eagerly. + # + # TODO(akshayka): Do the same for graphs as well, so that losses + # collected in a while_loop can be run outside its control flow + # context and so that losses won't be swallowed up by graph functions + # (i.e., `.losses()` should always create regularizers). + self._losses.append(lambda: regularizer(variable)) + + def _handle_activity_regularization(self, inputs, outputs): + # Apply activity regularization. + # Note that it should be applied every time the layer creates a new + # output, since it is output-specific. + if self._activity_regularizer: + output_list = nest.flatten(outputs) + for output in output_list: + with ops.name_scope('ActivityRegularizer'): + activity_regularization = self._activity_regularizer(output) + self.add_loss(activity_regularization, inputs=inputs) def call(self, inputs, **kwargs): # pylint: disable=unused-argument """This is where the layer's logic lives. @@ -218,6 +599,215 @@ class Layer(tf_base_layers.Layer): """ return inputs + def __call__(self, inputs, *args, **kwargs): + """Wraps `call`, applying pre- and post-processing steps. + + Arguments: + inputs: input tensor(s). + *args: additional positional arguments to be passed to `self.call`. + **kwargs: additional keyword arguments to be passed to `self.call`. + + Returns: + Output tensor(s). + + Note: + - The following optional keyword arguments are reserved for specific uses: + * `training`: Boolean scalar tensor of Python boolean indicating + whether the `call` is meant for training or inference. + * `mask`: Boolean input mask. + - If the layer's `call` method takes a `mask` argument (as some Keras + layers do), its default value will be set to the mask generated + for `inputs` by the previous layer (if `input` did come from + a layer that generated a corresponding mask, i.e. if it came from + a Keras layer with masking support. + + Raises: + ValueError: if the layer's `call` method returns None (an invalid value). + """ + input_list = nest.flatten(inputs) + + build_graph = not context.executing_eagerly() + # TODO(fchollet, allenl): Make deferred mode work with subclassed Models + # which don't use an "inputs" argument. + in_deferred_mode = isinstance(input_list[0], DeferredTensor) + + # Handle Keras mask propagation from previous layer to current layer. + previous_mask = None + if (not hasattr(self, '_compute_previous_mask') or + self._compute_previous_mask): + previous_mask = collect_previous_mask(inputs) + 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. + kwargs['mask'] = previous_mask + + input_shapes = None + + with ops.name_scope(self._name_scope()): + if not self.built: + if not build_graph: + # Activity regularization is currently unsupported in Eager mode. + if self._activity_regularizer: + 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: + for x in input_list: + if hasattr(x, '_keras_history'): + raise ValueError('_keras_history currently unsupported in ' + 'Eager mode. Found _keras_history in %s while ' + 'executing __call__ for %s(%s)' % + (x, self.__class_.__name__, self)) + + # Check input assumptions set before layer building, e.g. input rank. + self._assert_input_compatibility(inputs) + if input_list and self._dtype is None: + try: + self._dtype = input_list[0].dtype.base_dtype.name + except AttributeError: + pass + if all(hasattr(x, 'get_shape') for x in input_list): + input_shapes = nest.map_structure(lambda x: x.get_shape(), inputs) + self.build(input_shapes) + + # Check input assumptions set after layer building, e.g. input shape. + if build_graph or in_deferred_mode: + self._assert_input_compatibility(inputs) + + if not in_deferred_mode: + outputs = self.call(inputs, *args, **kwargs) + if outputs is None: + raise ValueError('A layer\'s `call` method should return a Tensor ' + 'or a list of Tensors, not None (layer: ' + + self.name + ').') + 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 = [ + # TODO(fchollet): name the deferred tensors? + DeferredTensor(shape=shape, dtype=self._dtype) + for shape in output_shapes + ] + if len(outputs) == 1: + outputs = outputs[0] + + if build_graph: + self._handle_activity_regularization(inputs, outputs) + # TODO(fchollet): consider enabling masking for Eager mode. + self._set_mask_metadata(inputs, outputs, previous_mask) + + if in_deferred_mode or build_graph and have_all_keras_metadata(inputs): + inputs, outputs = self._set_connectivity_metadata_( + inputs, outputs, args, kwargs) + + self.built = True + if context.executing_eagerly(): + return outputs + + if hasattr(self, '_symbolic_set_inputs') and not self.inputs: + # Subclassed network: explicitly set metadata normally set by a call to + # self._set_inputs(). This is not relevant in eager execution. + self._symbolic_set_inputs(inputs, outputs) + + if in_deferred_mode or build_graph: + self._set_learning_phase_metadata(inputs, outputs) + + # Optionally load weight values that were specified at layer instantiation. + # TODO(fchollet): consider enabling this with eager execution too. + if hasattr(self, '_initial_weights') and self._initial_weights is not None: + self.set_weights(self._initial_weights) + del self._initial_weights + return outputs + + def apply(self, inputs, *args, **kwargs): + """Apply the layer on a input. + + This simply wraps `self.__call__`. + + Arguments: + inputs: Input tensor(s). + *args: additional positional arguments to be passed to `self.call`. + **kwargs: additional keyword arguments to be passed to `self.call`. + + Returns: + Output tensor(s). + """ + return self.__call__(inputs, *args, **kwargs) + + def _set_learning_phase_metadata(self, inputs, outputs): + # Update learning phase info. To work with subclassed models, + # this should be done even if Keras metadata is absent. + output_tensors = generic_utils.to_list(outputs) + 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)): + try: + output_tensors[i]._uses_learning_phase = getattr( + output_tensors[i], '_uses_learning_phase', False) or uses_lp + except AttributeError: + # An output element happens to be a C type (such as tuple or dict). + # We don't track learning phase info in such edge cases. + pass + + def _set_mask_metadata(self, inputs, outputs, previous_mask): + if hasattr(self, 'compute_mask'): + output_mask = self.compute_mask(inputs, previous_mask) + 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): + try: + x._keras_mask = m # pylint: disable=protected-access + except AttributeError: + pass # C type such as dict. Masking not supported in this case. + else: + try: + outputs._keras_mask = output_mask # pylint: disable=protected-access + except AttributeError: + pass # C type such as dict. Masking not supported in this case. + + def _set_connectivity_metadata_(self, inputs, outputs, args, kwargs): + if args and getattr(self, '_uses_inputs_arg', True): + raise TypeError( + 'This Layer takes an `inputs` argument to call(), and only the ' + '`inputs` argument may be specified as a positional argument. ' + 'Pass everything else as a keyword argument (those arguments will' + ' not be tracked as inputs to the Layer).') + + # If the layer returns tensors from its inputs, unmodified, + # we copy them to avoid loss of tensor metadata. + output_ls = nest.flatten(outputs) + output_ls_copy = [] + for x in output_ls: + if x in nest.flatten(inputs): + with ops.name_scope(self.name): + x = array_ops.identity(x) + output_ls_copy.append(x) + if len(output_ls_copy) == 1: + outputs = output_ls_copy[0] + else: + outputs = output_ls_copy + + inputs, kwargs = self._inputs_from_call_args( + call_args=(inputs,) + args, call_kwargs=kwargs) + # 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). + kwargs.pop('mask', None) # `mask` should not be serialized. + self._add_inbound_node( + input_tensors=inputs, output_tensors=outputs, arguments=kwargs) + return inputs, outputs + def _inputs_from_call_args(self, call_args, call_kwargs): """Get Layer inputs from __call__ *args and **kwargs. @@ -282,71 +872,6 @@ class Layer(tf_base_layers.Layer): input_arg_values.extend(bound_args[call_arg_spec.varargs]) return input_arg_values, non_input_arg_values - def __call__(self, inputs, *args, **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. - *args: Additional positional arguments to be passed to `call()`. Only - allowed in subclassed Models with custom call() signatures. In other - cases, `Layer` inputs must be passed using the `inputs` argument and - non-inputs must be keyword arguments. - **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. - TypeError: If positional arguments are passed and this `Layer` is not a - subclassed `Model`. - """ - # Actually call the layer (optionally building it). - output = super(Layer, self).__call__(inputs, *args, **kwargs) - - if args and getattr(self, '_uses_inputs_arg', True): - raise TypeError( - 'This Layer takes an `inputs` argument to call(), and only the ' - '`inputs` argument may be specified as a positional argument. Pass ' - 'everything else as a keyword argument (those arguments will not be ' - 'tracked as inputs to the Layer).') - - if context.executing_eagerly(): - return output - - inputs, kwargs = self._inputs_from_call_args( - call_args=(inputs,) + args, call_kwargs=kwargs) - - 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. @@ -362,13 +887,7 @@ class Layer(tf_base_layers.Layer): 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 + raise NotImplementedError def compute_mask(self, inputs, mask=None): # pylint: disable=unused-argument """Computes an output mask tensor. @@ -396,6 +915,87 @@ class Layer(tf_base_layers.Layer): # carry over the input mask return mask + def _add_inbound_node(self, + input_tensors, + output_tensors, + arguments=None): + """Internal method to create an inbound node for the layer. + + Arguments: + input_tensors: list of input tensors. + output_tensors: list of output tensors. + arguments: dictionary of keyword arguments that were passed to the + `call` method of the layer at the call that created the node. + """ + input_tensors = nest.flatten(input_tensors) + output_tensors = nest.flatten(output_tensors) + + # Collect input tensor(s) coordinates. + inbound_layers = [] + node_indices = [] + tensor_indices = [] + for x in input_tensors: + assert hasattr(x, '_keras_history') + inbound_layer, node_index, tensor_index = x._keras_history # pylint: disable=protected-access + inbound_layers.append(inbound_layer) + node_indices.append(node_index) + tensor_indices.append(tensor_index) + + # Create node, add it to inbound nodes. + Node( + self, + inbound_layers=inbound_layers, + node_indices=node_indices, + tensor_indices=tensor_indices, + input_tensors=input_tensors, + output_tensors=output_tensors, + arguments=arguments) + + # Update tensor history metadata. + for i in range(len(output_tensors)): + # The metadata attribute consists of 1) a layer instance + # 2) a node index for the layer, 3) a tensor index for the node. + # The allows layer reuse (multiple nodes per layer) and multi-output + # or multi-input layers (e.g. a layer can return multiple tensors, + # and each can be sent to a different layer). + output_tensors[i]._keras_history = (self, len(self._inbound_nodes) - 1, i) # pylint: disable=protected-access + + def _get_node_attribute_at_index(self, node_index, attr, attr_name): + """Private utility to retrieves an attribute (e.g. inputs) from a node. + + This is used to implement the methods: + - get_input_shape_at + - get_output_shape_at + - get_input_at + etc... + + Arguments: + node_index: Integer index of the node from which + to retrieve the attribute. + attr: Exact node attribute name. + attr_name: Human-readable attribute name, for error messages. + + Returns: + The layer's attribute `attr` at the node of index `node_index`. + + Raises: + RuntimeError: If the layer has no inbound nodes, or if called in Eager + mode. + ValueError: If the index provided does not match any node. + """ + if not self._inbound_nodes: + raise RuntimeError('The layer has never been called ' + 'and thus has no defined ' + attr_name + '.') + if not len(self._inbound_nodes) > node_index: + raise ValueError('Asked to get ' + attr_name + ' at node ' + + str(node_index) + ', but the layer has only ' + + str(len(self._inbound_nodes)) + ' inbound nodes.') + values = getattr(self._inbound_nodes[node_index], attr) + if len(values) == 1: + return values[0] + else: + return values + def get_input_mask_at(self, node_index): """Retrieves the input mask tensor(s) of a layer at a given node. @@ -476,6 +1076,325 @@ class Layer(tf_base_layers.Layer): else: return getattr(output, '_keras_mask', None) + def get_input_shape_at(self, node_index): + """Retrieves the input shape(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 shape tuple + (or list of shape tuples if the layer has multiple inputs). + + Raises: + RuntimeError: If called in Eager mode. + """ + return self._get_node_attribute_at_index(node_index, 'input_shapes', + 'input shape') + + def get_output_shape_at(self, node_index): + """Retrieves the output shape(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 shape tuple + (or list of shape tuples if the layer has multiple outputs). + + Raises: + RuntimeError: If called in Eager mode. + """ + return self._get_node_attribute_at_index(node_index, 'output_shapes', + 'output shape') + + def get_input_at(self, node_index): + """Retrieves the input 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 tensor (or list of tensors if the layer has multiple inputs). + + Raises: + RuntimeError: If called in Eager mode. + """ + return self._get_node_attribute_at_index(node_index, 'input_tensors', + 'input') + + def get_output_at(self, node_index): + """Retrieves the output 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 tensor (or list of tensors if the layer has multiple outputs). + + Raises: + RuntimeError: If called in Eager mode. + """ + return self._get_node_attribute_at_index(node_index, 'output_tensors', + 'output') + + @property + def input(self): + """Retrieves the input tensor(s) of a layer. + + Only applicable if the layer has exactly one input, + i.e. if it is connected to one incoming layer. + + Returns: + Input tensor or list of input tensors. + + Raises: + AttributeError: if the layer is connected to + more than one incoming layers. + + Raises: + RuntimeError: If called in Eager mode. + AttributeError: If no inbound nodes are found. + """ + if not self._inbound_nodes: + raise AttributeError('Layer ' + self.name + + ' is not connected, no input to return.') + return self._get_node_attribute_at_index(0, 'input_tensors', 'input') + + @property + def output(self): + """Retrieves the output tensor(s) of a layer. + + Only applicable if the layer has exactly one output, + i.e. if it is connected to one incoming layer. + + Returns: + Output tensor or list of output tensors. + + Raises: + AttributeError: if the layer is connected to more than one incoming + layers. + RuntimeError: if called in Eager mode. + """ + if not self._inbound_nodes: + raise AttributeError('Layer ' + self.name + ' has no inbound nodes.') + return self._get_node_attribute_at_index(0, 'output_tensors', 'output') + + @property + def input_shape(self): + """Retrieves the input shape(s) of a layer. + + Only applicable if the layer has exactly one input, + i.e. if it is connected to one incoming layer, or if all inputs + have the same shape. + + Returns: + Input shape, as an integer shape tuple + (or list of shape tuples, one tuple per input tensor). + + Raises: + AttributeError: if the layer has no defined input_shape. + RuntimeError: if called in Eager mode. + """ + if not self._inbound_nodes: + raise AttributeError('The layer has never been called ' + 'and thus has no defined input shape.') + all_input_shapes = set( + [str(node.input_shapes) for node in self._inbound_nodes]) + if len(all_input_shapes) == 1: + input_shapes = self._inbound_nodes[0].input_shapes + if len(input_shapes) == 1: + return tuple(tensor_shape.TensorShape(input_shapes[0]).as_list()) + else: + return [ + tuple(tensor_shape.TensorShape(shape).as_list()) + for shape in input_shapes + ] + else: + raise AttributeError('The layer "' + str(self.name) + + ' has multiple inbound nodes, ' + 'with different input shapes. Hence ' + 'the notion of "input shape" is ' + 'ill-defined for the layer. ' + 'Use `get_input_shape_at(node_index)` ' + 'instead.') + + def count_params(self): + """Count the total number of scalars composing the weights. + + Returns: + An integer count. + + Raises: + ValueError: if the layer isn't yet built + (in which case its weights aren't yet defined). + """ + if not self.built: + if self.__class__.__name__ == 'Sequential': + self.build() # pylint: disable=no-value-for-parameter + else: + raise ValueError('You tried to call `count_params` on ' + self.name + + ', but the layer isn\'t built. ' + 'You can build it manually via: `' + self.name + + '.build(batch_input_shape)`.') + weight_shapes = [w.get_shape().as_list() for w in self.weights] + return int(sum([np.prod(w) for w in weight_shapes])) + + @property + def output_shape(self): + """Retrieves the output shape(s) of a layer. + + Only applicable if the layer has one output, + or if all outputs have the same shape. + + Returns: + Output shape, as an integer shape tuple + (or list of shape tuples, one tuple per output tensor). + + Raises: + AttributeError: if the layer has no defined output shape. + RuntimeError: if called in Eager mode. + """ + if not self._inbound_nodes: + raise AttributeError('The layer has never been called ' + 'and thus has no defined output shape.') + all_output_shapes = set( + [str(node.output_shapes) for node in self._inbound_nodes]) + if len(all_output_shapes) == 1: + output_shapes = self._inbound_nodes[0].output_shapes + if len(output_shapes) == 1: + return tuple(tensor_shape.TensorShape(output_shapes[0]).as_list()) + else: + return [ + tuple(tensor_shape.TensorShape(shape).as_list()) + for shape in output_shapes + ] + else: + raise AttributeError('The layer "%s"' + ' has multiple inbound nodes, ' + 'with different output shapes. Hence ' + 'the notion of "output shape" is ' + 'ill-defined for the layer. ' + 'Use `get_output_shape_at(node_index)` ' + 'instead.' % self.name) + + @property + def inbound_nodes(self): + """Deprecated, do NOT use! Only for compatibility with external Keras.""" + return self._inbound_nodes + + @property + def outbound_nodes(self): + """Deprecated, do NOT use! Only for compatibility with external Keras.""" + return self._outbound_nodes + + def _assert_input_compatibility(self, inputs): + """Checks compatibility between the layer and provided inputs. + + This checks that the tensor(s) `inputs` verify the input assumptions + of the layer (if any). If not, a clear and actional exception gets raised. + + Arguments: + inputs: input tensor or list of input tensors. + + Raises: + ValueError: in case of mismatch between + the provided inputs and the expectations of the layer. + """ + if not self.input_spec: + return + if not isinstance(self.input_spec, (list, tuple)): + input_spec = nest.flatten(self.input_spec) + else: + input_spec = self.input_spec + inputs = nest.flatten(inputs) + if len(inputs) != len(input_spec): + raise ValueError('Layer ' + self.name + ' expects ' + + str(len(input_spec)) + ' inputs, ' + 'but it received ' + str(len(inputs)) + + ' input tensors. Inputs received: ' + str(inputs)) + for input_index, (x, spec) in enumerate(zip(inputs, input_spec)): + if spec is None: + continue + + if (spec.ndim is not None or + spec.min_ndim is not None or + spec.max_ndim is not None): + if x.get_shape().ndims is None: + raise ValueError('Input ' + str(input_index) + ' of layer ' + + self.name + ' is incompatible with the layer: ' + 'its rank is undefined, but the layer requires a ' + 'defined rank.') + + # Check ndim. + if spec.ndim is not None: + ndim = x.get_shape().ndims + if ndim != spec.ndim: + raise ValueError('Input ' + str(input_index) + ' of layer ' + + self.name + ' is incompatible with the layer: ' + 'expected ndim=' + str(spec.ndim) + ', found ndim=' + + str(ndim) + '. Full shape received: ' + + str(x.get_shape().as_list())) + if spec.max_ndim is not None: + ndim = x.get_shape().ndims + if ndim is not None and ndim > spec.max_ndim: + raise ValueError('Input ' + str(input_index) + ' of layer ' + + self.name + ' is incompatible with the layer: ' + 'expected max_ndim=' + str(spec.max_ndim) + + ', found ndim=' + str(ndim)) + if spec.min_ndim is not None: + ndim = x.get_shape().ndims + if ndim is not None and ndim < spec.min_ndim: + raise ValueError('Input ' + str(input_index) + ' of layer ' + + self.name + ' is incompatible with the layer: ' + ': expected min_ndim=' + str(spec.min_ndim) + + ', found ndim=' + str(ndim) + + '. Full shape received: ' + + str(x.get_shape().as_list())) + # Check dtype. + if spec.dtype is not None: + if x.dtype != spec.dtype: + raise ValueError('Input ' + str(input_index) + ' of layer ' + + self.name + ' is incompatible with the layer: ' + 'expected dtype=' + str(spec.dtype) + + ', found dtype=' + str(x.dtype)) + # Check specific shape axes. + if spec.axes: + shape = x.get_shape().as_list() + if shape is not None: + for axis, value in spec.axes.items(): + if hasattr(value, 'value'): + value = value.value + if value is not None and shape[int(axis)] not in {value, None}: + raise ValueError( + 'Input ' + str(input_index) + ' of layer ' + self.name + ' is' + ' incompatible with the layer: expected axis ' + str(axis) + + ' of input shape to have value ' + str(value) + + ' but received input with shape ' + str(shape)) + # Check shape. + if spec.shape is not None: + shape = x.get_shape().as_list() + if shape is not None: + for spec_dim, dim in zip(spec.shape, shape): + if spec_dim is not None and dim is not None: + if spec_dim != dim: + raise ValueError('Input ' + str(input_index) + + ' is incompatible with layer ' + self.name + + ': expected shape=' + str(spec.shape) + + ', found shape=' + str(shape)) + def set_weights(self, weights): """Sets the weights of the layer, from Numpy arrays. @@ -500,14 +1419,14 @@ class Layer(tf_base_layers.Layer): if not params: return weight_value_tuples = [] - param_values = K.batch_get_value(params) + param_values = backend.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) + backend.batch_set_value(weight_value_tuples) def get_weights(self): """Returns the current weights of the layer. @@ -516,7 +1435,7 @@ class Layer(tf_base_layers.Layer): Weights values as a list of numpy arrays. """ params = self.weights - return K.batch_get_value(params) + return backend.batch_get_value(params) def get_config(self): """Returns the config of the layer. @@ -558,9 +1477,196 @@ class Layer(tf_base_layers.Layer): """ return cls(**config) - @tf_base_layers.Layer.activity_regularizer.setter - def activity_regularizer(self, activity_regularizer): - self._activity_regularizer = activity_regularizer + +@tf_export('keras.layers.InputSpec', 'layers.InputSpec') +class InputSpec(object): + """Specifies the ndim, dtype and shape of every input to a layer. + + Every layer should expose (if appropriate) an `input_spec` attribute: + a list of instances of InputSpec (one per input tensor). + + A None entry in a shape is compatible with any dimension, + a None shape is compatible with any shape. + + Arguments: + dtype: Expected DataType of the input. + shape: Shape tuple, expected shape of the input + (may include None for unchecked axes). + ndim: Integer, expected rank of the input. + max_ndim: Integer, maximum rank of the input. + min_ndim: Integer, minimum rank of the input. + axes: Dictionary mapping integer axes to + a specific dimension value. + """ + + def __init__(self, + dtype=None, + shape=None, + ndim=None, + max_ndim=None, + min_ndim=None, + axes=None): + self.dtype = dtype + self.shape = shape + if shape is not None: + self.ndim = len(shape) + else: + self.ndim = ndim + self.max_ndim = max_ndim + self.min_ndim = min_ndim + self.axes = axes or {} + + def __repr__(self): + spec = [('dtype=' + str(self.dtype)) if self.dtype else '', + ('shape=' + str(self.shape)) if self.shape else '', + ('ndim=' + str(self.ndim)) if self.ndim else '', + ('max_ndim=' + str(self.max_ndim)) if self.max_ndim else '', + ('min_ndim=' + str(self.min_ndim)) if self.min_ndim else '', + ('axes=' + str(self.axes)) if self.axes else ''] + return 'InputSpec(%s)' % ', '.join(x for x in spec if x) + + +class Node(object): + """A `Node` describes the connectivity between two layers. + + Each time a layer is connected to some new input, + a node is added to `layer._inbound_nodes`. + Each time the output of a layer is used by another layer, + a node is added to `layer._outbound_nodes`. + + Arguments: + outbound_layer: the layer that takes + `input_tensors` and turns them into `output_tensors` + (the node gets created when the `call` + method of the layer was called). + inbound_layers: a list of layers, the same length as `input_tensors`, + the layers from where `input_tensors` originate. + node_indices: a list of integers, the same length as `inbound_layers`. + `node_indices[i]` is the origin node of `input_tensors[i]` + (necessary since each inbound layer might have several nodes, + e.g. if the layer is being shared with a different data stream). + tensor_indices: a list of integers, + the same length as `inbound_layers`. + `tensor_indices[i]` is the index of `input_tensors[i]` within the + output of the inbound layer + (necessary since each inbound layer might + have multiple tensor outputs, with each one being + independently manipulable). + input_tensors: list of input tensors. + output_tensors: list of output tensors. + arguments: dictionary of keyword arguments that were passed to the + `call` method of the layer at the call that created the node. + + `node_indices` and `tensor_indices` are basically fine-grained coordinates + describing the origin of the `input_tensors`. + + A node from layer A to layer B is added to: + - A._outbound_nodes + - B._inbound_nodes + """ + + def __init__(self, + outbound_layer, + inbound_layers, + node_indices, + tensor_indices, + input_tensors, + output_tensors, + arguments=None): + # Layer instance (NOT a list). + if isinstance(outbound_layer, list): + raise ValueError( + '`outbound_layer` should be a layer instance, not a list.') + # this is the layer that takes a list of input tensors + # and turns them into a list of output tensors. + # the current node will be added to + # the inbound_nodes of outbound_layer. + self.outbound_layer = outbound_layer + + # The following 3 properties describe where + # the input tensors come from: which layers, + # and for each layer, which node and which + # tensor output of each node. + + # List of layer instances. + self.inbound_layers = inbound_layers + # List of integers, 1:1 mapping with inbound_layers. + self.node_indices = node_indices + # List of integers, 1:1 mapping with inbound_layers. + self.tensor_indices = tensor_indices + + # Following 2 properties: + # tensor inputs and outputs of outbound_layer. + + # List of tensors. 1:1 mapping with inbound_layers. + self.input_tensors = input_tensors + # List of tensors, created by outbound_layer.call(). + self.output_tensors = output_tensors + + # Following 2 properties: input and output shapes. + + # List of shape tuples, shapes of input_tensors. + self.input_shapes = [static_shape(x) for x in input_tensors] + # List of shape tuples, shapes of output_tensors. + self.output_shapes = [static_shape(x) for x in output_tensors] + + # Optional keyword arguments to layer's `call`. + self.arguments = arguments + + # Add nodes to all layers involved. + for layer in inbound_layers: + if layer is not None: + # For compatibility with external Keras, we use the deprecated + # accessor here. + layer.outbound_nodes.append(self) + # For compatibility with external Keras, we use the deprecated + # accessor here. + outbound_layer.inbound_nodes.append(self) + + def get_config(self): + inbound_names = [] + for layer in self.inbound_layers: + if layer: + inbound_names.append(layer.name) + else: + inbound_names.append(None) + return { + 'outbound_layer': self.outbound_layer.name, + 'inbound_layers': inbound_names, + 'node_indices': self.node_indices, + 'tensor_indices': self.tensor_indices + } + + +class DeferredTensor(object): + """Tensor-like object used to build graphs of layers in Eager mode. + + When calling a layer on a DeferredTensor, the layer will not perform any + computation and will simply perfom shape inference to return new + DeferredTensors with appropriate shape information. Thus DeferredTensor + behaves like a graph-mode Tensor when manipulated by layers. + """ + + def __init__(self, shape, dtype, name=None): + self.shape = tensor_shape.TensorShape(shape) + 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): + return self.shape + + def __str__(self): + return "DeferredTensor('%s', shape=%s, dtype=%s)" % (self.name, + self.get_shape(), + self.dtype.name) + + def __repr__(self): + return "" % (self.name, + self.get_shape(), + self.dtype.name) def shape_type_conversion(fn): @@ -589,3 +1695,251 @@ def shape_type_conversion(fn): return tensor_shape.TensorShape(output_shape) return wrapper + + +def object_list_uid(object_list): + """Creates a single string from object ids.""" + object_list = nest.flatten(object_list) + return ', '.join([str(abs(id(x))) for x in object_list]) + + +def static_shape(x): + """Get the static shape of a Tensor, or None if it is unavailable.""" + if x is None: + return None + try: + return tuple(x.get_shape().as_list()) + except ValueError: + return None + + +def get_reachable_from_inputs(inputs, targets=None): + """Returns the set of tensors/ops 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() + if isinstance(x, ops.Operation): + outputs = x.outputs[:] or [] + outputs += x._control_outputs + elif isinstance(x, ops.Tensor): + outputs = x.consumers() + elif isinstance(x, tf_variables.Variable): + outputs = [x.op] + else: + raise TypeError('Expected Operation, Variable, or Tensor, got ' + str(x)) + + 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 + + +def unique_layer_name(name, name_uid_map=None, avoid_names=None, namespace='', + zero_based=False): + """Makes a layer name (or arbitrary string) unique within a TensorFlow graph. + + Arguments: + name: String name to make unique. + name_uid_map: An optional defaultdict(int) to use when creating unique + names. If None (default), uses a per-Graph dictionary. + avoid_names: An optional set or dict with names which should not be used. If + None (default) does not avoid any names. + namespace: Gets a name which is unique within the (graph, namespace). Layers + which are not Networks use a blank namespace and so get graph-global + names. + zero_based: If True, name sequences start with no suffix (e.g. "dense", + "dense_1"). If False, naming is one-based ("dense_1", "dense_2"). + + Returns: + Unique string name. + + Example: + + ```python + _unique_layer_name('dense') # dense_1 + _unique_layer_name('dense') # dense_2 + ``` + """ + if name_uid_map is None: + name_uid_map = get_default_graph_uid_map() + if avoid_names is None: + avoid_names = set() + proposed_name = None + while proposed_name is None or proposed_name in avoid_names: + name_key = (namespace, name) + if zero_based: + number = name_uid_map[name_key] + if number: + proposed_name = name + '_' + str(number) + else: + proposed_name = name + name_uid_map[name_key] += 1 + else: + name_uid_map[name_key] += 1 + proposed_name = name + '_' + str(name_uid_map[name_key]) + return proposed_name + + +def to_snake_case(name): + intermediate = re.sub('(.)([A-Z][a-z0-9]+)', r'\1_\2', name) + insecure = re.sub('([a-z])([A-Z])', r'\1_\2', intermediate).lower() + # If the class is private the name starts with "_" which is not secure + # for creating scopes. We prefix the name with "private" in this case. + if insecure[0] != '_': + return insecure + return 'private' + insecure + + +def is_all_none(iterable_or_element): + if not isinstance(iterable_or_element, (list, tuple)): + iterable = [iterable_or_element] + else: + iterable = iterable_or_element + # We cannot use Python's `any` because the iterable may return Tensors. + for element in iterable: + if element is not None: + return False + return True + + +def have_all_keras_metadata(iterable_or_element): + if not isinstance(iterable_or_element, (list, tuple)): + iterable = [iterable_or_element] + else: + iterable = iterable_or_element + return all([hasattr(x, '_keras_history') for x in iterable]) + + +def collect_previous_mask(input_tensors): + """Retrieves the output mask(s) of the previous node. + + Arguments: + input_tensors: A tensor or list of tensors. + + Returns: + A mask tensor or list of mask tensors. + """ + input_tensors = nest.flatten(input_tensors) + masks = [] + for x in input_tensors: + if hasattr(x, '_keras_mask'): + mask = x._keras_mask # pylint: disable=protected-access + masks.append(mask) + else: + masks.append(None) + if len(masks) == 1: + return masks[0] + return masks + + +def is_tensor_or_tensor_list(v): + v = nest.flatten(v) + if v and isinstance(v[0], ops.Tensor): + return True + else: + return False + + +def get_default_graph_uid_map(): + # TODO(fchollet): refactor this into backend. + graph = ops.get_default_graph() + name_uid_map = backend.PER_GRAPH_LAYER_NAME_UIDS.get(graph, None) + if name_uid_map is None: + name_uid_map = collections.defaultdict(int) + backend.PER_GRAPH_LAYER_NAME_UIDS[graph] = name_uid_map + return name_uid_map + + +def make_variable(name, + shape=None, + dtype=dtypes.float32, + initializer=None, + partition_info=None, + trainable=True, + caching_device=None, + validate_shape=True, + constraint=None, + use_resource=None, + partitioner=None): # pylint: disable=unused-argument + """Temporary util to create a variable (relies on `variable_scope.variable`). + + Some reuse-related technicalities prevent us from using + `variable_scope.get_variable()` directly, so we use a subcomponent + that has fewer constraints (`variable_scope.variable()`). + + In the longer term, it seems like a similar "default variable creator" method + should exist in `CheckpointableBase` instead. When this happens, we can get + rid of this temporary solution. + + TODO(fchollet): remove this method when no longer needed. + TODO(fchollet): handle `partitioner` argument. + + Arguments: + name: Variable name. + shape: Variable shape. + dtype: The type of the variable. Defaults to `self.dtype` or `float32`. + initializer: Initializer instance (callable). + partition_info: Not handled at this time. + trainable: Whether the variable should be part of the layer's + "trainable_variables" (e.g. variables, biases) + or "non_trainable_variables" (e.g. BatchNorm mean, stddev). + Note, if the current variable scope is marked as non-trainable + then this parameter is ignored and any added variables are also + marked as non-trainable. + caching_device: Passed to `vs.variable`. + validate_shape: Passed to `vs.variable`. + constraint: Constraint instance (callable). + use_resource: Whether to use a `ResourceVariable`. + partitioner: Not handled at this time. + + Returns: + Variable instance. + """ + initializing_from_value = False + if initializer is not None and not callable(initializer): + initializing_from_value = True + + with ops.init_scope(): + if initializing_from_value: + init_val = initializer + variable_dtype = None + else: + # Instantiate initializer if provided initializer is a type object. + if isinstance(initializer, type(init_ops.Initializer)): + initializer = initializer(dtype=dtype) + init_val = lambda: initializer( # pylint: disable=g-long-lambda + shape, dtype=dtype, partition_info=partition_info) + variable_dtype = dtype.base_dtype + if use_resource is None: + use_resource = True + + v = vs.variable( + initial_value=init_val, + name=name, + trainable=trainable, + caching_device=caching_device, + dtype=variable_dtype, + validate_shape=validate_shape, + constraint=constraint, + use_resource=use_resource) + return v diff --git a/tensorflow/python/keras/_impl/keras/engine/input_layer.py b/tensorflow/python/keras/_impl/keras/engine/input_layer.py index b51dd8a2189d0c8542c84dfeac9be0d72b96ff1b..bd9dcbe3c576851123dfcabe3e36379019627ac5 100644 --- a/tensorflow/python/keras/_impl/keras/engine/input_layer.py +++ b/tensorflow/python/keras/_impl/keras/engine/input_layer.py @@ -23,7 +23,6 @@ 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 @@ -95,7 +94,7 @@ class InputLayer(base_layer.Layer): 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 + input_tensor = base_layer.DeferredTensor( # pylint: disable=protected-access shape=batch_input_shape, dtype=dtype, name=self.name) @@ -123,7 +122,7 @@ class InputLayer(base_layer.Layer): # 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( + base_layer.Node( self, inbound_layers=[], node_indices=[], diff --git a/tensorflow/python/keras/_impl/keras/engine/network.py b/tensorflow/python/keras/_impl/keras/engine/network.py index 9f1c7de1157a3659ccb27e4850e99e09016d0067..cc177c14a894040df37f75bbdc6b2651336fe869 100644 --- a/tensorflow/python/keras/_impl/keras/engine/network.py +++ b/tensorflow/python/keras/_impl/keras/engine/network.py @@ -35,8 +35,6 @@ 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 @@ -82,7 +80,7 @@ class Network(base_layer.Layer): # self.losses # self.updates - self._init_set_name(name) + self._init_set_name(name, zero_based=True) 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 @@ -132,14 +130,14 @@ class Network(base_layer.Layer): 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 + if not isinstance(tensor, base_layer.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 + if not isinstance(tensor, base_layer.DeferredTensor): # pylint: disable=protected-access raise TypeError('When eager execution is enabled, ' 'outputs must come from a call to ' 'a layer (called after ' @@ -230,7 +228,7 @@ class Network(base_layer.Layer): self._layers_by_depth = layers_by_depth # Create the node linking internal inputs to internal outputs. - tf_base_layers.Node( + base_layer.Node( outbound_layer=self, inbound_layers=[], node_indices=[], @@ -243,8 +241,8 @@ class Network(base_layer.Layer): 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)) + mask_cache_key = (base_layer.object_list_uid(self.inputs) + '_' + + base_layer.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 @@ -289,7 +287,7 @@ class Network(base_layer.Layer): self.built = False def __setattr__(self, name, value): - if isinstance(value, (tf_base_layers.Layer, Network)): + if isinstance(value, (base_layer.Layer, Network)): try: is_graph_network = self._is_graph_network except AttributeError: @@ -299,6 +297,10 @@ class Network(base_layer.Layer): if not is_graph_network: if value not in self._layers: self._layers.append(value) + if hasattr(value, '_use_resource_variables'): + # In subclassed models, legacy layers (tf.layers) must always use + # resource variables. + value._use_resource_variables = True if isinstance(value, checkpointable.CheckpointableBase): # Layer (and therefore Network/Model) inherit from CheckpointableBase # rather than Checkpointable, which means there is no Checkpointable @@ -387,8 +389,8 @@ class Network(base_layer.Layer): 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)) + cache_key = (base_layer.object_list_uid(inputs) + + '_' + base_layer.object_list_uid(masks)) if cache_key in self._output_mask_cache: return self._output_mask_cache[cache_key] else: @@ -502,8 +504,7 @@ class Network(base_layer.Layer): relevant_inputs += inputs else: relevant_inputs.append(inputs) - reachable = tf_layers_util.get_reachable_from_inputs(relevant_inputs, - updates) + reachable = base_layer.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 @@ -540,8 +541,7 @@ class Network(base_layer.Layer): relevant_inputs += inputs else: relevant_inputs.append(inputs) - reachable = tf_layers_util.get_reachable_from_inputs(relevant_inputs, - losses) + reachable = base_layer.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 @@ -623,8 +623,8 @@ class Network(base_layer.Layer): 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)) + cache_key = (base_layer.object_list_uid(inputs) + + '_' + base_layer.object_list_uid(masks)) if cache_key in self._output_tensor_cache: # Cache hit. return self._output_tensor_cache[cache_key] @@ -656,7 +656,7 @@ class Network(base_layer.Layer): ': model has ' + str(len(self._input_layers)) + ' tensor inputs.') - cache_key = tf_layers_util.object_list_uid(input_shapes) + cache_key = base_layer.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`). @@ -845,7 +845,7 @@ class Network(base_layer.Layer): 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_shapes.append(base_layer.static_shape(x)) output_tensors.append(tensor) output_masks.append(mask) @@ -859,14 +859,14 @@ class Network(base_layer.Layer): 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)) + cache_key = (base_layer.object_list_uid(inputs) + + '_' + base_layer.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) + input_shapes = [base_layer.static_shape(x) for x in inputs] + cache_key = base_layer.object_list_uid(input_shapes) self._output_shape_cache[cache_key] = output_shapes return output_tensors, output_masks diff --git a/tensorflow/python/keras/_impl/keras/engine/saving_test.py b/tensorflow/python/keras/_impl/keras/engine/saving_test.py index dde090120456f968267e1c572f22eda1bd6ed7c4..3b1578cddfd97b31cae8619cdd2d8e1997585f51 100644 --- a/tensorflow/python/keras/_impl/keras/engine/saving_test.py +++ b/tensorflow/python/keras/_impl/keras/engine/saving_test.py @@ -422,7 +422,7 @@ class TestWholeModelSaving(test.TestCase): 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) + f = keras.layers.Dense(2, name='nested_model_output' + ('x' * (2**14)))(f) nested_model = keras.Model(inputs=[x], outputs=[f], name='nested_model') x = keras.Input(shape=(2,), name='outer_model_input') diff --git a/tensorflow/python/keras/_impl/keras/engine/sequential.py b/tensorflow/python/keras/_impl/keras/engine/sequential.py index 2ef99d5ab3f432058fdf685b99b01aa0b5eeffdc..bd13ca671340551c3e96895951be360b15e55cfe 100644 --- a/tensorflow/python/keras/_impl/keras/engine/sequential.py +++ b/tensorflow/python/keras/_impl/keras/engine/sequential.py @@ -123,7 +123,7 @@ class Sequential(Model): multiple output tensors, or is already connected somewhere else (forbidden in `Sequential` models). """ - if not isinstance(layer, (base_layer.Layer, base_layer.TFBaseLayer)): + if not isinstance(layer, base_layer.Layer): raise TypeError('The added layer must be ' 'an instance of class Layer. ' 'Found: ' + str(layer)) diff --git a/tensorflow/python/keras/_impl/keras/engine/sequential_test.py b/tensorflow/python/keras/_impl/keras/engine/sequential_test.py index c9a47581df03e0fc1ad38552ba8634862435cd80..8aba16aef3e187e9e33bdb65c7d44b0e622730ef 100644 --- a/tensorflow/python/keras/_impl/keras/engine/sequential_test.py +++ b/tensorflow/python/keras/_impl/keras/engine/sequential_test.py @@ -151,6 +151,7 @@ class TestSequential(test.TestCase): with self.test_session(): model = keras.models.Sequential() model.add(keras.layers.BatchNormalization(input_shape=(4,))) + assert model.updates model.trainable = False assert not model.updates diff --git a/tensorflow/python/keras/_impl/keras/engine/topology_test.py b/tensorflow/python/keras/_impl/keras/engine/topology_test.py index 9ab4b6fdcf55cc6186b96dd4e747f3600a4f78f8..49cc1cd3b38325b4f42d5b26bac9442d7cc09b05 100644 --- a/tensorflow/python/keras/_impl/keras/engine/topology_test.py +++ b/tensorflow/python/keras/_impl/keras/engine/topology_test.py @@ -25,7 +25,7 @@ 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.keras._impl.keras.engine import base_layer from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import state_ops @@ -52,11 +52,13 @@ class TopologyConstructionTest(test.TestCase): (1, 1), 'float32', trainable=False) - self.add_update(state_ops.assign_add(self.a, [[1.]])) + self.add_update(state_ops.assign_add(self.a, [[1.]], + name='unconditional_update')) self.built = True def call(self, inputs): - self.add_update(state_ops.assign_add(self.a, inputs), + self.add_update(state_ops.assign_add(self.b, inputs, + name='conditional_update'), inputs=True) return inputs + 1 @@ -97,10 +99,20 @@ class TopologyConstructionTest(test.TestCase): 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) + network.add_update(state_ops.assign_add(layer.b, x4), inputs=True) self.assertEqual(len(network.updates), 5) self.assertEqual(len(network.get_updates_for(x4)), 2) + def test_get_updates_bn(self): + x1 = keras.Input(shape=(1,)) + layer = keras.layers.BatchNormalization() + _ = layer.apply(x1) + + print('BN updates', layer._updates) + self.assertEqual(len(layer.updates), 2) + self.assertEqual(len(layer.get_updates_for(x1)), 2) + self.assertEqual(len(layer.get_updates_for(None)), 0) + def test_get_losses(self): class MyLayer(keras.layers.Layer): @@ -875,25 +887,25 @@ class TopologyConstructionTest(test.TestCase): class DeferredModeTest(test.TestCase): def testDeferredTensorAttributes(self): - x = tf_base_layers._DeferredTensor(shape=(None, 2), - dtype='float32', - name='x') + x = base_layer.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.assertIsInstance(inputs, base_layer.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.assertIsInstance(x, base_layer.DeferredTensor) self.assertEqual(x.dtype.name, 'float32') self.assertEqual(x.shape.as_list(), [None, 2]) @@ -936,5 +948,34 @@ class DeferredModeTest(test.TestCase): self.assertEqual(outputs[0].shape.as_list(), [10, 16]) self.assertEqual(outputs[1].shape.as_list(), [10, 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( + keras.engine.base_layer.get_reachable_from_inputs([pl_1]), + {pl_1, x_1, x_4, x_5, x_1.op, x_4.op, x_5.op}) + self.assertEqual( + keras.engine.base_layer.get_reachable_from_inputs([pl_1, pl_2]), + {pl_1, pl_2, x_1, x_2, x_4, x_5, x_1.op, x_2.op, x_4.op, x_5.op}) + self.assertEqual( + keras.engine.base_layer.get_reachable_from_inputs([pl_3]), + {pl_3, x_3, x_5, x_3.op, x_5.op}) + self.assertEqual( + keras.engine.base_layer.get_reachable_from_inputs([x_3]), + {x_3, x_5, x_5.op}) + + 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 71de657da81b92a2fc6b1eef9041147be6ff307e..146e8fdac9a2dd51ad928958b5281485c1e732ae 100644 --- a/tensorflow/python/keras/_impl/keras/engine/training.py +++ b/tensorflow/python/keras/_impl/keras/engine/training.py @@ -20,6 +20,8 @@ from __future__ import print_function import numpy as np +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.data.ops import iterator_ops from tensorflow.python.eager import context from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_util @@ -31,10 +33,10 @@ 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 DeferredTensor 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 @@ -276,6 +278,8 @@ class Model(Network): self.metrics_names.append(self.output_names[i] + '_loss') self.nested_metrics = training_utils.collect_metrics(metrics, self.output_names) + with K.name_scope('metrics'): + training_utils.populate_metric_names(self) self._feed_sample_weight_modes = [] for i in range(len(self.outputs)): self._feed_sample_weight_modes.append(None) @@ -462,7 +466,6 @@ class Model(Network): output_weighted_metrics = nested_weighted_metrics[i] def handle_metrics(metrics, weights=None): - metric_name_prefix = 'weighted_' if weights is not None else '' for metric in metrics: if metric in ('accuracy', 'acc', 'crossentropy', 'ce'): @@ -489,39 +492,19 @@ class Model(Network): metric_fn = metrics_module.categorical_accuracy elif metric in ('crossentropy', 'ce'): metric_fn = metrics_module.categorical_crossentropy - if metric in ('accuracy', 'acc'): - suffix = 'acc' - elif metric in ('crossentropy', 'ce'): - suffix = 'ce' 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 = 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 - + metric_name = training_utils.get_base_metric_name( + metric, weighted=weights is not None) with K.name_scope(metric_name): metric_result = weighted_metric_fn( y_true, y_pred, weights=weights, mask=masks[i]) - # 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) + training_utils.add_metric_name(self, metric_name, i) self.metrics_tensors.append(metric_result) # Keep track of state updates created by @@ -653,12 +636,20 @@ class Model(Network): 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. + x: Input data. It could be: + - A Numpy array (or array-like), or a list of arrays + (in case the model has multiple inputs). + - A TensorFlow tensor, or a list of tensors + (in case the model has multiple inputs). + - A dict mapping input names to the corresponding array/tensors, + if the model has named inputs. + - A `tf.data` dataset iterator. + y: Target data. Like the input data `x`, + it could be either Numpy array(s) or TensorFlow tensor(s). + It should be consistent with `x` (you cannot have Numpy inputs and + tensor targets, or inversely). If `x` is a dataset iterator, + `y` should not be specified + (since targets will be obtained from the iterator). 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 @@ -678,6 +669,31 @@ class Model(Network): RuntimeError: If the model was never compiled. """ # First, we build/compile the model on the fly if necessary. + if isinstance(x, dataset_ops.Dataset): + raise ValueError('You passed a `Dataset` instance to your model (%s), ' + 'which is not supported. Instead, pass an `Iterator`, ' + 'which you can obtain e.g. via ' + '`dataset.make_one_shot_iterator()` (the exact method ' + 'to use will depend on your specific dataset).' % x) + if isinstance(x, iterator_ops.Iterator): + if y is not None: + raise ValueError('You passed a dataset iterator (%s) as input `x` to ' + 'your model. In that case, you should not specify ' + 'a target (`y`) argument, since the dataset iterator ' + 'generates both input data and target data. ' + 'Received: %s' % (x, y)) + if not context.executing_eagerly(): + x, y = x.get_next() + # TODO(fchollet): handle case of `get_next` not returning 2 tensors? + else: + # TODO(psv): implement this. The way to support it will be to typecheck + # for `iterator` before `_standardize_user_data` is called and redirect + # to new training/eval functions in `training_eager.py`. The model + # may need to get built using the specs of the data from the first batch + # drawn from the iterator. + raise ValueError('Dataset iterators are not supported ' + 'with eager execution yet.') + all_inputs = [] if not self.built: # We need to use `x` to set the model inputs. @@ -891,15 +907,6 @@ class Model(Network): 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) - def _eager_set_inputs(self, inputs): """Set model's input and output specs based on the input data received. @@ -933,11 +940,11 @@ class Model(Network): 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] + 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] + 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 = [ @@ -1044,22 +1051,26 @@ class Model(Network): """Trains the model for a fixed number of epochs (iterations on a dataset). Arguments: - x: Numpy array of training data (if the model has a single input), - or list of Numpy arrays (if the model has multiple inputs). - If input layers in the model are named, you can also pass a - dictionary mapping input names to Numpy arrays. - `x` can be `None` (default) if feeding from - TensorFlow data tensors. - y: Numpy array of target (label) data - (if the model has a single output), - or list of Numpy arrays (if the model has multiple outputs). - If output layers in the model are named, you can also pass a - dictionary mapping output names to Numpy arrays. - `y` can be `None` (default) if feeding from - TensorFlow data tensors. + x: Input data. It could be: + - A Numpy array (or array-like), or a list of arrays + (in case the model has multiple inputs). + - A TensorFlow tensor, or a list of tensors + (in case the model has multiple inputs). + - A dict mapping input names to the corresponding array/tensors, + if the model has named inputs. + - A `tf.data` dataset iterator. + y: Target data. Like the input data `x`, + it could be either Numpy array(s) or TensorFlow tensor(s). + It should be consistent with `x` (you cannot have Numpy inputs and + tensor targets, or inversely). If `x` is a dataset iterator, + `y` should not be specified + (since targets will be obtained from the iterator). batch_size: Integer or `None`. Number of samples per gradient update. If unspecified, `batch_size` will default to 32. + Do not specify the `batch_size` is your data is in the + form of symbolic tensors or dataset iterators (since they generate + batches). epochs: Integer. Number of epochs to train the model. An epoch is an iteration over the entire `x` and `y` data provided. @@ -1081,11 +1092,14 @@ class Model(Network): 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 + validation_data: Data 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. `validation_data` will override `validation_split`. + `validation_data` could be: + - tuple `(x_val, y_val)` of Numpy arrays or tensors + - tuple `(x_val, y_val, val_sample_weights)` of Numpy arrays + - dataset iterator shuffle: Boolean (whether to shuffle the training data before each epoch) or str (for 'batch'). 'batch' is a special option for dealing with the @@ -1162,17 +1176,22 @@ class Model(Network): batch_size=batch_size) # Prepare validation data. if validation_data: - if len(validation_data) == 2: + if isinstance(validation_data, iterator_ops.Iterator): + val_x = validation_data + val_y = None + val_sample_weight = None + elif 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( - 'When passing validation_data, ' - 'it must contain 2 (x_val, y_val) ' - 'or 3 (x_val, y_val, val_sample_weights) ' - 'items, however it contains %d items' % len(validation_data)) + 'When passing a `validation_data` argument, ' + 'it must contain either 2 items (x_val, y_val), ' + 'or 3 items (x_val, y_val, val_sample_weights), ' + 'or alternatively it could be a dataset iterator. However we ' + 'received `validation_data=%s`' % validation_data) val_x, val_y, val_sample_weights = self._standardize_user_data( val_x, @@ -1246,22 +1265,26 @@ class Model(Network): Computation is done in batches. Arguments: - x: Numpy array of test data (if the model has a single input), - or list of Numpy arrays (if the model has multiple inputs). - If input layers in the model are named, you can also pass a - dictionary mapping input names to Numpy arrays. - `x` can be `None` (default) if feeding from - TensorFlow data tensors. - y: Numpy array of target (label) data - (if the model has a single output), - or list of Numpy arrays (if the model has multiple outputs). - If output layers in the model are named, you can also pass a - dictionary mapping output names to Numpy arrays. - `y` can be `None` (default) if feeding from - TensorFlow data tensors. + x: Input data. It could be: + - A Numpy array (or array-like), or a list of arrays + (in case the model has multiple inputs). + - A TensorFlow tensor, or a list of tensors + (in case the model has multiple inputs). + - A dict mapping input names to the corresponding array/tensors, + if the model has named inputs. + - A `tf.data` dataset iterator. + y: Target data. Like the input data `x`, + it could be either Numpy array(s) or TensorFlow tensor(s). + It should be consistent with `x` (you cannot have Numpy inputs and + tensor targets, or inversely). If `x` is a dataset iterator, + `y` should not be specified + (since targets will be obtained from the iterator). batch_size: Integer or `None`. - Number of samples per evaluation step. + Number of samples per gradient update. If unspecified, `batch_size` will default to 32. + Do not specify the `batch_size` is your data is in the + form of symbolic tensors or dataset iterators (since they generate + batches). verbose: 0 or 1. Verbosity mode. 0 = silent, 1 = progress bar. sample_weight: Optional Numpy array of weights for @@ -1319,9 +1342,13 @@ class Model(Network): Computation is done in batches. Arguments: - x: The input data, as a Numpy array - (or list of Numpy arrays if the model has multiple outputs). - batch_size: Integer. If unspecified, it will default to 32. + x: Input samples, as Numpy array(s) or tensor(s). + batch_size: Integer or `None`. + Number of samples per gradient update. + If unspecified, `batch_size` will default to 32. + Do not specify the `batch_size` is your data is in the + form of symbolic tensors or dataset iterators (since they generate + batches). verbose: Verbosity mode, 0 or 1. steps: Total number of steps (batches of samples) before declaring the prediction round finished. @@ -1352,20 +1379,24 @@ class Model(Network): 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): + def train_on_batch(self, x, y=None, sample_weight=None, class_weight=None): """Runs a single gradient update on a single batch of data. Arguments: - x: Numpy array of training data, - or list of Numpy arrays if the model has multiple inputs. - If all inputs in the model are named, - you can also pass a dictionary - mapping input names to Numpy arrays. - y: Numpy array of target data, - or list of Numpy arrays if the model has multiple outputs. - If all outputs in the model are named, - you can also pass a dictionary - mapping output names to Numpy arrays. + x: Input data. It could be: + - A Numpy array (or array-like), or a list of arrays + (in case the model has multiple inputs). + - A TensorFlow tensor, or a list of tensors + (in case the model has multiple inputs). + - A dict mapping input names to the corresponding array/tensors, + if the model has named inputs. + - A `tf.data` dataset iterator. + y: Target data. Like the input data `x`, + it could be either Numpy array(s) or TensorFlow tensor(s). + It should be consistent with `x` (you cannot have Numpy inputs and + tensor targets, or inversely). If `x` is a dataset iterator, + `y` should not be specified + (since targets will be obtained from the iterator). sample_weight: Optional array of the same length as x, containing weights to apply to the model's loss for each sample. In the case of temporal data, you can pass a 2D array @@ -1412,20 +1443,24 @@ class Model(Network): return outputs[0] return outputs - def test_on_batch(self, x, y, sample_weight=None): + def test_on_batch(self, x, y=None, sample_weight=None): """Test the model on a single batch of samples. Arguments: - x: Numpy array of test data, - or list of Numpy arrays if the model has multiple inputs. - If all inputs in the model are named, - you can also pass a dictionary - mapping input names to Numpy arrays. - y: Numpy array of target data, - or list of Numpy arrays if the model has multiple outputs. - If all outputs in the model are named, - you can also pass a dictionary - mapping output names to Numpy arrays. + x: Input data. It could be: + - A Numpy array (or array-like), or a list of arrays + (in case the model has multiple inputs). + - A TensorFlow tensor, or a list of tensors + (in case the model has multiple inputs). + - A dict mapping input names to the corresponding array/tensors, + if the model has named inputs. + - A `tf.data` dataset iterator. + y: Target data. Like the input data `x`, + it could be either Numpy array(s) or TensorFlow tensor(s). + It should be consistent with `x` (you cannot have Numpy inputs and + tensor targets, or inversely). If `x` is a dataset iterator, + `y` should not be specified + (since targets will be obtained from the iterator). sample_weight: Optional array of the same length as x, containing weights to apply to the model's loss for each sample. In the case of temporal data, you can pass a 2D array @@ -1465,7 +1500,7 @@ class Model(Network): """Returns predictions for a single batch of samples. Arguments: - x: Input samples, as a Numpy array. + x: Input samples, as Numpy array(s) or tensor(s). Returns: Numpy array(s) of predictions. diff --git a/tensorflow/python/keras/_impl/keras/engine/training_arrays.py b/tensorflow/python/keras/_impl/keras/engine/training_arrays.py index 18116e3a14d6b1365f1a9db1a23243cd07763a62..4164cae864c7f8aa8f59bb66d585aaa682fc8ff8 100644 --- a/tensorflow/python/keras/_impl/keras/engine/training_arrays.py +++ b/tensorflow/python/keras/_impl/keras/engine/training_arrays.py @@ -23,6 +23,7 @@ import copy import numpy as np +from tensorflow.python.framework import errors 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 @@ -30,6 +31,7 @@ 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 +from tensorflow.python.platform import tf_logging as logging try: from scipy.sparse import issparse # pylint: disable=g-import-not-at-top @@ -190,7 +192,15 @@ def fit_loop(model, batch_logs['batch'] = step_index batch_logs['size'] = 1 callbacks.on_batch_begin(step_index, batch_logs) - outs = f(ins) + try: + outs = f(ins) + except errors.OutOfRangeError: + logging.warning('Your dataset iterator ran out of data; ' + 'interrupting training. Make sure that your dataset ' + 'can generate at least `steps_per_epoch * epochs` ' + 'batches (in this case, %d batches).' % + steps_per_epoch * epochs) + break if not isinstance(outs, list): outs = [outs] diff --git a/tensorflow/python/keras/_impl/keras/engine/training_eager.py b/tensorflow/python/keras/_impl/keras/engine/training_eager.py index 4cdb5f108a05bb88ed328ca20351914160906e86..ad239d6151e02a0c1b2753908cc775db8498d244 100644 --- a/tensorflow/python/keras/_impl/keras/engine/training_eager.py +++ b/tensorflow/python/keras/_impl/keras/engine/training_eager.py @@ -100,7 +100,7 @@ def _eager_metrics_fn(model, outputs, targets): metric_names.append(metric_name) metric_results.append(backend.mean(metric_result)) - return metric_names, metric_results + return metric_results def _model_loss(model, inputs, targets, sample_weights=None, training=False): @@ -150,8 +150,13 @@ def _model_loss(model, inputs, targets, sample_weights=None, training=False): 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)) + targets[i], outs[i], weights, mask=mask) + # If the number of outputs is 1 then we don't append the loss metric + # associated with each model output. When there are multiple outputs + # associated with a model, each output's loss is calculated and returned + # as part of the loss_metrics. + if len(model.outputs) > 1: + loss_metrics.append(backend.mean(output_loss)) loss_weight = model.loss_weights_list[i] if total_loss is None: @@ -274,7 +279,7 @@ def train_on_batch(model, inputs, targets, sample_weights=None): model, inputs, targets, sample_weights=sample_weights, training=True) if not isinstance(outs, list): outs = [outs] - _, metrics_results = _eager_metrics_fn( + metrics_results = _eager_metrics_fn( model, outs, targets) if not isinstance(loss, list): loss = [loss] @@ -304,7 +309,7 @@ def test_on_batch(model, inputs, targets, sample_weights=None): model, inputs, targets, sample_weights=sample_weights, training=False) if not isinstance(outs, list): outs = [outs] - _, metrics_results = _eager_metrics_fn( + metrics_results = _eager_metrics_fn( model, outs, targets) if not isinstance(loss, list): loss = [loss] @@ -498,34 +503,12 @@ def fit_loop( 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) + 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 @@ -611,7 +594,7 @@ def test_loop(model, inputs, targets, targets_batch, sample_weights=sample_weights_batch, training=False) - _, metrics_results = _eager_metrics_fn(model, loss_outs, targets_batch) + 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): diff --git a/tensorflow/python/keras/_impl/keras/engine/training_eager_test.py b/tensorflow/python/keras/_impl/keras/engine/training_eager_test.py index 6cdb6b0753fce1bebec0060524e76d32929d0228..deaf1d13064e285e9b85befc699074cf33dd218f 100644 --- a/tensorflow/python/keras/_impl/keras/engine/training_eager_test.py +++ b/tensorflow/python/keras/_impl/keras/engine/training_eager_test.py @@ -21,6 +21,7 @@ from __future__ import print_function import numpy as np from tensorflow.python.framework import ops +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 @@ -211,7 +212,7 @@ class TrainingTest(test.TestCase): optimizer = RMSPropOptimizer(learning_rate=0.001) loss = 'mse' loss_weights = [1., 0.5] - metrics = ['mae'] + metrics = ['acc', 'mae'] model.compile( optimizer, loss, @@ -230,20 +231,20 @@ class TrainingTest(test.TestCase): [input_a_np, input_b_np], [output_d_np, output_e_np], batch_size=5, verbose=0) - self.assertEqual(len(out), 5) + self.assertEqual(len(out), 7) 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) + self.assertEqual(len(out), 7) 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) + self.assertEqual(len(out), 7) out = model.test_on_batch([input_a_np, input_b_np], [output_d_np, output_e_np]) - self.assertEqual(len(out), 5) + self.assertEqual(len(out), 7) # Test evaluate with dictionary inputs model.evaluate( @@ -624,6 +625,28 @@ class LossWeightingTest(test.TestCase): 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}) +class CorrectnessTest(test.TestCase): + + @tf_test_util.run_in_graph_and_eager_modes() + def test_loss_correctness(self): + # Test that training loss is the same in eager and graph + # (by comparing it to a reference value in a deterministic case) + model = keras.Sequential() + model.add(keras.layers.Dense(3, + activation='relu', + input_dim=4, + kernel_initializer='ones')) + model.add(keras.layers.Dense(2, + activation='softmax', + kernel_initializer='ones')) + model.compile(loss='sparse_categorical_crossentropy', + optimizer=RMSPropOptimizer(learning_rate=0.001)) + x = np.ones((100, 4)) + np.random.seed(123) + y = np.random.randint(0, 1, size=(100, 1)) + history = model.fit(x, y, epochs=1, batch_size=10) + self.assertEqual( + np.around(history.history['loss'][-1], decimals=4), 0.6173) if __name__ == '__main__': ops.enable_eager_execution() diff --git a/tensorflow/python/keras/_impl/keras/engine/training_test.py b/tensorflow/python/keras/_impl/keras/engine/training_test.py index 08fd26dd18d5bc1b171d780be133f02f51b9c248..58011a141268e588e6b746afca982f36f4d4d176 100644 --- a/tensorflow/python/keras/_impl/keras/engine/training_test.py +++ b/tensorflow/python/keras/_impl/keras/engine/training_test.py @@ -23,11 +23,17 @@ import unittest import numpy as np +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.framework import tensor_shape +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.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.ops import array_ops from tensorflow.python.platform import test +from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.training.rmsprop import RMSPropOptimizer try: import scipy.sparse as scipy_sparse # pylint: disable=g-import-not-at-top @@ -1140,6 +1146,21 @@ class TestTrainingWithDataTensors(test.TestCase): epochs=1, steps_per_epoch=2, verbose=0, validation_data=(inputs, targets), validation_steps=2) + # Test with dynamic shape + inputs = array_ops.placeholder_with_default( + np.zeros((2, 3)), shape=tensor_shape.TensorShape([None, 3])) + targets = array_ops.placeholder_with_default( + np.zeros((2, 4)), shape=tensor_shape.TensorShape([None, 4])) + self.assertEqual(inputs.shape[0].value, None) + model.fit(inputs, targets, epochs=1, steps_per_epoch=2, verbose=0) + model.evaluate(inputs, targets, steps=2, verbose=0) + model.predict(inputs, steps=2) + model.train_on_batch(inputs, targets) + model.test_on_batch(inputs, targets) + model.fit(inputs, targets, + epochs=1, steps_per_epoch=2, verbose=0, + validation_data=(inputs, targets), validation_steps=2) + def test_training_and_eval_methods_on_symbolic_tensors_multi_io(self): with self.test_session(): a = keras.layers.Input(shape=(3,), name='input_a') @@ -1667,15 +1688,101 @@ class TestTrainingWithDataTensors(test.TestCase): model.train_on_batch([input_a_np, input_b_np], [output_a_np, output_b_np]) + @tf_test_util.run_in_graph_and_eager_modes() + def test_metric_names_are_identical_in_graph_and_eager(self): + a = keras.layers.Input(shape=(3,), name='input_a') + b = keras.layers.Input(shape=(3,), name='input_b') -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] + 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', 'acc'] + model.compile(optimizer, loss, metrics=metrics, loss_weights=loss_weights) + reference_metric_names = ['loss', 'dense_loss', 'dropout_loss', + 'dense_mean_absolute_error', + 'dense_acc', + 'dropout_mean_absolute_error', + 'dropout_acc'] + self.assertEqual(reference_metric_names, model.metrics_names) + + +class TestTrainingWithDatasetIterators(test.TestCase): + + def test_training_and_eval_methods_on_iterators_single_io(self): + with self.test_session(): + x = keras.layers.Input(shape=(3,), name='input') + y = keras.layers.Dense(4, name='dense')(x) + model = keras.Model(x, y) + + optimizer = 'rmsprop' + loss = 'mse' + metrics = ['mae'] + model.compile(optimizer, loss, metrics=metrics) + inputs = np.zeros((10, 3)) + targets = np.zeros((10, 4)) + dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets)) + dataset = dataset.repeat(100) + dataset = dataset.batch(10) + iterator = dataset.make_one_shot_iterator() + + model.fit(iterator, epochs=1, steps_per_epoch=2, verbose=0) + model.evaluate(iterator, steps=2, verbose=0) + model.predict(iterator, steps=2) + model.train_on_batch(iterator) + model.test_on_batch(iterator) + # Test with validation data + model.fit(iterator, + epochs=1, steps_per_epoch=2, verbose=0, + validation_data=iterator, validation_steps=2) + # Test with validation split + with self.assertRaisesRegexp(ValueError, + 'you cannot use `validation_split`'): + model.fit(iterator, + epochs=1, steps_per_epoch=2, verbose=0, + validation_split=0.5, validation_steps=2) + + # Test invalid usage + with self.assertRaisesRegexp(ValueError, + 'Instead, pass an `Iterator`'): + model.fit(dataset, + epochs=1, steps_per_epoch=2, verbose=0) + with self.assertRaisesRegexp(ValueError, + 'you should not specify a target'): + model.fit(iterator, iterator, + epochs=1, steps_per_epoch=2, verbose=0) + + def test_iterators_running_out_of_data(self): + with self.test_session(): + x = keras.layers.Input(shape=(3,), name='input') + y = keras.layers.Dense(4, name='dense')(x) + model = keras.Model(x, y) + + optimizer = 'rmsprop' + loss = 'mse' + metrics = ['mae'] + model.compile(optimizer, loss, metrics=metrics) + + inputs = np.zeros((10, 3)) + targets = np.zeros((10, 4)) + dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets)) + dataset = dataset.repeat(2) + dataset = dataset.batch(10) + iterator = dataset.make_one_shot_iterator() + + with test.mock.patch.object(logging, 'warning') as mock_log: + model.fit(iterator, epochs=1, steps_per_epoch=3, verbose=0) + self.assertRegexpMatches( + str(mock_log.call_args), + 'dataset iterator ran out of data') + + +if __name__ == '__main__': test.main() diff --git a/tensorflow/python/keras/_impl/keras/engine/training_utils.py b/tensorflow/python/keras/_impl/keras/engine/training_utils.py index a3fc8ef2a0359c527a2757c1888d61822e35d7a9..662938f421b3a338d6720b2911347664711c1ac1 100644 --- a/tensorflow/python/keras/_impl/keras/engine/training_utils.py +++ b/tensorflow/python/keras/_impl/keras/engine/training_utils.py @@ -26,6 +26,7 @@ from tensorflow.python.eager import context 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 +from tensorflow.python.keras._impl.keras import metrics as metrics_module from tensorflow.python.ops import math_ops @@ -61,22 +62,21 @@ def check_num_samples(ins, 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.') - if has_symbolic_tensors(ins) and steps is None: - raise ValueError('If your data is in the form of symbolic tensors, ' - 'you should specify the `' + steps_name + '` argument ' - '(instead of the `batch_size` argument).') - if ins and hasattr(ins[0], 'shape'): - num_samples = int(ins[0].shape[0]) - elif steps is None: + if steps is not None and batch_size is not None: raise ValueError( - 'Either the input data should have ' - 'a defined shape, or ' + steps_name + ' should be specified.') - return num_samples + 'If ' + steps_name + ' is set, the `batch_size` must be None.') + + if not ins or has_symbolic_tensors(ins): + if steps is None: + raise ValueError('If your data is in the form of symbolic tensors, ' + 'you should specify the `' + steps_name + '` argument ' + '(instead of the `batch_size` argument, ' + 'because symbolic tensors are expected to produce ' + 'batches of input data).') + return None + if hasattr(ins[0], 'shape'): + return int(ins[0].shape[0]) + return None # Edge case where ins == [static_learning_phase] def standardize_single_array(x): @@ -553,3 +553,64 @@ def standardize_weights(y, def has_symbolic_tensors(ls): return (any(tensor_util.is_tensor(v) for v in ls) and not context.executing_eagerly()) + + +def populate_metric_names(model): + for i in range(len(model.outputs)): + metrics = model.nested_metrics[i] + for metric in metrics: + base_metric_name = get_base_metric_name(metric) + add_metric_name(model, base_metric_name, i) + + +def get_base_metric_name(metric, weighted=False): + """Returns the metric name given the metric function. + + Arguments: + metric: Metric function name or reference. + weighted: Boolean indicating if the metric for which we are adding + names is weighted. + + Returns: + a metric name. + """ + metric_name_prefix = 'weighted_' if weighted else '' + if metric in ('accuracy', 'acc', 'crossentropy', 'ce'): + if metric in ('accuracy', 'acc'): + suffix = 'acc' + elif metric in ('crossentropy', 'ce'): + suffix = 'ce' + metric_name = metric_name_prefix + suffix + else: + metric_fn = metrics_module.get(metric) + # 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 + + return metric_name + + +def add_metric_name(model, metric_name, index): + """Makes the metric name unique and adds it to the model's metric name list. + + If there are multiple outputs for which the metrics are calculated, the + metric names have to be made unique by appending an integer. + + Arguments: + model: Model to which we are adding metric names. + metric_name: Metric name that corresponds to the metric specified by the + user. For example: 'acc' + index: The index of the model output for which the metric name is being + added. + """ + if len(model.output_names) > 1: + metric_name = '%s_%s' % (model.output_names[index], metric_name) + j = 1 + base_metric_name = metric_name + while metric_name in model.metrics_names: + metric_name = '%s_%d' % (base_metric_name, j) + j += 1 + model.metrics_names.append(metric_name) diff --git a/tensorflow/python/keras/_impl/keras/estimator.py b/tensorflow/python/keras/_impl/keras/estimator.py index b922a6c68399e6c0c43b98c54ff0f550326c2199..c3c3fceb454773f18400bd31fe8b649c557fc013 100644 --- a/tensorflow/python/keras/_impl/keras/estimator.py +++ b/tensorflow/python/keras/_impl/keras/estimator.py @@ -29,12 +29,14 @@ 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.framework import tensor_util 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 check_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import metrics as metrics_module from tensorflow.python.ops import variables as variables_module @@ -55,6 +57,17 @@ def _cast_tensor_to_floatx(x): return math_ops.cast(x, K.floatx()) +def _convert_tensor(x): + """Create or cast tensor if needed.""" + if not tensor_util.is_tensor(x): + # x is a numpy array + x = sparse_tensor_lib.convert_to_tensor_or_sparse_tensor(x) + if check_ops.is_numeric_tensor(x): + # is_numeric_tensor returns False if provided with a numpy array + x = _cast_tensor_to_floatx(x) + return x + + def _any_variable_initalized(): """Check if any variable has been initialized in the Keras model. @@ -86,7 +99,7 @@ def _create_ordered_io(keras_model, estimator_io, is_input=True): 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] + return [_convert_tensor(x) for x in estimator_io] elif isinstance(estimator_io, dict): if is_input: if keras_model._is_graph_network: @@ -108,12 +121,12 @@ def _create_ordered_io(keras_model, estimator_io, is_input=True): '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]) + tensors = [_convert_tensor(estimator_io[io_name]) for io_name in keras_io_names] return tensors else: # Plain array. - return _cast_tensor_to_floatx(estimator_io) + return _convert_tensor(estimator_io) def _in_place_subclassed_model_reset(model): @@ -274,8 +287,7 @@ def _clone_and_build_model(mode, is_input=False) else: target_tensors = [ - _cast_tensor_to_floatx( - sparse_tensor_lib.convert_to_tensor_or_sparse_tensor(labels)) + _convert_tensor(labels) ] if keras_model._is_graph_network: diff --git a/tensorflow/python/keras/_impl/keras/estimator_test.py b/tensorflow/python/keras/_impl/keras/estimator_test.py index 653cdc01e245666f2b8fc1e742bb8ef71512d279..80fa87d0410871c30b8a7c46e7ff02bc81c96f3b 100644 --- a/tensorflow/python/keras/_impl/keras/estimator_test.py +++ b/tensorflow/python/keras/_impl/keras/estimator_test.py @@ -30,6 +30,7 @@ from tensorflow.python.estimator.inputs import numpy_io from tensorflow.python.framework import ops from tensorflow.python.framework import test_util from tensorflow.python.keras._impl import keras +from tensorflow.python.keras._impl.keras import backend as K from tensorflow.python.keras._impl.keras import testing_utils from tensorflow.python.keras._impl.keras.applications import mobilenet from tensorflow.python.keras._impl.keras.optimizers import SGD @@ -142,16 +143,20 @@ def randomize_io_type(array, name): def multi_inputs_multi_outputs_model(): - # test multi-input layer a = keras.layers.Input(shape=(16,), name='input_a') b = keras.layers.Input(shape=(16,), name='input_b') + m = keras.layers.Input(shape=(8,), dtype='bool', name='input_m') dense = keras.layers.Dense(8, name='dense_1') + a_2 = dense(a) + # Apply a mask + s_2 = keras.layers.Lambda(lambda k: + K.switch(k[0], k[1], K.zeros_like(k[1])))([m, a_2]) b_2 = dense(b) - merged = keras.layers.concatenate([a_2, b_2], name='merge') + merged = keras.layers.concatenate([s_2, b_2], name='merge') c = keras.layers.Dense(3, activation='softmax', name='dense_2')(merged) d = keras.layers.Dense(2, activation='softmax', name='dense_3')(merged) - model = keras.models.Model(inputs=[a, b], outputs=[c, d]) + model = keras.models.Model(inputs=[a, b, m], outputs=[c, d]) model.compile( loss='categorical_crossentropy', optimizer='rmsprop', @@ -352,18 +357,27 @@ class TestKerasEstimator(test_util.TensorFlowTestCase): test_samples=50, input_shape=(16,), num_classes=2) + np.random.seed(_RANDOM_SEED) + (input_m_train, _), (input_m_test, _) = testing_utils.get_test_data( + train_samples=_TRAIN_SIZE, + test_samples=50, + input_shape=(8,), + num_classes=2) + c_train = keras.utils.to_categorical(c_train) c_test = keras.utils.to_categorical(c_test) d_train = keras.utils.to_categorical(d_train) d_test = keras.utils.to_categorical(d_test) def train_input_fn(): - input_dict = {'input_a': a_train, 'input_b': b_train} + input_dict = {'input_a': a_train, 'input_b': b_train, + 'input_m': input_m_train > 0} output_dict = {'dense_2': c_train, 'dense_3': d_train} return input_dict, output_dict def eval_input_fn(): - input_dict = {'input_a': a_test, 'input_b': b_test} + input_dict = {'input_a': a_test, 'input_b': b_test, + 'input_m': input_m_test > 0} output_dict = {'dense_2': c_test, 'dense_3': d_test} return input_dict, output_dict diff --git a/tensorflow/python/keras/_impl/keras/initializers.py b/tensorflow/python/keras/_impl/keras/initializers.py index 300bed5e1437074d010760c427c14f68e58ac363..ecb71d00e2c78ced6095aaa3a0180b454b04917a 100644 --- a/tensorflow/python/keras/_impl/keras/initializers.py +++ b/tensorflow/python/keras/_impl/keras/initializers.py @@ -201,6 +201,8 @@ def deserialize(config, custom_objects=None): @tf_export('keras.initializers.get') def get(identifier): + if identifier is None: + return None if isinstance(identifier, dict): return deserialize(identifier) elif isinstance(identifier, six.string_types): diff --git a/tensorflow/python/keras/_impl/keras/integration_test.py b/tensorflow/python/keras/_impl/keras/integration_test.py index 53dd56ba9a5ba7e973ddca227d0d06a99c8ecddd..43aff67ef93c8ec495beafdd17c5557b6398671f 100644 --- a/tensorflow/python/keras/_impl/keras/integration_test.py +++ b/tensorflow/python/keras/_impl/keras/integration_test.py @@ -29,16 +29,15 @@ from tensorflow.python.platform import test class KerasIntegrationTest(test.TestCase): - def test_vector_classification_declarative(self): + def test_vector_classification_sequential(self): with self.test_session(): np.random.seed(1337) - (x_train, y_train), (x_test, y_test) = testing_utils.get_test_data( - train_samples=200, - test_samples=100, + (x_train, y_train), _ = testing_utils.get_test_data( + train_samples=100, + test_samples=0, input_shape=(10,), num_classes=2) y_train = keras.utils.to_categorical(y_train) - y_test = keras.utils.to_categorical(y_test) model = keras.models.Sequential([ keras.layers.Dense(16, @@ -48,23 +47,22 @@ class KerasIntegrationTest(test.TestCase): keras.layers.Dense(y_train.shape[-1], activation='softmax') ]) model.compile(loss='categorical_crossentropy', - optimizer='rmsprop', + optimizer=keras.optimizers.Adam(lr=0.1), metrics=['accuracy']) history = model.fit(x_train, y_train, epochs=10, batch_size=16, - validation_data=(x_test, y_test), + validation_data=(x_train, y_train), verbose=2) - self.assertGreater(history.history['val_acc'][-1], 0.85) + self.assertGreater(history.history['val_acc'][-1], 0.7) def test_vector_classification_functional(self): with self.test_session(): np.random.seed(1337) - (x_train, y_train), (x_test, y_test) = testing_utils.get_test_data( - train_samples=200, - test_samples=100, - input_shape=(10,), + (x_train, y_train), _ = testing_utils.get_test_data( + train_samples=100, + test_samples=0, + input_shape=(20,), num_classes=2) y_train = keras.utils.to_categorical(y_train) - y_test = keras.utils.to_categorical(y_test) inputs = keras.layers.Input(shape=x_train.shape[1:]) x = keras.layers.Dense(16, activation='relu')(inputs) @@ -73,77 +71,78 @@ class KerasIntegrationTest(test.TestCase): model = keras.models.Model(inputs, outputs) model.compile(loss='categorical_crossentropy', - optimizer='rmsprop', + optimizer=keras.optimizers.Adam(lr=0.1), metrics=['accuracy']) history = model.fit(x_train, y_train, epochs=10, batch_size=16, - validation_data=(x_test, y_test), + validation_data=(x_train, y_train), verbose=2) - self.assertGreater(history.history['val_acc'][-1], 0.85) + self.assertGreater(history.history['val_acc'][-1], 0.7) - def test_temporal_classification_declarative(self): + def test_temporal_classification_sequential(self): with self.test_session(): - np.random.seed(1336) - (x_train, y_train), (x_test, y_test) = testing_utils.get_test_data( - train_samples=200, - test_samples=100, - input_shape=(4, 8), + np.random.seed(1337) + (x_train, y_train), _ = testing_utils.get_test_data( + train_samples=100, + test_samples=0, + input_shape=(4, 10), num_classes=2) y_train = keras.utils.to_categorical(y_train) - y_test = keras.utils.to_categorical(y_test) model = keras.models.Sequential() model.add(keras.layers.LSTM(5, return_sequences=True, input_shape=x_train.shape[1:])) model.add(keras.layers.GRU(y_train.shape[-1], activation='softmax')) model.compile(loss='categorical_crossentropy', - optimizer='adam', + optimizer=keras.optimizers.Adam(lr=0.1), metrics=['accuracy']) - history = model.fit(x_train, y_train, epochs=10, batch_size=16, - validation_data=(x_test, y_test), + history = model.fit(x_train, y_train, epochs=15, batch_size=16, + validation_data=(x_train, y_train), verbose=2) - self.assertGreater(history.history['val_acc'][-1], 0.80) + self.assertGreater(history.history['val_acc'][-1], 0.7) - def test_image_classification_declarative(self): + def test_image_classification_sequential(self): with self.test_session(): np.random.seed(1337) - (x_train, y_train), (x_test, y_test) = testing_utils.get_test_data( - train_samples=200, - test_samples=100, - input_shape=(8, 8, 3), + (x_train, y_train), _ = testing_utils.get_test_data( + train_samples=100, + test_samples=0, + input_shape=(12, 12, 3), num_classes=2) y_train = keras.utils.to_categorical(y_train) - y_test = keras.utils.to_categorical(y_test) model = keras.models.Sequential() model.add(keras.layers.Conv2D( - 8, 3, + 4, 3, + padding='same', activation='relu', input_shape=x_train.shape[1:])) - model.add(keras.layers.BatchNormalization()) model.add(keras.layers.Conv2D( 8, 3, padding='same', activation='relu')) - model.add(keras.layers.GlobalMaxPooling2D()) + model.add(keras.layers.Conv2D( + 16, 3, + padding='same', + activation='relu')) + model.add(keras.layers.Flatten()) model.add(keras.layers.Dense(y_train.shape[-1], activation='softmax')) model.compile(loss='categorical_crossentropy', - optimizer='adam', + optimizer=keras.optimizers.SGD(lr=0.01, momentum=0.8), metrics=['accuracy']) history = model.fit(x_train, y_train, epochs=10, batch_size=16, - validation_data=(x_test, y_test), + validation_data=(x_train, y_train), verbose=2) - self.assertGreater(history.history['val_acc'][-1], 0.85) + self.assertGreater(history.history['val_acc'][-1], 0.7) def test_video_classification_functional(self): with self.test_session(): np.random.seed(1337) - (x_train, y_train), (x_test, y_test) = testing_utils.get_test_data( - train_samples=200, - test_samples=100, + (x_train, y_train), _ = testing_utils.get_test_data( + train_samples=100, + test_samples=0, input_shape=(4, 8, 8, 3), num_classes=3) y_train = keras.utils.to_categorical(y_train) - y_test = keras.utils.to_categorical(y_test) inputs = keras.layers.Input(shape=x_train.shape[1:]) x = keras.layers.TimeDistributed( @@ -159,22 +158,21 @@ class KerasIntegrationTest(test.TestCase): optimizer=keras.optimizers.SGD(lr=0.01, momentum=0.8), metrics=['accuracy']) history = model.fit(x_train, y_train, epochs=10, batch_size=16, - validation_data=(x_test, y_test), + validation_data=(x_train, y_train), verbose=2) - self.assertGreater(history.history['val_acc'][-1], 0.70) + self.assertGreater(history.history['val_acc'][-1], 0.7) def test_vector_classification_shared_sequential(self): # Test that Sequential models that feature internal updates # and internal losses can be shared. with self.test_session(): np.random.seed(1337) - (x_train, y_train), (x_test, y_test) = testing_utils.get_test_data( - train_samples=200, - test_samples=100, + (x_train, y_train), _ = testing_utils.get_test_data( + train_samples=100, + test_samples=0, input_shape=(10,), num_classes=2) y_train = keras.utils.to_categorical(y_train) - y_test = keras.utils.to_categorical(y_test) base_model = keras.models.Sequential([ keras.layers.Dense(16, @@ -189,27 +187,26 @@ class KerasIntegrationTest(test.TestCase): y = keras.layers.Dense(y_train.shape[-1], activation='softmax')(y) model = keras.models.Model(x, y) model.compile(loss='categorical_crossentropy', - optimizer='rmsprop', + optimizer=keras.optimizers.Adam(lr=0.1), metrics=['accuracy']) self.assertEqual(len(model.losses), 2) self.assertEqual(len(model.updates), 2) history = model.fit(x_train, y_train, epochs=10, batch_size=16, - validation_data=(x_test, y_test), + validation_data=(x_train, y_train), verbose=2) - self.assertGreater(history.history['val_acc'][-1], 0.84) + self.assertGreater(history.history['val_acc'][-1], 0.7) def test_vector_classification_shared_model(self): # Test that functional models that feature internal updates # and internal losses can be shared. with self.test_session(): np.random.seed(1337) - (x_train, y_train), (x_test, y_test) = testing_utils.get_test_data( - train_samples=200, - test_samples=100, + (x_train, y_train), _ = testing_utils.get_test_data( + train_samples=100, + test_samples=0, input_shape=(10,), num_classes=2) y_train = keras.utils.to_categorical(y_train) - y_test = keras.utils.to_categorical(y_test) inputs = keras.layers.Input(x_train.shape[1:]) x = keras.layers.Dense(16, @@ -225,12 +222,12 @@ class KerasIntegrationTest(test.TestCase): y = keras.layers.Dense(y_train.shape[-1], activation='softmax')(y) model = keras.models.Model(x, y) model.compile(loss='categorical_crossentropy', - optimizer='rmsprop', + optimizer=keras.optimizers.Adam(lr=0.1), metrics=['accuracy']) history = model.fit(x_train, y_train, epochs=10, batch_size=16, - validation_data=(x_test, y_test), + validation_data=(x_train, y_train), verbose=2) - self.assertGreater(history.history['val_acc'][-1], 0.85) + self.assertGreater(history.history['val_acc'][-1], 0.7) def test_embedding_with_clipnorm(self): with self.test_session(): @@ -242,9 +239,9 @@ class KerasIntegrationTest(test.TestCase): def test_using_tf_layers_in_keras_sequential_model(self): with self.test_session(): np.random.seed(1337) - (x_train, y_train), (x_test, y_test) = testing_utils.get_test_data( - train_samples=200, - test_samples=100, + (x_train, y_train), _ = testing_utils.get_test_data( + train_samples=100, + test_samples=0, input_shape=(10,), num_classes=2) @@ -254,25 +251,23 @@ class KerasIntegrationTest(test.TestCase): model.summary() y_train = keras.utils.to_categorical(y_train) - y_test = keras.utils.to_categorical(y_test) model.compile(loss='categorical_crossentropy', - optimizer='adam', + optimizer=keras.optimizers.Adam(lr=0.1), metrics=['accuracy']) history = model.fit(x_train, y_train, epochs=10, batch_size=16, - validation_data=(x_test, y_test), + validation_data=(x_train, y_train), verbose=0) - self.assertGreater(history.history['val_acc'][-1], 0.85) + self.assertGreater(history.history['val_acc'][-1], 0.7) def test_using_tf_layers_in_keras_functional_model(self): with self.test_session(): np.random.seed(1337) - (x_train, y_train), (x_test, y_test) = testing_utils.get_test_data( - train_samples=200, - test_samples=100, + (x_train, y_train), _ = testing_utils.get_test_data( + train_samples=100, + test_samples=0, input_shape=(10,), num_classes=2) y_train = keras.utils.to_categorical(y_train) - y_test = keras.utils.to_categorical(y_test) inputs = keras.Input(shape=(10,)) x = tf_core_layers.Dense(32, activation=nn.relu)(inputs) @@ -281,12 +276,12 @@ class KerasIntegrationTest(test.TestCase): model.summary() model.compile(loss='categorical_crossentropy', - optimizer='adam', + optimizer=keras.optimizers.Adam(lr=0.1), metrics=['accuracy']) history = model.fit(x_train, y_train, epochs=10, batch_size=16, - validation_data=(x_test, y_test), + validation_data=(x_train, y_train), verbose=0) - self.assertGreater(history.history['val_acc'][-1], 0.85) + self.assertGreater(history.history['val_acc'][-1], 0.7) if __name__ == '__main__': diff --git a/tensorflow/python/keras/_impl/keras/layers/convolutional.py b/tensorflow/python/keras/_impl/keras/layers/convolutional.py index 7cdebc6aa4f8460d14bef3edb9edf7e4db5066c5..12b965587f5695fe410d2363956f54386c9fa8cf 100644 --- a/tensorflow/python/keras/_impl/keras/layers/convolutional.py +++ b/tensorflow/python/keras/_impl/keras/layers/convolutional.py @@ -19,9 +19,10 @@ 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 activations -from tensorflow.python.keras._impl.keras import backend as K +from tensorflow.python.keras._impl.keras import backend from tensorflow.python.keras._impl.keras import constraints from tensorflow.python.keras._impl.keras import initializers from tensorflow.python.keras._impl.keras import regularizers @@ -38,12 +39,232 @@ from tensorflow.python.keras._impl.keras.layers.pooling import MaxPooling2D 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.ops import array_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(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. If `use_bias` is True (and a `bias_initializer` is provided), + a bias vector is created and added to the outputs. Finally, if + `activation` is not `None`, it is applied to the outputs as well. + + 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. + use_bias: Boolean, whether the layer uses a bias. + kernel_initializer: An initializer for the convolution kernel. + bias_initializer: An initializer for the bias vector. If None, the default + initializer will be used. + kernel_regularizer: Optional regularizer for the convolution kernel. + bias_regularizer: Optional regularizer for the bias vector. + activity_regularizer: Optional regularizer function for the output. + kernel_constraint: Optional projection function to be applied to the + kernel after being updated by an `Optimizer` (e.g. used to implement + norm constraints or value constraints for layer weights). 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. + bias_constraint: Optional projection function to be applied to the + bias after being updated by an `Optimizer`. + trainable: Boolean, if `True` also add variables to the graph collection + `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`). + name: A string, the name of the layer. + """ + + def __init__(self, rank, + filters, + kernel_size, + strides=1, + padding='valid', + data_format=None, + dilation_rate=1, + activation=None, + use_bias=True, + kernel_initializer='glorot_uniform', + bias_initializer='zeros', + kernel_regularizer=None, + bias_regularizer=None, + activity_regularizer=None, + kernel_constraint=None, + bias_constraint=None, + trainable=True, + name=None, + **kwargs): + super(Conv, self).__init__( + trainable=trainable, + name=name, + activity_regularizer=regularizers.get(activity_regularizer), + **kwargs) + self.rank = rank + self.filters = filters + self.kernel_size = conv_utils.normalize_tuple( + kernel_size, rank, 'kernel_size') + self.strides = conv_utils.normalize_tuple(strides, rank, 'strides') + self.padding = conv_utils.normalize_padding(padding) + self.data_format = conv_utils.normalize_data_format(data_format) + self.dilation_rate = conv_utils.normalize_tuple( + dilation_rate, rank, 'dilation_rate') + self.activation = activations.get(activation) + self.use_bias = use_bias + self.kernel_initializer = initializers.get(kernel_initializer) + self.bias_initializer = initializers.get(bias_initializer) + self.kernel_regularizer = regularizers.get(kernel_regularizer) + self.bias_regularizer = regularizers.get(bias_regularizer) + self.kernel_constraint = constraints.get(kernel_constraint) + self.bias_constraint = constraints.get(bias_constraint) + self.input_spec = InputSpec(ndim=self.rank + 2) + + 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 = int(input_shape[channel_axis]) + kernel_shape = self.kernel_size + (input_dim, self.filters) + + self.kernel = self.add_variable(name='kernel', + shape=kernel_shape, + initializer=self.kernel_initializer, + regularizer=self.kernel_regularizer, + constraint=self.kernel_constraint, + trainable=True, + dtype=self.dtype) + if self.use_bias: + self.bias = self.add_variable(name='bias', + shape=(self.filters,), + initializer=self.bias_initializer, + regularizer=self.bias_regularizer, + constraint=self.bias_constraint, + trainable=True, + dtype=self.dtype) + else: + self.bias = None + self.input_spec = InputSpec(ndim=self.rank + 2, + axes={channel_axis: input_dim}) + self._convolution_op = nn_ops.Convolution( + input_shape, + filter_shape=self.kernel.get_shape(), + dilation_rate=self.dilation_rate, + strides=self.strides, + padding=self.padding.upper(), + data_format=conv_utils.convert_data_format(self.data_format, + self.rank + 2)) + self.built = True + + def call(self, inputs): + outputs = self._convolution_op(inputs, self.kernel) + + if self.use_bias: + 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, (1, self.filters, 1)) + outputs += bias + if self.rank == 2: + outputs = nn.bias_add(outputs, self.bias, 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() + if outputs_shape[0] is None: + outputs_shape[0] = -1 + 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, data_format='NCHW') + outputs = array_ops.reshape(outputs_4d, outputs_shape) + else: + outputs = nn.bias_add(outputs, self.bias, data_format='NHWC') + + if self.activation is not None: + return self.activation(outputs) + return outputs + + 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 = conv_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 = conv_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) + + def get_config(self): + config = { + 'filters': self.filters, + 'kernel_size': self.kernel_size, + 'strides': self.strides, + 'padding': self.padding, + 'data_format': self.data_format, + 'dilation_rate': self.dilation_rate, + 'activation': activations.serialize(self.activation), + 'use_bias': self.use_bias, + 'kernel_initializer': initializers.serialize(self.kernel_initializer), + 'bias_initializer': initializers.serialize(self.bias_initializer), + 'kernel_regularizer': regularizers.serialize(self.kernel_regularizer), + 'bias_regularizer': regularizers.serialize(self.bias_regularizer), + 'activity_regularizer': + regularizers.serialize(self.activity_regularizer), + 'kernel_constraint': constraints.serialize(self.kernel_constraint), + 'bias_constraint': constraints.serialize(self.bias_constraint) + } + base_config = super(Conv, self).get_config() + return dict(list(base_config.items()) + list(config.items())) + + @tf_export('keras.layers.Conv1D', 'keras.layers.Convolution1D') -class Conv1D(tf_convolutional_layers.Conv1D, Layer): +class Conv1D(Conv): """1D convolution layer (e.g. temporal convolution). This layer creates a convolution kernel that is convolved @@ -74,6 +295,8 @@ class Conv1D(tf_convolutional_layers.Conv1D, Layer): where the model should not violate the temporal order. See [WaveNet: A Generative Model for Raw Audio, section 2.1](https://arxiv.org/abs/1609.03499). + data_format: A string, + one of `channels_last` (default) or `channels_first`. 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 @@ -105,6 +328,7 @@ class Conv1D(tf_convolutional_layers.Conv1D, Layer): kernel_size, strides=1, padding='valid', + data_format='channels_last', dilation_rate=1, activation=None, use_bias=True, @@ -117,11 +341,12 @@ class Conv1D(tf_convolutional_layers.Conv1D, Layer): bias_constraint=None, **kwargs): super(Conv1D, self).__init__( + rank=1, filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, - data_format='channels_last', + data_format=data_format, dilation_rate=dilation_rate, activation=activations.get(activation), use_bias=use_bias, @@ -134,30 +359,9 @@ class Conv1D(tf_convolutional_layers.Conv1D, Layer): bias_constraint=constraints.get(bias_constraint), **kwargs) - def get_config(self): - config = { - 'filters': self.filters, - 'kernel_size': self.kernel_size, - 'strides': self.strides, - 'padding': self.padding, - 'dilation_rate': self.dilation_rate, - 'activation': activations.serialize(self.activation), - 'use_bias': self.use_bias, - 'kernel_initializer': initializers.serialize(self.kernel_initializer), - 'bias_initializer': initializers.serialize(self.bias_initializer), - 'kernel_regularizer': regularizers.serialize(self.kernel_regularizer), - 'bias_regularizer': regularizers.serialize(self.bias_regularizer), - 'activity_regularizer': - regularizers.serialize(self.activity_regularizer), - 'kernel_constraint': constraints.serialize(self.kernel_constraint), - 'bias_constraint': constraints.serialize(self.bias_constraint) - } - base_config = super(Conv1D, self).get_config() - return dict(list(base_config.items()) + list(config.items())) - @tf_export('keras.layers.Conv2D', 'keras.layers.Convolution2D') -class Conv2D(tf_convolutional_layers.Conv2D, Layer): +class Conv2D(Conv): """2D convolution layer (e.g. spatial convolution over images). This layer creates a convolution kernel that is convolved @@ -247,9 +451,8 @@ class Conv2D(tf_convolutional_layers.Conv2D, Layer): kernel_constraint=None, bias_constraint=None, **kwargs): - if data_format is None: - data_format = K.image_data_format() super(Conv2D, self).__init__( + rank=2, filters=filters, kernel_size=kernel_size, strides=strides, @@ -267,31 +470,9 @@ class Conv2D(tf_convolutional_layers.Conv2D, Layer): bias_constraint=constraints.get(bias_constraint), **kwargs) - def get_config(self): - config = { - 'filters': self.filters, - 'kernel_size': self.kernel_size, - 'strides': self.strides, - 'padding': self.padding, - 'data_format': self.data_format, - 'dilation_rate': self.dilation_rate, - 'activation': activations.serialize(self.activation), - 'use_bias': self.use_bias, - 'kernel_initializer': initializers.serialize(self.kernel_initializer), - 'bias_initializer': initializers.serialize(self.bias_initializer), - 'kernel_regularizer': regularizers.serialize(self.kernel_regularizer), - 'bias_regularizer': regularizers.serialize(self.bias_regularizer), - 'activity_regularizer': - regularizers.serialize(self.activity_regularizer), - 'kernel_constraint': constraints.serialize(self.kernel_constraint), - 'bias_constraint': constraints.serialize(self.bias_constraint) - } - base_config = super(Conv2D, self).get_config() - return dict(list(base_config.items()) + list(config.items())) - @tf_export('keras.layers.Conv3D', 'keras.layers.Convolution3D') -class Conv3D(tf_convolutional_layers.Conv3D, Layer): +class Conv3D(Conv): """3D convolution layer (e.g. spatial convolution over volumes). This layer creates a convolution kernel that is convolved @@ -388,9 +569,8 @@ class Conv3D(tf_convolutional_layers.Conv3D, Layer): kernel_constraint=None, bias_constraint=None, **kwargs): - if data_format is None: - data_format = K.image_data_format() super(Conv3D, self).__init__( + rank=3, filters=filters, kernel_size=kernel_size, strides=strides, @@ -408,32 +588,10 @@ class Conv3D(tf_convolutional_layers.Conv3D, Layer): bias_constraint=constraints.get(bias_constraint), **kwargs) - def get_config(self): - config = { - 'filters': self.filters, - 'kernel_size': self.kernel_size, - 'strides': self.strides, - 'padding': self.padding, - 'data_format': self.data_format, - 'dilation_rate': self.dilation_rate, - 'activation': activations.serialize(self.activation), - 'use_bias': self.use_bias, - 'kernel_initializer': initializers.serialize(self.kernel_initializer), - 'bias_initializer': initializers.serialize(self.bias_initializer), - 'kernel_regularizer': regularizers.serialize(self.kernel_regularizer), - 'bias_regularizer': regularizers.serialize(self.bias_regularizer), - 'activity_regularizer': - regularizers.serialize(self.activity_regularizer), - 'kernel_constraint': constraints.serialize(self.kernel_constraint), - 'bias_constraint': constraints.serialize(self.bias_constraint) - } - base_config = super(Conv3D, self).get_config() - return dict(list(base_config.items()) + list(config.items())) - @tf_export('keras.layers.Conv2DTranspose', 'keras.layers.Convolution2DTranspose') -class Conv2DTranspose(tf_convolutional_layers.Conv2DTranspose, Layer): +class Conv2DTranspose(Conv2D): """Transposed convolution layer (sometimes called Deconvolution). The need for transposed convolutions generally arises @@ -529,8 +687,6 @@ class Conv2DTranspose(tf_convolutional_layers.Conv2DTranspose, Layer): kernel_constraint=None, bias_constraint=None, **kwargs): - if data_format is None: - data_format = K.image_data_format() super(Conv2DTranspose, self).__init__( filters=filters, kernel_size=kernel_size, @@ -548,31 +704,124 @@ class Conv2DTranspose(tf_convolutional_layers.Conv2DTranspose, Layer): bias_constraint=constraints.get(bias_constraint), **kwargs) - def get_config(self): - config = { - 'filters': self.filters, - 'kernel_size': self.kernel_size, - 'strides': self.strides, - 'padding': self.padding, - 'data_format': self.data_format, - 'activation': activations.serialize(self.activation), - 'use_bias': self.use_bias, - 'kernel_initializer': initializers.serialize(self.kernel_initializer), - 'bias_initializer': initializers.serialize(self.bias_initializer), - 'kernel_regularizer': regularizers.serialize(self.kernel_regularizer), - 'bias_regularizer': regularizers.serialize(self.bias_regularizer), - 'activity_regularizer': - regularizers.serialize(self.activity_regularizer), - 'kernel_constraint': constraints.serialize(self.kernel_constraint), - 'bias_constraint': constraints.serialize(self.bias_constraint) - } - base_config = super(Conv2DTranspose, self).get_config() - return dict(list(base_config.items()) + list(config.items())) + def build(self, input_shape): + input_shape = tensor_shape.TensorShape(input_shape) + if len(input_shape) != 4: + raise ValueError('Inputs should have rank 4. Received input shape: ' + + str(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 = int(input_shape[channel_axis]) + self.input_spec = InputSpec(ndim=4, axes={channel_axis: input_dim}) + kernel_shape = self.kernel_size + (self.filters, input_dim) + + self.kernel = self.add_variable(name='kernel', + shape=kernel_shape, + initializer=self.kernel_initializer, + regularizer=self.kernel_regularizer, + constraint=self.kernel_constraint, + trainable=True, + dtype=self.dtype) + if self.use_bias: + self.bias = self.add_variable(name='bias', + shape=(self.filters,), + initializer=self.bias_initializer, + regularizer=self.bias_regularizer, + constraint=self.bias_constraint, + trainable=True, + dtype=self.dtype) + else: + self.bias = None + self.built = True + + def call(self, inputs): + inputs_shape = array_ops.shape(inputs) + batch_size = inputs_shape[0] + if self.data_format == 'channels_first': + c_axis, h_axis, w_axis = 1, 2, 3 + else: + c_axis, h_axis, w_axis = 3, 1, 2 + + height, width = inputs_shape[h_axis], inputs_shape[w_axis] + kernel_h, kernel_w = self.kernel_size + stride_h, stride_w = self.strides + + # Infer the dynamic output shape: + out_height = conv_utils.deconv_output_length(height, + kernel_h, + self.padding, + stride_h) + out_width = conv_utils.deconv_output_length(width, + kernel_w, + self.padding, + stride_w) + if self.data_format == 'channels_first': + output_shape = (batch_size, self.filters, out_height, out_width) + strides = (1, 1, stride_h, stride_w) + else: + output_shape = (batch_size, out_height, out_width, self.filters) + strides = (1, stride_h, stride_w, 1) + + output_shape_tensor = array_ops.stack(output_shape) + outputs = nn.conv2d_transpose( + inputs, + self.kernel, + output_shape_tensor, + strides, + padding=self.padding.upper(), + data_format=conv_utils.convert_data_format(self.data_format, ndim=4)) + + if not context.executing_eagerly(): + # Infer the static output shape: + out_shape = inputs.get_shape().as_list() + out_shape[c_axis] = self.filters + out_shape[h_axis] = conv_utils.deconv_output_length(out_shape[h_axis], + kernel_h, + self.padding, + stride_h) + out_shape[w_axis] = conv_utils.deconv_output_length(out_shape[w_axis], + kernel_w, + self.padding, + stride_w) + outputs.set_shape(out_shape) + + if self.use_bias: + outputs = nn.bias_add( + outputs, + self.bias, + data_format=conv_utils.convert_data_format(self.data_format, ndim=4)) + + if self.activation is not None: + return self.activation(outputs) + return outputs + + def compute_output_shape(self, input_shape): + input_shape = tensor_shape.TensorShape(input_shape).as_list() + output_shape = list(input_shape) + if self.data_format == 'channels_first': + c_axis, h_axis, w_axis = 1, 2, 3 + else: + c_axis, h_axis, w_axis = 3, 1, 2 + + kernel_h, kernel_w = self.kernel_size + stride_h, stride_w = self.strides + + output_shape[c_axis] = self.filters + output_shape[h_axis] = conv_utils.deconv_output_length( + output_shape[h_axis], kernel_h, self.padding, stride_h) + output_shape[w_axis] = conv_utils.deconv_output_length( + output_shape[w_axis], kernel_w, self.padding, stride_w) + return tensor_shape.TensorShape(output_shape) @tf_export('keras.layers.Conv3DTranspose', 'keras.layers.Convolution3DTranspose') -class Conv3DTranspose(tf_convolutional_layers.Conv3DTranspose, Layer): +class Conv3DTranspose(Conv3D): """Transposed convolution layer (sometimes called Deconvolution). The need for transposed convolutions generally arises @@ -679,8 +928,6 @@ class Conv3DTranspose(tf_convolutional_layers.Conv3DTranspose, Layer): kernel_constraint=None, bias_constraint=None, **kwargs): - if data_format is None: - data_format = K.image_data_format() super(Conv3DTranspose, self).__init__( filters=filters, kernel_size=kernel_size, @@ -698,6 +945,314 @@ class Conv3DTranspose(tf_convolutional_layers.Conv3DTranspose, Layer): bias_constraint=constraints.get(bias_constraint), **kwargs) + def build(self, input_shape): + input_shape = tensor_shape.TensorShape(input_shape) + if len(input_shape) != 5: + raise ValueError('Inputs should have rank 5, received input shape:', + str(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: ' + str(input_shape)) + input_dim = int(input_shape[channel_axis]) + kernel_shape = self.kernel_size + (self.filters, input_dim) + self.input_spec = InputSpec(ndim=5, axes={channel_axis: input_dim}) + + self.kernel = self.add_variable( + 'kernel', + shape=kernel_shape, + initializer=self.kernel_initializer, + regularizer=self.kernel_regularizer, + constraint=self.kernel_constraint, + trainable=True, + dtype=self.dtype) + if self.use_bias: + self.bias = self.add_variable( + 'bias', + shape=(self.filters,), + initializer=self.bias_initializer, + regularizer=self.bias_regularizer, + constraint=self.bias_constraint, + trainable=True, + dtype=self.dtype) + else: + self.bias = None + self.built = True + + def call(self, inputs): + inputs_shape = array_ops.shape(inputs) + batch_size = inputs_shape[0] + if self.data_format == 'channels_first': + c_axis, d_axis, h_axis, w_axis = 1, 2, 3, 4 + else: + c_axis, d_axis, h_axis, w_axis = 4, 1, 2, 3 + + self.input_spec = InputSpec(ndim=5, axes={c_axis: inputs_shape[c_axis]}) + + depth = inputs_shape[d_axis] + height = inputs_shape[h_axis] + width = inputs_shape[w_axis] + + kernel_d, kernel_h, kernel_w = self.kernel_size + stride_d, stride_h, stride_w = self.strides + + # Infer the dynamic output shape: + out_depth = conv_utils.deconv_output_length(depth, + kernel_d, + self.padding, + stride_d) + out_height = conv_utils.deconv_output_length(height, + kernel_h, + self.padding, + stride_h) + out_width = conv_utils.deconv_output_length(width, + kernel_w, + self.padding, + stride_w) + if self.data_format == 'channels_first': + output_shape = (batch_size, self.filters, out_depth, out_height, + out_width) + strides = (1, 1, stride_d, stride_h, stride_w) + else: + output_shape = (batch_size, out_depth, out_height, out_width, + self.filters) + strides = (1, stride_d, stride_h, stride_w, 1) + + output_shape_tensor = array_ops.stack(output_shape) + outputs = nn.conv3d_transpose( + inputs, + self.kernel, + output_shape_tensor, + strides, + data_format=conv_utils.convert_data_format(self.data_format, ndim=5), + padding=self.padding.upper()) + + if not context.executing_eagerly(): + # Infer the static output shape: + out_shape = inputs.get_shape().as_list() + out_shape[c_axis] = self.filters + out_shape[d_axis] = conv_utils.deconv_output_length(out_shape[d_axis], + kernel_d, + self.padding, + stride_d) + out_shape[h_axis] = conv_utils.deconv_output_length(out_shape[h_axis], + kernel_h, + self.padding, + stride_h) + out_shape[w_axis] = conv_utils.deconv_output_length(out_shape[w_axis], + kernel_w, + self.padding, + stride_w) + outputs.set_shape(out_shape) + + if self.use_bias: + outputs_shape = outputs.shape.as_list() + if outputs_shape[0] is None: + outputs_shape[0] = -1 + if self.data_format == 'channels_first': + outputs_4d = array_ops.reshape(outputs, [ + outputs_shape[0], outputs_shape[1], + outputs_shape[2] * outputs_shape[3], outputs_shape[4] + ]) + else: + 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, + data_format=conv_utils.convert_data_format(self.data_format, ndim=4)) + outputs = array_ops.reshape(outputs_4d, outputs_shape) + + if self.activation is not None: + return self.activation(outputs) + return outputs + + def compute_output_shape(self, input_shape): + input_shape = tensor_shape.TensorShape(input_shape).as_list() + output_shape = list(input_shape) + if self.data_format == 'channels_first': + c_axis, d_axis, h_axis, w_axis = 1, 2, 3, 4 + else: + c_axis, d_axis, h_axis, w_axis = 4, 1, 2, 3 + + kernel_d, kernel_h, kernel_w = self.kernel_size + stride_d, stride_h, stride_w = self.strides + + output_shape[c_axis] = self.filters + output_shape[d_axis] = conv_utils.deconv_output_length( + output_shape[d_axis], kernel_d, self.padding, stride_d) + output_shape[h_axis] = conv_utils.deconv_output_length( + output_shape[h_axis], kernel_h, self.padding, stride_h) + output_shape[w_axis] = conv_utils.deconv_output_length( + output_shape[w_axis], kernel_w, self.padding, stride_w) + return tensor_shape.TensorShape(output_shape) + + +class SeparableConv(Conv): + """Abstract base layer for separable nD convolution. + + This layer performs a depthwise convolution that acts separately on + channels, followed by a pointwise convolution that mixes channels. + If `use_bias` is True and a bias initializer is provided, + it adds a bias vector to the output. + It then optionally applies an activation function to produce the final output. + + 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: A tuple or list of integers specifying the spatial + dimensions of the filters. Can be a single integer to specify the same + value for all spatial dimensions. + strides: A tuple or list of integers specifying the strides + of the convolution. 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, ..., channels)` while `channels_first` corresponds to + inputs with shape `(batch, channels, ...)`. + 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. + depth_multiplier: The number of depthwise convolution output channels for + each input channel. The total number of depthwise convolution output + channels will be equal to `num_filters_in * depth_multiplier`. + activation: Activation function. Set it to None to maintain a + linear activation. + use_bias: Boolean, whether the layer uses a bias. + depthwise_initializer: An initializer for the depthwise convolution kernel. + pointwise_initializer: An initializer for the pointwise convolution kernel. + bias_initializer: An initializer for the bias vector. If None, the default + initializer will be used. + depthwise_regularizer: Optional regularizer for the depthwise + convolution kernel. + pointwise_regularizer: Optional regularizer for the pointwise + convolution kernel. + bias_regularizer: Optional regularizer for the bias vector. + activity_regularizer: Optional regularizer function for the output. + depthwise_constraint: Optional projection function to be applied to the + depthwise kernel after being updated by an `Optimizer` (e.g. used for + norm constraints or value constraints for layer weights). 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. + pointwise_constraint: Optional projection function to be applied to the + pointwise kernel after being updated by an `Optimizer`. + bias_constraint: Optional projection function to be applied to the + bias after being updated by an `Optimizer`. + trainable: Boolean, if `True` also add variables to the graph collection + `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`). + name: A string, the name of the layer. + """ + + def __init__(self, + rank, + filters, + kernel_size, + strides=1, + padding='valid', + data_format=None, + dilation_rate=1, + depth_multiplier=1, + activation=None, + use_bias=True, + depthwise_initializer='glorot_uniform', + pointwise_initializer='glorot_uniform', + bias_initializer='zeros', + depthwise_regularizer=None, + pointwise_regularizer=None, + bias_regularizer=None, + activity_regularizer=None, + depthwise_constraint=None, + pointwise_constraint=None, + bias_constraint=None, + trainable=True, + name=None, + **kwargs): + super(SeparableConv, self).__init__( + rank=rank, + filters=filters, + kernel_size=kernel_size, + strides=strides, + padding=padding, + data_format=data_format, + dilation_rate=dilation_rate, + activation=activations.get(activation), + use_bias=use_bias, + bias_regularizer=regularizers.get(bias_regularizer), + activity_regularizer=regularizers.get(activity_regularizer), + bias_constraint=bias_constraint, + trainable=trainable, + name=name, + **kwargs) + self.depth_multiplier = depth_multiplier + self.depthwise_initializer = initializers.get(depthwise_initializer) + self.pointwise_initializer = initializers.get(pointwise_initializer) + self.depthwise_regularizer = regularizers.get(depthwise_regularizer) + self.pointwise_regularizer = regularizers.get(pointwise_regularizer) + self.depthwise_constraint = constraints.get(depthwise_constraint) + self.pointwise_constraint = constraints.get(pointwise_constraint) + + 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 = int(input_shape[channel_axis]) + self.input_spec = InputSpec(ndim=self.rank + 2, + axes={channel_axis: input_dim}) + depthwise_kernel_shape = self.kernel_size + (input_dim, + self.depth_multiplier) + pointwise_kernel_shape = ( + 1,) * self.rank + (self.depth_multiplier * input_dim, self.filters) + + self.depthwise_kernel = self.add_variable( + name='depthwise_kernel', + shape=depthwise_kernel_shape, + initializer=self.depthwise_initializer, + regularizer=self.depthwise_regularizer, + constraint=self.depthwise_constraint, + trainable=True, + dtype=self.dtype) + self.pointwise_kernel = self.add_variable( + name='pointwise_kernel', + shape=pointwise_kernel_shape, + initializer=self.pointwise_initializer, + regularizer=self.pointwise_regularizer, + constraint=self.pointwise_constraint, + trainable=True, + dtype=self.dtype) + if self.use_bias: + self.bias = self.add_variable(name='bias', + shape=(self.filters,), + initializer=self.bias_initializer, + regularizer=self.bias_regularizer, + constraint=self.bias_constraint, + trainable=True, + dtype=self.dtype) + else: + self.bias = None + self.built = True + + def call(self, inputs): + raise NotImplementedError + def get_config(self): config = { 'filters': self.filters, @@ -705,24 +1260,34 @@ class Conv3DTranspose(tf_convolutional_layers.Conv3DTranspose, Layer): 'strides': self.strides, 'padding': self.padding, 'data_format': self.data_format, + 'dilation_rate': self.dilation_rate, 'activation': activations.serialize(self.activation), 'use_bias': self.use_bias, - 'kernel_initializer': initializers.serialize(self.kernel_initializer), + 'depthwise_initializer': + initializers.serialize(self.depthwise_initializer), + 'pointwise_initializer': + initializers.serialize(self.pointwise_initializer), 'bias_initializer': initializers.serialize(self.bias_initializer), - 'kernel_regularizer': regularizers.serialize(self.kernel_regularizer), + 'depthwise_regularizer': + regularizers.serialize(self.depthwise_regularizer), + 'pointwise_regularizer': + regularizers.serialize(self.pointwise_regularizer), 'bias_regularizer': regularizers.serialize(self.bias_regularizer), 'activity_regularizer': regularizers.serialize(self.activity_regularizer), - 'kernel_constraint': constraints.serialize(self.kernel_constraint), + 'depthwise_constraint': + constraints.serialize(self.depthwise_constraint), + 'pointwise_constraint': + constraints.serialize(self.pointwise_constraint), 'bias_constraint': constraints.serialize(self.bias_constraint) } - base_config = super(Conv3DTranspose, self).get_config() + base_config = super(SeparableConv, self).get_config() return dict(list(base_config.items()) + list(config.items())) @tf_export('keras.layers.SeparableConv1D', 'keras.layers.SeparableConvolution1D') -class SeparableConv1D(tf_convolutional_layers.SeparableConv1D, Layer): +class SeparableConv1D(SeparableConv): """Depthwise separable 1D convolution. This layer performs a depthwise convolution that acts separately on @@ -802,15 +1367,15 @@ class SeparableConv1D(tf_convolutional_layers.SeparableConv1D, Layer): pointwise_constraint=None, bias_constraint=None, **kwargs): - if data_format is None: - data_format = K.image_data_format() super(SeparableConv1D, self).__init__( + rank=1, filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, data_format=data_format, dilation_rate=dilation_rate, + depth_multiplier=depth_multiplier, activation=activations.get(activation), use_bias=use_bias, depthwise_initializer=initializers.get(depthwise_initializer), @@ -825,44 +1390,46 @@ class SeparableConv1D(tf_convolutional_layers.SeparableConv1D, Layer): bias_constraint=constraints.get(bias_constraint), **kwargs) - def get_config(self): - config = { - 'filters': self.filters, - 'kernel_size': self.kernel_size, - 'strides': self.strides, - 'padding': self.padding, - 'data_format': self.data_format, - 'dilation_rate': self.dilation_rate, - 'activation': activations.serialize(self.activation), - 'use_bias': self.use_bias, - 'depthwise_initializer': - initializers.serialize(self.depthwise_initializer), - 'pointwise_initializer': - initializers.serialize(self.pointwise_initializer), - 'bias_initializer': - initializers.serialize(self.bias_initializer), - 'depthwise_regularizer': - regularizers.serialize(self.depthwise_regularizer), - 'pointwise_regularizer': - regularizers.serialize(self.pointwise_regularizer), - 'bias_regularizer': - regularizers.serialize(self.bias_regularizer), - 'activity_regularizer': - regularizers.serialize(self.activity_regularizer), - 'depthwise_constraint': - constraints.serialize(self.depthwise_constraint), - 'pointwise_constraint': - constraints.serialize(self.pointwise_constraint), - 'bias_constraint': - constraints.serialize(self.bias_constraint) - } - base_config = super(SeparableConv1D, self).get_config() - return dict(list(base_config.items()) + list(config.items())) + def call(self, inputs): + if self.data_format == 'channels_last': + strides = (1,) + self.strides * 2 + (1,) + spatial_start_dim = 1 + else: + strides = (1, 1) + self.strides * 2 + spatial_start_dim = 2 + + # Explicitly broadcast inputs and kernels to 4D. + # TODO(fchollet): refactor when a native separable_conv1d op is available. + inputs = array_ops.expand_dims(inputs, spatial_start_dim) + depthwise_kernel = array_ops.expand_dims(self.depthwise_kernel, 0) + pointwise_kernel = array_ops.expand_dims(self.pointwise_kernel, 0) + dilation_rate = (1,) + self.dilation_rate + + outputs = nn.separable_conv2d( + inputs, + depthwise_kernel, + pointwise_kernel, + strides=strides, + padding=self.padding.upper(), + rate=dilation_rate, + data_format=conv_utils.convert_data_format(self.data_format, ndim=4)) + + if self.use_bias: + outputs = nn.bias_add( + outputs, + self.bias, + data_format=conv_utils.convert_data_format(self.data_format, ndim=4)) + + outputs = array_ops.squeeze(outputs, [spatial_start_dim]) + + if self.activation is not None: + return self.activation(outputs) + return outputs @tf_export('keras.layers.SeparableConv2D', 'keras.layers.SeparableConvolution2D') -class SeparableConv2D(tf_convolutional_layers.SeparableConv2D, Layer): +class SeparableConv2D(SeparableConv): """Depthwise separable 2D convolution. Separable convolutions consist in first performing @@ -959,15 +1526,15 @@ class SeparableConv2D(tf_convolutional_layers.SeparableConv2D, Layer): pointwise_constraint=None, bias_constraint=None, **kwargs): - if data_format is None: - data_format = K.image_data_format() super(SeparableConv2D, self).__init__( + rank=2, filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, data_format=data_format, dilation_rate=dilation_rate, + depth_multiplier=depth_multiplier, activation=activations.get(activation), use_bias=use_bias, depthwise_initializer=initializers.get(depthwise_initializer), @@ -982,47 +1549,30 @@ class SeparableConv2D(tf_convolutional_layers.SeparableConv2D, Layer): bias_constraint=constraints.get(bias_constraint), **kwargs) - def get_config(self): - config = { - 'filters': - self.filters, - 'kernel_size': - self.kernel_size, - 'strides': - self.strides, - 'padding': - self.padding, - 'data_format': - self.data_format, - 'dilation_rate': - self.dilation_rate, - 'activation': - activations.serialize(self.activation), - 'use_bias': - self.use_bias, - 'depthwise_initializer': - initializers.serialize(self.depthwise_initializer), - 'pointwise_initializer': - initializers.serialize(self.pointwise_initializer), - 'bias_initializer': - initializers.serialize(self.bias_initializer), - 'depthwise_regularizer': - regularizers.serialize(self.depthwise_regularizer), - 'pointwise_regularizer': - regularizers.serialize(self.pointwise_regularizer), - 'bias_regularizer': - regularizers.serialize(self.bias_regularizer), - 'activity_regularizer': - regularizers.serialize(self.activity_regularizer), - 'depthwise_constraint': - constraints.serialize(self.depthwise_constraint), - 'pointwise_constraint': - constraints.serialize(self.pointwise_constraint), - 'bias_constraint': - constraints.serialize(self.bias_constraint) - } - base_config = super(SeparableConv2D, self).get_config() - return dict(list(base_config.items()) + list(config.items())) + def call(self, inputs): + # Apply the actual ops. + if self.data_format == 'channels_last': + strides = (1,) + self.strides + (1,) + else: + strides = (1, 1) + self.strides + outputs = nn.separable_conv2d( + inputs, + self.depthwise_kernel, + self.pointwise_kernel, + strides=strides, + padding=self.padding.upper(), + rate=self.dilation_rate, + data_format=conv_utils.convert_data_format(self.data_format, ndim=4)) + + if self.use_bias: + outputs = nn.bias_add( + outputs, + self.bias, + data_format=conv_utils.convert_data_format(self.data_format, ndim=4)) + + if self.activation is not None: + return self.activation(outputs) + return outputs @tf_export('keras.layers.DepthwiseConv2D') @@ -1162,7 +1712,7 @@ class DepthwiseConv2D(Conv2D): self.built = True def call(self, inputs, training=None): - outputs = K.depthwise_conv2d( + outputs = backend.depthwise_conv2d( inputs, self.depthwise_kernel, strides=self.strides, @@ -1171,7 +1721,7 @@ class DepthwiseConv2D(Conv2D): data_format=self.data_format) if self.bias: - outputs = K.bias_add( + outputs = backend.bias_add( outputs, self.bias, data_format=self.data_format) @@ -1246,7 +1796,7 @@ class UpSampling1D(Layer): return tensor_shape.TensorShape([input_shape[0], size, input_shape[2]]) def call(self, inputs): - output = K.repeat_elements(inputs, self.size, axis=1) + output = backend.repeat_elements(inputs, self.size, axis=1) return output def get_config(self): @@ -1315,7 +1865,8 @@ class UpSampling2D(Layer): [input_shape[0], height, width, input_shape[3]]) def call(self, inputs): - return K.resize_images(inputs, self.size[0], self.size[1], self.data_format) + return backend.resize_images( + inputs, self.size[0], self.size[1], self.data_format) def get_config(self): config = {'size': self.size, 'data_format': self.data_format} @@ -1387,8 +1938,8 @@ class UpSampling3D(Layer): [input_shape[0], dim1, dim2, dim3, input_shape[4]]) def call(self, inputs): - return K.resize_volumes(inputs, self.size[0], self.size[1], self.size[2], - self.data_format) + return backend.resize_volumes( + inputs, self.size[0], self.size[1], self.size[2], self.data_format) def get_config(self): config = {'size': self.size, 'data_format': self.data_format} @@ -1429,7 +1980,7 @@ class ZeroPadding1D(Layer): return tensor_shape.TensorShape([input_shape[0], length, input_shape[2]]) def call(self, inputs): - return K.temporal_padding(inputs, padding=self.padding) + return backend.temporal_padding(inputs, padding=self.padding) def get_config(self): config = {'padding': self.padding} @@ -1530,7 +2081,7 @@ class ZeroPadding2D(Layer): [input_shape[0], rows, cols, input_shape[3]]) def call(self, inputs): - return K.spatial_2d_padding( + return backend.spatial_2d_padding( inputs, padding=self.padding, data_format=self.data_format) def get_config(self): @@ -1648,7 +2199,7 @@ class ZeroPadding3D(Layer): [input_shape[0], dim1, dim2, dim3, input_shape[4]]) def call(self, inputs): - return K.spatial_3d_padding( + return backend.spatial_3d_padding( inputs, padding=self.padding, data_format=self.data_format) def get_config(self): diff --git a/tensorflow/python/keras/_impl/keras/layers/core.py b/tensorflow/python/keras/_impl/keras/layers/core.py index c74fc1e4c0a764b4cc0d09129be4e5287a9bdd05..9c4cb0f4fda681ce3236222460cd87439ea67810 100644 --- a/tensorflow/python/keras/_impl/keras/layers/core.py +++ b/tensorflow/python/keras/_impl/keras/layers/core.py @@ -24,6 +24,7 @@ import types as python_types import numpy as np 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 activations from tensorflow.python.keras._impl.keras import backend as K @@ -32,13 +33,14 @@ 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.utils.generic_utils import deserialize_keras_object -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.keras._impl.keras.utils import generic_utils +from tensorflow.python.keras._impl.keras.utils import tf_utils from tensorflow.python.ops import array_ops +from tensorflow.python.ops import gen_math_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 standard_ops from tensorflow.python.util.tf_export import tf_export @@ -94,7 +96,7 @@ class Masking(Layer): @tf_export('keras.layers.Dropout') -class Dropout(tf_core_layers.Dropout, Layer): +class Dropout(Layer): """Applies Dropout to the input. Dropout consists in randomly setting @@ -113,23 +115,40 @@ class Dropout(tf_core_layers.Dropout, Layer): """ def __init__(self, rate, noise_shape=None, seed=None, **kwargs): - # Inheritance call order: - # 1) tf.layers.Dropout, 2) keras.layers.Layer, 3) tf.layers.Layer - super(Dropout, self).__init__(rate=rate, - noise_shape=noise_shape, - seed=seed, - **kwargs) + super(Dropout, self).__init__(**kwargs) + self.rate = rate + self.noise_shape = noise_shape + self.seed = seed self.supports_masking = True + def _get_noise_shape(self, inputs): + # Subclasses of `Dropout` may implement `_get_noise_shape(self, inputs)`, + # which will override `self.noise_shape`, and allows for custom noise + # shapes with dynamically sized inputs. + if self.noise_shape is None: + return self.noise_shape + return nn_ops._get_noise_shape(inputs, self.noise_shape) # pylint: disable=protected-access + def call(self, inputs, training=None): + original_training_value = training if training is None: training = K.learning_phase() - output = super(Dropout, self).call(inputs, training=training) + + def dropped_inputs(): + return nn.dropout(inputs, 1 - self.rate, + noise_shape=self._get_noise_shape(inputs), + seed=self.seed) + output = tf_utils.smart_cond(training, + dropped_inputs, + lambda: array_ops.identity(inputs)) # EagerTensor object has no attribute _uses_learning_phase - if not context.executing_eagerly() and training is K.learning_phase(): + if not context.executing_eagerly() and original_training_value is None: output._uses_learning_phase = True # pylint: disable=protected-access return output + def compute_output_shape(self, input_shape): + return input_shape + def get_config(self): config = { 'rate': self.rate, @@ -479,7 +498,7 @@ class Permute(Layer): @tf_export('keras.layers.Flatten') -class Flatten(tf_core_layers.Flatten, Layer): +class Flatten(Layer): """Flattens the input. Does not affect the batch size. Example: @@ -495,7 +514,25 @@ class Flatten(tf_core_layers.Flatten, Layer): # now: model.output_shape == (None, 65536) ``` """ - pass + + def __init__(self, **kwargs): + super(Flatten, self).__init__(**kwargs) + self.input_spec = InputSpec(min_ndim=2) + + def call(self, inputs): + outputs = array_ops.reshape(inputs, (array_ops.shape(inputs)[0], -1)) + if not context.executing_eagerly(): + outputs.set_shape(self.compute_output_shape(inputs.get_shape())) + return outputs + + def compute_output_shape(self, input_shape): + input_shape = tensor_shape.TensorShape(input_shape).as_list() + output_shape = [input_shape[0]] + if all(input_shape[1:]): + output_shape += [np.prod(input_shape[1:])] + else: + output_shape += [None] + return tensor_shape.TensorShape(output_shape) @tf_export('keras.layers.RepeatVector') @@ -611,10 +648,12 @@ class Lambda(Layer): 'must be a list, a tuple, or a function.') self._output_shape = output_shape - def _compute_output_shape(self, input_shape): + def compute_output_shape(self, input_shape): input_shape = tuple(tensor_shape.TensorShape(input_shape).as_list()) if self._output_shape is None: + if context.executing_eagerly(): + raise NotImplementedError x = K.placeholder(shape=input_shape) x = self.call(x) if isinstance(x, list): @@ -640,7 +679,7 @@ class Lambda(Layer): def call(self, inputs, mask=None): arguments = self.arguments - if has_arg(self.function, 'mask'): + if generic_utils.has_arg(self.function, 'mask'): arguments['mask'] = mask return self.function(inputs, **arguments) @@ -651,14 +690,14 @@ class Lambda(Layer): def get_config(self): if isinstance(self.function, python_types.LambdaType): - function = func_dump(self.function) + function = generic_utils.func_dump(self.function) function_type = 'lambda' else: function = self.function.__name__ function_type = 'function' if isinstance(self._output_shape, python_types.LambdaType): - output_shape = func_dump(self._output_shape) + output_shape = generic_utils.func_dump(self._output_shape) output_shape_type = 'lambda' elif callable(self._output_shape): output_shape = self._output_shape.__name__ @@ -686,26 +725,27 @@ class Lambda(Layer): function_type = config.pop('function_type') if function_type == 'function': # Simple lookup in custom objects - function = deserialize_keras_object( + function = generic_utils.deserialize_keras_object( config['function'], custom_objects=custom_objects, printable_module_name='function in Lambda layer') elif function_type == 'lambda': # Unsafe deserialization from bytecode - function = func_load(config['function'], globs=globs) + function = generic_utils.func_load(config['function'], globs=globs) else: raise TypeError('Unknown function type:', function_type) output_shape_type = config.pop('output_shape_type') if output_shape_type == 'function': # Simple lookup in custom objects - output_shape = deserialize_keras_object( + output_shape = generic_utils.deserialize_keras_object( config['output_shape'], custom_objects=custom_objects, printable_module_name='output_shape function in Lambda layer') elif output_shape_type == 'lambda': # Unsafe deserialization from bytecode - output_shape = func_load(config['output_shape'], globs=globs) + output_shape = generic_utils.func_load(config['output_shape'], + globs=globs) else: output_shape = config['output_shape'] @@ -725,7 +765,7 @@ class Lambda(Layer): @tf_export('keras.layers.Dense') -class Dense(tf_core_layers.Dense, Layer): +class Dense(Layer): """Just your regular densely-connected NN layer. `Dense` implements the operation: @@ -795,21 +835,74 @@ class Dense(tf_core_layers.Dense, Layer): if 'input_shape' not in kwargs and 'input_dim' in kwargs: kwargs['input_shape'] = (kwargs.pop('input_dim'),) - # Inheritance call order: - # 1) tf.layers.Dense, 2) keras.layers.Layer, 3) tf.layers.Layer super(Dense, self).__init__( - units, - activation=activations.get(activation), - use_bias=use_bias, - kernel_initializer=initializers.get(kernel_initializer), - bias_initializer=initializers.get(bias_initializer), - kernel_regularizer=regularizers.get(kernel_regularizer), - bias_regularizer=regularizers.get(bias_regularizer), - activity_regularizer=regularizers.get(activity_regularizer), - kernel_constraint=constraints.get(kernel_constraint), - bias_constraint=constraints.get(bias_constraint), - **kwargs) + activity_regularizer=regularizers.get(activity_regularizer), **kwargs) + self.units = int(units) + self.activation = activations.get(activation) + self.use_bias = use_bias + self.kernel_initializer = initializers.get(kernel_initializer) + self.bias_initializer = initializers.get(bias_initializer) + self.kernel_regularizer = regularizers.get(kernel_regularizer) + self.bias_regularizer = regularizers.get(bias_regularizer) + self.kernel_constraint = constraints.get(kernel_constraint) + self.bias_constraint = constraints.get(bias_constraint) + self.supports_masking = True + self.input_spec = InputSpec(min_ndim=2) + + def build(self, input_shape): + input_shape = tensor_shape.TensorShape(input_shape) + if input_shape[-1].value is None: + raise ValueError('The last dimension of the inputs to `Dense` ' + 'should be defined. Found `None`.') + self.input_spec = InputSpec(min_ndim=2, + axes={-1: input_shape[-1].value}) + self.kernel = self.add_variable('kernel', + shape=[input_shape[-1].value, self.units], + initializer=self.kernel_initializer, + regularizer=self.kernel_regularizer, + constraint=self.kernel_constraint, + dtype=self.dtype, + trainable=True) + if self.use_bias: + self.bias = self.add_variable('bias', + shape=[self.units,], + initializer=self.bias_initializer, + regularizer=self.bias_regularizer, + constraint=self.bias_constraint, + dtype=self.dtype, + trainable=True) + else: + self.bias = None + self.built = True + + def call(self, inputs): + inputs = ops.convert_to_tensor(inputs, dtype=self.dtype) + shape = inputs.get_shape().as_list() + if len(shape) > 2: + # Broadcasting is required for the inputs. + outputs = standard_ops.tensordot(inputs, self.kernel, [[len(shape) - 1], + [0]]) + # Reshape the output back to the original ndim of the input. + if not context.executing_eagerly(): + output_shape = shape[:-1] + [self.units] + outputs.set_shape(output_shape) + else: + 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: + return self.activation(outputs) # pylint: disable=not-callable + return outputs + + def compute_output_shape(self, input_shape): + input_shape = tensor_shape.TensorShape(input_shape) + input_shape = input_shape.with_rank_at_least(2) + if input_shape[-1].value is None: + raise ValueError( + 'The innermost dimension of input_shape must be defined, but saw: %s' + % input_shape) + return input_shape[:-1].concatenate(self.units) def get_config(self): config = { diff --git a/tensorflow/python/keras/_impl/keras/layers/core_test.py b/tensorflow/python/keras/_impl/keras/layers/core_test.py index 551d1b1c3a0a80ed81ad03afc4c9510a231e33ef..d22d8d12dc4e76998c177dbe96fb87e3fffa5175 100644 --- a/tensorflow/python/keras/_impl/keras/layers/core_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/core_test.py @@ -129,7 +129,6 @@ class CoreLayersTest(test.TestCase): 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): testing_utils.layer_test( keras.layers.Lambda, diff --git a/tensorflow/python/keras/_impl/keras/layers/embeddings.py b/tensorflow/python/keras/_impl/keras/layers/embeddings.py index 540e2d945c986aebbd7028e4a1f2e4566747320f..07b8726b8595f7c5bbf36e296df79c2e82afdbaf 100644 --- a/tensorflow/python/keras/_impl/keras/layers/embeddings.py +++ b/tensorflow/python/keras/_impl/keras/layers/embeddings.py @@ -24,7 +24,7 @@ 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.base_layer import shape_type_conversion -from tensorflow.python.ops import array_ops +from tensorflow.python.ops import embedding_ops from tensorflow.python.ops import math_ops from tensorflow.python.util.tf_export import tf_export @@ -102,7 +102,8 @@ class Embedding(Layer): kwargs['input_shape'] = (input_length,) else: kwargs['input_shape'] = (None,) - super(Embedding, self).__init__(**kwargs) + dtype = kwargs.pop('dtype', K.floatx()) + super(Embedding, self).__init__(dtype=dtype, **kwargs) self.input_dim = input_dim self.output_dim = output_dim @@ -120,8 +121,7 @@ class Embedding(Layer): initializer=self.embeddings_initializer, name='embeddings', regularizer=self.embeddings_regularizer, - constraint=self.embeddings_constraint, - dtype=self.dtype) + constraint=self.embeddings_constraint) self.built = True def compute_mask(self, inputs, mask=None): @@ -155,7 +155,7 @@ class Embedding(Layer): def call(self, inputs): if K.dtype(inputs) != 'int32': inputs = math_ops.cast(inputs, 'int32') - out = array_ops.gather(self.embeddings, inputs) + out = embedding_ops.embedding_lookup(self.embeddings, inputs) return out def get_config(self): diff --git a/tensorflow/python/keras/_impl/keras/layers/embeddings_test.py b/tensorflow/python/keras/_impl/keras/layers/embeddings_test.py index 26fd1f1c114587c2f1b3e0155f1259dd5f0dcf60..6ebf5dc94adb423abae7ec9e6910fb86439410f1 100644 --- a/tensorflow/python/keras/_impl/keras/layers/embeddings_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/embeddings_test.py @@ -18,6 +18,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import numpy as np + from tensorflow.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 @@ -26,7 +28,7 @@ from tensorflow.python.platform import test class EmbeddingTest(test.TestCase): - @tf_test_util.run_in_graph_and_eager_modes() + @tf_test_util.run_in_graph_and_eager_modes(use_gpu=False) def test_embedding(self): testing_utils.layer_test( keras.layers.Embedding, @@ -65,6 +67,17 @@ class EmbeddingTest(test.TestCase): input_dtype='int32', expected_output_dtype='float32') + def test_embedding_correctness(self): + with self.test_session(): + layer = keras.layers.Embedding(output_dim=2, input_dim=2) + layer.build((None, 2)) + matrix = np.array([[1, 1], [2, 2]]) + layer.set_weights([matrix]) + + inputs = keras.backend.constant([[0, 1, 0]], dtype='int32') + outputs = keras.backend.eval(layer(inputs)) + self.assertAllClose(outputs, [[[1, 1], [2, 2], [1, 1]]]) + if __name__ == '__main__': test.main() diff --git a/tensorflow/python/keras/_impl/keras/layers/normalization.py b/tensorflow/python/keras/_impl/keras/layers/normalization.py index 3b44b20bf822429351002c0f81fe8f9596d595d3..c16fc07fb4ecda66bd8bcc70dce5d753c73f5dd9 100644 --- a/tensorflow/python/keras/_impl/keras/layers/normalization.py +++ b/tensorflow/python/keras/_impl/keras/layers/normalization.py @@ -19,17 +19,28 @@ from __future__ import division from __future__ import print_function from tensorflow.python.eager import context +from tensorflow.python.framework import dtypes +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 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 import Layer -from tensorflow.python.layers import normalization as tf_normalization_layers +from tensorflow.python.keras._impl.keras.utils import tf_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 nn +from tensorflow.python.ops import state_ops +from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.training import distribute as distribute_lib from tensorflow.python.util.tf_export import tf_export @tf_export('keras.layers.BatchNormalization') -class BatchNormalization(tf_normalization_layers.BatchNormalization, Layer): +class BatchNormalization(Layer): """Batch normalization layer (Ioffe and Szegedy, 2014). Normalize the activations of the previous layer at each batch, @@ -37,28 +48,63 @@ class BatchNormalization(tf_normalization_layers.BatchNormalization, Layer): close to 0 and the activation standard deviation close to 1. Arguments: - axis: Integer, the axis that should be normalized - (typically the features axis). - For instance, after a `Conv2D` layer with - `data_format="channels_first"`, - set `axis=1` in `BatchNormalization`. - momentum: Momentum for the moving average. - epsilon: Small float added to variance to avoid dividing by zero. - center: If True, add offset of `beta` to normalized tensor. - If False, `beta` is ignored. - scale: If True, multiply by `gamma`. - If False, `gamma` is not used. - When the next layer is linear (also e.g. `nn.relu`), - this can be disabled since the scaling - will be done by the next layer. - beta_initializer: Initializer for the beta weight. - gamma_initializer: Initializer for the gamma weight. - moving_mean_initializer: Initializer for the moving mean. - moving_variance_initializer: Initializer for the moving variance. - beta_regularizer: Optional regularizer for the beta weight. - gamma_regularizer: Optional regularizer for the gamma weight. - beta_constraint: Optional constraint for the beta weight. - gamma_constraint: Optional constraint for the gamma weight. + axis: Integer, the axis that should be normalized + (typically the features axis). + For instance, after a `Conv2D` layer with + `data_format="channels_first"`, + set `axis=1` in `BatchNormalization`. + momentum: Momentum for the moving average. + epsilon: Small float added to variance to avoid dividing by zero. + center: If True, add offset of `beta` to normalized tensor. + If False, `beta` is ignored. + scale: If True, multiply by `gamma`. + If False, `gamma` is not used. + When the next layer is linear (also e.g. `nn.relu`), + this can be disabled since the scaling + will be done by the next layer. + beta_initializer: Initializer for the beta weight. + gamma_initializer: Initializer for the gamma weight. + moving_mean_initializer: Initializer for the moving mean. + moving_variance_initializer: Initializer for the moving variance. + beta_regularizer: Optional regularizer for the beta weight. + gamma_regularizer: Optional regularizer for the gamma weight. + beta_constraint: Optional constraint for the beta weight. + gamma_constraint: Optional constraint for the gamma weight. + renorm: Whether to use Batch Renormalization + (https://arxiv.org/abs/1702.03275). This adds extra variables during + training. The inference is the same for either value of this parameter. + renorm_clipping: A dictionary that may map keys 'rmax', 'rmin', 'dmax' to + scalar `Tensors` used to clip the renorm correction. The correction + `(r, d)` is used as `corrected_value = normalized_value * r + d`, with + `r` clipped to [rmin, rmax], and `d` to [-dmax, dmax]. Missing rmax, rmin, + dmax are set to inf, 0, inf, respectively. + renorm_momentum: Momentum used to update the moving means and standard + deviations with renorm. Unlike `momentum`, this affects training + 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 `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`, + which means batch normalization is performed across the whole batch. When + `virtual_batch_size` is not `None`, instead perform "Ghost Batch + Normalization", which creates virtual sub-batches which are each + normalized separately (with shared gamma, beta, and moving statistics). + Must divide the actual batch size during execution. + adjustment: A function taking the `Tensor` containing the (dynamic) shape of + the input tensor and returning a pair (scale, bias) to apply to the + normalized values (before gamma and beta), only during training. For + example, if axis==-1, + `adjustment = lambda shape: ( + tf.random_uniform(shape[-1:], 0.93, 1.07), + tf.random_uniform(shape[-1:], -0.1, 0.1))` + will scale the normalized value by up to 7% up or down, then shift the + result by up to 0.1 (with independent scaling and bias for each feature + but shared across all examples), and finally apply gamma and/or beta. If + `None`, no adjustment is applied. Cannot be specified if + virtual_batch_size is specified. Input shape: Arbitrary. Use the keyword argument `input_shape` @@ -87,33 +133,508 @@ class BatchNormalization(tf_normalization_layers.BatchNormalization, Layer): gamma_regularizer=None, beta_constraint=None, gamma_constraint=None, + renorm=False, + renorm_clipping=None, + renorm_momentum=0.99, + fused=None, + trainable=True, + virtual_batch_size=None, + adjustment=None, + name=None, **kwargs): - self.supports_masking = True super(BatchNormalization, self).__init__( - axis=axis, - momentum=momentum, - epsilon=epsilon, - center=center, - scale=scale, - beta_initializer=initializers.get(beta_initializer), - gamma_initializer=initializers.get(gamma_initializer), - moving_mean_initializer=initializers.get(moving_mean_initializer), - moving_variance_initializer=initializers.get( - moving_variance_initializer), - beta_regularizer=regularizers.get(beta_regularizer), - gamma_regularizer=regularizers.get(gamma_regularizer), - beta_constraint=constraints.get(beta_constraint), - gamma_constraint=constraints.get(gamma_constraint), - **kwargs - ) + name=name, trainable=trainable, **kwargs) + if isinstance(axis, list): + self.axis = axis[:] + else: + self.axis = axis + self.momentum = momentum + self.epsilon = epsilon + self.center = center + self.scale = scale + self.beta_initializer = initializers.get(beta_initializer) + self.gamma_initializer = initializers.get(gamma_initializer) + self.moving_mean_initializer = initializers.get(moving_mean_initializer) + self.moving_variance_initializer = initializers.get( + moving_variance_initializer) + self.beta_regularizer = regularizers.get(beta_regularizer) + self.gamma_regularizer = regularizers.get(gamma_regularizer) + self.beta_constraint = constraints.get(beta_constraint) + self.gamma_constraint = constraints.get(gamma_constraint) + self.renorm = renorm + self.virtual_batch_size = virtual_batch_size + self.adjustment = adjustment + if fused is None: + fused = True + self.supports_masking = True + + self.fused = fused + self._bessels_correction_test_only = True + + if renorm: + renorm_clipping = renorm_clipping or {} + keys = ['rmax', 'rmin', 'dmax'] + if set(renorm_clipping) - set(keys): + raise ValueError('renorm_clipping %s contains keys not in %s' % + (renorm_clipping, keys)) + self.renorm_clipping = renorm_clipping + self.renorm_momentum = renorm_momentum + + def _add_tower_local_variable(self, *args, **kwargs): + tower_context = distribute_lib.get_tower_context() + with tower_context.tower_local_var_scope('mean'): + return self.add_variable(*args, **kwargs) + + def build(self, input_shape): + input_shape = tensor_shape.TensorShape(input_shape) + if not input_shape.ndims: + raise ValueError('Input has undefined rank:', input_shape) + ndims = len(input_shape) + + # Convert axis to list and resolve negatives + if isinstance(self.axis, int): + self.axis = [self.axis] + + if not isinstance(self.axis, list): + raise TypeError('axis must be int or list, type given: %s' + % type(self.axis)) + + for idx, x in enumerate(self.axis): + if x < 0: + self.axis[idx] = ndims + x + + # Validate axes + for x in self.axis: + if x < 0 or x >= ndims: + raise ValueError('Invalid axis: %d' % x) + if len(self.axis) != len(set(self.axis)): + raise ValueError('Duplicate axis: %s' % self.axis) + + if self.virtual_batch_size is not None: + if self.virtual_batch_size <= 0: + raise ValueError('virtual_batch_size must be a positive integer that ' + 'divides the true batch size of the input Tensor') + # If using virtual batches, the first dimension must be the batch + # dimension and cannot be the batch norm axis + if 0 in self.axis: + raise ValueError('When using virtual_batch_size, the batch dimension ' + 'must be 0 and thus axis cannot include 0') + if self.adjustment is not None: + raise ValueError('When using virtual_batch_size, adjustment cannot ' + 'be specified') + + if self.fused: + # Currently fused batch norm doesn't support renorm. It also only supports + # an input tensor of rank 4 and a channel dimension on axis 1 or 3. + # TODO(yaozhang): if input is not 4D, reshape it to 4D and reshape the + # output back to its original shape accordingly. + self.fused = (not self.renorm and + ndims == 4 and + self.axis in [[1], [3]] and + self.virtual_batch_size is None and + self.adjustment is None) + # TODO(chrisying): fused batch norm is currently not supported for + # multi-axis batch norm and by extension virtual batches. In some cases, + # it might be possible to use fused batch norm but would require reshaping + # the Tensor to 4D with the axis in 1 or 3 (preferred 1) which is + # particularly tricky. A compromise might be to just support the most + # common use case (turning 5D w/ virtual batch to NCHW) + + if self.fused: + if self.axis == [1]: + self._data_format = 'NCHW' + elif self.axis == [3]: + self._data_format = 'NHWC' + else: + raise ValueError('Unsupported axis, fused batch norm only supports ' + 'axis == [1] or axis == [3]') + + # Raise parameters of fp16 batch norm to fp32 + if self.dtype == dtypes.float16 or self.dtype == dtypes.bfloat16: + param_dtype = dtypes.float32 + else: + param_dtype = self.dtype or dtypes.float32 + + axis_to_dim = {x: input_shape[x].value for x in self.axis} + for x in axis_to_dim: + if axis_to_dim[x] is None: + raise ValueError('Input has undefined `axis` dimension. Input shape: ', + input_shape) + self.input_spec = InputSpec(ndim=ndims, axes=axis_to_dim) + + if len(axis_to_dim) == 1 and self.virtual_batch_size is None: + # Single axis batch norm (most common/default use-case) + param_shape = (list(axis_to_dim.values())[0],) + else: + # Parameter shape is the original shape but with 1 in all non-axis dims + param_shape = [axis_to_dim[i] if i in axis_to_dim + else 1 for i in range(ndims)] + if self.virtual_batch_size is not None: + # When using virtual batches, add an extra dim at index 1 + param_shape.insert(1, 1) + for idx, x in enumerate(self.axis): + self.axis[idx] = x + 1 # Account for added dimension + + if self.scale: + self.gamma = self.add_variable( + name='gamma', + shape=param_shape, + dtype=param_dtype, + initializer=self.gamma_initializer, + regularizer=self.gamma_regularizer, + constraint=self.gamma_constraint, + trainable=True) + else: + self.gamma = None + if self.fused: + self._gamma_const = array_ops.constant( + 1.0, dtype=param_dtype, shape=param_shape) + + if self.center: + self.beta = self.add_variable( + name='beta', + shape=param_shape, + dtype=param_dtype, + initializer=self.beta_initializer, + regularizer=self.beta_regularizer, + constraint=self.beta_constraint, + trainable=True) + else: + self.beta = None + if self.fused: + self._beta_const = array_ops.constant( + 0.0, dtype=param_dtype, shape=param_shape) + + try: + # Disable variable partitioning when creating the moving mean and variance + if hasattr(self, '_scope') and self._scope: + partitioner = self._scope.partitioner + self._scope.set_partitioner(None) + else: + partitioner = None + self.moving_mean = self._add_tower_local_variable( + name='moving_mean', + shape=param_shape, + dtype=param_dtype, + initializer=self.moving_mean_initializer, + trainable=False) + + self.moving_variance = self._add_tower_local_variable( + name='moving_variance', + shape=param_shape, + dtype=param_dtype, + initializer=self.moving_variance_initializer, + trainable=False) + + if self.renorm: + # Create variables to maintain the moving mean and standard deviation. + # These are used in training and thus are different from the moving + # averages above. The renorm variables are colocated with moving_mean + # and moving_variance. + # NOTE: below, the outer `with device` block causes the current device + # stack to be cleared. The nested ones use a `lambda` to set the desired + # device and ignore any devices that may be set by the custom getter. + def _renorm_variable(name, shape): + var = self._add_tower_local_variable( + name=name, + shape=shape, + dtype=param_dtype, + initializer=init_ops.zeros_initializer(), + trainable=False) + return var + + with distribute_lib.get_distribution_strategy().colocate_vars_with( + self.moving_mean): + self.renorm_mean = _renorm_variable('renorm_mean', param_shape) + self.renorm_mean_weight = _renorm_variable('renorm_mean_weight', ()) + # We initialize renorm_stddev to 0, and maintain the (0-initialized) + # 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. + with distribute_lib.get_distribution_strategy().colocate_vars_with( + self.moving_variance): + self.renorm_stddev = _renorm_variable('renorm_stddev', param_shape) + self.renorm_stddev_weight = _renorm_variable('renorm_stddev_weight', + ()) + finally: + if partitioner: + self._scope.set_partitioner(partitioner) + self.built = True + + def _assign_moving_average(self, variable, value, momentum): + with ops.name_scope(None, 'AssignMovingAvg', + [variable, value, momentum]) as scope: + decay = ops.convert_to_tensor(1.0 - momentum, name='decay') + if decay.dtype != variable.dtype.base_dtype: + decay = math_ops.cast(decay, variable.dtype.base_dtype) + update_delta = (variable - value) * decay + return state_ops.assign_sub(variable, update_delta, name=scope) + + def _fused_batch_norm(self, inputs, training): + """Returns the output of fused batch norm.""" + beta = self.beta if self.center else self._beta_const + gamma = self.gamma if self.scale else self._gamma_const + + def _fused_batch_norm_training(): + return nn.fused_batch_norm( + inputs, + gamma, + beta, + epsilon=self.epsilon, + data_format=self._data_format) + + def _fused_batch_norm_inference(): + return nn.fused_batch_norm( + inputs, + gamma, + beta, + mean=self.moving_mean, + variance=self.moving_variance, + epsilon=self.epsilon, + is_training=False, + data_format=self._data_format) + + output, mean, variance = tf_utils.smart_cond( + training, _fused_batch_norm_training, _fused_batch_norm_inference) + if not self._bessels_correction_test_only: + # Remove Bessel's correction to be consistent with non-fused batch norm. + # Note that the variance computed by fused batch norm is + # with Bessel's correction. + sample_size = math_ops.cast( + array_ops.size(inputs) / array_ops.size(variance), variance.dtype) + factor = (sample_size - math_ops.cast(1.0, variance.dtype)) / sample_size + variance *= factor + + training_value = tf_utils.constant_value(training) + if training_value is None: + momentum = tf_utils.smart_cond(training, + lambda: self.momentum, + lambda: 1.0) + else: + momentum = ops.convert_to_tensor(self.momentum) + if training_value or training_value is None: + mean_update = self._assign_moving_average(self.moving_mean, mean, + momentum) + variance_update = self._assign_moving_average(self.moving_variance, + variance, momentum) + self.add_update(mean_update, inputs=True) + self.add_update(variance_update, inputs=True) + + return output + + def _renorm_correction_and_moments(self, mean, variance, training): + """Returns the correction and update values for renorm.""" + stddev = math_ops.sqrt(variance + self.epsilon) + # Compute the average mean and standard deviation, as if they were + # initialized with this batch's moments. + mixed_renorm_mean = (self.renorm_mean + + (1. - self.renorm_mean_weight) * mean) + mixed_renorm_stddev = (self.renorm_stddev + + (1. - self.renorm_stddev_weight) * stddev) + # Compute the corrections for batch renorm. + r = stddev / mixed_renorm_stddev + d = (mean - mixed_renorm_mean) / mixed_renorm_stddev + # Ensure the corrections use pre-update moving averages. + with ops.control_dependencies([r, d]): + mean = array_ops.identity(mean) + stddev = array_ops.identity(stddev) + rmin, rmax, dmax = [self.renorm_clipping.get(key) + for key in ['rmin', 'rmax', 'dmax']] + if rmin is not None: + r = math_ops.maximum(r, rmin) + if rmax is not None: + r = math_ops.minimum(r, rmax) + if dmax is not None: + d = math_ops.maximum(d, -dmax) + d = math_ops.minimum(d, dmax) + # When not training, use r=1, d=0. + r = tf_utils.smart_cond(training, lambda: r, lambda: array_ops.ones_like(r)) + d = tf_utils.smart_cond(training, + lambda: d, + lambda: array_ops.zeros_like(d)) + + def _update_renorm_variable(var, weight, value): + """Updates a moving average and weight, returns the unbiased value.""" + value = array_ops.identity(value) + def _do_update(): + """Updates the var and weight, returns their updated ratio.""" + # Update the variables without zero debiasing. The debiasing will be + # accomplished by dividing the exponential moving average by the weight. + # For example, after a single update, the moving average would be + # (1-decay) * value. and the weight will be 1-decay, with their ratio + # giving the value. + # Make sure the weight is not updated until before r and d computation. + with ops.control_dependencies([value]): + weight_value = array_ops.constant(1., dtype=weight.dtype) + new_var = self._assign_moving_average(var, value, self.renorm_momentum) + new_weight = self._assign_moving_average(weight, weight_value, + self.renorm_momentum) + # TODO(yuefengz): the updates to var and weighted can not be batched + # together if we fetch their updated values here. Consider calculating + # new values and delaying the updates. + return new_var / new_weight + + def _fake_update(): + return array_ops.identity(var) + return tf_utils.smart_cond(training, _do_update, _fake_update) + + # TODO(yuefengz): colocate the operations + new_mean = _update_renorm_variable(self.renorm_mean, + self.renorm_mean_weight, mean) + new_stddev = _update_renorm_variable(self.renorm_stddev, + self.renorm_stddev_weight, stddev) + # Make sqrt(moving_variance + epsilon) = new_stddev. + new_variance = math_ops.square(new_stddev) - self.epsilon + + return (r, d, new_mean, new_variance) def call(self, inputs, training=None): + original_training_value = training if training is None: training = K.learning_phase() - output = super(BatchNormalization, self).call(inputs, training=training) - if not context.executing_eagerly() and training is K.learning_phase(): - output._uses_learning_phase = True # pylint: disable=protected-access - return output + + 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 + original_shape = [-1] + inputs.shape.as_list()[1:] + expanded_shape = [self.virtual_batch_size, -1] + original_shape[1:] + + # Will cause errors if virtual_batch_size does not divide the batch size + inputs = array_ops.reshape(inputs, expanded_shape) + + def undo_virtual_batching(outputs): + outputs = array_ops.reshape(outputs, original_shape) + return outputs + + if self.fused: + outputs = self._fused_batch_norm(inputs, training=training) + if self.virtual_batch_size is not None: + # Currently never reaches here since fused_batch_norm does not support + # virtual batching + outputs = undo_virtual_batching(outputs) + if not context.executing_eagerly() and original_training_value is None: + outputs._uses_learning_phase = True # pylint: disable=protected-access + return outputs + + # Compute the axes along which to reduce the mean / variance + input_shape = inputs.get_shape() + ndims = len(input_shape) + reduction_axes = [i for i in range(ndims) if i not in self.axis] + if self.virtual_batch_size is not None: + del reduction_axes[1] # Do not reduce along virtual batch dim + + # Broadcasting only necessary for single-axis batch norm where the axis is + # not the last dimension + broadcast_shape = [1] * ndims + broadcast_shape[self.axis[0]] = input_shape[self.axis[0]].value + 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 + + scale, offset = _broadcast(self.gamma), _broadcast(self.beta) + + def _compose_transforms(scale, offset, then_scale, then_offset): + if then_scale is not None: + scale *= then_scale + offset *= then_scale + if then_offset is not None: + offset += then_offset + return (scale, offset) + + # Determine a boolean value for `training`: could be True, False, or None. + training_value = tf_utils.constant_value(training) + if training_value is not False: + if self.adjustment: + adj_scale, adj_bias = self.adjustment(array_ops.shape(inputs)) + # Adjust only during training. + adj_scale = tf_utils.smart_cond(training, + lambda: adj_scale, + lambda: array_ops.ones_like(adj_scale)) + adj_bias = tf_utils.smart_cond(training, + lambda: adj_bias, + lambda: array_ops.zeros_like(adj_bias)) + scale, offset = _compose_transforms(adj_scale, adj_bias, scale, offset) + + # Some of the computations here are not necessary when training==False + # but not a constant. However, this makes the code simpler. + keep_dims = self.virtual_batch_size is not None or len(self.axis) > 1 + mean, variance = nn.moments(inputs, reduction_axes, keep_dims=keep_dims) + + moving_mean = self.moving_mean + moving_variance = self.moving_variance + + mean = tf_utils.smart_cond(training, + lambda: mean, + lambda: moving_mean) + variance = tf_utils.smart_cond(training, + lambda: variance, + lambda: moving_variance) + + if self.renorm: + r, d, new_mean, new_variance = self._renorm_correction_and_moments( + mean, variance, training) + # When training, the normalized values (say, x) will be transformed as + # x * gamma + beta without renorm, and (x * r + d) * gamma + beta + # = x * (r * gamma) + (d * gamma + beta) with renorm. + r = _broadcast(array_ops.stop_gradient(r, name='renorm_r')) + d = _broadcast(array_ops.stop_gradient(d, name='renorm_d')) + scale, offset = _compose_transforms(r, d, scale, offset) + else: + new_mean, new_variance = mean, variance + + if self.virtual_batch_size is not None: + # This isn't strictly correct since in ghost batch norm, you are + # supposed to sequentially update the moving_mean and moving_variance + # with each sub-batch. However, since the moving statistics are only + # used during evaluation, it is more efficient to just update in one + # step and should not make a significant difference in the result. + new_mean = math_ops.reduce_mean(new_mean, + axis=1, keepdims=True) + new_variance = math_ops.reduce_mean(new_variance, + axis=1, keepdims=True) + + def _do_update(var, value): + if in_eager_mode and not self.trainable: + return + + return self._assign_moving_average(var, value, self.momentum) + + mean_update = tf_utils.smart_cond( + training, + lambda: _do_update(self.moving_mean, new_mean), + lambda: self.moving_mean) + variance_update = tf_utils.smart_cond( + training, + lambda: _do_update(self.moving_variance, new_variance), + lambda: self.moving_variance) + if not context.executing_eagerly(): + self.add_update(mean_update, inputs=True) + self.add_update(variance_update, inputs=True) + + else: + mean, variance = self.moving_mean, self.moving_variance + + outputs = nn.batch_normalization(inputs, + _broadcast(mean), + _broadcast(variance), + offset, + scale, + self.epsilon) + # If some components of the shape got lost due to adjustments, fix that. + outputs.set_shape(input_shape) + + if self.virtual_batch_size is not None: + outputs = undo_virtual_batching(outputs) + if not context.executing_eagerly() and original_training_value is None: + outputs._uses_learning_phase = True # pylint: disable=protected-access + return outputs + + def compute_output_shape(self, input_shape): + return input_shape def get_config(self): config = { @@ -133,5 +654,19 @@ class BatchNormalization(tf_normalization_layers.BatchNormalization, Layer): 'beta_constraint': constraints.serialize(self.beta_constraint), 'gamma_constraint': constraints.serialize(self.gamma_constraint) } + # Only add TensorFlow-specific parameters if they are set, so as to preserve + # model compatibility with external Keras. + if self.renorm: + config['renorm'] = True + config['renorm_clipping'] = self.renorm_clipping + config['renorm_momentum'] = self.renorm_momentum + if self.virtual_batch_size is not None: + config['virtual_batch_size'] = self.virtual_batch_size + # Note: adjustment is not serializable. + if self.adjustment is not None: + logging.warning('The `adjustment` function of this `BatchNormalization` ' + 'layer cannot be serialized and has been omitted from ' + 'the layer config. It will not be included when ' + 're-creating the layer from the saved config.') base_config = super(BatchNormalization, self).get_config() return dict(list(base_config.items()) + list(config.items())) diff --git a/tensorflow/python/keras/_impl/keras/layers/normalization_test.py b/tensorflow/python/keras/_impl/keras/layers/normalization_test.py index 2b3628c3f1023612297465bdf3286246261992a2..fa9277e3d1e5bb0b9633abc46a96a11816dddb2d 100644 --- a/tensorflow/python/keras/_impl/keras/layers/normalization_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/normalization_test.py @@ -114,6 +114,26 @@ class NormalizationLayersTest(test.TestCase): np.testing.assert_allclose(np.mean(out, axis=(0, 2, 3)), 0.0, atol=1e-1) np.testing.assert_allclose(np.std(out, axis=(0, 2, 3)), 1.0, atol=1e-1) + def test_batchnorm_convnet_channel_last(self): + with self.test_session(): + # keras.backend.set_learning_phase(True) + + model = keras.models.Sequential() + norm = keras.layers.BatchNormalization( + axis=-1, input_shape=(4, 4, 3), momentum=0.8) + model.add(norm) + model.compile(loss='mse', optimizer='sgd') + + # centered on 5.0, variance 10.0 + x = np.random.normal(loc=5.0, scale=10.0, size=(1000, 4, 4, 3)) + model.fit(x, x, epochs=4, verbose=0) + out = model.predict(x) + out -= np.reshape(keras.backend.eval(norm.beta), (1, 1, 1, 3)) + out /= np.reshape(keras.backend.eval(norm.gamma), (1, 1, 1, 3)) + + np.testing.assert_allclose(np.mean(out, axis=(0, 1, 2)), 0.0, atol=1e-1) + np.testing.assert_allclose(np.std(out, axis=(0, 1, 2)), 1.0, atol=1e-1) + def test_shared_batchnorm(self): """Test that a BN layer can be shared across different data streams. """ diff --git a/tensorflow/python/keras/_impl/keras/layers/pooling.py b/tensorflow/python/keras/_impl/keras/layers/pooling.py index 15d53379769d8142f5b2755a07479f60751346d2..86bc8a680a529a9ea17592a42207fab58adeebce 100644 --- a/tensorflow/python/keras/_impl/keras/layers/pooling.py +++ b/tensorflow/python/keras/_impl/keras/layers/pooling.py @@ -19,16 +19,98 @@ from __future__ import division from __future__ import print_function from tensorflow.python.framework import tensor_shape -from tensorflow.python.keras._impl.keras import backend as K +from tensorflow.python.keras._impl.keras import backend 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.ops import array_ops +from tensorflow.python.ops import nn from tensorflow.python.util.tf_export import tf_export +class Pooling1D(Layer): + """Pooling layer for arbitrary pooling functions, for 1D inputs. + + This class only exists for code reuse. It will never be an exposed API. + + Arguments: + pool_function: The pooling function to apply, e.g. `tf.nn.max_pool`. + pool_size: An integer or tuple/list of a single integer, + representing the size of the pooling window. + strides: An integer or tuple/list of a single integer, specifying the + strides of the pooling operation. + padding: A string. The padding method, either '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)`. + name: A string, the name of the layer. + """ + + def __init__(self, pool_function, pool_size, strides, + padding='valid', data_format=None, + name=None, **kwargs): + super(Pooling1D, self).__init__(name=name, **kwargs) + if data_format is None: + data_format = backend.image_data_format() + if strides is None: + strides = pool_size + self.pool_function = pool_function + self.pool_size = conv_utils.normalize_tuple(pool_size, 1, 'pool_size') + self.strides = conv_utils.normalize_tuple(strides, 1, 'strides') + self.padding = conv_utils.normalize_padding(padding) + self.data_format = conv_utils.normalize_data_format(data_format) + self.input_spec = InputSpec(ndim=3) + + def call(self, inputs): + # There is no TF op for 1D pooling, hence we make the inputs 4D. + if self.data_format == 'channels_last': + # input is NWC, make it NHWC + inputs = array_ops.expand_dims(inputs, 1) + # pool on the W dim + pool_shape = (1, 1) + self.pool_size + (1,) + strides = (1, 1) + self.strides + (1,) + data_format = 'NHWC' + else: + # input is NCW, make it NCHW + inputs = array_ops.expand_dims(inputs, 2) + # pool on the W dim + pool_shape = (1, 1, 1) + self.pool_size + strides = (1, 1, 1) + self.strides + data_format = 'NCHW' + + outputs = self.pool_function( + inputs, + ksize=pool_shape, + strides=strides, + padding=self.padding.upper(), + data_format=data_format) + + if self.data_format == 'channels_last': + return array_ops.squeeze(outputs, 1) + else: + return array_ops.squeeze(outputs, 2) + + def compute_output_shape(self, input_shape): + input_shape = tensor_shape.TensorShape(input_shape).as_list() + length = conv_utils.conv_output_length(input_shape[1], self.pool_size[0], + self.padding, self.strides[0]) + return tensor_shape.TensorShape([input_shape[0], length, input_shape[2]]) + + def get_config(self): + config = { + 'strides': self.strides, + 'pool_size': self.pool_size, + 'padding': self.padding + } + base_config = super(Pooling1D, self).get_config() + return dict(list(base_config.items()) + list(config.items())) + + @tf_export('keras.layers.MaxPool1D', 'keras.layers.MaxPooling1D') -class MaxPooling1D(tf_pooling_layers.MaxPooling1D, Layer): +class MaxPooling1D(Pooling1D): """Max pooling operation for temporal data. Arguments: @@ -45,23 +127,20 @@ class MaxPooling1D(tf_pooling_layers.MaxPooling1D, Layer): 3D tensor with shape: `(batch_size, downsampled_steps, features)`. """ - def __init__(self, pool_size=2, strides=None, padding='valid', **kwargs): - if strides is None: - strides = pool_size - super(MaxPooling1D, self).__init__(pool_size, strides, padding, **kwargs) + def __init__(self, pool_size=2, strides=None, + padding='valid', data_format=None, **kwargs): - def get_config(self): - config = { - 'strides': self.strides, - 'pool_size': self.pool_size, - 'padding': self.padding - } - base_config = super(MaxPooling1D, self).get_config() - return dict(list(base_config.items()) + list(config.items())) + super(MaxPooling1D, self).__init__( + nn.max_pool, + pool_size=pool_size, + strides=strides, + padding=padding, + data_format=data_format, + **kwargs) @tf_export('keras.layers.AveragePooling1D', 'keras.layers.AvgPool1D') -class AveragePooling1D(tf_pooling_layers.AveragePooling1D, Layer): +class AveragePooling1D(Pooling1D): """Average pooling for temporal data. Arguments: @@ -78,24 +157,104 @@ class AveragePooling1D(tf_pooling_layers.AveragePooling1D, Layer): 3D tensor with shape: `(batch_size, downsampled_steps, features)`. """ - def __init__(self, pool_size=2, strides=None, padding='valid', **kwargs): + def __init__(self, pool_size=2, strides=None, + padding='valid', data_format=None, **kwargs): + super(AveragePooling1D, self).__init__( + nn.avg_pool, + pool_size=pool_size, + strides=strides, + padding=padding, + data_format=data_format, + **kwargs) + + +class Pooling2D(Layer): + """Pooling layer for arbitrary pooling functions, for 2D inputs (e.g. images). + + This class only exists for code reuse. It will never be an exposed API. + + Arguments: + pool_function: The pooling function to apply, e.g. `tf.nn.max_pool`. + pool_size: An integer or tuple/list of 2 integers: (pool_height, pool_width) + specifying the size of the pooling 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 pooling operation. + Can be a single integer to specify the same value for + all spatial dimensions. + padding: A string. The padding method, either '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)`. + name: A string, the name of the layer. + """ + + def __init__(self, pool_function, pool_size, strides, + padding='valid', data_format=None, + name=None, **kwargs): + super(Pooling2D, self).__init__(name=name, **kwargs) + if data_format is None: + data_format = backend.image_data_format() if strides is None: strides = pool_size - super(AveragePooling1D, self).__init__(pool_size, strides, padding, - **kwargs) + self.pool_function = pool_function + self.pool_size = conv_utils.normalize_tuple(pool_size, 2, 'pool_size') + self.strides = conv_utils.normalize_tuple(strides, 2, 'strides') + self.padding = conv_utils.normalize_padding(padding) + self.data_format = conv_utils.normalize_data_format(data_format) + self.input_spec = InputSpec(ndim=4) + + def call(self, inputs): + if self.data_format == 'channels_last': + pool_shape = (1,) + self.pool_size + (1,) + strides = (1,) + self.strides + (1,) + else: + pool_shape = (1, 1) + self.pool_size + strides = (1, 1) + self.strides + outputs = self.pool_function( + inputs, + ksize=pool_shape, + strides=strides, + padding=self.padding.upper(), + data_format=conv_utils.convert_data_format(self.data_format, 4)) + return outputs + + def compute_output_shape(self, input_shape): + input_shape = tensor_shape.TensorShape(input_shape).as_list() + if self.data_format == 'channels_first': + rows = input_shape[2] + cols = input_shape[3] + else: + rows = input_shape[1] + cols = input_shape[2] + rows = conv_utils.conv_output_length(rows, self.pool_size[0], self.padding, + self.strides[0]) + cols = conv_utils.conv_output_length(cols, self.pool_size[1], self.padding, + self.strides[1]) + if self.data_format == 'channels_first': + return tensor_shape.TensorShape( + [input_shape[0], input_shape[1], rows, cols]) + else: + return tensor_shape.TensorShape( + [input_shape[0], rows, cols, input_shape[3]]) def get_config(self): config = { - 'strides': self.strides, 'pool_size': self.pool_size, - 'padding': self.padding + 'padding': self.padding, + 'strides': self.strides, + 'data_format': self.data_format } - base_config = super(AveragePooling1D, self).get_config() + base_config = super(Pooling2D, self).get_config() return dict(list(base_config.items()) + list(config.items())) @tf_export('keras.layers.MaxPool2D', 'keras.layers.MaxPooling2D') -class MaxPooling2D(tf_pooling_layers.MaxPooling2D, Layer): +class MaxPooling2D(Pooling2D): """Max pooling operation for spatial data. Arguments: @@ -142,26 +301,14 @@ class MaxPooling2D(tf_pooling_layers.MaxPooling2D, Layer): padding='valid', data_format=None, **kwargs): - if data_format is None: - data_format = K.image_data_format() - if strides is None: - strides = pool_size - super(MaxPooling2D, self).__init__(pool_size, strides, padding, data_format, - **kwargs) - - def get_config(self): - config = { - 'pool_size': self.pool_size, - 'padding': self.padding, - 'strides': self.strides, - 'data_format': self.data_format - } - base_config = super(MaxPooling2D, self).get_config() - return dict(list(base_config.items()) + list(config.items())) + super(MaxPooling2D, self).__init__( + nn.max_pool, + pool_size=pool_size, strides=strides, + padding=padding, data_format=data_format, **kwargs) @tf_export('keras.layers.AveragePooling2D', 'keras.layers.AvgPool2D') -class AveragePooling2D(tf_pooling_layers.AveragePooling2D, Layer): +class AveragePooling2D(Pooling2D): """Average pooling operation for spatial data. Arguments: @@ -208,12 +355,96 @@ class AveragePooling2D(tf_pooling_layers.AveragePooling2D, Layer): padding='valid', data_format=None, **kwargs): + super(AveragePooling2D, self).__init__( + nn.avg_pool, + pool_size=pool_size, strides=strides, + padding=padding, data_format=data_format, **kwargs) + + +class Pooling3D(Layer): + """Pooling layer for arbitrary pooling functions, for 3D inputs. + + This class only exists for code reuse. It will never be an exposed API. + + Arguments: + pool_function: The pooling function to apply, e.g. `tf.nn.max_pool`. + pool_size: An integer or tuple/list of 3 integers: + (pool_depth, pool_height, pool_width) + specifying the size of the pooling 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 pooling operation. + Can be a single integer to specify the same value for + all spatial dimensions. + padding: A string. The padding method, either '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)`. + name: A string, the name of the layer. + """ + + def __init__(self, pool_function, pool_size, strides, + padding='valid', data_format='channels_last', + name=None, **kwargs): + super(Pooling3D, self).__init__(name=name, **kwargs) if data_format is None: - data_format = K.image_data_format() + data_format = backend.image_data_format() if strides is None: strides = pool_size - super(AveragePooling2D, self).__init__(pool_size, strides, padding, - data_format, **kwargs) + self.pool_function = pool_function + self.pool_size = conv_utils.normalize_tuple(pool_size, 3, 'pool_size') + self.strides = conv_utils.normalize_tuple(strides, 3, 'strides') + self.padding = conv_utils.normalize_padding(padding) + self.data_format = conv_utils.normalize_data_format(data_format) + self.input_spec = InputSpec(ndim=5) + + def call(self, inputs): + pool_shape = (1,) + self.pool_size + (1,) + strides = (1,) + self.strides + (1,) + + if self.data_format == 'channels_first': + # TF does not support `channels_first` with 3D pooling operations, + # so we must handle this case manually. + # TODO(fchollet): remove this when TF pooling is feature-complete. + inputs = array_ops.transpose(inputs, (0, 2, 3, 4, 1)) + + outputs = self.pool_function( + inputs, + ksize=pool_shape, + strides=strides, + padding=self.padding.upper()) + + if self.data_format == 'channels_first': + outputs = array_ops.transpose(outputs, (0, 4, 1, 2, 3)) + return outputs + + def compute_output_shape(self, input_shape): + input_shape = tensor_shape.TensorShape(input_shape).as_list() + if self.data_format == 'channels_first': + len_dim1 = input_shape[2] + len_dim2 = input_shape[3] + len_dim3 = input_shape[4] + else: + len_dim1 = input_shape[1] + len_dim2 = input_shape[2] + len_dim3 = input_shape[3] + len_dim1 = conv_utils.conv_output_length(len_dim1, self.pool_size[0], + self.padding, self.strides[0]) + len_dim2 = conv_utils.conv_output_length(len_dim2, self.pool_size[1], + self.padding, self.strides[1]) + len_dim3 = conv_utils.conv_output_length(len_dim3, self.pool_size[2], + self.padding, self.strides[2]) + if self.data_format == 'channels_first': + return tensor_shape.TensorShape( + [input_shape[0], input_shape[1], len_dim1, len_dim2, len_dim3]) + else: + return tensor_shape.TensorShape( + [input_shape[0], len_dim1, len_dim2, len_dim3, input_shape[4]]) def get_config(self): config = { @@ -222,12 +453,12 @@ class AveragePooling2D(tf_pooling_layers.AveragePooling2D, Layer): 'strides': self.strides, 'data_format': self.data_format } - base_config = super(AveragePooling2D, self).get_config() + base_config = super(Pooling3D, self).get_config() return dict(list(base_config.items()) + list(config.items())) @tf_export('keras.layers.MaxPool3D', 'keras.layers.MaxPooling3D') -class MaxPooling3D(tf_pooling_layers.MaxPooling3D, Layer): +class MaxPooling3D(Pooling3D): """Max pooling operation for 3D data (spatial or spatio-temporal). Arguments: @@ -270,26 +501,14 @@ class MaxPooling3D(tf_pooling_layers.MaxPooling3D, Layer): padding='valid', data_format=None, **kwargs): - if data_format is None: - data_format = K.image_data_format() - if strides is None: - strides = pool_size - super(MaxPooling3D, self).__init__(pool_size, strides, padding, data_format, - **kwargs) - - def get_config(self): - config = { - 'pool_size': self.pool_size, - 'padding': self.padding, - 'strides': self.strides, - 'data_format': self.data_format - } - base_config = super(MaxPooling3D, self).get_config() - return dict(list(base_config.items()) + list(config.items())) + super(MaxPooling3D, self).__init__( + nn.max_pool3d, + pool_size=pool_size, strides=strides, + padding=padding, data_format=data_format, **kwargs) @tf_export('keras.layers.AveragePooling3D', 'keras.layers.AvgPool3D') -class AveragePooling3D(tf_pooling_layers.AveragePooling3D, Layer): +class AveragePooling3D(Pooling3D): """Average pooling operation for 3D data (spatial or spatio-temporal). Arguments: @@ -332,30 +551,18 @@ class AveragePooling3D(tf_pooling_layers.AveragePooling3D, Layer): padding='valid', data_format=None, **kwargs): - if data_format is None: - data_format = K.image_data_format() - if strides is None: - strides = pool_size - super(AveragePooling3D, self).__init__(pool_size, strides, padding, - data_format, **kwargs) - - def get_config(self): - config = { - 'pool_size': self.pool_size, - 'padding': self.padding, - 'strides': self.strides, - 'data_format': self.data_format - } - base_config = super(AveragePooling3D, self).get_config() - return dict(list(base_config.items()) + list(config.items())) + super(AveragePooling3D, self).__init__( + nn.avg_pool3d, + pool_size=pool_size, strides=strides, + padding=padding, data_format=data_format, **kwargs) -class _GlobalPooling1D(Layer): +class GlobalPooling1D(Layer): """Abstract class for different global pooling 1D layers. """ def __init__(self, **kwargs): - super(_GlobalPooling1D, self).__init__(**kwargs) + super(GlobalPooling1D, self).__init__(**kwargs) self.input_spec = InputSpec(ndim=3) def compute_output_shape(self, input_shape): @@ -368,7 +575,7 @@ class _GlobalPooling1D(Layer): @tf_export('keras.layers.GlobalAveragePooling1D', 'keras.layers.GlobalAvgPool1D') -class GlobalAveragePooling1D(_GlobalPooling1D): +class GlobalAveragePooling1D(GlobalPooling1D): """Global average pooling operation for temporal data. Input shape: @@ -380,11 +587,11 @@ class GlobalAveragePooling1D(_GlobalPooling1D): """ def call(self, inputs): - return K.mean(inputs, axis=1) + return backend.mean(inputs, axis=1) @tf_export('keras.layers.GlobalMaxPool1D', 'keras.layers.GlobalMaxPooling1D') -class GlobalMaxPooling1D(_GlobalPooling1D): +class GlobalMaxPooling1D(GlobalPooling1D): """Global max pooling operation for temporal data. Input shape: @@ -396,15 +603,15 @@ class GlobalMaxPooling1D(_GlobalPooling1D): """ def call(self, inputs): - return K.max(inputs, axis=1) + return backend.max(inputs, axis=1) -class _GlobalPooling2D(Layer): +class GlobalPooling2D(Layer): """Abstract class for different global pooling 2D layers. """ def __init__(self, data_format=None, **kwargs): - super(_GlobalPooling2D, self).__init__(**kwargs) + super(GlobalPooling2D, self).__init__(**kwargs) self.data_format = conv_utils.normalize_data_format(data_format) self.input_spec = InputSpec(ndim=4) @@ -420,13 +627,13 @@ class _GlobalPooling2D(Layer): def get_config(self): config = {'data_format': self.data_format} - base_config = super(_GlobalPooling2D, self).get_config() + base_config = super(GlobalPooling2D, self).get_config() return dict(list(base_config.items()) + list(config.items())) @tf_export('keras.layers.GlobalAveragePooling2D', 'keras.layers.GlobalAvgPool2D') -class GlobalAveragePooling2D(_GlobalPooling2D): +class GlobalAveragePooling2D(GlobalPooling2D): """Global average pooling operation for spatial data. Arguments: @@ -456,13 +663,13 @@ class GlobalAveragePooling2D(_GlobalPooling2D): def call(self, inputs): if self.data_format == 'channels_last': - return K.mean(inputs, axis=[1, 2]) + return backend.mean(inputs, axis=[1, 2]) else: - return K.mean(inputs, axis=[2, 3]) + return backend.mean(inputs, axis=[2, 3]) @tf_export('keras.layers.GlobalMaxPool2D', 'keras.layers.GlobalMaxPooling2D') -class GlobalMaxPooling2D(_GlobalPooling2D): +class GlobalMaxPooling2D(GlobalPooling2D): """Global max pooling operation for spatial data. Arguments: @@ -492,17 +699,17 @@ class GlobalMaxPooling2D(_GlobalPooling2D): def call(self, inputs): if self.data_format == 'channels_last': - return K.max(inputs, axis=[1, 2]) + return backend.max(inputs, axis=[1, 2]) else: - return K.max(inputs, axis=[2, 3]) + return backend.max(inputs, axis=[2, 3]) -class _GlobalPooling3D(Layer): +class GlobalPooling3D(Layer): """Abstract class for different global pooling 3D layers. """ def __init__(self, data_format=None, **kwargs): - super(_GlobalPooling3D, self).__init__(**kwargs) + super(GlobalPooling3D, self).__init__(**kwargs) self.data_format = conv_utils.normalize_data_format(data_format) self.input_spec = InputSpec(ndim=5) @@ -518,13 +725,13 @@ class _GlobalPooling3D(Layer): def get_config(self): config = {'data_format': self.data_format} - base_config = super(_GlobalPooling3D, self).get_config() + base_config = super(GlobalPooling3D, self).get_config() return dict(list(base_config.items()) + list(config.items())) @tf_export('keras.layers.GlobalAveragePooling3D', 'keras.layers.GlobalAvgPool3D') -class GlobalAveragePooling3D(_GlobalPooling3D): +class GlobalAveragePooling3D(GlobalPooling3D): """Global Average pooling operation for 3D data. Arguments: @@ -554,13 +761,13 @@ class GlobalAveragePooling3D(_GlobalPooling3D): def call(self, inputs): if self.data_format == 'channels_last': - return K.mean(inputs, axis=[1, 2, 3]) + return backend.mean(inputs, axis=[1, 2, 3]) else: - return K.mean(inputs, axis=[2, 3, 4]) + return backend.mean(inputs, axis=[2, 3, 4]) @tf_export('keras.layers.GlobalMaxPool3D', 'keras.layers.GlobalMaxPooling3D') -class GlobalMaxPooling3D(_GlobalPooling3D): +class GlobalMaxPooling3D(GlobalPooling3D): """Global Max pooling operation for 3D data. Arguments: @@ -590,9 +797,9 @@ class GlobalMaxPooling3D(_GlobalPooling3D): def call(self, inputs): if self.data_format == 'channels_last': - return K.max(inputs, axis=[1, 2, 3]) + return backend.max(inputs, axis=[1, 2, 3]) else: - return K.max(inputs, axis=[2, 3, 4]) + return backend.max(inputs, axis=[2, 3, 4]) # Aliases diff --git a/tensorflow/python/keras/_impl/keras/layers/pooling_test.py b/tensorflow/python/keras/_impl/keras/layers/pooling_test.py index bb003c1dddf80e2a745c1268a3a7d045f4e8b036..2c08b647ea0fafb7519240b0c81e8fa77f034f7f 100644 --- a/tensorflow/python/keras/_impl/keras/layers/pooling_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/pooling_test.py @@ -27,14 +27,14 @@ from tensorflow.python.platform import test class GlobalPoolingTest(test.TestCase): - @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) + @tf_test_util.run_in_graph_and_eager_modes() def test_globalpooling_1d(self): 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) + @tf_test_util.run_in_graph_and_eager_modes() def test_globalpooling_2d(self): testing_utils.layer_test( keras.layers.pooling.GlobalMaxPooling2D, @@ -53,7 +53,7 @@ class GlobalPoolingTest(test.TestCase): kwargs={'data_format': 'channels_last'}, input_shape=(3, 5, 6, 4)) - @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) + @tf_test_util.run_in_graph_and_eager_modes() def test_globalpooling_3d(self): testing_utils.layer_test( keras.layers.pooling.GlobalMaxPooling3D, @@ -75,7 +75,7 @@ class GlobalPoolingTest(test.TestCase): class Pooling2DTest(test.TestCase): - @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) + @tf_test_util.run_in_graph_and_eager_modes() def test_maxpooling_2d(self): pool_size = (3, 3) for strides in [(1, 1), (2, 2)]: @@ -88,7 +88,7 @@ class Pooling2DTest(test.TestCase): }, input_shape=(3, 5, 6, 4)) - @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) + @tf_test_util.run_in_graph_and_eager_modes() def test_averagepooling_2d(self): testing_utils.layer_test( keras.layers.AveragePooling2D, @@ -122,7 +122,7 @@ class Pooling2DTest(test.TestCase): class Pooling3DTest(test.TestCase): - @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) + @tf_test_util.run_in_graph_and_eager_modes() def test_maxpooling_3d(self): pool_size = (3, 3, 3) testing_utils.layer_test( @@ -141,7 +141,7 @@ class Pooling3DTest(test.TestCase): }, input_shape=(3, 4, 11, 12, 10)) - @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) + @tf_test_util.run_in_graph_and_eager_modes() def test_averagepooling_3d(self): pool_size = (3, 3, 3) testing_utils.layer_test( @@ -163,7 +163,7 @@ class Pooling3DTest(test.TestCase): class Pooling1DTest(test.TestCase): - @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) + @tf_test_util.run_in_graph_and_eager_modes() def test_maxpooling_1d(self): for padding in ['valid', 'same']: for stride in [1, 2]: @@ -173,7 +173,7 @@ class Pooling1DTest(test.TestCase): 'padding': padding}, input_shape=(3, 5, 4)) - @tf_test_util.run_in_graph_and_eager_modes(use_gpu=True) + @tf_test_util.run_in_graph_and_eager_modes() def test_averagepooling_1d(self): for padding in ['valid', 'same']: for stride in [1, 2]: diff --git a/tensorflow/python/keras/_impl/keras/layers/recurrent_test.py b/tensorflow/python/keras/_impl/keras/layers/recurrent_test.py index 641b563a2594126f0bb6267cda52131ccbc766af..4c68c18825a47d87806a7a09d4054f974d569e00 100644 --- a/tensorflow/python/keras/_impl/keras/layers/recurrent_test.py +++ b/tensorflow/python/keras/_impl/keras/layers/recurrent_test.py @@ -435,8 +435,8 @@ class RNNTest(test.TestCase): 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]) + self.assertEqual(len(layer.get_updates_for(None)), 1) + self.assertEqual(len(layer.get_updates_for(x)), 1) def test_rnn_dynamic_trainability(self): layer_class = keras.layers.SimpleRNN diff --git a/tensorflow/python/keras/_impl/keras/layers/wrappers.py b/tensorflow/python/keras/_impl/keras/layers/wrappers.py index c510e464ae54b2dbe87533f02971a4af6c9c7c45..9aee5f03b6d79f0b363f79d2b7a18c0b20a2883f 100644 --- a/tensorflow/python/keras/_impl/keras/layers/wrappers.py +++ b/tensorflow/python/keras/_impl/keras/layers/wrappers.py @@ -23,11 +23,11 @@ import copy 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 InputSpec from tensorflow.python.keras._impl.keras.engine import Layer 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.ops import array_ops from tensorflow.python.util.tf_export import tf_export @@ -213,7 +213,7 @@ class TimeDistributed(Wrapper): input_length = array_ops.shape(inputs)[1] # Shape: (num_samples * timesteps, ...). And track the # transformation in self._input_map. - input_uid = tf_layers_util.object_list_uid(inputs) + input_uid = base_layer.object_list_uid(inputs) inputs = array_ops.reshape(inputs, (-1,) + input_shape[2:]) self._input_map[input_uid] = inputs # (num_samples * timesteps, ...) diff --git a/tensorflow/python/keras/_impl/keras/model_subclassing_test.py b/tensorflow/python/keras/_impl/keras/model_subclassing_test.py index 444590033042d915b12645fb0239833b666a02f7..bc8698f235aac0f5fb0c3303cc4c70aa1efa08bc 100644 --- a/tensorflow/python/keras/_impl/keras/model_subclassing_test.py +++ b/tensorflow/python/keras/_impl/keras/model_subclassing_test.py @@ -607,12 +607,6 @@ class CustomCallSignatureTests(test.TestCase): self.assertAllClose(10. * expected_output, self.evaluate(output)) output = model(first, second=second, training=False) self.assertAllClose(expected_output, self.evaluate(output)) - if not context.executing_eagerly(): - six.assertCountEqual(self, [first, second], model.inputs) - with self.assertRaises(TypeError): - # tf.layers.Layer expects an "inputs" argument, so all-keywords doesn't - # work at the moment. - model(first=first, second=second, fiddle_with_output='yes') @test_util.run_in_graph_and_eager_modes() def test_inputs_in_signature(self): @@ -622,10 +616,14 @@ class CustomCallSignatureTests(test.TestCase): def call(self, inputs, some_other_arg, training=False): return inputs + def compute_output_shape(self, input_shape): + return input_shape + model = HasInputsAndOtherPositional() with self.assertRaisesRegexp( TypeError, 'everything else as a keyword argument'): - model(array_ops.ones([]), array_ops.ones([])) + x1, x2 = keras.Input((1, 1)), keras.Input((1, 1)) + model(x1, x2) @test_util.run_in_graph_and_eager_modes() def test_kwargs_in_signature(self): @@ -649,13 +647,14 @@ class CustomCallSignatureTests(test.TestCase): def call(self, x, *args, **kwargs): return [x] + list(args) + def compute_output_shape(self, input_shape): + return input_shape + model = HasArgs() - arg1 = array_ops.ones([]) - arg2 = array_ops.ones([]) - arg3 = array_ops.ones([]) - model(arg1, arg2, arg3, a=3) + x1, x2, x3 = keras.Input((1, 1)), keras.Input((1, 1)), keras.Input((1, 1)) + model(x1, x2, x3, a=3) if not context.executing_eagerly(): - six.assertCountEqual(self, [arg1, arg2, arg3], model.inputs) + six.assertCountEqual(self, [x1, x2, x3], model.inputs) def test_args_and_keywords_in_signature(self): @@ -666,11 +665,9 @@ class CustomCallSignatureTests(test.TestCase): with context.graph_mode(): model = HasArgs() - arg1 = array_ops.ones([]) - arg2 = array_ops.ones([]) - arg3 = array_ops.ones([]) + x1, x2, x3 = keras.Input((1, 1)), keras.Input((1, 1)), keras.Input((1, 1)) with self.assertRaisesRegexp(TypeError, 'args and arguments with'): - model(arg1, arg2, arg3, a=3) + model(x1, x2, x3, a=3) def test_training_no_default(self): @@ -694,11 +691,9 @@ class CustomCallSignatureTests(test.TestCase): with context.graph_mode(): model = TrainingNoDefaultWithPositional() - arg1 = array_ops.ones([]) - arg2 = array_ops.ones([]) - arg3 = array_ops.ones([]) + x1, x2, x3 = keras.Input((1, 1)), keras.Input((1, 1)), keras.Input((1, 1)) with self.assertRaisesRegexp(TypeError, 'after a non-input'): - model(arg1, arg2, arg3) + model(x1, x2, x3) if __name__ == '__main__': test.main() diff --git a/tensorflow/python/keras/_impl/keras/utils/conv_utils.py b/tensorflow/python/keras/_impl/keras/utils/conv_utils.py index 583079d9626361eb594f16a57af86f103e5ee74d..8882a3a46bcb9de7283a67f001e67ed8644a0cf7 100644 --- a/tensorflow/python/keras/_impl/keras/utils/conv_utils.py +++ b/tensorflow/python/keras/_impl/keras/utils/conv_utils.py @@ -21,17 +21,146 @@ from __future__ import print_function import numpy as np from six.moves import range # pylint: disable=redefined-builtin -# pylint: disable=unused-import -from tensorflow.python.keras._impl.keras import backend as K -from tensorflow.python.layers.utils import conv_input_length -from tensorflow.python.layers.utils import conv_output_length -from tensorflow.python.layers.utils import deconv_output_length as deconv_length -from tensorflow.python.layers.utils import normalize_tuple +from tensorflow.python.keras._impl.keras import backend + + +def convert_data_format(data_format, ndim): + if data_format == 'channels_last': + if ndim == 3: + return 'NWC' + elif ndim == 4: + return 'NHWC' + elif ndim == 5: + return 'NDHWC' + else: + raise ValueError('Input rank not supported:', ndim) + elif data_format == 'channels_first': + if ndim == 3: + return 'NCW' + elif ndim == 4: + return 'NCHW' + elif ndim == 5: + return 'NCDHW' + else: + raise ValueError('Input rank not supported:', ndim) + else: + raise ValueError('Invalid data_format:', data_format) + + +def normalize_tuple(value, n, name): + """Transforms a single integer or iterable of integers into an integer tuple. + + Arguments: + value: The value to validate and convert. Could an int, or any iterable + of ints. + n: The size of the tuple to be returned. + name: The name of the argument being validated, e.g. "strides" or + "kernel_size". This is only used to format error messages. + + Returns: + A tuple of n integers. + + Raises: + ValueError: If something else than an int/long or iterable thereof was + passed. + """ + if isinstance(value, int): + return (value,) * n + else: + try: + value_tuple = tuple(value) + except TypeError: + raise ValueError('The `' + name + '` argument must be a tuple of ' + + str(n) + ' integers. Received: ' + str(value)) + if len(value_tuple) != n: + raise ValueError('The `' + name + '` argument must be a tuple of ' + + str(n) + ' integers. Received: ' + str(value)) + for single_value in value_tuple: + try: + int(single_value) + except (ValueError, TypeError): + raise ValueError('The `' + name + '` argument must be a tuple of ' + + str(n) + ' integers. Received: ' + str(value) + ' ' + 'including element ' + str(single_value) + ' of type' + + ' ' + str(type(single_value))) + return value_tuple + + +def conv_output_length(input_length, filter_size, padding, stride, dilation=1): + """Determines output length of a convolution given input length. + + Arguments: + input_length: integer. + filter_size: integer. + padding: one of "same", "valid", "full". + stride: integer. + dilation: dilation rate, integer. + + Returns: + The output length (integer). + """ + if input_length is None: + return None + assert padding in {'same', 'valid', 'full'} + dilated_filter_size = filter_size + (filter_size - 1) * (dilation - 1) + if padding == 'same': + output_length = input_length + elif padding == 'valid': + output_length = input_length - dilated_filter_size + 1 + elif padding == 'full': + output_length = input_length + dilated_filter_size - 1 + return (output_length + stride - 1) // stride + + +def conv_input_length(output_length, filter_size, padding, stride): + """Determines input length of a convolution given output length. + + Arguments: + output_length: integer. + filter_size: integer. + padding: one of "same", "valid", "full". + stride: integer. + + Returns: + The input length (integer). + """ + if output_length is None: + return None + assert padding in {'same', 'valid', 'full'} + if padding == 'same': + pad = filter_size // 2 + elif padding == 'valid': + pad = 0 + elif padding == 'full': + pad = filter_size - 1 + return (output_length - 1) * stride - 2 * pad + filter_size + + +def deconv_output_length(input_length, filter_size, padding, stride): + """Determines output length of a transposed convolution given input length. + + Arguments: + input_length: integer. + filter_size: integer. + padding: one of "same", "valid", "full". + stride: integer. + + Returns: + The output length (integer). + """ + if input_length is None: + return None + input_length *= stride + if padding == 'valid': + input_length += max(filter_size - stride, 0) + elif padding == 'full': + input_length -= (stride + filter_size - 2) + return input_length def normalize_data_format(value): if value is None: - value = K.image_data_format() + value = backend.image_data_format() data_format = value.lower() if data_format not in {'channels_first', 'channels_last'}: raise ValueError('The `data_format` argument must be one of ' diff --git a/tensorflow/python/keras/_impl/keras/utils/io_utils.py b/tensorflow/python/keras/_impl/keras/utils/io_utils.py index bbf1d2a3d9c3948271780ec3fad3316b4e6d53c3..f82e3277de70a631c93f0ef3c240f41ddb3390a7 100644 --- a/tensorflow/python/keras/_impl/keras/utils/io_utils.py +++ b/tensorflow/python/keras/_impl/keras/utils/io_utils.py @@ -19,9 +19,9 @@ from __future__ import division from __future__ import print_function from collections import defaultdict -import sys import numpy as np +import six from tensorflow.python.util.tf_export import tf_export @@ -160,13 +160,11 @@ def ask_to_proceed_with_overwrite(filepath): Returns: True if we can proceed with overwrite, False otherwise. """ - get_input = input - if sys.version_info[:2] <= (2, 7): - get_input = raw_input - overwrite = get_input('[WARNING] %s already exists - overwrite? ' - '[y/n]' % (filepath)) - while overwrite not in ['y', 'n']: - overwrite = get_input('Enter "y" (overwrite) or "n" (cancel).') + overwrite = six.moves.input('[WARNING] %s already exists - overwrite? ' + '[y/n]' % (filepath)).strip().lower() + while overwrite not in ('y', 'n'): + overwrite = six.moves.input('Enter "y" (overwrite) or "n" ' + '(cancel).').strip().lower() if overwrite == 'n': return False print('[TIP] Next time specify overwrite=True!') diff --git a/tensorflow/python/keras/_impl/keras/utils/tf_utils.py b/tensorflow/python/keras/_impl/keras/utils/tf_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8da5f7777733767f31fad205a23c2f08f9ffbb1c --- /dev/null +++ b/tensorflow/python/keras/_impl/keras/utils/tf_utils.py @@ -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. +# ============================================================================== +"""TensorFlow-related utilities.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.framework import smart_cond as smart_module +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import variables + + +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 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. + + Returns: + True or False if `pred` has a constant boolean value, None otherwise. + + Raises: + TypeError: If `pred` is not a Variable, Tensor or bool, or Python + integer 1 or 0. + """ + # Allow integer booleans. + 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) diff --git a/tensorflow/python/kernel_tests/BUILD b/tensorflow/python/kernel_tests/BUILD index 3033b489774177bfa06dcf0ae056b39e8d5b59b6..ba8f1fd3caf4def5f6648298d882b75165861359 100644 --- a/tensorflow/python/kernel_tests/BUILD +++ b/tensorflow/python/kernel_tests/BUILD @@ -592,7 +592,7 @@ cuda_py_test( cuda_py_test( name = "matrix_solve_op_test", - size = "small", + size = "medium", srcs = ["matrix_solve_op_test.py"], additional_deps = [ "//third_party/py/numpy", @@ -917,6 +917,20 @@ tf_py_test( ], ) +tf_py_test( + name = "string_strip_op_test", + size = "small", + srcs = ["string_strip_op_test.py"], + additional_deps = [ + "//third_party/py/numpy", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:errors", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:string_ops", + ], +) + tf_py_test( name = "substr_op_test", size = "small", @@ -1029,12 +1043,14 @@ tf_py_test( "//tensorflow/python:framework_for_generated_wrappers", "//tensorflow/python:framework_test_lib", "//tensorflow/python:init_ops", + "//tensorflow/python:layers", "//tensorflow/python:math_ops", "//tensorflow/python:variable_scope", "//tensorflow/python:resource_variable_ops", "//tensorflow/python:state_ops", "//tensorflow/python:variables", "//tensorflow/python/eager:context", + "//tensorflow/python/eager:function", ], tags = ["no_windows"], ) @@ -1188,6 +1204,21 @@ cuda_py_test( "//tensorflow/python/eager:context", ], shard_count = 10, + tags = [ + "noasan", # times out + ], +) + +cuda_py_test( + name = "broadcast_to_ops_test", + size = "small", + srcs = ["broadcast_to_ops_test.py"], + additional_deps = [ + "//third_party/py/numpy", + "//tensorflow/python:array_ops", + "//tensorflow/python:client", + "//tensorflow/python:client_testlib", + ], ) cuda_py_test( @@ -1601,11 +1632,16 @@ cuda_py_test( "//tensorflow/python:variable_scope", "//tensorflow/python:variables", ], + shard_count = 4, + tags = [ + "noasan", + "notap", + ], ) cuda_py_test( name = "linalg_ops_test", - size = "small", + size = "medium", srcs = ["linalg_ops_test.py"], additional_deps = [ "//third_party/py/numpy", @@ -2667,10 +2703,6 @@ cuda_py_test( "//tensorflow/python:variables", ], shard_count = 50, - tags = [ - "manual", - "notap", # b/30226163 - ], ) cuda_py_test( @@ -2820,7 +2852,7 @@ sycl_py_test( tf_py_test( name = "sets_test", - size = "small", + size = "medium", srcs = ["sets_test.py"], additional_deps = [ "//third_party/py/numpy", @@ -2872,7 +2904,10 @@ tf_py_test( "//tensorflow/python:variables", ], shard_count = 10, - tags = ["no_windows_gpu"], + tags = [ + "no_windows_gpu", + "noasan", + ], ) tf_py_test( diff --git a/tensorflow/python/kernel_tests/broadcast_to_ops_test.py b/tensorflow/python/kernel_tests/broadcast_to_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..6a1bd958ba89080ff38e461646b07edbc6daec21 --- /dev/null +++ b/tensorflow/python/kernel_tests/broadcast_to_ops_test.py @@ -0,0 +1,85 @@ +# 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 broadcast_to ops.""" +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 test_util +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import test as test_lib + + +class BroadcastToTest(test_util.TensorFlowTestCase): + + def testBroadcastToBasic(self): + for dtype in [np.uint8, np.uint16, np.int8, np.int16, np.int32, np.int64]: + with self.test_session(use_gpu=True): + x = np.array([1, 2, 3], dtype=dtype) + v_tf = array_ops.broadcast_to(constant_op.constant(x), [3, 3]) + v_np = np.broadcast_to(x, [3, 3]) + self.assertAllEqual(v_tf.eval(), v_np) + + def testBroadcastToString(self): + with self.test_session(use_gpu=True): + x = np.array([b"1", b"2", b"3"]) + v_tf = array_ops.broadcast_to(constant_op.constant(x), [3, 3]) + v_np = np.broadcast_to(x, [3, 3]) + self.assertAllEqual(v_tf.eval(), v_np) + + def testBroadcastToBool(self): + with self.test_session(use_gpu=True): + x = np.array([True, False, True], dtype=np.bool) + v_tf = array_ops.broadcast_to(constant_op.constant(x), [3, 3]) + v_np = np.broadcast_to(x, [3, 3]) + self.assertAllEqual(v_tf.eval(), v_np) + + def testBroadcastToShape(self): + for input_dim in range(1, 6): + for output_dim in range(input_dim, 6): + with self.test_session(use_gpu=True): + input_shape = [2] * input_dim + output_shape = [2] * output_dim + x = np.array(np.random.randint(5, size=input_shape), dtype=np.int32) + v_tf = array_ops.broadcast_to(constant_op.constant(x), output_shape) + v_np = np.broadcast_to(x, output_shape) + self.assertAllEqual(v_tf.eval(), v_np) + + def testBroadcastToScalar(self): + with self.test_session(use_gpu=True): + x = np.array(1, dtype=np.int32) + v_tf = array_ops.broadcast_to(constant_op.constant(x), [3, 3]) + v_np = np.broadcast_to(x, [3, 3]) + self.assertAllEqual(v_tf.eval(), v_np) + + def testBroadcastToShapeTypeAndInference(self): + for dtype in [dtypes.int32, dtypes.int64]: + with self.test_session(use_gpu=True): + x = np.array([1, 2, 3]) + v_tf = array_ops.broadcast_to( + constant_op.constant(x), + constant_op.constant([3, 3], dtype=dtype)) + shape = v_tf.get_shape().as_list() + v_np = np.broadcast_to(x, [3, 3]) + self.assertAllEqual(v_tf.eval(), v_np) + # check shape inference when shape input is constant + self.assertAllEqual(shape, v_np.shape) + +if __name__ == "__main__": + test_lib.main() diff --git a/tensorflow/python/kernel_tests/clip_ops_test.py b/tensorflow/python/kernel_tests/clip_ops_test.py index cb1359be159f5d8983f149cf42b2723dc0581ea8..e08123b0417912c479476d8147d832d1715b8882 100644 --- a/tensorflow/python/kernel_tests/clip_ops_test.py +++ b/tensorflow/python/kernel_tests/clip_ops_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.framework import ops from tensorflow.python.ops import clip_ops from tensorflow.python.ops import gradient_checker @@ -29,19 +28,19 @@ from tensorflow.python.platform import test class ClipTest(test.TestCase): - def testClipByValueGradient(self): + def DISABLED_testClipByValueGradient(self): inputs = constant_op.constant([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32) outputs_1 = clip_ops.clip_by_value(inputs, 0.5, 3.5) min_val = constant_op.constant([0.5, 0.5, 0.5, 0.5], dtype=dtypes.float32) max_val = constant_op.constant([3.5, 3.5, 3.5, 3.5], dtype=dtypes.float32) outputs_2 = clip_ops.clip_by_value(inputs, min_val, max_val) with self.test_session(): - error_1 = gradient_checker.compute_gradient_error(inputs, [4], - outputs_1, [4]) + error_1 = gradient_checker.compute_gradient_error(inputs, [4], outputs_1, + [4]) self.assertLess(error_1, 1e-4) - error_2 = gradient_checker.compute_gradient_error(inputs, [4], - outputs_2, [4]) + error_2 = gradient_checker.compute_gradient_error(inputs, [4], outputs_2, + [4]) self.assertLess(error_2, 1e-4) # ClipByValue test @@ -56,10 +55,11 @@ class ClipTest(test.TestCase): self.assertAllClose(np_ans, tf_ans) # [Tensor, Scalar, Scalar] - def testClipByValue0Type(self): - for dtype in [dtypes.float16, dtypes.float32, dtypes.float64, - dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64, - dtypes.uint8, dtypes.uint16]: + def DISABLED_testClipByValue0Type(self): + for dtype in [ + dtypes.float16, dtypes.float32, dtypes.float64, dtypes.int8, + dtypes.int16, dtypes.int32, dtypes.int64, dtypes.uint8, dtypes.uint16 + ]: with self.test_session(use_gpu=True): x = constant_op.constant([1, 2, 3, 4, 5, 6], shape=[2, 3], dtype=dtype) np_ans = [[2, 2, 3], [4, 4, 4]] @@ -71,15 +71,16 @@ class ClipTest(test.TestCase): self.assertAllClose(np_ans, tf_ans) # [Tensor, Tensor, Scalar] - def testClipByValue1Type(self): - for dtype in [dtypes.float16, dtypes.float32, dtypes.float64, - dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64, - dtypes.uint8, dtypes.uint16]: + def DISABLED_testClipByValue1Type(self): + for dtype in [ + dtypes.float16, dtypes.float32, dtypes.float64, dtypes.int8, + dtypes.int16, dtypes.int32, dtypes.int64, dtypes.uint8, dtypes.uint16 + ]: with self.test_session(use_gpu=True): x = constant_op.constant([1, 2, 3, 4, 5, 6], shape=[2, 3], dtype=dtype) np_ans = [[2, 2, 3], [4, 4, 4]] - clip_value_min = constant_op.constant([2, 2, 2, 3, 3, 3], shape=[2, 3], - dtype=dtype) + clip_value_min = constant_op.constant( + [2, 2, 2, 3, 3, 3], shape=[2, 3], dtype=dtype) clip_value_max = 4 ans = clip_ops.clip_by_value(x, clip_value_min, clip_value_max) tf_ans = ans.eval() @@ -87,33 +88,35 @@ class ClipTest(test.TestCase): self.assertAllClose(np_ans, tf_ans) # [Tensor, Scalar, Tensor] - def testClipByValue2Type(self): - for dtype in [dtypes.float16, dtypes.float32, dtypes.float64, - dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64, - dtypes.uint8, dtypes.uint16]: + def DISABLED_testClipByValue2Type(self): + for dtype in [ + dtypes.float16, dtypes.float32, dtypes.float64, dtypes.int8, + dtypes.int16, dtypes.int32, dtypes.int64, dtypes.uint8, dtypes.uint16 + ]: with self.test_session(use_gpu=True): x = constant_op.constant([1, 2, 3, 4, 5, 6], shape=[2, 3], dtype=dtype) np_ans = [[4, 4, 4], [4, 5, 6]] clip_value_min = 4 - clip_value_max = constant_op.constant([6, 6, 6, 6, 6, 6], shape=[2, 3], - dtype=dtype) + clip_value_max = constant_op.constant( + [6, 6, 6, 6, 6, 6], shape=[2, 3], dtype=dtype) ans = clip_ops.clip_by_value(x, clip_value_min, clip_value_max) tf_ans = ans.eval() self.assertAllClose(np_ans, tf_ans) # [Tensor, Tensor, Tensor] - def testClipByValue3Type(self): - for dtype in [dtypes.float16, dtypes.float32, dtypes.float64, - dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64, - dtypes.uint8, dtypes.uint16]: + def DISABLED_testClipByValue3Type(self): + for dtype in [ + dtypes.float16, dtypes.float32, dtypes.float64, dtypes.int8, + dtypes.int16, dtypes.int32, dtypes.int64, dtypes.uint8, dtypes.uint16 + ]: with self.test_session(use_gpu=True): x = constant_op.constant([1, 2, 3, 4, 5, 6], shape=[2, 3], dtype=dtype) np_ans = [[2, 2, 3], [5, 5, 6]] - clip_value_min = constant_op.constant([2, 2, 2, 5, 5, 5], shape=[2, 3], - dtype=dtype) - clip_value_max = constant_op.constant([5, 5, 5, 7, 7, 7], shape=[2, 3], - dtype=dtype) + clip_value_min = constant_op.constant( + [2, 2, 2, 5, 5, 5], shape=[2, 3], dtype=dtype) + clip_value_max = constant_op.constant( + [5, 5, 5, 7, 7, 7], shape=[2, 3], dtype=dtype) ans = clip_ops.clip_by_value(x, clip_value_min, clip_value_max) tf_ans = ans.eval() @@ -124,15 +127,14 @@ class ClipTest(test.TestCase): x = constant_op.constant([-5.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3, 1]) # Use a nonsensical shape. clip = constant_op.constant([1.0, 2.0]) - with self.assertRaises(errors_impl.InvalidArgumentError): - ans = clip_ops.clip_by_value(x, -clip, clip) - tf_ans = ans.eval() - with self.assertRaises(errors_impl.InvalidArgumentError): - ans = clip_ops.clip_by_value(x, 1.0, clip) - tf_ans = ans.eval() + with self.assertRaises(ValueError): + _ = clip_ops.clip_by_value(x, -clip, clip) + with self.assertRaises(ValueError): + _ = clip_ops.clip_by_value(x, 1.0, clip) def testClipByValueNonFinite(self): - with self.test_session(use_gpu=True): + # TODO(b/78016351): Enable test on GPU once the bug is fixed. + with self.test_session(): x = constant_op.constant([float('NaN'), float('Inf'), -float('Inf')]) np_ans = [float('NaN'), 4.0, -4.0] clip_value = 4.0 diff --git a/tensorflow/python/kernel_tests/confusion_matrix_test.py b/tensorflow/python/kernel_tests/confusion_matrix_test.py index 670a625f0f1dd84c523de8acb17f9a410d184ad5..79e419867d70071280b7c88b6bfa820b935b24cd 100644 --- a/tensorflow/python/kernel_tests/confusion_matrix_test.py +++ b/tensorflow/python/kernel_tests/confusion_matrix_test.py @@ -19,6 +19,7 @@ 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 @@ -104,11 +105,7 @@ class ConfusionMatrixTest(test.TestCase): d, l, cm_out = sess.run([data, lab, cm], {m_neg: 0.0, m_pos: 1.0, s: 1.0}) truth = np.zeros([2, 2], dtype=np_dtype) - try: - range_builder = xrange - except NameError: # In Python 3. - range_builder = range - for i in range_builder(len(d)): + for i in xrange(len(d)): truth[l[i], d[i]] += 1 self.assertEqual(cm_out.dtype, np_dtype) diff --git a/tensorflow/python/kernel_tests/constant_op_test.py b/tensorflow/python/kernel_tests/constant_op_test.py index 749313b00d8b01eac821662ee4cfa61cff4e34b6..107ee37fabbae56c5bf715e1e7953b62ac3c526b 100644 --- a/tensorflow/python/kernel_tests/constant_op_test.py +++ b/tensorflow/python/kernel_tests/constant_op_test.py @@ -65,6 +65,11 @@ class ConstantTest(test.TestCase): self._testCpu(x) self._testGpu(x) + def testInvalidDType(self): + # Test case for GitHub issue 18474 + with self.assertRaises(TypeError): + constant_op.constant(dtypes_lib.string, "[,]") + def testBFloat16(self): bfloat16 = dtypes_lib.bfloat16.as_numpy_dtype self._testAll(np.arange(-15, 15).reshape([2, 3, 5]).astype(bfloat16)) 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 75f8644f694c4cebb7dbdac4599244dda427bc05..e27eb00818a9f4b8dd8b8c9caf5bd48ee7400928 100644 --- a/tensorflow/python/kernel_tests/control_flow_ops_py_test.py +++ b/tensorflow/python/kernel_tests/control_flow_ops_py_test.py @@ -664,6 +664,23 @@ class ControlFlowTest(test.TestCase): self.assertAllEqual(42.0, grad.eval(feed_dict={c: 1})) self.assertAllEqual(3.0, grad.eval(feed_dict={c: 3})) + def testCondGrad_3(self): + with self.test_session(): + c = array_ops.placeholder(dtypes.int32, shape=[]) + ox = constant_op.constant(10.0) + pred = math_ops.less(c, 2) + + def fn1(x): + m = x * x + return gradients_impl.gradients(m, [ox])[0] + + fn2 = lambda: math_ops.multiply(ox, 3.0) + y = math_ops.multiply(7.0, ox) + r = control_flow_ops.cond(pred, lambda: fn1(y), fn2) + + self.assertAllEqual(980.0, r.eval(feed_dict={c: 1})) + self.assertAllEqual(30.0, r.eval(feed_dict={c: 3})) + def testNestedCond_Simple(self): with self.test_session(): x = constant_op.constant(0., name="X") diff --git a/tensorflow/python/kernel_tests/conv3d_transpose_test.py b/tensorflow/python/kernel_tests/conv3d_transpose_test.py index a8b3af509622a871f0afe78c510f9ce994078866..8973a450fa246e3c924f4352d72a1bbd4f7851ea 100644 --- a/tensorflow/python/kernel_tests/conv3d_transpose_test.py +++ b/tensorflow/python/kernel_tests/conv3d_transpose_test.py @@ -119,6 +119,18 @@ class Conv3DTransposeTest(test.TestCase): target = 3.0 self.assertAllClose(target, value[n, d, h, w, k]) + def testConv3DTransposeShapeMismatch(self): + # Test case for GitHub issue 18460 + x_shape = [2, 2, 3, 4, 3] + f_shape = [3, 3, 3, 2, 2] + y_shape = [2, 2, 6, 8, 6] + strides = [1, 1, 2, 2, 2] + np.random.seed(1) + x_value = np.random.random_sample(x_shape).astype(np.float64) + f_value = np.random.random_sample(f_shape).astype(np.float64) + nn_ops.conv3d_transpose( + x_value, f_value, y_shape, strides, data_format='NCDHW') + def testConv3DTransposeValid(self): with self.test_session(): strides = [1, 2, 2, 2, 1] diff --git a/tensorflow/python/kernel_tests/conv_ops_3d_test.py b/tensorflow/python/kernel_tests/conv_ops_3d_test.py index f4616fd661f989c1c3e4939a3d062b0260f8572e..0b531125f36c6d01268081aefe3815d38720a23f 100644 --- a/tensorflow/python/kernel_tests/conv_ops_3d_test.py +++ b/tensorflow/python/kernel_tests/conv_ops_3d_test.py @@ -28,6 +28,7 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import tensor_shape from tensorflow.python.framework import test_util from tensorflow.python.ops import gradient_checker +from tensorflow.python.ops import gradients_impl from tensorflow.python.ops import nn_ops import tensorflow.python.ops.nn_grad # pylint: disable=unused-import from tensorflow.python.platform import test @@ -61,18 +62,18 @@ class Conv3DTest(test.TestCase): def _SetupValuesForDevice(self, tensor_in_sizes, filter_in_sizes, stride, padding, data_format, dtype, use_gpu): - total_size_1 = 1 - total_size_2 = 1 + total_size_tensor = 1 + total_size_filter = 1 for s in tensor_in_sizes: - total_size_1 *= s + total_size_tensor *= s for s in filter_in_sizes: - total_size_2 *= s + total_size_filter *= s # Initializes the input tensor with array containing numbers from 0 to 1. # We keep the input tensor values fairly small to avoid overflowing float16 # during the conv3d. - x1 = [f * 1.0 / total_size_1 for f in range(1, total_size_1 + 1)] - x2 = [f * 1.0 / total_size_2 for f in range(1, total_size_2 + 1)] + x1 = [f * 1.0 / total_size_tensor for f in range(1, total_size_tensor + 1)] + x2 = [f * 1.0 / total_size_filter for f in range(1, total_size_filter + 1)] with self.test_session(use_gpu=use_gpu): t1 = constant_op.constant(x1, shape=tensor_in_sizes, dtype=dtype) t2 = constant_op.constant(x2, shape=filter_in_sizes, dtype=dtype) @@ -118,6 +119,79 @@ class Conv3DTest(test.TestCase): self.assertAllClose(expected, value.flatten(), atol=tol, rtol=tol) + def _ComputeReferenceDilatedConv(self, tensor_in_sizes, filter_in_sizes, + stride, dilation, padding, data_format, + use_gpu): + total_size_tensor = 1 + total_size_filter = 1 + for s in tensor_in_sizes: + total_size_tensor *= s + for s in filter_in_sizes: + total_size_filter *= s + + # Initializes the input tensor with array containing incrementing + # numbers from 1. + x1 = [f * 1.0 for f in range(1, total_size_tensor + 1)] + x2 = [f * 1.0 for f in range(1, total_size_filter + 1)] + with self.test_session(use_gpu=use_gpu): + t1 = constant_op.constant(x1, shape=tensor_in_sizes) + t2 = constant_op.constant(x2, shape=filter_in_sizes) + if isinstance(stride, collections.Iterable): + strides = list(stride) + else: + strides = [stride, stride, stride] + if data_format == "NCDHW": + t1 = test_util.NHWCToNCHW(t1) + full_strides = [1, 1] + strides + full_dilation = [1, 1] + dilation + else: + full_strides = [1] + strides + [1] + full_dilation = [1] + dilation + [1] + expected = nn_ops.convolution( + t1, + t2, + padding=padding, + strides=strides, + dilation_rate=dilation, + data_format=data_format) + computed = nn_ops.conv3d( + t1, + t2, + strides=full_strides, + dilations=full_dilation, + padding=padding, + data_format=data_format) + if data_format == "NCDHW": + expected = test_util.NCHWToNHWC(expected) + computed = test_util.NCHWToNHWC(computed) + return expected, computed + + def _VerifyDilatedConvValues(self, tensor_in_sizes, filter_in_sizes, stride, + padding, dilations): + expected_results = [] + computed_results = [] + default_dilations = ( + dilations[0] == 1 and dilations[1] == 1 and dilations[2] == 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, stride, dilations, padding, + data_format, use_gpu) + expected_results.append(expected) + computed_results.append(computed) + tolerance = 1e-2 if use_gpu else 1e-5 + with self.test_session() as sess: + expected_values = sess.run(expected_results) + computed_values = sess.run(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-6) + def testConv3D1x1x1Filter(self): expected_output = [ 0.18518519, 0.22222222, 0.25925926, 0.40740741, 0.5, 0.59259259, @@ -145,6 +219,15 @@ class Conv3DTest(test.TestCase): padding="VALID", expected=expected_output) + def testConv3D1x1x1Filter2x1x1Dilation(self): + if test.is_gpu_available(cuda_only=True): + self._VerifyDilatedConvValues( + tensor_in_sizes=[1, 3, 6, 1, 1], + filter_in_sizes=[1, 1, 1, 1, 1], + stride=1, + padding="VALID", + dilations=[2, 1, 1]) + # Expected values computed using scipy's correlate function. def testConv3D2x2x2Filter(self): expected_output = [ @@ -161,6 +244,15 @@ class Conv3DTest(test.TestCase): padding="VALID", expected=expected_output) + def testConv3D2x2x2Filter1x2x1Dilation(self): + if test.is_gpu_available(cuda_only=True): + self._VerifyDilatedConvValues( + tensor_in_sizes=[1, 4, 6, 3, 1], + filter_in_sizes=[2, 2, 2, 1, 1], + stride=1, + padding="VALID", + dilations=[1, 2, 1]) + def testConv3DStrides(self): expected_output = [ 0.06071429, 0.08988095, 0.10238095, 0.11488095, 0.12738095, 0.13988095, @@ -546,6 +638,98 @@ class Conv3DTest(test.TestCase): padding="SAME", test_input=False) + # Testing for backprops + def _RunAndVerifyBackprop(self, input_sizes, filter_sizes, output_sizes, + strides, dilations, padding, data_format, use_gpu, + err, mode): + total_input_size = 1 + total_filter_size = 1 + for s in input_sizes: + total_input_size *= s + for s in filter_sizes: + total_filter_size *= s + # Initializes the input tensor with array containing incrementing + # numbers from 1. + x1 = [f * 1.0 for f in range(1, total_input_size + 1)] + x2 = [f * 1.0 for f in range(1, total_filter_size + 1)] + default_dilations = ( + dilations[0] == 1 and dilations[1] == 1 and dilations[2] == 1) + + # 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: + with self.test_session(use_gpu=use_gpu) as sess: + if data_format == "NCDHW": + input_sizes = test_util.NHWCToNCHW(input_sizes) + t1 = constant_op.constant(x1, shape=input_sizes) + t2 = constant_op.constant(x2, shape=filter_sizes) + full_strides = [1] + strides + [1] + full_dilations = [1] + dilations + [1] + if data_format == "NCDHW": + full_strides = test_util.NHWCToNCHW(full_strides) + full_dilations = test_util.NHWCToNCHW(full_dilations) + actual = nn_ops.conv3d( + t1, + t2, + strides=full_strides, + dilations=full_dilations, + padding=padding, + data_format=data_format) + expected = nn_ops.convolution( + t1, + t2, + padding=padding, + strides=strides, + dilation_rate=dilations, + data_format=data_format) + if data_format == "NCDHW": + actual = test_util.NCHWToNHWC(actual) + expected = test_util.NCHWToNHWC(expected) + actual_grad = gradients_impl.gradients(actual, t1 + if mode == "input" else t2)[0] + expected_grad = gradients_impl.gradients(expected, t1 + if mode == "input" else t2)[0] + # "values" consists of two tensors for two backprops + actual_value = sess.run(actual_grad) + expected_value = sess.run(expected_grad) + self.assertShapeEqual(actual_value, actual_grad) + self.assertShapeEqual(expected_value, expected_grad) + print("expected = ", expected_value) + print("actual = ", actual_value) + self.assertArrayNear(expected_value.flatten(), actual_value.flatten(), + err) + + def testConv3D2x2Depth3ValidBackpropFilterStride1x1Dilation2x1(self): + if test.is_gpu_available(cuda_only=True): + for (data_format, use_gpu) in GetTestConfigs(): + self._RunAndVerifyBackprop( + input_sizes=[1, 3, 6, 1, 1], + filter_sizes=[2, 2, 1, 1, 1], + output_sizes=[1, 1, 5, 1, 1], + strides=[1, 1, 1], + dilations=[2, 1, 1], + padding="VALID", + data_format=data_format, + use_gpu=use_gpu, + err=1e-5, + mode="filter") + + def testConv3D2x2Depth3ValidBackpropInputStride1x1Dilation2x1(self): + if test.is_gpu_available(cuda_only=True): + for (data_format, use_gpu) in GetTestConfigs(): + self._RunAndVerifyBackprop( + input_sizes=[1, 3, 6, 1, 1], + filter_sizes=[2, 2, 1, 1, 1], + output_sizes=[1, 1, 5, 1, 1], + strides=[1, 1, 1], + dilations=[2, 1, 1], + padding="VALID", + data_format=data_format, + use_gpu=use_gpu, + err=1e-5, + mode="input") + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/kernel_tests/cwise_ops_test.py b/tensorflow/python/kernel_tests/cwise_ops_test.py index 34e77512434ea26d2693f0bb1600ff81dd15e84f..87da89831c8ded9b8382c7bb251948b6d202300e 100644 --- a/tensorflow/python/kernel_tests/cwise_ops_test.py +++ b/tensorflow/python/kernel_tests/cwise_ops_test.py @@ -398,14 +398,17 @@ class UnaryOpTest(test.TestCase): self._compareCpu(x, np.abs, _ABS) self._compareCpu(x, np.negative, math_ops.negative) self._compareCpu(x, np.negative, _NEG) - self._compareCpu(x, np.square, math_ops.square) self._compareCpu(x, np.sign, math_ops.sign) self._compareBothSparse(x, np.abs, math_ops.abs) self._compareBothSparse(x, np.negative, math_ops.negative) - self._compareBothSparse(x, np.square, math_ops.square) self._compareBothSparse(x, np.sign, math_ops.sign) + def testInt64Square(self): + x = np.arange(-6 << 20, 6 << 20, 2 << 20).reshape(1, 3, 2).astype(np.int64) + self._compareCpu(x, np.square, math_ops.square) + self._compareBothSparse(x, np.square, math_ops.square) + def testComplex64Basic(self): x = np.complex(1, 1) * np.arange(-3, 3).reshape(1, 3, 2).astype( np.complex64) diff --git a/tensorflow/python/kernel_tests/decode_csv_op_test.py b/tensorflow/python/kernel_tests/decode_csv_op_test.py index fec52fa9cc7bcab1da67e797c2e121edac8c9345..4f49d726765e6019715a9b40f531b82df7f33126 100644 --- a/tensorflow/python/kernel_tests/decode_csv_op_test.py +++ b/tensorflow/python/kernel_tests/decode_csv_op_test.py @@ -78,9 +78,11 @@ class DecodeCSVOpTest(test.TestCase): self._test(args, expected_out) def test2DNoQuoteDelimiter(self): - args = {"records": [["1", "2"], ['""', '"']], - "record_defaults": [[""]], - "use_quote_delim": False} + args = { + "records": [["1", "2"], ['""', '"']], + "record_defaults": [[""]], + "use_quote_delim": False + } expected_out = [[[b"1", b"2"], [b'""', b'"']]] self._test(args, expected_out) @@ -88,8 +90,7 @@ class DecodeCSVOpTest(test.TestCase): def testDouble(self): args = { "records": ["1.0", "-1.79e+308", '"1.79e+308"'], - "record_defaults": [np.array( - [], dtype=np.double)], + "record_defaults": [np.array([], dtype=np.double)], } expected_out = [[1.0, -1.79e+308, 1.79e+308]] @@ -99,8 +100,7 @@ class DecodeCSVOpTest(test.TestCase): def testInt64(self): args = { "records": ["1", "2", '"2147483648"'], - "record_defaults": [np.array( - [], dtype=np.int64)], + "record_defaults": [np.array([], dtype=np.int64)], } expected_out = [[1, 2, 2147483648]] @@ -173,8 +173,7 @@ class DecodeCSVOpTest(test.TestCase): def testWithoutDefaultsError(self): args = { "records": [",1", "0.2,3", "3.0,"], - "record_defaults": [[1.0], np.array( - [], dtype=np.int32)] + "record_defaults": [[1.0], np.array([], dtype=np.int32)] } self._test( @@ -183,8 +182,7 @@ class DecodeCSVOpTest(test.TestCase): def testWrongFieldIntError(self): args = { "records": [",1", "0.2,234a", "3.0,2"], - "record_defaults": [[1.0], np.array( - [], dtype=np.int32)] + "record_defaults": [[1.0], np.array([], dtype=np.int32)] } self._test( @@ -202,8 +200,7 @@ class DecodeCSVOpTest(test.TestCase): def testWrongFieldFloatError(self): args = { "records": [",1", "0.2,2", "3.0adf,3"], - "record_defaults": [[1.0], np.array( - [], dtype=np.int32)] + "record_defaults": [[1.0], np.array([], dtype=np.int32)] } self._test( @@ -229,6 +226,73 @@ class DecodeCSVOpTest(test.TestCase): self._test( args, expected_err_re="Quoted field has to end with quote followed.*") + def testSelectCols(self): + args = { + "records": [",,", "4,5,6"], + "record_defaults": [[1], [2]], + "select_cols": [0, 1] + } + expected_out = [[1, 4], [2, 5]] + self._test(args, expected_out) + + def testSelectColsInclLast(self): + # The last col is a edge-casey; add test for that + args = { + "records": [",,", "4,5,6"], + "record_defaults": [[0], [1], [2]], + "select_cols": [0, 1, 2] + } + expected_out = [[0, 4], [1, 5], [2, 6]] + self._test(args, expected_out) + + def testWrongSelectColsInclLast(self): + # The last col is a edge-casey; add test for that + args = { + "records": [",,", "4,5,6"], + "record_defaults": [[0], [1], [2]], + "select_cols": [0, 1, 3] + } + self._test(args, expected_err_re="Expect 3 fields but have 2 in record 0") + + def testWrongSelectColsLen(self): + args = { + "records": ["1,2,3", "4,5,6"], + "record_defaults": [[0], [0], [0]], + "select_cols": [0] + } + with self.assertRaisesWithPredicateMatch( + ValueError, "Length of select_cols and record_defaults do not match."): + self._test(args) + + def testWrongSelectColsSorting(self): + args = { + "records": ["1,2,3"], + "record_defaults": [[0], [1]], + "select_cols": [1, 0] + } + with self.assertRaisesWithPredicateMatch( + ValueError, "select_cols is not strictly increasing."): + self._test(args) + + def testWrongSelectColsIndicesNegative(self): + args = { + "records": ["1,2,3"], + "record_defaults": [[0], [1]], + "select_cols": [-1, 0] # -1 is not a valid index + } + with self.assertRaisesWithPredicateMatch( + ValueError, "select_cols contains negative values."): + self._test(args) + + def testWrongSelectColsIndicesTooHigh(self): + args = { + "records": ["1,2,3"], + "record_defaults": [[0], [1]], + "select_cols": [0, 3] # 3 is not a valid index + } + # Only successfully parses one of the columns + self._test(args, expected_err_re="Expect 2 fields but have 1 in record 0") + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/kernel_tests/distributions/bijector_test.py b/tensorflow/python/kernel_tests/distributions/bijector_test.py index 9f9fb5c0bb4c0e9d68ddf6034a8649ad5a6bd8e9..18582241e2fb69dffc0b66aa361aa77fbb97944f 100644 --- a/tensorflow/python/kernel_tests/distributions/bijector_test.py +++ b/tensorflow/python/kernel_tests/distributions/bijector_test.py @@ -20,6 +20,7 @@ from __future__ import print_function import abc +import numpy as np import six from tensorflow.python.framework import constant_op @@ -43,11 +44,10 @@ class BaseBijectorTest(test.TestCase): """Minimal specification of a `Bijector`.""" def __init__(self): - super(_BareBonesBijector, self).__init__() + super(_BareBonesBijector, self).__init__(forward_min_event_ndims=0) with self.test_session() as sess: bij = _BareBonesBijector() - self.assertEqual(None, bij.event_ndims) self.assertEqual([], bij.graph_parents) self.assertEqual(False, bij.is_constant_jacobian) self.assertEqual(False, bij.validate_args) @@ -67,13 +67,21 @@ class BaseBijectorTest(test.TestCase): self.assertAllEqual(shape, inverse_event_shape_) self.assertAllEqual(shape, bij.inverse_event_shape(shape)) - for fn in ["forward", - "inverse", - "inverse_log_det_jacobian", - "forward_log_det_jacobian"]: - with self.assertRaisesRegexp( - NotImplementedError, fn + " not implemented"): - getattr(bij, fn)(0) + with self.assertRaisesRegexp( + NotImplementedError, "inverse not implemented"): + bij.inverse(0) + + with self.assertRaisesRegexp( + NotImplementedError, "forward not implemented"): + bij.forward(0) + + with self.assertRaisesRegexp( + NotImplementedError, "inverse_log_det_jacobian not implemented"): + bij.inverse_log_det_jacobian(0, event_ndims=0) + + with self.assertRaisesRegexp( + NotImplementedError, "forward_log_det_jacobian not implemented"): + bij.forward_log_det_jacobian(0, event_ndims=0) class IntentionallyMissingError(Exception): @@ -85,7 +93,7 @@ class BrokenBijector(bijector.Bijector): def __init__(self, forward_missing=False, inverse_missing=False): super(BrokenBijector, self).__init__( - event_ndims=0, validate_args=False, name="broken") + validate_args=False, forward_min_event_ndims=0, name="broken") self._forward_missing = forward_missing self._inverse_missing = inverse_missing @@ -120,35 +128,42 @@ class BijectorCachingTestBase(object): def testCachingOfForwardResults(self): broken_bijector = self.broken_bijector_cls(inverse_missing=True) - with self.test_session(): - x = constant_op.constant(1.1) + x = constant_op.constant(1.1) + + # Call forward and forward_log_det_jacobian one-by-one (not together). + y = broken_bijector.forward(x) + _ = broken_bijector.forward_log_det_jacobian(x, event_ndims=0) - # Call forward and forward_log_det_jacobian one-by-one (not together). - y = broken_bijector.forward(x) - _ = broken_bijector.forward_log_det_jacobian(x) + # Now, everything should be cached if the argument is y. + broken_bijector.inverse_log_det_jacobian(y, event_ndims=0) + try: + broken_bijector.inverse(y) + broken_bijector.inverse_log_det_jacobian(y, event_ndims=0) + except IntentionallyMissingError: + raise AssertionError("Tests failed! Cached values not used.") - # Now, everything should be cached if the argument is y. - try: - broken_bijector.inverse(y) - broken_bijector.inverse_log_det_jacobian(y) - except IntentionallyMissingError: - raise AssertionError("Tests failed! Cached values not used.") + # Different event_ndims should not be cached. + with self.assertRaises(IntentionallyMissingError): + broken_bijector.inverse_log_det_jacobian(y, event_ndims=1) def testCachingOfInverseResults(self): broken_bijector = self.broken_bijector_cls(forward_missing=True) - with self.test_session(): - y = constant_op.constant(1.1) + y = constant_op.constant(1.1) - # Call inverse and inverse_log_det_jacobian one-by-one (not together). - x = broken_bijector.inverse(y) - _ = broken_bijector.inverse_log_det_jacobian(y) + # Call inverse and inverse_log_det_jacobian one-by-one (not together). + x = broken_bijector.inverse(y) + _ = broken_bijector.inverse_log_det_jacobian(y, event_ndims=0) - # Now, everything should be cached if the argument is x. - try: - broken_bijector.forward(x) - broken_bijector.forward_log_det_jacobian(x) - except IntentionallyMissingError: - raise AssertionError("Tests failed! Cached values not used.") + # Now, everything should be cached if the argument is x. + try: + broken_bijector.forward(x) + broken_bijector.forward_log_det_jacobian(x, event_ndims=0) + except IntentionallyMissingError: + raise AssertionError("Tests failed! Cached values not used.") + + # Different event_ndims should not be cached. + with self.assertRaises(IntentionallyMissingError): + broken_bijector.forward_log_det_jacobian(x, event_ndims=1) class BijectorCachingTest(BijectorCachingTestBase, test.TestCase): @@ -159,5 +174,107 @@ class BijectorCachingTest(BijectorCachingTestBase, test.TestCase): return BrokenBijector +class ExpOnlyJacobian(bijector.Bijector): + """Only used for jacobian calculations.""" + + def __init__(self, forward_min_event_ndims=0): + super(ExpOnlyJacobian, self).__init__( + validate_args=False, + is_constant_jacobian=False, + forward_min_event_ndims=forward_min_event_ndims, + name="exp") + + def _inverse_log_det_jacobian(self, y): + return -math_ops.log(y) + + def _forward_log_det_jacobian(self, x): + return math_ops.log(x) + + +class ConstantJacobian(bijector.Bijector): + """Only used for jacobian calculations.""" + + def __init__(self, forward_min_event_ndims=0): + super(ConstantJacobian, self).__init__( + validate_args=False, + is_constant_jacobian=True, + forward_min_event_ndims=forward_min_event_ndims, + name="c") + + def _inverse_log_det_jacobian(self, y): + return constant_op.constant(2., y.dtype) + + def _forward_log_det_jacobian(self, x): + return constant_op.constant(-2., x.dtype) + + +class BijectorReduceEventDimsTest(test.TestCase): + """Test caching with BrokenBijector.""" + + def testReduceEventNdimsForward(self): + x = [[[1., 2.], [3., 4.]]] + bij = ExpOnlyJacobian() + self.assertAllClose( + np.log(x), + self.evaluate(bij.forward_log_det_jacobian(x, event_ndims=0))) + self.assertAllClose( + np.sum(np.log(x), axis=-1), + self.evaluate(bij.forward_log_det_jacobian(x, event_ndims=1))) + self.assertAllClose( + np.sum(np.log(x), axis=(-1, -2)), + self.evaluate(bij.forward_log_det_jacobian(x, event_ndims=2))) + + def testReduceEventNdimsForwardRaiseError(self): + x = [[[1., 2.], [3., 4.]]] + bij = ExpOnlyJacobian(forward_min_event_ndims=1) + with self.assertRaisesRegexp(ValueError, "must be larger than"): + bij.forward_log_det_jacobian(x, event_ndims=0) + + def testReduceEventNdimsInverse(self): + x = [[[1., 2.], [3., 4.]]] + bij = ExpOnlyJacobian() + self.assertAllClose( + -np.log(x), + self.evaluate(bij.inverse_log_det_jacobian(x, event_ndims=0))) + self.assertAllClose( + np.sum(-np.log(x), axis=-1), + self.evaluate(bij.inverse_log_det_jacobian(x, event_ndims=1))) + self.assertAllClose( + np.sum(-np.log(x), axis=(-1, -2)), + self.evaluate(bij.inverse_log_det_jacobian(x, event_ndims=2))) + + def testReduceEventNdimsInverseRaiseError(self): + x = [[[1., 2.], [3., 4.]]] + bij = ExpOnlyJacobian(forward_min_event_ndims=1) + with self.assertRaisesRegexp(ValueError, "must be larger than"): + bij.inverse_log_det_jacobian(x, event_ndims=0) + + def testReduceEventNdimsForwardConstJacobian(self): + x = [[[1., 2.], [3., 4.]]] + bij = ConstantJacobian() + self.assertAllClose( + -2., + self.evaluate(bij.forward_log_det_jacobian(x, event_ndims=0))) + self.assertAllClose( + -4., + self.evaluate(bij.forward_log_det_jacobian(x, event_ndims=1))) + self.assertAllClose( + -8., + self.evaluate(bij.forward_log_det_jacobian(x, event_ndims=2))) + + def testReduceEventNdimsInverseConstJacobian(self): + x = [[[1., 2.], [3., 4.]]] + bij = ConstantJacobian() + self.assertAllClose( + 2., + self.evaluate(bij.inverse_log_det_jacobian(x, event_ndims=0))) + self.assertAllClose( + 4., + self.evaluate(bij.inverse_log_det_jacobian(x, event_ndims=1))) + self.assertAllClose( + 8., + self.evaluate(bij.inverse_log_det_jacobian(x, event_ndims=2))) + + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/kernel_tests/distributions/identity_bijector_test.py b/tensorflow/python/kernel_tests/distributions/identity_bijector_test.py index e8f9d0b728d8f831becc82cdba0ae2bf3d5da52a..b347c20db25df6dc0f278d9b34b4588277104850 100644 --- a/tensorflow/python/kernel_tests/distributions/identity_bijector_test.py +++ b/tensorflow/python/kernel_tests/distributions/identity_bijector_test.py @@ -27,14 +27,19 @@ class IdentityBijectorTest(test.TestCase): """Tests correctness of the Y = g(X) = X transformation.""" def testBijector(self): - with self.test_session(): - bijector = identity_bijector.Identity() - self.assertEqual("identity", bijector.name) - x = [[[0.], [1.]]] - self.assertAllEqual(x, bijector.forward(x).eval()) - self.assertAllEqual(x, bijector.inverse(x).eval()) - self.assertAllEqual(0., bijector.inverse_log_det_jacobian(x).eval()) - self.assertAllEqual(0., bijector.forward_log_det_jacobian(x).eval()) + bijector = identity_bijector.Identity(validate_args=True) + self.assertEqual("identity", bijector.name) + x = [[[0.], [1.]]] + self.assertAllEqual(x, self.evaluate(bijector.forward(x))) + self.assertAllEqual(x, self.evaluate(bijector.inverse(x))) + self.assertAllEqual( + 0., + self.evaluate( + bijector.inverse_log_det_jacobian(x, event_ndims=3))) + self.assertAllEqual( + 0., + self.evaluate( + bijector.forward_log_det_jacobian(x, event_ndims=3))) def testScalarCongruency(self): with self.test_session(): diff --git a/tensorflow/python/kernel_tests/gather_op_test.py b/tensorflow/python/kernel_tests/gather_op_test.py index 9a946925693370912613f4dde33bbbda176060e4..a2fcd751dfa94605d271587640815fae6ac1c360 100644 --- a/tensorflow/python/kernel_tests/gather_op_test.py +++ b/tensorflow/python/kernel_tests/gather_op_test.py @@ -149,6 +149,15 @@ class GatherTest(test.TestCase): self.assertAllEqual([b"asdf", b"qwer"], array_ops.gather(params, 0, axis=1).eval()) + def testUInt32AndUInt64(self): + for unsigned_type in (dtypes.uint32, dtypes.uint64): + params = self._buildParams( + np.array([[1, 2, 3], [7, 8, 9]]), unsigned_type) + with self.test_session(): + self.assertAllEqual([7, 8, 9], + array_ops.gather(params, 1, axis=0).eval()) + self.assertAllEqual([1, 7], array_ops.gather(params, 0, axis=1).eval()) + def testUnknownIndices(self): params = constant_op.constant([[0, 1, 2]]) indices = array_ops.placeholder(dtypes.int32) diff --git a/tensorflow/python/kernel_tests/init_ops_test.py b/tensorflow/python/kernel_tests/init_ops_test.py index 1e5c118cbc3573af0a2ce95239f499a5e52a0c86..a9b55854f1b4a3dfc49f05397ca32bc7b2ccb88e 100644 --- a/tensorflow/python/kernel_tests/init_ops_test.py +++ b/tensorflow/python/kernel_tests/init_ops_test.py @@ -551,7 +551,6 @@ class OrthogonalInitializerTest(test.TestCase): init2 = init_ops.orthogonal_initializer(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) @@ -610,15 +609,16 @@ class ConvolutionDeltaOrthogonalInitializerTest(test.TestCase): 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): + gain = 3.14 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 outputs. + # Check orthogonality by computing ratio between + # the 2-norms of the inputs and outputs. if len(kernel_size) == 1: shape = [4, 32, 64] convolution = convolutional.conv1d @@ -634,9 +634,10 @@ class ConvolutionDeltaOrthogonalInitializerTest(test.TestCase): inputs, padding="same", filters=128, kernel_size=kernel_size, use_bias=False, kernel_initializer=init_ops.convolutional_delta_orthogonal( - gain=3.14)) + gain=gain)) outputs_shape = shape[0:-1] + [128] outputs_2norm = linalg_ops.norm(outputs) + ratio = outputs_2norm / inputs_2norm my_ops = variables.global_variables_initializer() with self.test_session(use_gpu=True) as sess: sess.run(my_ops) @@ -644,10 +645,8 @@ class ConvolutionDeltaOrthogonalInitializerTest(test.TestCase): 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) + self.assertAllClose(sess.run(ratio), np.sqrt(gain), + rtol=tol, atol=tol) def testNonuniformity(self): value = 0 @@ -655,7 +654,7 @@ class ConvolutionDeltaOrthogonalInitializerTest(test.TestCase): shape = [3, 3, 10, 10] count = 70 tol = 1e-5 - with self.test_session(use_gpu=True): # as sess: + with self.test_session(use_gpu=True): for i in range(count): x = variable_scope.get_variable("{}".format(i), shape=shape, initializer= @@ -674,6 +673,340 @@ class ConvolutionDeltaOrthogonalInitializerTest(test.TestCase): self.assertAllClose(abs_value, count, rtol=tol, atol=tol) +class ConvolutionOrthogonal1dInitializerTest(test.TestCase): + + def testInitializerIdentical(self): + for dtype in [dtypes.float32, dtypes.float64]: + init1 = init_ops.convolutional_orthogonal_1d(seed=1, dtype=dtype) + init2 = init_ops.convolutional_orthogonal_1d(seed=1, dtype=dtype) + self.assertTrue(identicaltest(self, init1, init2, (3, 10, 10))) + + def testInitializerDifferent(self): + for dtype in [dtypes.float32, dtypes.float64]: + init1 = init_ops.convolutional_orthogonal_1d(seed=1, dtype=dtype) + init2 = init_ops.convolutional_orthogonal_1d(seed=2, dtype=dtype) + self.assertFalse(identicaltest(self, init1, init2, (3, 10, 10))) + + def testDuplicatedInitializer(self): + init = init_ops.convolutional_orthogonal_1d() + self.assertFalse(duplicated_initializer(self, init, 1, (3, 10, 10))) + + def testInvalidDataType(self): + self.assertRaises( + ValueError, init_ops.convolutional_orthogonal_1d, + dtype=dtypes.string) + + def testInvalidShape(self): + init1 = init_ops.convolutional_orthogonal_1d() + with self.test_session(graph=ops.Graph(), use_gpu=True): + self.assertRaises(ValueError, init1, shape=[3, 6, 5]) + + def testGain(self): + shape = (3, 10, 10) + for dtype in [dtypes.float32, dtypes.float64]: + init1 = init_ops.convolutional_orthogonal_1d(seed=1, dtype=dtype) + init2 = init_ops.convolutional_orthogonal_1d(gain=3.14, + seed=1, dtype=dtype) + with self.test_session(graph=ops.Graph(), use_gpu=True): + t1 = init1(shape).eval() + t2 = init2(shape).eval() + return np.allclose(t1, t2 / 3.14, rtol=1e-15, atol=1e-15) + + def testNonuniformity(self): + value = 0 + abs_value = 0 + shape = [3, 10, 10] + count = 70 + tol = 1e-5 + with self.test_session(use_gpu=True): + for i in range(count): + x = variable_scope.get_variable("{}".format(i), shape=shape, + initializer= + init_ops.convolutional_orthogonal_1d) + x.initializer.run() + y = np.sum(x.eval(), axis=0) + determinant = np.linalg.det(y) + value += determinant + abs_value += np.abs(determinant) + + # Check there is some variation in the signs of the determinants. + self.assertLess(value, count - tol) + self.assertLess(-count + tol, value) + # Check all determinants have absolute value 1 + # Compute the sum of the absolute values of 'count' determinants + self.assertAllClose(abs_value, count, rtol=tol, atol=tol) + + def testShapesValues(self): + def circular_pad(input_, width, kernel_size): + """Pad input_ for computing (circular) convolution. + + Args: + input_: the input tensor + width: the width of the tensor. + kernel_size: the kernel size of the filter. + Returns: + a tensor whose width is (width + kernel_size - 1). + """ + + beginning = kernel_size // 2 + end = kernel_size - 1 - beginning + + tmp_up = array_ops.slice(input_, [0, width - beginning, 0], + [-1, beginning, -1]) + tmp_down = array_ops.slice(input_, [0, 0, 0], [-1, end, -1]) + tmp = array_ops.concat([tmp_up, input_, tmp_down], 1) + + return tmp + + cout = 64 + shape = [10, 20, 32] + outputs_shape = shape[0:-1] + [cout] + dtype = dtypes.float32 + tol = 1e-3 + gain = 3.14 + # Check orthogonality/isometry by computing the ratio between + # the 2-norms of the inputs and ouputs. + for kernel_size in [[1], [2], [3], [4], [5], [6]]: + convolution = convolutional.conv1d + inputs = random_ops.random_normal(shape, dtype=dtype) + inputs_2norm = linalg_ops.norm(inputs) + input_with_circular_pad = circular_pad(inputs, shape[1], kernel_size[0]) + outputs = convolution( + input_with_circular_pad, padding="valid", filters=cout, + kernel_size=kernel_size[0], use_bias=False, + kernel_initializer=init_ops.convolutional_orthogonal_1d(gain=gain)) + outputs_2norm = linalg_ops.norm(outputs) + ratio = outputs_2norm / inputs_2norm + 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 orthogonal kernel. + self.assertAllClose(sess.run(ratio), np.sqrt(gain), rtol=tol, atol=tol) + + +class ConvolutionOrthogonal2dInitializerTest(test.TestCase): + + def testInitializerIdentical(self): + for dtype in [dtypes.float32, dtypes.float64]: + init1 = init_ops.convolutional_orthogonal_2d(seed=1, dtype=dtype) + init2 = init_ops.convolutional_orthogonal_2d(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_orthogonal_2d(seed=1, dtype=dtype) + init2 = init_ops.convolutional_orthogonal_2d(seed=2, dtype=dtype) + self.assertFalse(identicaltest(self, init1, init2, (3, 3, 10, 10))) + + def testDuplicatedInitializer(self): + init = init_ops.convolutional_orthogonal_2d() + self.assertFalse(duplicated_initializer(self, init, 1, (3, 3, 10, 10))) + + def testInvalidDataType(self): + self.assertRaises( + ValueError, init_ops.convolutional_orthogonal_2d, + dtype=dtypes.string) + + def testInvalidShape(self): + init1 = init_ops.convolutional_orthogonal_2d() + 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_orthogonal_2d(seed=1, dtype=dtype) + init2 = init_ops.convolutional_orthogonal_2d(gain=3.14, + seed=1, dtype=dtype) + with self.test_session(graph=ops.Graph(), use_gpu=True): + t1 = init1(shape).eval() + t2 = init2(shape).eval() + return np.allclose(t1, t2 / 3.14, rtol=1e-15, atol=1e-15) + + def testShapesValues(self): + def circular_pad(input_, width, kernel_size): + """Pad input_ for computing (circular) convolution. + + Args: + input_: the input tensor + width: the width of the tensor. + kernel_size: the kernel size of the filter. + Returns: + a tensor whose width is (width + kernel_size - 1). + """ + beginning = kernel_size // 2 + end = kernel_size - 1 - beginning + + tmp_up = array_ops.slice(input_, [0, width - beginning, 0, 0], + [-1, beginning, width, -1]) + tmp_down = array_ops.slice(input_, [0, 0, 0, 0], [-1, end, width, -1]) + tmp = array_ops.concat([tmp_up, input_, tmp_down], 1) + + new_width = width + kernel_size - 1 + tmp_left = array_ops.slice(tmp, [0, 0, width - beginning, 0], + [-1, new_width, beginning, -1]) + tmp_right = array_ops.slice(tmp, [0, 0, 0, 0], [-1, new_width, end, -1]) + + final = array_ops.concat([tmp_left, tmp, tmp_right], 2) + return final + + cout = 45 + shape = [64, 28, 28, 32] + outputs_shape = shape[0:-1] + [cout] + dtype = dtypes.float32 + tol = 1e-3 + gain = 3.14 + # Check orthogonality/isometry by computing the ratio between + # the 2-norms of the inputs and ouputs. + for kernel_size in [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]]: + convolution = convolutional.conv2d + inputs = random_ops.random_normal(shape, dtype=dtype) + inputs_2norm = linalg_ops.norm(inputs) + input_with_circular_pad = circular_pad(inputs, shape[1], kernel_size[0]) + outputs = convolution( + input_with_circular_pad, padding="valid", filters=cout, + kernel_size=kernel_size, use_bias=False, + kernel_initializer=init_ops.convolutional_orthogonal_2d(gain=gain)) + outputs_2norm = linalg_ops.norm(outputs) + ratio = outputs_2norm / inputs_2norm + 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 orthogonal kernel. + self.assertAllClose(sess.run(ratio), np.sqrt(gain), rtol=tol, atol=tol) + + +class ConvolutionOrthogonal3dInitializerTest(test.TestCase): + + def testInitializerIdentical(self): + for dtype in [dtypes.float32, dtypes.float64]: + init1 = init_ops.convolutional_orthogonal_3d(seed=1, dtype=dtype) + init2 = init_ops.convolutional_orthogonal_3d(seed=1, dtype=dtype) + self.assertTrue(identicaltest(self, init1, init2, (3, 3, 3, 10, 10))) + + def testInitializerDifferent(self): + for dtype in [dtypes.float32, dtypes.float64]: + init1 = init_ops.convolutional_orthogonal_3d(seed=1, dtype=dtype) + init2 = init_ops.convolutional_orthogonal_3d(seed=2, dtype=dtype) + self.assertFalse(identicaltest(self, init1, init2, (3, 3, 3, 10, 10))) + + def testDuplicatedInitializer(self): + init = init_ops.convolutional_orthogonal_3d() + self.assertFalse(duplicated_initializer(self, init, 1, (3, 3, 3, 10, 10))) + + def testInvalidDataType(self): + self.assertRaises( + ValueError, init_ops.convolutional_orthogonal_3d, + dtype=dtypes.string) + + def testInvalidShape(self): + init1 = init_ops.convolutional_orthogonal_3d() + with self.test_session(graph=ops.Graph(), use_gpu=True): + self.assertRaises(ValueError, init1, shape=[3, 3, 3, 6, 5]) + + def testGain(self): + shape = (3, 3, 3, 10, 10) + for dtype in [dtypes.float32, dtypes.float64]: + init1 = init_ops.convolutional_orthogonal_3d(seed=1, dtype=dtype) + init2 = init_ops.convolutional_orthogonal_3d(gain=3.14, + seed=1, dtype=dtype) + with self.test_session(graph=ops.Graph(), use_gpu=True): + t1 = init1(shape).eval() + t2 = init2(shape).eval() + return np.allclose(t1, t2 / 3.14, rtol=1e-15, atol=1e-15) + + def testNonuniformity(self): + value = 0 + abs_value = 0 + shape = [3, 3, 3, 5, 5] + count = 20 + tol = 1e-5 + with self.test_session(use_gpu=True): + for i in range(count): + x = variable_scope.get_variable("{}".format(i), shape=shape, + initializer= + init_ops.convolutional_orthogonal_3d) + x.initializer.run() + y = np.sum(x.eval(), axis=(0, 1, 2)) + determinant = np.linalg.det(y) + value += determinant + abs_value += np.abs(determinant) + + # Check there is some variation in the signs of the determinants + self.assertLess(value, count - tol) + self.assertLess(-count + tol, value) + # Check all determinants have absolute value 1 + # Compute the sum of the absolute values of 'count' determinants + self.assertAllClose(abs_value, count, rtol=tol, atol=tol) + + def testShapesValues(self): + def circular_pad(input_, width, kernel_size): + """Padding input_ for computing circular convolution. + + Args: + input_: the input tensor + width: the width of the tensor. + kernel_size: the kernel size of the filter. + + Returns: + a tensor whose width is (width + kernel_size - 1). + """ + + beginning = kernel_size // 2 + end = kernel_size - 1 - beginning + + tmp_up = array_ops.slice(input_, [0, width - beginning, 0, 0, 0], + [-1, beginning, -1, -1, -1]) + tmp_down = array_ops.slice(input_, [0, 0, 0, 0, 0], + [-1, end, -1, -1, -1]) + tmp = array_ops.concat([tmp_up, input_, tmp_down], 1) + + tmp_left = array_ops.slice(tmp, [0, 0, width - beginning, 0, 0], + [-1, -1, beginning, -1, -1]) + tmp_right = array_ops.slice(tmp, [0, 0, 0, 0, 0], + [-1, -1, end, -1, -1]) + tmp = array_ops.concat([tmp_left, tmp, tmp_right], 2) + + tmp_front = array_ops.slice(tmp, [0, 0, 0, width - beginning, 0], + [-1, -1, -1, beginning, -1]) + tmp_back = array_ops.slice(tmp, [0, 0, 0, 0, 0], [-1, -1, -1, end, -1]) + return array_ops.concat([tmp_front, tmp, tmp_back], 3) + + cout = 32 + shape = [1, 7, 7, 7, 16] + outputs_shape = shape[0:-1] + [cout] + dtype = dtypes.float32 + tol = 1e-3 + gain = 3.14 + # Check orthogonality/isometry by computing the ratio between + # the 2-norms of the inputs and ouputs. + for kernel_size in [[1, 1, 1], [2, 2, 2], [3, 3, 3]]: + convolution = convolutional.conv3d + inputs = random_ops.random_normal(shape, dtype=dtype) + inputs_2norm = linalg_ops.norm(inputs) + input_with_circular_pad = circular_pad(inputs, shape[1], kernel_size[0]) + outputs = convolution( + input_with_circular_pad, padding="valid", filters=cout, + kernel_size=kernel_size[0], use_bias=False, + kernel_initializer=init_ops.convolutional_orthogonal_3d(gain=gain)) + outputs_2norm = linalg_ops.norm(outputs) + ratio = outputs_2norm / inputs_2norm + 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 orthogonal kernel. + self.assertAllClose(sess.run(ratio), np.sqrt(gain), rtol=tol, atol=tol) + + class IdentityInitializerTest(test.TestCase): def testInvalidDataType(self): diff --git a/tensorflow/python/kernel_tests/linalg/BUILD b/tensorflow/python/kernel_tests/linalg/BUILD index 9555e510997a6aa07797dffa1a6e4810b0b4e5d2..7ffa48b65308b50d57a51d06fea1f88928cb2554 100644 --- a/tensorflow/python/kernel_tests/linalg/BUILD +++ b/tensorflow/python/kernel_tests/linalg/BUILD @@ -123,6 +123,10 @@ cuda_py_test( "//tensorflow/python:platform_test", ], shard_count = 5, + tags = [ + "noasan", # times out + "optonly", + ], ) cuda_py_test( diff --git a/tensorflow/python/kernel_tests/list_ops_test.py b/tensorflow/python/kernel_tests/list_ops_test.py index 6173a1def3365f455e9fd5066a2b40620ec50a93..098f9724a2a65f544005e1799f6329a29bb02abe 100644 --- a/tensorflow/python/kernel_tests/list_ops_test.py +++ b/tensorflow/python/kernel_tests/list_ops_test.py @@ -318,6 +318,108 @@ class ListOpsTest(test_util.TensorFlowTestCase): [[1.0, 2.0]] * 4) self.assertAllEqual(self.evaluate(updated_v_stacked), expected) + @test_util.run_in_graph_and_eager_modes() + def testConcat(self): + c = constant_op.constant([1.0, 2.0], dtype=dtypes.float32) + l0 = list_ops.tensor_list_from_tensor(c, element_shape=scalar_shape()) + l1 = list_ops.tensor_list_from_tensor([-1.0], element_shape=scalar_shape()) + l_batch_0 = array_ops.stack([l0, l1]) + l_batch_1 = array_ops.stack([l1, l0]) + + l_concat_01 = list_ops.tensor_list_concat_lists( + l_batch_0, l_batch_1, element_dtype=dtypes.float32) + l_concat_10 = list_ops.tensor_list_concat_lists( + l_batch_1, l_batch_0, element_dtype=dtypes.float32) + l_concat_00 = list_ops.tensor_list_concat_lists( + l_batch_0, l_batch_0, element_dtype=dtypes.float32) + l_concat_11 = list_ops.tensor_list_concat_lists( + l_batch_1, l_batch_1, element_dtype=dtypes.float32) + + expected_00 = [[1.0, 2.0, 1.0, 2.0], [-1.0, -1.0]] + expected_01 = [[1.0, 2.0, -1.0], [-1.0, 1.0, 2.0]] + expected_10 = [[-1.0, 1.0, 2.0], [1.0, 2.0, -1.0]] + expected_11 = [[-1.0, -1.0], [1.0, 2.0, 1.0, 2.0]] + + for i, (concat, expected) in enumerate(zip( + [l_concat_00, l_concat_01, l_concat_10, l_concat_11], + [expected_00, expected_01, expected_10, expected_11])): + splitted = array_ops.unstack(concat) + splitted_stacked_ret = self.evaluate( + (list_ops.tensor_list_stack(splitted[0], dtypes.float32), + list_ops.tensor_list_stack(splitted[1], dtypes.float32))) + print("Test concat %d: %s, %s, %s, %s" + % (i, expected[0], splitted_stacked_ret[0], + expected[1], splitted_stacked_ret[1])) + self.assertAllClose(expected[0], splitted_stacked_ret[0]) + self.assertAllClose(expected[1], splitted_stacked_ret[1]) + + # Concatenating mismatched shapes fails. + with self.assertRaises((errors.InvalidArgumentError, ValueError)): + self.evaluate( + list_ops.tensor_list_concat_lists( + l_batch_0, + list_ops.empty_tensor_list(scalar_shape(), dtypes.float32), + element_dtype=dtypes.float32)) + + with self.assertRaisesRegexp(errors.InvalidArgumentError, + "element shapes are not identical at index 0"): + l_batch_of_vec_tls = array_ops.stack( + [list_ops.tensor_list_from_tensor([[1.0]], element_shape=[1])] * 2) + self.evaluate( + list_ops.tensor_list_concat_lists(l_batch_0, l_batch_of_vec_tls, + element_dtype=dtypes.float32)) + + with self.assertRaisesRegexp(errors.InvalidArgumentError, + r"input_b\[0\].dtype != element_dtype."): + l_batch_of_int_tls = array_ops.stack( + [list_ops.tensor_list_from_tensor([1], element_shape=scalar_shape())] + * 2) + self.evaluate( + list_ops.tensor_list_concat_lists(l_batch_0, l_batch_of_int_tls, + element_dtype=dtypes.float32)) + + @test_util.run_in_graph_and_eager_modes() + def testPushBackBatch(self): + c = constant_op.constant([1.0, 2.0], dtype=dtypes.float32) + l0 = list_ops.tensor_list_from_tensor(c, element_shape=scalar_shape()) + l1 = list_ops.tensor_list_from_tensor([-1.0], element_shape=scalar_shape()) + l_batch = array_ops.stack([l0, l1]) + l_push = list_ops.tensor_list_push_back_batch(l_batch, [3.0, 4.0]) + l_unstack = array_ops.unstack(l_push) + l0_ret = list_ops.tensor_list_stack(l_unstack[0], dtypes.float32) + l1_ret = list_ops.tensor_list_stack(l_unstack[1], dtypes.float32) + self.assertAllClose([1.0, 2.0, 3.0], self.evaluate(l0_ret)) + self.assertAllClose([-1.0, 4.0], self.evaluate(l1_ret)) + + with ops.control_dependencies([l_push]): + l_unstack_orig = array_ops.unstack(l_batch) + l0_orig_ret = list_ops.tensor_list_stack(l_unstack_orig[0], + dtypes.float32) + l1_orig_ret = list_ops.tensor_list_stack(l_unstack_orig[1], + dtypes.float32) + + # Check that without aliasing, push_back_batch still works; and + # that it doesn't modify the input. + l0_r_v, l1_r_v, l0_orig_v, l1_orig_v = self.evaluate( + (l0_ret, l1_ret, l0_orig_ret, l1_orig_ret)) + self.assertAllClose([1.0, 2.0, 3.0], l0_r_v) + self.assertAllClose([-1.0, 4.0], l1_r_v) + self.assertAllClose([1.0, 2.0], l0_orig_v) + self.assertAllClose([-1.0], l1_orig_v) + + # Pushing back mismatched shapes fails. + with self.assertRaises((errors.InvalidArgumentError, ValueError)): + self.evaluate(list_ops.tensor_list_push_back_batch(l_batch, [])) + + with self.assertRaisesRegexp(errors.InvalidArgumentError, + "incompatible shape to a list at index 0"): + self.evaluate( + list_ops.tensor_list_push_back_batch(l_batch, [[3.0], [4.0]])) + + with self.assertRaisesRegexp(errors.InvalidArgumentError, + "Invalid data type at index 0"): + self.evaluate(list_ops.tensor_list_push_back_batch(l_batch, [3, 4])) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/kernel_tests/manip_ops_test.py b/tensorflow/python/kernel_tests/manip_ops_test.py index b8200ac0cb1e4315a56181779c70da1126d8fc15..f31426713c49ba03b978ae7820e22c5c17319139 100644 --- a/tensorflow/python/kernel_tests/manip_ops_test.py +++ b/tensorflow/python/kernel_tests/manip_ops_test.py @@ -20,8 +20,10 @@ 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 errors_impl 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 manip_ops from tensorflow.python.platform import test as test_lib @@ -88,41 +90,78 @@ 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 testNegativeAxis(self): + self._testAll(np.random.randint(-100, 100, (5)).astype(np.int32), 3, -1) + self._testAll(np.random.randint(-100, 100, (4, 4)).astype(np.int32), 3, -2) + # Make sure negative axis shoudl be 0 <= axis + dims < dims + with self.test_session(): + with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, + "is out of range"): + manip_ops.roll(np.random.randint(-100, 100, (4, 4)).astype(np.int32), + 3, -10).eval() + + def testInvalidInputShape(self): + # The input should be 1-D or higher, checked in shape function. + with self.assertRaisesRegexp( + ValueError, "Shape must be at least rank 1 but is rank 0"): + manip_ops.roll(7, 1, 0) + def testRollInputMustVectorHigherRaises(self): - tensor = 7 + # The input should be 1-D or higher, checked in kernel. + tensor = array_ops.placeholder(dtype=dtypes.int32) shift = 1 axis = 0 with self.test_session(): with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, "input must be 1-D or higher"): - manip_ops.roll(tensor, shift, axis).eval() + manip_ops.roll(tensor, shift, axis).eval(feed_dict={tensor: 7}) + + def testInvalidAxisShape(self): + # The axis should be a scalar or 1-D, checked in shape function. + with self.assertRaisesRegexp( + ValueError, "Shape must be at most rank 1 but is rank 2"): + manip_ops.roll([[1, 2], [3, 4]], 1, [[0, 1]]) def testRollAxisMustBeScalarOrVectorRaises(self): + # The axis should be a scalar or 1-D, checked in kernel. tensor = [[1, 2], [3, 4]] shift = 1 - axis = [[0, 1]] + axis = array_ops.placeholder(dtype=dtypes.int32) with self.test_session(): with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, "axis must be a scalar or a 1-D vector"): - manip_ops.roll(tensor, shift, axis).eval() + manip_ops.roll(tensor, shift, axis).eval(feed_dict={axis: [[0, 1]]}) + + def testInvalidShiftShape(self): + # The shift should be a scalar or 1-D, checked in shape function. + with self.assertRaisesRegexp( + ValueError, "Shape must be at most rank 1 but is rank 2"): + manip_ops.roll([[1, 2], [3, 4]], [[0, 1]], 1) def testRollShiftMustBeScalarOrVectorRaises(self): + # The shift should be a scalar or 1-D, checked in kernel. tensor = [[1, 2], [3, 4]] - shift = [[0, 1]] + shift = array_ops.placeholder(dtype=dtypes.int32) axis = 1 with self.test_session(): with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, "shift must be a scalar or a 1-D vector"): - manip_ops.roll(tensor, shift, axis).eval() + manip_ops.roll(tensor, shift, axis).eval(feed_dict={shift: [[0, 1]]}) + + def testInvalidShiftAndAxisNotEqualShape(self): + # The shift and axis must be same size, checked in shape function. + with self.assertRaisesRegexp(ValueError, "both shapes must be equal"): + manip_ops.roll([[1, 2], [3, 4]], [1], [0, 1]) def testRollShiftAndAxisMustBeSameSizeRaises(self): + # The shift and axis must be same size, checked in kernel. tensor = [[1, 2], [3, 4]] - shift = [1] + shift = array_ops.placeholder(dtype=dtypes.int32) axis = [0, 1] with self.test_session(): with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, "shift and axis must have the same size"): - manip_ops.roll(tensor, shift, axis).eval() + manip_ops.roll(tensor, shift, axis).eval(feed_dict={shift: [1]}) def testRollAxisOutOfRangeRaises(self): tensor = [1, 2] diff --git a/tensorflow/python/kernel_tests/norm_op_test.py b/tensorflow/python/kernel_tests/norm_op_test.py index d85512fae69801a0c14789f2ed2122559815f80d..3f71b326a2fcb8accfd3182ce5d42f30aa2c74b4 100644 --- a/tensorflow/python/kernel_tests/norm_op_test.py +++ b/tensorflow/python/kernel_tests/norm_op_test.py @@ -37,17 +37,17 @@ class NormOpTest(test_lib.TestCase): def testBadOrder(self): matrix = [[0., 1.], [2., 3.]] - for ord_ in "foo", -7, -1.1, 0: + for ord_ in "fro", -7, -1.1, 0: with self.assertRaisesRegexp(ValueError, "'ord' must be a supported vector norm"): - linalg_ops.norm(matrix, ord="fro") + linalg_ops.norm(matrix, ord=ord_) - for ord_ in "foo", -7, -1.1, 0: + for ord_ in "fro", -7, -1.1, 0: with self.assertRaisesRegexp(ValueError, "'ord' must be a supported vector norm"): linalg_ops.norm(matrix, ord=ord_, axis=-1) - for ord_ in 1.1, 2: + for ord_ in "foo", -7, -1.1, 1.1: with self.assertRaisesRegexp(ValueError, "'ord' must be a supported matrix norm"): linalg_ops.norm(matrix, ord=ord_, axis=[-2, -1]) @@ -69,14 +69,14 @@ def _GetNormOpTest(dtype_, shape_, ord_, axis_, keep_dims_, use_static_shape_): if use_static_shape_: tf_matrix = constant_op.constant(matrix) tf_norm = linalg_ops.norm( - tf_matrix, ord=ord_, axis=axis_, keep_dims=keep_dims_) + tf_matrix, ord=ord_, axis=axis_, keepdims=keep_dims_) tf_norm_val = sess.run(tf_norm) else: tf_matrix = array_ops.placeholder(dtype_) tf_norm = linalg_ops.norm( - tf_matrix, ord=ord_, axis=axis_, keep_dims=keep_dims_) + tf_matrix, ord=ord_, axis=axis_, keepdims=keep_dims_) tf_norm_val = sess.run(tf_norm, feed_dict={tf_matrix: matrix}) - self.assertAllClose(np_norm, tf_norm_val) + self.assertAllClose(np_norm, tf_norm_val, rtol=1e-5, atol=1e-5) def Test(self): is_matrix_norm = (isinstance(axis_, tuple) or @@ -85,8 +85,6 @@ def _GetNormOpTest(dtype_, shape_, ord_, axis_, keep_dims_, use_static_shape_): if ((not is_matrix_norm and ord_ == "fro") or (is_matrix_norm and is_fancy_p_norm)): self.skipTest("Not supported by neither numpy.linalg.norm nor tf.norm") - if is_matrix_norm and ord_ == 2: - self.skipTest("Not supported by tf.norm") if ord_ == 'euclidean' or (axis_ is None and len(shape) > 2): self.skipTest("Not supported by numpy.linalg.norm") matrix = np.random.randn(*shape_).astype(dtype_) diff --git a/tensorflow/python/kernel_tests/py_func_test.py b/tensorflow/python/kernel_tests/py_func_test.py index 5b508b7c0e72180194fa1a4c95bc4282d4694605..b9f44d728a1d9843df1e836594f9caa7010d8a94 100644 --- a/tensorflow/python/kernel_tests/py_func_test.py +++ b/tensorflow/python/kernel_tests/py_func_test.py @@ -52,6 +52,38 @@ class PyFuncTest(test.TestCase): """Encapsulates tests for py_func and eager_py_func.""" # ----- Tests for py_func ----- + def testRealDataTypes(self): + def sum_func(x, y): + return x + y + for dtype in [dtypes.float16, dtypes.float32, dtypes.float64, + dtypes.uint8, dtypes.int8, dtypes.uint16, dtypes.int16, + dtypes.int32, dtypes.int64]: + with self.test_session(): + x = constant_op.constant(1, dtype=dtype) + y = constant_op.constant(2, dtype=dtype) + z = self.evaluate(script_ops.py_func(sum_func, [x, y], dtype)) + self.assertEqual(z, 3) + + def testComplexDataTypes(self): + def sub_func(x, y): + return x - y + for dtype in [dtypes.complex64, dtypes.complex128]: + with self.test_session(): + x = constant_op.constant(1 + 1j, dtype=dtype) + y = constant_op.constant(2 - 2j, dtype=dtype) + z = self.evaluate(script_ops.py_func(sub_func, [x, y], dtype)) + self.assertEqual(z, -1 + 3j) + + def testBoolDataTypes(self): + def and_func(x, y): + return x and y + dtype = dtypes.bool + with self.test_session(): + x = constant_op.constant(True, dtype=dtype) + y = constant_op.constant(False, dtype=dtype) + z = self.evaluate(script_ops.py_func(and_func, [x, y], dtype)) + self.assertEqual(z, False) + def testSingleType(self): with self.test_session(): x = constant_op.constant(1.0, dtypes.float32) diff --git a/tensorflow/python/kernel_tests/random/multinomial_op_test.py b/tensorflow/python/kernel_tests/random/multinomial_op_test.py index a9dc7b7de000024f23b88406bf0c1c2f32ac4fac..051c7d86bf2342f15b587fc350bfbede7fae2285 100644 --- a/tensorflow/python/kernel_tests/random/multinomial_op_test.py +++ b/tensorflow/python/kernel_tests/random/multinomial_op_test.py @@ -46,7 +46,7 @@ def composed_sampler(logits, num_samples): logits = array_ops.expand_dims(logits, -1) # [batch size, num samples] - return math_ops.argmax(logits + noise, dimension=1) + return math_ops.argmax(logits + noise, axis=1) native_sampler = random_ops.multinomial diff --git a/tensorflow/python/kernel_tests/random/random_ops_test.py b/tensorflow/python/kernel_tests/random/random_ops_test.py index df37dd98ece57ae7c3835ab63b720b29fc19c975..e4b5c3832a2252aedc8820a650b022cd30b7f285 100644 --- a/tensorflow/python/kernel_tests/random/random_ops_test.py +++ b/tensorflow/python/kernel_tests/random/random_ops_test.py @@ -228,6 +228,17 @@ class RandomUniformTest(test.TestCase): print("count = ", count) self.assertTrue(count < count_limit) + def testUniformIntsWithInvalidShape(self): + for dtype in dtypes.int32, dtypes.int64: + with self.assertRaisesRegexp( + ValueError, "Shape must be rank 0 but is rank 1"): + random_ops.random_uniform( + [1000], minval=[1, 2], maxval=3, dtype=dtype) + with self.assertRaisesRegexp( + ValueError, "Shape must be rank 0 but is rank 1"): + random_ops.random_uniform( + [1000], minval=1, maxval=[2, 3], dtype=dtype) + # Check that uniform ints actually follow a uniform distribution. def testUniformInts(self): minv = -2 diff --git a/tensorflow/python/kernel_tests/reduce_join_op_test.py b/tensorflow/python/kernel_tests/reduce_join_op_test.py index 7f3049b9f841c55759e78c8bc66301f325d2a7c6..fb9e5cc2a3727b70f731ca8ae946b6d5ef325833 100644 --- a/tensorflow/python/kernel_tests/reduce_join_op_test.py +++ b/tensorflow/python/kernel_tests/reduce_join_op_test.py @@ -160,7 +160,7 @@ class ReduceJoinTest(UnicodeTestCase): separator=separator) if not reduction_indices: truth = constant_op.constant(truth) - truth_squeezed = array_ops.squeeze(truth, squeeze_dims=reduction_indices) + truth_squeezed = array_ops.squeeze(truth, axis=reduction_indices) output_array = output.eval() output_keep_dims_array = output_keep_dims.eval() truth_array = truth.eval() diff --git a/tensorflow/python/kernel_tests/reduction_ops_test.py b/tensorflow/python/kernel_tests/reduction_ops_test.py index 589ea54973c10902c461f552d5c54b6fad6ecf67..ea78b58d88f7ff04a0a5d2272d2c94e1c97009da 100644 --- a/tensorflow/python/kernel_tests/reduction_ops_test.py +++ b/tensorflow/python/kernel_tests/reduction_ops_test.py @@ -889,9 +889,9 @@ class AnyReductionTest(test.TestCase): class CountNonzeroReductionTest(test.TestCase): - def _compare(self, x, reduction_axes, keepdims, use_gpu=False, + def _compare(self, x, reduction_axes, keepdims, use_gpu=False, zero=0, feed_dict=None): - np_ans = (x != 0).astype(np.int32) + np_ans = (x != zero).astype(np.int32) if reduction_axes is None: np_ans = np.sum(np_ans, keepdims=keepdims) else: @@ -958,6 +958,37 @@ class CountNonzeroReductionTest(test.TestCase): y = math_ops.count_nonzero(x, [0]) self.assertAllEqual(y.eval(), np.zeros(9938)) + def testStringReduce(self): + # Test case for GitHub issue 18712 + with self.test_session() as sess: + v = math_ops.count_nonzero(constant_op.constant(["test"])) + self.assertAllClose(sess.run(v), 1) + + def testStringReduce1D(self): + # Create a 1D array of strings + x = np.asarray(["", "", "a", "", "", "b"]) + self._compare(x, None, keepdims=False, zero=np.str("")) + self._compare(x, [], keepdims=False, zero=np.str("")) + self._compare(x, [0], keepdims=False, zero=np.str("")) + self._compare(x, None, keepdims=True, zero=np.str("")) + self._compare(x, [], keepdims=True, zero=np.str("")) + self._compare(x, [0], keepdims=True, zero=np.str("")) + + def testStringReduce2D(self): + # Create a 2D array of strings + x = np.asarray([["", "", "a", "", "", "b"], + ["", "c", "", "d", "", ""], + ["e", "", "f", "", "", ""]]) + self._compare(x, None, keepdims=False, zero=np.str("")) + self._compare(x, [], keepdims=False, zero=np.str("")) + self._compare(x, [0], keepdims=False, zero=np.str("")) + self._compare(x, [1], keepdims=False, zero=np.str("")) + self._compare(x, [0, 1], keepdims=False, zero=np.str("")) + self._compare(x, None, keepdims=True, zero=np.str("")) + self._compare(x, [], keepdims=True, zero=np.str("")) + self._compare(x, [0], keepdims=True, zero=np.str("")) + self._compare(x, [0, 1], keepdims=True, zero=np.str("")) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/kernel_tests/resource_variable_ops_test.py b/tensorflow/python/kernel_tests/resource_variable_ops_test.py index edc63264a3549e91f2d6278a935be29eda5c99be..984192258c9724dd9d73105c65177786def98e83 100644 --- a/tensorflow/python/kernel_tests/resource_variable_ops_test.py +++ b/tensorflow/python/kernel_tests/resource_variable_ops_test.py @@ -36,6 +36,9 @@ 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 momentum +from tensorflow.python.training import saver +from tensorflow.python.training import training_util from tensorflow.python.util import compat @@ -174,73 +177,59 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase): resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32)) self.assertEqual(read, 2) - @test_util.run_in_graph_and_eager_modes(use_gpu=True) + @test_util.run_in_graph_and_eager_modes() 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]]) + 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]]) - @test_util.run_in_graph_and_eager_modes(use_gpu=True) + @test_util.run_in_graph_and_eager_modes() def testScatterSub(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_sub(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), [[-1]]) + 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_sub( + 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), [[-1]]) - @test_util.run_in_graph_and_eager_modes(use_gpu=True) + @test_util.run_in_graph_and_eager_modes() def testScatterMul(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_mul(handle, [0], - constant_op.constant( - [[5]], - dtype=dtypes.int32))) - read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(self.evaluate(read), [[5]]) + 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_mul( + handle, [0], constant_op.constant([[5]], dtype=dtypes.int32))) + read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) + self.assertEqual(self.evaluate(read), [[5]]) - @test_util.run_in_graph_and_eager_modes(use_gpu=True) + @test_util.run_in_graph_and_eager_modes() def testScatterDiv(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( - [[6]], - dtype=dtypes.int32))) - self.evaluate( - resource_variable_ops.resource_scatter_div(handle, [0], - constant_op.constant( - [[3]], - dtype=dtypes.int32))) - read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(self.evaluate(read), [[2]]) + 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([[6]], dtype=dtypes.int32))) + self.evaluate( + resource_variable_ops.resource_scatter_div( + handle, [0], constant_op.constant([[3]], dtype=dtypes.int32))) + read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) + self.assertEqual(self.evaluate(read), [[2]]) - @test_util.run_in_graph_and_eager_modes(use_gpu=True) + @test_util.run_in_graph_and_eager_modes() def testScatterMin(self): with ops.device("cpu:0"): handle = resource_variable_ops.var_handle_op( @@ -258,131 +247,115 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase): read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) self.assertEqual(self.evaluate(read), [[3]]) - @test_util.run_in_graph_and_eager_modes(use_gpu=True) + def testMetagraph(self): + with ops.Graph().as_default(): + with variable_scope.variable_scope("foo", use_resource=True): + a = variable_scope.get_variable("a", initializer=10.0) + + momentum.MomentumOptimizer( + learning_rate=0.001, momentum=0.1).minimize( + a, + colocate_gradients_with_ops=True, + global_step=training_util.get_or_create_global_step()) + + graph = ops.get_default_graph() + meta_graph_def = saver.export_meta_graph(graph=graph) + + with ops.Graph().as_default(): + saver.import_meta_graph(meta_graph_def, import_scope="") + meta_graph_two = saver.export_meta_graph(graph=graph) + self.assertEqual(meta_graph_def, meta_graph_two) + + @test_util.run_in_graph_and_eager_modes() def testScatterMax(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( - [[6]], - dtype=dtypes.int32))) - self.evaluate( - resource_variable_ops.resource_scatter_max(handle, [0], - constant_op.constant( - [[3]], - dtype=dtypes.int32))) - read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(self.evaluate(read), [[6]]) + 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([[6]], dtype=dtypes.int32))) + self.evaluate( + resource_variable_ops.resource_scatter_max( + handle, [0], constant_op.constant([[3]], dtype=dtypes.int32))) + read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) + self.assertEqual(self.evaluate(read), [[6]]) - @test_util.run_in_graph_and_eager_modes(use_gpu=True) + @test_util.run_in_graph_and_eager_modes() def testScatterAddScalar(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]]) + 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]]) - @test_util.run_in_graph_and_eager_modes(use_gpu=True) + @test_util.run_in_graph_and_eager_modes() def testScatterSubScalar(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_sub(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), [[-1]]) + 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_sub( + 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), [[-1]]) - @test_util.run_in_graph_and_eager_modes(use_gpu=True) + @test_util.run_in_graph_and_eager_modes() def testScatterMulScalar(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_mul(handle, [0], - constant_op.constant( - 5, - dtype=dtypes.int32))) - read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(self.evaluate(read), [[5]]) + 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_mul( + handle, [0], constant_op.constant(5, dtype=dtypes.int32))) + read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) + self.assertEqual(self.evaluate(read), [[5]]) - @test_util.run_in_graph_and_eager_modes(use_gpu=True) + @test_util.run_in_graph_and_eager_modes() def testScatterDivScalar(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( - [[6]], - dtype=dtypes.int32))) - self.evaluate( - resource_variable_ops.resource_scatter_div(handle, [0], - constant_op.constant( - 3, - dtype=dtypes.int32))) - read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(self.evaluate(read), [[2]]) + 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([[6]], dtype=dtypes.int32))) + self.evaluate( + resource_variable_ops.resource_scatter_div( + handle, [0], constant_op.constant(3, dtype=dtypes.int32))) + read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) + self.assertEqual(self.evaluate(read), [[2]]) - @test_util.run_in_graph_and_eager_modes(use_gpu=True) + @test_util.run_in_graph_and_eager_modes() def testScatterMinScalar(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( - [[6]], - dtype=dtypes.int32))) - self.evaluate( - resource_variable_ops.resource_scatter_min(handle, [0], - constant_op.constant( - 3, - dtype=dtypes.int32))) - read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(self.evaluate(read), [[3]]) + 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([[6]], dtype=dtypes.int32))) + self.evaluate( + resource_variable_ops.resource_scatter_min( + handle, [0], constant_op.constant(3, dtype=dtypes.int32))) + read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) + self.assertEqual(self.evaluate(read), [[3]]) - @test_util.run_in_graph_and_eager_modes(use_gpu=True) + @test_util.run_in_graph_and_eager_modes() def testScatterMaxScalar(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( - [[6]], - dtype=dtypes.int32))) - self.evaluate( - resource_variable_ops.resource_scatter_max(handle, [0], - constant_op.constant( - 3, - dtype=dtypes.int32))) - read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(self.evaluate(read), [[6]]) + 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([[6]], dtype=dtypes.int32))) + self.evaluate( + resource_variable_ops.resource_scatter_max( + handle, [0], constant_op.constant(3, dtype=dtypes.int32))) + read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) + self.assertEqual(self.evaluate(read), [[6]]) def testScatterUpdateString(self): handle = resource_variable_ops.var_handle_op( diff --git a/tensorflow/python/kernel_tests/rnn_test.py b/tensorflow/python/kernel_tests/rnn_test.py index 9a0409c796ab60da3d47cf7d46ef6fbd5bd82394..fe5ad84c104502f0e09d3a963b406f49d6b97b71 100644 --- a/tensorflow/python/kernel_tests/rnn_test.py +++ b/tensorflow/python/kernel_tests/rnn_test.py @@ -206,6 +206,28 @@ class RNNTest(test.TestCase): self.assertAllEqual(4, state[0]) self.assertAllEqual([[[1]], [[2]], [[3]], [[4]]], state[1]) + def _assert_cell_builds(self, cell_class, dtype, batch_size, in_size, + out_size): + cell = cell_class(out_size, dtype=dtype) + in_shape = tensor_shape.TensorShape((batch_size, in_size)) + cell.build(in_shape) + state_output = cell.zero_state(batch_size, dtype) + cell_output, _ = cell(array_ops.zeros(in_shape, dtype), state_output) + self.assertAllEqual([batch_size, out_size], cell_output.shape.as_list()) + + @test_util.run_in_graph_and_eager_modes() + def testCellsBuild(self): + f32 = dtypes.float32 + f64 = dtypes.float64 + self._assert_cell_builds(rnn_cell_impl.BasicRNNCell, f32, 5, 7, 3) + self._assert_cell_builds(rnn_cell_impl.BasicRNNCell, f64, 5, 7, 3) + self._assert_cell_builds(rnn_cell_impl.BasicLSTMCell, f32, 5, 7, 3) + self._assert_cell_builds(rnn_cell_impl.BasicLSTMCell, f64, 5, 7, 3) + self._assert_cell_builds(rnn_cell_impl.GRUCell, f32, 5, 7, 3) + self._assert_cell_builds(rnn_cell_impl.GRUCell, f64, 5, 7, 3) + self._assert_cell_builds(rnn_cell_impl.LSTMCell, f32, 5, 7, 3) + self._assert_cell_builds(rnn_cell_impl.LSTMCell, f64, 5, 7, 3) + ######### Benchmarking RNN code diff --git a/tensorflow/python/kernel_tests/scatter_nd_ops_test.py b/tensorflow/python/kernel_tests/scatter_nd_ops_test.py index 9f5794951524b2689daa5fc4eefb19703262b8f0..b7477a768ab718c5a57293cc6f774298e1c9f891 100644 --- a/tensorflow/python/kernel_tests/scatter_nd_ops_test.py +++ b/tensorflow/python/kernel_tests/scatter_nd_ops_test.py @@ -364,6 +364,42 @@ class ScatterNdTest(test.TestCase): del input_ # input_ is not used in scatter_nd return array_ops.scatter_nd(indices, updates, shape) + def testString(self): + indices = constant_op.constant([[4], [3], [1], [7]], + dtype=dtypes.int32) + updates = constant_op.constant(["four", "three", "one", "seven"], + dtype=dtypes.string) + expected = np.array([b"", b"one", b"", b"three", b"four", + b"", b"", b"seven"]) + scatter = self.scatter_nd(indices, updates, shape=(8,)) + with self.test_session() as sess: + result = sess.run(scatter) + self.assertAllEqual(expected, result) + + # Same indice is updated twice by same value. + indices = constant_op.constant([[4], [3], [3], [7]], + dtype=dtypes.int32) + updates = constant_op.constant(["a", "b", "b", "c"], + dtype=dtypes.string) + expected = np.array([b"", b"", b"", b"bb", b"a", b"", b"", b"c"]) + scatter = self.scatter_nd(indices, updates, shape=(8,)) + with self.test_session() as sess: + result = sess.run(scatter) + self.assertAllEqual(expected, result) + + # Same indice is updated twice by different value. + indices = constant_op.constant([[4], [3], [3], [7]], + dtype=dtypes.int32) + updates = constant_op.constant(["a", "b", "c", "d"], + dtype=dtypes.string) + expected = [np.array([b"", b"", b"", b"bc", b"a", b"", b"", b"d"]), + np.array([b"", b"", b"", b"cb", b"a", b"", b"", b"d"])] + scatter = self.scatter_nd(indices, updates, shape=(8,)) + with self.test_session() as sess: + result = sess.run(scatter) + self.assertTrue(np.array_equal(result, expected[0]) or + np.array_equal(result, expected[1])) + def testRank3ValidShape(self): indices = array_ops.zeros([2, 2, 2], dtypes.int32) updates = array_ops.zeros([2, 2, 2], dtypes.int32) @@ -584,6 +620,10 @@ class ScatterNdNonAliasingAddTest(ScatterNdTest): shape, dtype=updates.dtype)) return array_ops.scatter_nd_non_aliasing_add(input_, indices, updates) + def testString(self): + # Not supported yet. + pass + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/kernel_tests/string_strip_op_test.py b/tensorflow/python/kernel_tests/string_strip_op_test.py new file mode 100644 index 0000000000000000000000000000000000000000..30fd477ff42cf1e6c96b80936226df7dc15997d4 --- /dev/null +++ b/tensorflow/python/kernel_tests/string_strip_op_test.py @@ -0,0 +1,56 @@ +# 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 string_strip_op.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.ops import string_ops +from tensorflow.python.platform import test + + +class StringStripOpTest(test.TestCase): + """ Test cases for tf.string_strip.""" + + def test_string_strip(self): + strings = ["pigs on the wing", "animals"] + + with self.test_session() as sess: + output = string_ops.string_strip(strings) + output = sess.run(output) + self.assertAllEqual(output, [b"pigs on the wing", b"animals"]) + + def test_string_strip_2d(self): + strings = [["pigs on the wing", "animals"], + [" hello ", "\n\tworld \r \n"]] + + with self.test_session() as sess: + output = string_ops.string_strip(strings) + output = sess.run(output) + self.assertAllEqual(output, [[b"pigs on the wing", b"animals"], + [b"hello", b"world"]]) + + def test_string_strip_with_empty_strings(self): + strings = [" hello ", "", "world ", " \t \r \n "] + + with self.test_session() as sess: + output = string_ops.string_strip(strings) + output = sess.run(output) + self.assertAllEqual(output, [b"hello", b"", b"world", b""]) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/python/kernel_tests/variable_scope_test.py b/tensorflow/python/kernel_tests/variable_scope_test.py index 86ab9fbb70b5efcf06cc064617df14deb18c1f98..51aa671098905e840b7c96cd5a984887d347adf9 100644 --- a/tensorflow/python/kernel_tests/variable_scope_test.py +++ b/tensorflow/python/kernel_tests/variable_scope_test.py @@ -24,11 +24,13 @@ import threading import numpy from tensorflow.python.eager import context +from tensorflow.python.eager import function 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 test_util +from tensorflow.python.layers import core as core_layers from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import init_ops @@ -118,6 +120,16 @@ class VariableScopeTest(test.TestCase): w = variable_scope.get_variable("w", []) self.assertEqual(w.dtype.base_dtype, dtypes.float16) + def testGetVariableInGraphNestedUnderEagerContext(self): + with context.eager_mode(): + + @function.defun + def f(): + v = variable_scope.get_variable("should_be_resource", []) + self.assertEqual(type(v), resource_variable_ops.ResourceVariable) + + f() + def testEagerVariableStore(self): with context.eager_mode(): store = variable_scope.EagerVariableStore() @@ -156,6 +168,28 @@ class VariableScopeTest(test.TestCase): for v in new_store.variables(): self.assertEqual(v.numpy(), 1) + def testEagerVariableStoreWithEagerDefun(self): + with context.eager_mode(): + + @function.defun + def f(): + x = constant_op.constant([[2.0]]) + d1 = core_layers.Dense( + 1, name="my_dense", kernel_initializer=init_ops.ones_initializer()) + _ = d1(x) # create variables + self.assertEqual(len(d1.variables), 2) + v1, v2 = d1.variables + d2 = core_layers.Dense( + 1, + name="my_dense", + kernel_initializer=init_ops.ones_initializer(), + _reuse=True) + _ = d2(x) + self.assertEqual(len(d2.variables), 2) + v3, v4 = d2.variables + self.assertAllEqual([v1, v2], [v3, v4]) + f() + @test_util.run_in_graph_and_eager_modes() def testInitFromNonTensorValue(self): v = variable_scope.get_variable("v4", initializer=4, dtype=dtypes.int32) @@ -209,15 +243,15 @@ class VariableScopeTest(test.TestCase): with variable_scope.variable_scope("not_cached", caching_device=""): v2_not_cached = variable_scope.get_variable("v", []) - self.assertFalse(v2_not_cached.value().device.startswith( - caching_device)) + self.assertFalse( + v2_not_cached.value().device.startswith(caching_device)) with variable_scope.variable_scope( "not_cached_identity_device", caching_device=lambda op: op.device): v2_identity_device = variable_scope.get_variable("v", []) - self.assertFalse(v2_identity_device.value().device.startswith( - caching_device)) + self.assertFalse( + v2_identity_device.value().device.startswith(caching_device)) with variable_scope.variable_scope("we_will_do_it_live") as vs_live: vs_live.set_caching_device("/job:live") @@ -484,15 +518,19 @@ class VariableScopeTest(test.TestCase): def testVarScopeGetOrCreateReuse(self): with self.test_session(): + def test_value(value): x = constant_op.constant(value) - with variable_scope.variable_scope("testVarScopeGetOrCreateReuse_bar", - reuse=variable_scope.AUTO_REUSE): + with variable_scope.variable_scope( + "testVarScopeGetOrCreateReuse_bar", + reuse=variable_scope.AUTO_REUSE): _ = state_ops.assign(variable_scope.get_variable("var", []), x) - with variable_scope.variable_scope("testVarScopeGetOrCreateReuse_bar", - reuse=variable_scope.AUTO_REUSE): + with variable_scope.variable_scope( + "testVarScopeGetOrCreateReuse_bar", + reuse=variable_scope.AUTO_REUSE): _ = variable_scope.get_variable("var", []) self.assertEqual(value, x.eval()) + test_value(42.) # Variable is created. test_value(13.) # Variable is reused hereafter. test_value(17.) @@ -551,19 +589,16 @@ class VariableScopeTest(test.TestCase): with variable_scope.variable_scope("default") as default: with variable_scope.variable_scope(None, "layer"): self.assertEqual( - variable_scope.get_variable("w", []).name, - "default/layer/w:0") + variable_scope.get_variable("w", []).name, "default/layer/w:0") with variable_scope.variable_scope(None, "layer"): self.assertEqual( - variable_scope.get_variable("w", []).name, - "default/layer_1/w:0") + variable_scope.get_variable("w", []).name, "default/layer_1/w:0") with variable_scope.variable_scope(default): pass # No matter the jump in the middle, unique numbering continues. with variable_scope.variable_scope(None, "layer"): self.assertEqual( - variable_scope.get_variable("w", []).name, - "default/layer_2/w:0") + variable_scope.get_variable("w", []).name, "default/layer_2/w:0") def testVarOpScopeReuse(self): with self.test_session(): @@ -935,12 +970,12 @@ class VariableScopeTest(test.TestCase): def testGetCollection(self): with self.test_session(): _ = variable_scope.get_variable("testGetCollection_a", []) - _ = variable_scope.get_variable("testGetCollection_b", [], - trainable=False) + _ = variable_scope.get_variable( + "testGetCollection_b", [], trainable=False) with variable_scope.variable_scope("testGetCollection_foo_") as scope1: _ = variable_scope.get_variable("testGetCollection_a", []) - _ = variable_scope.get_variable("testGetCollection_b", [], - trainable=False) + _ = variable_scope.get_variable( + "testGetCollection_b", [], trainable=False) self.assertEqual([ v.name for v in scope1.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES) @@ -954,8 +989,8 @@ class VariableScopeTest(test.TestCase): ]) with variable_scope.variable_scope("testGetCollection_foo") as scope2: _ = variable_scope.get_variable("testGetCollection_a", []) - _ = variable_scope.get_variable("testGetCollection_b", [], - trainable=False) + _ = variable_scope.get_variable( + "testGetCollection_b", [], trainable=False) self.assertEqual([ v.name for v in scope2.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES) @@ -992,22 +1027,22 @@ class VariableScopeTest(test.TestCase): with variable_scope.variable_scope( "testGetTrainableVariables_foo") as scope: _ = variable_scope.get_variable("testGetTrainableVariables_b", []) - _ = variable_scope.get_variable("testGetTrainableVariables_c", [], - trainable=False) - self.assertEqual([v.name - for v in scope.trainable_variables()], - ["testGetTrainableVariables_foo/" - "testGetTrainableVariables_b:0"]) + _ = variable_scope.get_variable( + "testGetTrainableVariables_c", [], trainable=False) + self.assertEqual( + [v.name for v in scope.trainable_variables()], + ["testGetTrainableVariables_foo/" + "testGetTrainableVariables_b:0"]) def testGetGlobalVariables(self): with self.test_session(): _ = variable_scope.get_variable("testGetGlobalVariables_a", []) with variable_scope.variable_scope("testGetGlobalVariables_foo") as scope: _ = variable_scope.get_variable("testGetGlobalVariables_b", []) - self.assertEqual([v.name - for v in scope.global_variables()], - ["testGetGlobalVariables_foo/" - "testGetGlobalVariables_b:0"]) + self.assertEqual( + [v.name for v in scope.global_variables()], + ["testGetGlobalVariables_foo/" + "testGetGlobalVariables_b:0"]) def testGetLocalVariables(self): with self.test_session(): @@ -1016,10 +1051,8 @@ class VariableScopeTest(test.TestCase): with variable_scope.variable_scope("foo") as scope: _ = variable_scope.get_variable( "b", [], collections=[ops.GraphKeys.LOCAL_VARIABLES]) - _ = variable_scope.get_variable( - "c", []) - self.assertEqual([v.name - for v in scope.local_variables()], ["foo/b:0"]) + _ = variable_scope.get_variable("c", []) + self.assertEqual([v.name for v in scope.local_variables()], ["foo/b:0"]) def testGetVariableWithRefDtype(self): v = variable_scope.get_variable("v", shape=[3, 4], dtype=dtypes.float32) @@ -1242,10 +1275,8 @@ class VariableScopeWithCustomGetterTest(test.TestCase): with ops.name_scope("prod_getter"): return g_0 * g_1 - with variable_scope.variable_scope( - "prod_scope", custom_getter=prod_getter): - with variable_scope.variable_scope( - "sum_scope", custom_getter=sum_getter): + with variable_scope.variable_scope("prod_scope", custom_getter=prod_getter): + with variable_scope.variable_scope("sum_scope", custom_getter=sum_getter): with variable_scope.variable_scope( "inner_sum_scope", custom_getter=sum_getter): # take sums of sums of products @@ -1270,9 +1301,8 @@ class VariableScopeWithCustomGetterTest(test.TestCase): np_vars, np_v = sess.run([true_vars, v]) # take products of sums of products self.assertAllClose( - np_v, - (((np_vars[0] * np_vars[1]) + (np_vars[2] * np_vars[3])) - + ((np_vars[4] * np_vars[5]) + (np_vars[6] * np_vars[7])))) + np_v, (((np_vars[0] * np_vars[1]) + (np_vars[2] * np_vars[3])) + ( + (np_vars[4] * np_vars[5]) + (np_vars[6] * np_vars[7])))) def testVariableCreator(self): @@ -1368,7 +1398,11 @@ class VariableScopeMultithreadedTest(test.TestCase): graph = ops.get_default_graph() threads = [ - threading.Thread(target=thread_fn, args=(i, graph,)) for i in range(2)] + threading.Thread(target=thread_fn, args=( + i, + graph, + )) for i in range(2) + ] threads[0].start() # Allow thread 0 to finish before starting thread 1. diff --git a/tensorflow/python/layers/base.py b/tensorflow/python/layers/base.py index ec741d3265b4216bd962280b0b927d6ad8a51fe4..64db49c900c21d60ba2337f920d6fa2cb9ab7b5f 100644 --- a/tensorflow/python/layers/base.py +++ b/tensorflow/python/layers/base.py @@ -12,148 +12,91 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================= - -# pylint: disable=unused-import,g-bad-import-order """Contains the base Layer class, from which all layers inherit.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import collections import copy -import re -import weakref -import numpy as np 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 utils as layers_util -from tensorflow.python.framework import tensor_util -from tensorflow.python.ops import array_ops +from tensorflow.python.keras._impl.keras.engine import base_layer 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 -@tf_export('layers.Layer') -class Layer(checkpointable.CheckpointableBase): - """Base layer class. +InputSpec = base_layer.InputSpec # pylint: disable=invalid-name - This is the class from which all layers inherit, implementing common - infrastructure functionality. - A layer is a class implementing common neural networks operations, such - as convolution, batch norm, etc. These operations require managing variables, - losses, and updates, as well as applying TensorFlow ops to input tensors. +@tf_export('layers.Layer') +class Layer(base_layer.Layer): + """Base layer class. - Users will just instantiate it and then treat it as a callable. + It is considered legacy, and we recommend the use of `tf.keras.layers.Layer` + instead. - We recommend that descendants of Layer implement the following methods: - * `__init__()`: Save configuration in member variables - * `build()`: Called once from `__call__`, when we know the shapes of inputs - and `dtype`. Should have the calls to `add_variable()`, and then - call the super's `build()` (which sets `self.built = True`, which is - nice in case the user wants to call `build()` manually before the - first `__call__`). - * `call()`: Called in `__call__` after making sure `build()` has been called - once. Should actually perform the logic of applying the layer to the - input tensors (which should be passed in as the first argument). + Arguments: + trainable: Boolean, whether the layer's variables should be trainable. + name: String name of the layer. + dtype: Default dtype of the layer's weights (default of `None` means use the + type of the first input). Read-only properties: - `name`: The name of the layer (string). - `dtype`: Default dtype of the layer (default of `None` means use the + name: The name of the layer (string). + dtype: Default dtype of the layer's weights (default of `None` means use the type of the first input). - `trainable_variables`: List of trainable variables. - `non_trainable_variables`: List of non-trainable variables. - `variables`: List of all variables of this layer, trainable and + trainable_variables: List of trainable variables. + non_trainable_variables: List of non-trainable variables. + variables: List of all variables of this layer, trainable and non-trainable. - `updates`: List of update ops of this layer. - `losses`: List of losses added by this layer. + updates: List of update ops of this layer. + losses: List of losses added by this layer. + trainable_weights: List of variables to be included in backprop. + non_trainable_weights: List of variables that should not be + included in backprop. + weights: The concatenation of the lists trainable_weights and + non_trainable_weights (in this order). Mutable properties: - `trainable`: Whether the layer should be trained (boolean). - `input_spec`: Optional (list of) `InputSpec` object(s) specifying the + trainable: Whether the layer should be trained (boolean). + input_spec: Optional (list of) `InputSpec` object(s) specifying the constraints on inputs that can be accepted by the layer. """ def __init__(self, trainable=True, name=None, dtype=None, - activity_regularizer=None, **kwargs): - # We use a kwargs dict here because these kwargs only exist - # for compatibility reasons. - # The list of kwargs is subject to changes in the future. - # We do not want to commit to it or to expose the list to users at all. - # Note this is exactly as safe as defining kwargs in the function signature, - # the only difference being that the list of valid kwargs is defined - # below rather rather in the signature, and default values are defined - # in calls to kwargs.get(). - allowed_kwargs = { - '_scope', - '_reuse', - 'input_shape', # For compatibility with Keras `Sequential` model. - 'batch_size', # For compatibility with Keras `Sequential` model. - } - for kwarg in kwargs: - if kwarg not in allowed_kwargs: - raise TypeError('Keyword argument not understood:', kwarg) - - # Mutable properties - # Indicates whether the layer's weights are updated during training - # and whether the layer's updates are run during training - self.trainable = trainable - # A stateful layer is a layer whose updates are run during inference too, - # for instance stateful RNNs. - self.stateful = False - # Indicates whether `build` needs to be called upon layer call, to create - # the layer's weights. - self.built = False - # Provides information about which inputs are compatible with the layer. - self.input_spec = None - - if activity_regularizer and context.executing_eagerly(): - raise ValueError( - ('Activity regularization is not supported when executing eagerly. ' - 'Got activity_regularizer=%s') % (activity_regularizer,)) - self._activity_regularizer = activity_regularizer + **kwargs): + # For backwards compatibility, legacy layers do not use `ResourceVariable` + # by default. + self._use_resource_variables = False + scope = kwargs.pop('_scope', None) + self._reuse = kwargs.pop('_reuse', None) + + # Avoid an incorrect lint error self._trainable_weights = [] - self._non_trainable_weights = [] - self._updates = [] - # When executing eagerly, _losses is a list of zero-argument lambdas which - # return tensors. When using graph execution, _losses is a list of ops. - self._losses = [] - self._reuse = kwargs.get('_reuse') - self._graph = None # Will be set at build time. - self._dtype = None if dtype is None else dtypes.as_dtype(dtype).name - 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 self._call_fn_args - - # These lists will be filled via successive calls - # to self._add_inbound_node(). - self._inbound_nodes = [] - self._outbound_nodes = [] + self.built = False - self._init_set_name(name) + super(Layer, self).__init__(trainable=trainable, name=name, dtype=dtype, + **kwargs) - # Determine variable scope. - scope = kwargs.get('_scope') + self._graph = None + self._call_has_scope_arg = 'scope' in self._call_fn_args if scope: with vs.variable_scope(scope) as captured_scope: self._scope = captured_scope else: self._scope = None + self._current_scope = None - # Set `_batch_input_shape` attribute - # for compatibility with Keras `Sequential` model. - if 'input_shape' in kwargs: - batch_size = kwargs.get('batch_size') - self._batch_input_shape = (batch_size,) + tuple(kwargs['input_shape']) + @property + def graph(self): + if context.executing_eagerly(): + raise RuntimeError('Layer.graph not supported when executing eagerly.') + return self._graph def _init_set_name(self, name): # Determine layer name (non-unique). @@ -166,18 +109,15 @@ class Layer(checkpointable.CheckpointableBase): self._name, base_name = self._make_unique_name() self._base_name = base_name - @property - def dtype(self): - return self._dtype - - @property - def name(self): - return self._name - - @property - def activity_regularizer(self): - """Optional regularizer function for the output of this layer.""" - return self._activity_regularizer + def _make_unique_name(self, name_uid_map=None, avoid_names=None, + namespace='', zero_based=False): + base_name = base_layer.to_snake_case(self.__class__.__name__) + name = base_layer.unique_layer_name(base_name, + name_uid_map=name_uid_map, + avoid_names=avoid_names, + namespace=namespace, + zero_based=zero_based) + return (name, base_name) @property def scope_name(self): @@ -189,271 +129,16 @@ class Layer(checkpointable.CheckpointableBase): 'querying `scope_name`.') return self._scope.name - @property - def trainable_weights(self): - return self._trainable_weights if self.trainable else [] - - @property - def non_trainable_weights(self): - if self.trainable: - return self._non_trainable_weights - else: - return self._trainable_weights + self._non_trainable_weights - - @property - def trainable_variables(self): - return self.trainable_weights - - @property - def non_trainable_variables(self): - return self.non_trainable_weights - - @property - def weights(self): - """Returns the list of all layer variables/weights. - - Returns: - A list of variables. - """ - return self.trainable_weights + self.non_trainable_weights - - @property - def variables(self): - """Returns the list of all layer variables/weights. - - Returns: - A list of variables. - """ - return self.weights - - @property - def updates(self): - if context.executing_eagerly(): - raise RuntimeError('Layer.updates not supported in Eager mode.') - if not self.trainable and not self.stateful: - return [] - return self._updates - - def add_update(self, updates, inputs=None): - """Add update op(s), potentially dependent on layer inputs. - - Weight updates (for instance, the updates of the moving mean and variance - in a BatchNormalization layer) may be dependent on the inputs passed - when calling a layer. Hence, when reusing the same layer on - different inputs `a` and `b`, some entries in `layer.updates` may be - dependent on `a` and some on `b`. This method automatically keeps track - of dependencies. - - The `get_updates_for` method allows to retrieve the updates relevant to a - specific set of inputs. - - This call is ignored in Eager mode. - - Arguments: - updates: Update op, or list/tuple of update ops. - 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.executing_eagerly(): - return # Updates already applied when in eager mode. - - updates = _to_list(updates) - updates = [x if isinstance(x, ops.Operation) - else ops.convert_to_tensor(x) for x in updates] - self._updates += updates - if inputs is None: - for u in updates: - u._unconditional_update = True # pylint: disable=protected-access - else: - 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. - - Returns: - List of update ops of the layer that depend on `inputs`. - - Raises: - RuntimeError: If called 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 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): - """Losses which are associated with this `Layer`. - - Note that when executing eagerly, getting this property evaluates - regularizers. When using graph execution, variable regularization ops have - already been created and are simply returned here. - - Returns: - A list of tensors. - """ - 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. - return [regularizer() for regularizer in self._losses] - else: - return self._losses - def add_loss(self, losses, inputs=None): - """Add loss tensor(s), potentially dependent on layer inputs. - - Some losses (for instance, activity regularization losses) may be dependent - on the inputs passed when calling a layer. Hence, when reusing the same - layer on different inputs `a` and `b`, some entries in `layer.losses` may - be dependent on `a` and some on `b`. This method automatically keeps track - of dependencies. - - The `get_losses_for` method allows to retrieve the losses relevant to a - specific set of inputs. - - Note that `add_loss` is not supported when executing eagerly. Instead, - variable regularizers may be added through `add_variable`. Activity - regularization is not supported directly (but such losses may be returned - from `Layer.call()`). - - Arguments: - losses: Loss tensor, or list/tuple of tensors. - 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.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) - self._losses += losses - if inputs is None: - for loss in losses: - loss._unconditional_loss = True # pylint: disable=protected-access - else: - for loss in losses: - loss._unconditional_loss = False # pylint: disable=protected-access + previous_losses_length = len(self._losses) + super(Layer, self).add_loss(losses, inputs=inputs) # TODO(fchollet): deprecate collection below. - _add_elements_to_collection(losses, ops.GraphKeys.REGULARIZATION_LOSSES) - - def get_losses_for(self, inputs): - """Retrieves losses relevant to a specific set of inputs. - - Arguments: - inputs: Input tensor or list/tuple of input tensors. - - Returns: - List of loss tensors of the layer that depend on `inputs`. - - Raises: - RuntimeError: If called in Eager mode. - """ - if context.executing_eagerly(): - raise RuntimeError('Layer.get_losses_for not supported in Eager mode.') - - 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.""" - self.built = True - - def call(self, inputs, **kwargs): # pylint: disable=unused-argument - """The logic of the layer lives here. + new_losses = self._losses[previous_losses_length:] + _add_elements_to_collection(new_losses, ops.GraphKeys.REGULARIZATION_LOSSES) - Arguments: - inputs: input tensor(s). - **kwargs: additional keyword arguments. - - Returns: - Output tensor(s). - """ - return inputs - - def _name_scope_name(self, current_variable_scope): + def _name_scope(self): """Determines op naming for the Layer.""" - return current_variable_scope.original_name_scope - - def compute_output_shape(self, input_shape): - """Computes the output shape of the layer given the input shape. - - Args: - input_shape: A (possibly nested tuple of) `TensorShape`. It need not - be fully defined (e.g. the batch size may be unknown). - - Returns: - A (possibly nested tuple of) `TensorShape`. - - Raises: - TypeError: if `input_shape` is not a (possibly nested tuple of) - `TensorShape`. - ValueError: if `input_shape` is incomplete or is incompatible with the - the layer. - """ - raise NotImplementedError - - def _make_unique_name(self, name_uid_map=None, avoid_names=None, - namespace='', zero_based=False): - base_name = _to_snake_case(self.__class__.__name__) - name = _unique_layer_name(base_name, name_uid_map=name_uid_map, - avoid_names=avoid_names, namespace=namespace, - zero_based=zero_based) - return (name, base_name) + return self._current_scope.original_name_scope def _set_scope(self, scope=None): if self._scope is None: @@ -467,10 +152,11 @@ class Layer(checkpointable.CheckpointableBase): scope, default_name=self._base_name) as captured_scope: self._scope = captured_scope - def add_variable(self, name, shape, dtype=None, - initializer=None, regularizer=None, - trainable=True, constraint=None, - partitioner=None): + def add_weight(self, name, shape, dtype=None, + initializer=None, regularizer=None, + trainable=True, constraint=None, + use_resource=None, + partitioner=None): """Adds a new variable to the layer, or gets an existing one; returns it. Arguments: @@ -486,6 +172,7 @@ class Layer(checkpointable.CheckpointableBase): then this parameter is ignored and any added variables are also marked as non-trainable. constraint: constraint instance (callable). + use_resource: Whether to use `ResourceVariable`. partitioner: (optional) partitioner instance (callable). If provided, when the requested variable is created it will be split into multiple partitions according to `partitioner`. In this case, @@ -504,10 +191,6 @@ class Layer(checkpointable.CheckpointableBase): RuntimeError: If called with partioned variable regularization and eager execution is enabled. """ - - # `init_graph` should point to the graph in which variable initialization - # will occur; it should be None if and only if initialization will take - # place in the eager context. init_graph = None if not context.executing_eagerly(): default_graph = ops.get_default_graph() @@ -530,71 +213,43 @@ class Layer(checkpointable.CheckpointableBase): self._set_scope(None) reuse = self.built or self._reuse + prev_len_trainable = len(self._trainable_weights) with vs.variable_scope( self._scope, reuse=reuse, auxiliary_name_scope=False) as scope: - with ops.name_scope(self._name_scope_name(scope)): - 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, + self._current_scope = scope + with ops.name_scope(self._name_scope()): + use_resource = (use_resource or + self._use_resource_variables or + scope.use_resource) + variable = super(Layer, self).add_weight( + name, + shape, dtype=dtypes.as_dtype(dtype), + initializer=initializer or scope.initializer, + trainable=trainable, 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. - - if variable in existing_variables: - # To match the behavior of tf.get_variable(), we only apply - # regularization if the variable is newly created. - return variable - + partitioner=partitioner, + use_resource=use_resource, + getter=vs.get_variable) + + if regularizer: + if context.executing_eagerly() or variable not in existing_variables: + self._handle_weight_regularization(name, variable, regularizer) + + if init_graph is not None: + # Handle edge case where a custom getter has overridden `trainable`. + # There is one known occurrence of this, in unit test + # testBasicRNNCellNotTrainable in + # contrib.rnn.python.kernel_tests.core_rnn_cell_test with init_graph.as_default(): trainable_variables = tf_variables.trainable_variables() if (trainable and self.trainable and variable not in trainable_variables): # A custom getter / variable scope overrode the trainable flag. - trainable = False - - if regularizer: - if isinstance(variable, tf_variables.PartitionedVariable): - for v in variable: - with ops.colocate_with(v.op): - with ops.name_scope(name + '/Regularizer'): - regularization = regularizer(v) - if regularization is not None: - self.add_loss(regularization) - else: - with ops.colocate_with(variable.op): - with ops.name_scope(name + '/Regularizer'): - regularization = regularizer(variable) - if regularization is not None: - self.add_loss(regularization) - elif regularizer: # and initialization took place in an eager context - if isinstance(variable, tf_variables.PartitionedVariable): - raise RuntimeError( - 'Partitioned variable regularization is not yet ' - '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. - # This makes losses responsive to variable updates when executing - # eagerly. - # - # TODO(akshayka): Do the same for graphs as well, so that losses - # collected in a while_loop can be run outside its control flow - # context and so that losses won't be swallowed up by graph functions - # (i.e., `.losses()` should always create regularizers). - self._losses.append(lambda: regularizer(variable)) - - if trainable: - self._trainable_weights.append(variable) - else: - self._non_trainable_weights.append(variable) + extra_trainable_vars = self._trainable_weights[prev_len_trainable:] + self._trainable_weights = self._trainable_weights[ + :prev_len_trainable] + self._non_trainable_weights += extra_trainable_vars return variable def __call__(self, inputs, *args, **kwargs): @@ -622,35 +277,14 @@ class Layer(checkpointable.CheckpointableBase): ValueError: if the layer's `call` method returns None (an invalid value). """ self._set_scope(kwargs.pop('scope', None)) - input_list = nest.flatten(inputs) - build_graph = not context.executing_eagerly() - # TODO(fchollet, allenl): Make deferred mode work with subclassed Models - # which don't use an "inputs" argument. - 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 build_graph: + if not context.executing_eagerly(): try: # Set layer's "graph" at build time - self._graph = ops._get_graph_from_inputs(input_list, graph=self._graph) # pylint: disable=protected-access + self._graph = ops._get_graph_from_inputs(nest.flatten(inputs), # pylint: disable=protected-access + graph=self._graph) except ValueError as e: raise ValueError('Input graph and Layer graph are not the same: %s' % e) - if build_graph or in_deferred_mode: - user_kwargs = copy.copy(kwargs) - - # Handle Keras mask propagation from previous layer to current layer. - previous_mask = None - if (not hasattr(self, '_compute_previous_mask') or - self._compute_previous_mask): - previous_mask = _collect_previous_mask(inputs) - 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. - kwargs['mask'] = previous_mask if self.built: try: @@ -667,134 +301,27 @@ class Layer(checkpointable.CheckpointableBase): 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 build_graph: - # Activity regularization is currently unsupported in Eager mode. - if self._activity_regularizer: - 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'): - raise ValueError('_keras_history currently unsupported in ' - 'Eager mode. Found _keras_history in %s while ' - 'executing __call__ for %s(%s)' % - (x, self.__class_.__name__, self)) - - # Check input assumptions set before layer building, e.g. input rank. - self._assert_input_compatibility(inputs) - if input_list and self._dtype is None: - try: - self._dtype = input_list[0].dtype.base_dtype.name - except AttributeError: - pass - if all(hasattr(x, 'get_shape') for x in input_list): - input_shapes = nest.map_structure(lambda x: x.get_shape(), inputs) - self.build(input_shapes) - try: - # Note: not all sub-classes of Layer call Layer.__init__ (especially - # the ones under tensorflow/python/keras). Hence we recompute this - # attribute here if it is not set. - # TODO(agarwal): Fix the sub-classes and avoid this complexity. - call_has_scope_arg = self._call_has_scope_arg - except AttributeError: - 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 build_graph or in_deferred_mode: - self._assert_input_compatibility(inputs) - - if not in_deferred_mode: - outputs = self.call(inputs, *args, **kwargs) - if outputs is None: - raise ValueError('A layer\'s `call` method should return a Tensor ' - 'or a list of Tensors, not None.') - 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 = [ - # TODO(fchollet): name the deferred tensors? - _DeferredTensor(shape=shape, dtype=self._dtype) - for shape in output_shapes - ] - if len(outputs) == 1: - outputs = outputs[0] - 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. - if self._activity_regularizer: - output_list = nest.flatten(outputs) - for output in output_list: - with ops.name_scope('ActivityRegularizer'): - activity_regularization = self._activity_regularizer(output) - self.add_loss(activity_regularization, inputs=inputs) + with scope_context_manager as scope: + self._current_scope = scope - # 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, tuple)): - if output_mask is None: - output_mask = [None for _ in range(len(outputs))] - for x, m in zip(outputs, output_mask): - x._keras_mask = m # pylint: disable=protected-access - else: - outputs._keras_mask = output_mask # pylint: disable=protected-access + try: + call_has_scope_arg = self._call_has_scope_arg + except AttributeError: + 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 - if build_graph: - # If all input tensors have history metadata, - # we update the output tensors - # with corresponding history metadata, thus eventually allowing to use - # these tensors to instantiate a Network. - if _have_all_keras_metadata(inputs): - # If the layer returns tensors from its inputs, unmodified, - # we copy them to avoid loss of tensor metadata. - output_ls = nest.flatten(outputs) - output_ls_copy = [] - for x in output_ls: - if x in input_list: - with ops.name_scope(scope.original_name_scope): - x = array_ops.identity(x) - output_ls_copy.append(x) - if len(output_ls_copy) == 1: - outputs = output_ls_copy[0] - else: - outputs = output_ls_copy + # Actually call layer + outputs = super(Layer, self).__call__(inputs, *args, **kwargs) + if not context.executing_eagerly(): # Update global default collections. _add_elements_to_collection(self.updates, ops.GraphKeys.UPDATE_OPS) - - 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). - self._add_inbound_node( - input_tensors=inputs, output_tensors=outputs, arguments=user_kwargs) - - self.built = True return outputs - @property - def graph(self): - if context.executing_eagerly(): - raise RuntimeError('Layer.graph not supported in Eager mode.') - return self._graph - def __deepcopy__(self, memo): no_copy = set(['_graph']) shallow_copy = set(['_scope', '_always_reuse_variable_scope']) @@ -806,658 +333,12 @@ class Layer(checkpointable.CheckpointableBase): setattr(result, k, v) elif k in shallow_copy: setattr(result, k, copy.copy(v)) - elif _is_tensor_or_tensor_list(v): + elif base_layer.is_tensor_or_tensor_list(v): setattr(result, k, v) else: setattr(result, k, copy.deepcopy(v, memo)) return result - def apply(self, inputs, *args, **kwargs): - """Apply the layer on a input. - - This simply wraps `self.__call__`. - - Arguments: - inputs: Input tensor(s). - *args: additional positional arguments to be passed to `self.call`. - **kwargs: additional keyword arguments to be passed to `self.call`. - - Returns: - Output tensor(s). - """ - return self.__call__(inputs, *args, **kwargs) - - def _add_inbound_node(self, - input_tensors, - output_tensors, - arguments=None): - """Internal method to create an inbound node for the layer. - - Arguments: - input_tensors: list of input tensors. - output_tensors: list of output tensors. - arguments: dictionary of keyword arguments that were passed to the - `call` method of the layer at the call that created the node. - """ - input_tensors = nest.flatten(input_tensors) - output_tensors = nest.flatten(output_tensors) - - # Collect input tensor(s) coordinates. - inbound_layers = [] - node_indices = [] - tensor_indices = [] - for x in input_tensors: - assert hasattr(x, '_keras_history') - inbound_layer, node_index, tensor_index = x._keras_history # pylint: disable=protected-access - inbound_layers.append(inbound_layer) - node_indices.append(node_index) - tensor_indices.append(tensor_index) - - # Create node, add it to inbound nodes. - Node( - self, - inbound_layers=inbound_layers, - node_indices=node_indices, - tensor_indices=tensor_indices, - input_tensors=input_tensors, - output_tensors=output_tensors, - arguments=arguments) - - # Update tensor history metadata. - for i in range(len(output_tensors)): - # The metadata attribute consists of 1) a layer instance - # 2) a node index for the layer, 3) a tensor index for the node. - # The allows layer reuse (multiple nodes per layer) and multi-output - # or multi-input layers (e.g. a layer can return multiple tensors, - # and each can be sent to a different layer). - output_tensors[i]._keras_history = (self, len(self._inbound_nodes) - 1, i) # pylint: disable=protected-access - - def _get_node_attribute_at_index(self, node_index, attr, attr_name): - """Private utility to retrieves an attribute (e.g. inputs) from a node. - - This is used to implement the methods: - - get_input_shape_at - - get_output_shape_at - - get_input_at - etc... - - Arguments: - node_index: Integer index of the node from which - to retrieve the attribute. - attr: Exact node attribute name. - attr_name: Human-readable attribute name, for error messages. - - Returns: - The layer's attribute `attr` at the node of index `node_index`. - - Raises: - RuntimeError: If the layer has no inbound nodes, or if called in Eager - mode. - ValueError: If the index provided does not match any node. - """ - if not self._inbound_nodes: - raise RuntimeError('The layer has never been called ' - 'and thus has no defined ' + attr_name + '.') - if not len(self._inbound_nodes) > node_index: - raise ValueError('Asked to get ' + attr_name + ' at node ' + - str(node_index) + ', but the layer has only ' + - str(len(self._inbound_nodes)) + ' inbound nodes.') - values = getattr(self._inbound_nodes[node_index], attr) - if len(values) == 1: - return values[0] - else: - return values - - def get_input_shape_at(self, node_index): - """Retrieves the input shape(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 shape tuple - (or list of shape tuples if the layer has multiple inputs). - - Raises: - RuntimeError: If called in Eager mode. - """ - return self._get_node_attribute_at_index(node_index, 'input_shapes', - 'input shape') - - def get_output_shape_at(self, node_index): - """Retrieves the output shape(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 shape tuple - (or list of shape tuples if the layer has multiple outputs). - - Raises: - RuntimeError: If called 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', - 'output shape') - - def get_input_at(self, node_index): - """Retrieves the input 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 tensor (or list of tensors if the layer has multiple inputs). - - Raises: - RuntimeError: If called 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') - - def get_output_at(self, node_index): - """Retrieves the output 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 tensor (or list of tensors if the layer has multiple outputs). - - Raises: - RuntimeError: If called in Eager mode. - """ - return self._get_node_attribute_at_index(node_index, 'output_tensors', - 'output') - - @property - def input(self): - """Retrieves the input tensor(s) of a layer. - - Only applicable if the layer has exactly one input, - i.e. if it is connected to one incoming layer. - - Returns: - Input tensor or list of input tensors. - - Raises: - AttributeError: if the layer is connected to - more than one incoming layers. - - Raises: - RuntimeError: If called in Eager mode. - AttributeError: If no inbound nodes are found. - """ - if not self._inbound_nodes: - raise AttributeError('Layer ' + self.name + - ' is not connected, no input to return.') - return self._get_node_attribute_at_index(0, 'input_tensors', 'input') - - @property - def output(self): - """Retrieves the output tensor(s) of a layer. - - Only applicable if the layer has exactly one output, - i.e. if it is connected to one incoming layer. - - Returns: - Output tensor or list of output tensors. - - Raises: - AttributeError: if the layer is connected to more than one incoming - layers. - RuntimeError: if called in Eager mode. - """ - if not self._inbound_nodes: - raise AttributeError('Layer ' + self.name + ' has no inbound nodes.') - return self._get_node_attribute_at_index(0, 'output_tensors', 'output') - - @property - def input_shape(self): - """Retrieves the input shape(s) of a layer. - - Only applicable if the layer has exactly one input, - i.e. if it is connected to one incoming layer, or if all inputs - have the same shape. - - Returns: - Input shape, as an integer shape tuple - (or list of shape tuples, one tuple per input tensor). - - Raises: - AttributeError: if the layer has no defined input_shape. - RuntimeError: if called in Eager mode. - """ - if not self._inbound_nodes: - raise AttributeError('The layer has never been called ' - 'and thus has no defined input shape.') - all_input_shapes = set( - [str(node.input_shapes) for node in self._inbound_nodes]) - if len(all_input_shapes) == 1: - input_shapes = self._inbound_nodes[0].input_shapes - if len(input_shapes) == 1: - return tuple(tensor_shape.TensorShape(input_shapes[0]).as_list()) - else: - return [ - tuple(tensor_shape.TensorShape(shape).as_list()) - for shape in input_shapes - ] - else: - raise AttributeError('The layer "' + str(self.name) + - ' has multiple inbound nodes, ' - 'with different input shapes. Hence ' - 'the notion of "input shape" is ' - 'ill-defined for the layer. ' - 'Use `get_input_shape_at(node_index)` ' - 'instead.') - - def count_params(self): - """Count the total number of scalars composing the weights. - - Returns: - An integer count. - - Raises: - ValueError: if the layer isn't yet built - (in which case its weights aren't yet defined). - """ - if not self.built: - if self.__class__.__name__ == 'Sequential': - self.build() # pylint: disable=no-value-for-parameter - else: - raise ValueError('You tried to call `count_params` on ' + self.name + - ', but the layer isn\'t built. ' - 'You can build it manually via: `' + self.name + - '.build(batch_input_shape)`.') - weight_shapes = [w.get_shape().as_list() for w in self.weights] - return int(sum([np.prod(w) for w in weight_shapes])) - - @property - def output_shape(self): - """Retrieves the output shape(s) of a layer. - - Only applicable if the layer has one output, - or if all outputs have the same shape. - - Returns: - Output shape, as an integer shape tuple - (or list of shape tuples, one tuple per output tensor). - - Raises: - AttributeError: if the layer has no defined output shape. - RuntimeError: if called in Eager mode. - """ - if not self._inbound_nodes: - raise AttributeError('The layer has never been called ' - 'and thus has no defined output shape.') - all_output_shapes = set( - [str(node.output_shapes) for node in self._inbound_nodes]) - if len(all_output_shapes) == 1: - output_shapes = self._inbound_nodes[0].output_shapes - if len(output_shapes) == 1: - return tuple(tensor_shape.TensorShape(output_shapes[0]).as_list()) - else: - return [ - tuple(tensor_shape.TensorShape(shape).as_list()) - for shape in output_shapes - ] - else: - raise AttributeError('The layer "%s"' - ' has multiple inbound nodes, ' - 'with different output shapes. Hence ' - 'the notion of "output shape" is ' - 'ill-defined for the layer. ' - 'Use `get_output_shape_at(node_index)` ' - 'instead.' % self.name) - - @property - def inbound_nodes(self): - """Deprecated, do NOT use! Only for compatibility with external Keras.""" - return self._inbound_nodes - - @property - def outbound_nodes(self): - """Deprecated, do NOT use! Only for compatibility with external Keras.""" - return self._outbound_nodes - - def _assert_input_compatibility(self, inputs): - """Checks compatibility between the layer and provided inputs. - - This checks that the tensor(s) `inputs` verify the input assumptions - of the layer (if any). If not, a clear and actional exception gets raised. - - Arguments: - inputs: input tensor or list of input tensors. - - Raises: - ValueError: in case of mismatch between - the provided inputs and the expectations of the layer. - """ - if not self.input_spec: - return - if not isinstance(self.input_spec, (list, tuple)): - input_spec = nest.flatten(self.input_spec) - else: - input_spec = self.input_spec - inputs = nest.flatten(inputs) - if len(inputs) != len(input_spec): - raise ValueError('Layer ' + self.name + ' expects ' + - str(len(input_spec)) + ' inputs, ' - 'but it received ' + str(len(inputs)) + - ' input tensors. Inputs received: ' + str(inputs)) - for input_index, (x, spec) in enumerate(zip(inputs, input_spec)): - if spec is None: - continue - - if (spec.ndim is not None or - spec.min_ndim is not None or - spec.max_ndim is not None): - if x.get_shape().ndims is None: - raise ValueError('Input ' + str(input_index) + ' of layer ' + - self.name + ' is incompatible with the layer: ' - 'its rank is undefined, but the layer requires a ' - 'defined rank.') - - # Check ndim. - if spec.ndim is not None: - ndim = x.get_shape().ndims - if ndim != spec.ndim: - raise ValueError('Input ' + str(input_index) + ' of layer ' + - self.name + ' is incompatible with the layer: ' - 'expected ndim=' + str(spec.ndim) + ', found ndim=' + - str(ndim) + '. Full shape received: ' + - str(x.get_shape().as_list())) - if spec.max_ndim is not None: - ndim = x.get_shape().ndims - if ndim is not None and ndim > spec.max_ndim: - raise ValueError('Input ' + str(input_index) + ' of layer ' + - self.name + ' is incompatible with the layer: ' - 'expected max_ndim=' + str(spec.max_ndim) + - ', found ndim=' + str(ndim)) - if spec.min_ndim is not None: - ndim = x.get_shape().ndims - if ndim is not None and ndim < spec.min_ndim: - raise ValueError('Input ' + str(input_index) + ' of layer ' + - self.name + ' is incompatible with the layer: ' - ': expected min_ndim=' + str(spec.min_ndim) + - ', found ndim=' + str(ndim) + - '. Full shape received: ' + - str(x.get_shape().as_list())) - # Check dtype. - if spec.dtype is not None: - if x.dtype != spec.dtype: - raise ValueError('Input ' + str(input_index) + ' of layer ' + - self.name + ' is incompatible with the layer: ' - 'expected dtype=' + str(spec.dtype) + - ', found dtype=' + str(x.dtype)) - # Check specific shape axes. - if spec.axes: - shape = x.get_shape().as_list() - if shape is not None: - for axis, value in spec.axes.items(): - if hasattr(value, 'value'): - value = value.value - if value is not None and shape[int(axis)] not in {value, None}: - raise ValueError( - 'Input ' + str(input_index) + ' of layer ' + self.name + ' is' - ' incompatible with the layer: expected axis ' + str(axis) + - ' of input shape to have value ' + str(value) + - ' but received input with shape ' + str(shape)) - # Check shape. - if spec.shape is not None: - shape = x.get_shape().as_list() - if shape is not None: - for spec_dim, dim in zip(spec.shape, shape): - if spec_dim is not None and dim is not None: - if spec_dim != dim: - raise ValueError('Input ' + str(input_index) + - ' is incompatible with layer ' + self.name + - ': expected shape=' + str(spec.shape) + - ', 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. - - Every layer should expose (if appropriate) an `input_spec` attribute: - a list of instances of InputSpec (one per input tensor). - - A None entry in a shape is compatible with any dimension, - a None shape is compatible with any shape. - - Arguments: - dtype: Expected DataType of the input. - shape: Shape tuple, expected shape of the input - (may include None for unchecked axes). - ndim: Integer, expected rank of the input. - max_ndim: Integer, maximum rank of the input. - min_ndim: Integer, minimum rank of the input. - axes: Dictionary mapping integer axes to - a specific dimension value. - """ - - def __init__(self, - dtype=None, - shape=None, - ndim=None, - max_ndim=None, - min_ndim=None, - axes=None): - self.dtype = dtype - self.shape = shape - if shape is not None: - self.ndim = len(shape) - else: - self.ndim = ndim - self.max_ndim = max_ndim - self.min_ndim = min_ndim - self.axes = axes or {} - - def __repr__(self): - spec = [('dtype=' + str(self.dtype)) if self.dtype else '', - ('shape=' + str(self.shape)) if self.shape else '', - ('ndim=' + str(self.ndim)) if self.ndim else '', - ('max_ndim=' + str(self.max_ndim)) if self.max_ndim else '', - ('min_ndim=' + str(self.min_ndim)) if self.min_ndim else '', - ('axes=' + str(self.axes)) if self.axes else ''] - return 'InputSpec(%s)' % ', '.join(x for x in spec if x) - - -class Node(object): - """A `Node` describes the connectivity between two layers. - - Each time a layer is connected to some new input, - a node is added to `layer._inbound_nodes`. - Each time the output of a layer is used by another layer, - a node is added to `layer._outbound_nodes`. - - Arguments: - outbound_layer: the layer that takes - `input_tensors` and turns them into `output_tensors` - (the node gets created when the `call` - method of the layer was called). - inbound_layers: a list of layers, the same length as `input_tensors`, - the layers from where `input_tensors` originate. - node_indices: a list of integers, the same length as `inbound_layers`. - `node_indices[i]` is the origin node of `input_tensors[i]` - (necessary since each inbound layer might have several nodes, - e.g. if the layer is being shared with a different data stream). - tensor_indices: a list of integers, - the same length as `inbound_layers`. - `tensor_indices[i]` is the index of `input_tensors[i]` within the - output of the inbound layer - (necessary since each inbound layer might - have multiple tensor outputs, with each one being - independently manipulable). - input_tensors: list of input tensors. - output_tensors: list of output tensors. - arguments: dictionary of keyword arguments that were passed to the - `call` method of the layer at the call that created the node. - - `node_indices` and `tensor_indices` are basically fine-grained coordinates - describing the origin of the `input_tensors`. - - A node from layer A to layer B is added to: - - A._outbound_nodes - - B._inbound_nodes - """ - - def __init__(self, - outbound_layer, - inbound_layers, - node_indices, - tensor_indices, - input_tensors, - output_tensors, - arguments=None): - # Layer instance (NOT a list). - if isinstance(outbound_layer, list): - raise ValueError( - '`outbound_layer` should be a layer instance, not a list.') - # this is the layer that takes a list of input tensors - # and turns them into a list of output tensors. - # the current node will be added to - # the inbound_nodes of outbound_layer. - self.outbound_layer = outbound_layer - - # The following 3 properties describe where - # the input tensors come from: which layers, - # and for each layer, which node and which - # tensor output of each node. - - # List of layer instances. - self.inbound_layers = inbound_layers - # List of integers, 1:1 mapping with inbound_layers. - self.node_indices = node_indices - # List of integers, 1:1 mapping with inbound_layers. - self.tensor_indices = tensor_indices - - # Following 2 properties: - # tensor inputs and outputs of outbound_layer. - - # List of tensors. 1:1 mapping with inbound_layers. - self.input_tensors = input_tensors - # List of tensors, created by outbound_layer.call(). - self.output_tensors = output_tensors - - # Following 2 properties: input and output shapes. - - # List of shape tuples, shapes of input_tensors. - self.input_shapes = [layers_util.static_shape(x) for x in input_tensors] - # List of shape tuples, shapes of output_tensors. - self.output_shapes = [layers_util.static_shape(x) for x in output_tensors] - - # Optional keyword arguments to layer's `call`. - self.arguments = arguments - - # Add nodes to all layers involved. - for layer in inbound_layers: - if layer is not None: - # For compatibility with external Keras, we use the deprecated - # accessor here. - layer.outbound_nodes.append(self) - # For compatibility with external Keras, we use the deprecated - # accessor here. - outbound_layer.inbound_nodes.append(self) - - def get_config(self): - inbound_names = [] - for layer in self.inbound_layers: - if layer: - inbound_names.append(layer.name) - else: - inbound_names.append(None) - return { - 'outbound_layer': self.outbound_layer.name, - 'inbound_layers': inbound_names, - 'node_indices': self.node_indices, - 'tensor_indices': self.tensor_indices - } - - -class _DeferredTensor(object): - """Tensor-like object used to build graphs of layers in Eager mode. - - When calling a layer on a DeferredTensor, the layer will not perform any - computation and will simply perfom shape inference to return new - DeferredTensors with appropriate shape information. Thus DeferredTensor - behaves like a graph-mode Tensor when manipulated by layers. - """ - - def __init__(self, shape, dtype, name=None): - self.shape = tensor_shape.TensorShape(shape) - 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): - return self.shape - - def __str__(self): - return "DeferredTensor('%s', shape=%s, dtype=%s)" % (self.name, - self.get_shape(), - self.dtype.name) - - def __repr__(self): - return "<_DeferredTensor '%s' shape=%s dtype=%s>" % (self.name, - self.get_shape(), - self.dtype.name) - - -def _is_tensor_or_tensor_list(v): - v = nest.flatten(v) - if v and isinstance(v[0], ops.Tensor): - return True - else: - return False - - -def _to_snake_case(name): - intermediate = re.sub('(.)([A-Z][a-z0-9]+)', r'\1_\2', name) - insecure = re.sub('([a-z])([A-Z])', r'\1_\2', intermediate).lower() - # If the class is private the name starts with "_" which is not secure - # for creating scopes. We prefix the name with "private" in this case. - if insecure[0] != '_': - return insecure - return 'private' + insecure - - -def _to_list(x): - """This normalizes a list/tuple or single element into a list. - - If a single element is passed, we return - a list of size 1 containing the element. - - Arguments: - x: list or tuple or single element. - - Returns: - A list. - """ - if isinstance(x, (list, tuple)): - return list(x) - return [x] - def _add_elements_to_collection(elements, collection_list): if context.executing_eagerly(): @@ -1473,105 +354,3 @@ def _add_elements_to_collection(elements, collection_list): if element not in collection_set: collection.append(element) - -def _is_all_none(iterable_or_element): - if not isinstance(iterable_or_element, (list, tuple)): - iterable = [iterable_or_element] - else: - iterable = iterable_or_element - # We cannot use Python's `any` because the iterable may return Tensors. - for element in iterable: - if element is not None: - return False - return True - - -def _have_all_keras_metadata(iterable_or_element): - if not isinstance(iterable_or_element, (list, tuple)): - iterable = [iterable_or_element] - else: - iterable = iterable_or_element - return all([hasattr(x, '_keras_history') for x in iterable]) - - -def _collect_previous_mask(input_tensors): - """Retrieves the output mask(s) of the previous node. - - Arguments: - input_tensors: A tensor or list of tensors. - - Returns: - A mask tensor or list of mask tensors. - """ - input_tensors = nest.flatten(input_tensors) - masks = [] - for x in input_tensors: - if hasattr(x, '_keras_mask'): - mask = x._keras_mask # pylint: disable=protected-access - masks.append(mask) - else: - masks.append(None) - if len(masks) == 1: - return masks[0] - return masks - - -# A global dictionary mapping graph objects to an index of counters used -# for various layer names in each graph. -# Allows to give unique autogenerated names to layers, in a graph-specific way. -PER_GRAPH_LAYER_NAME_UIDS = weakref.WeakKeyDictionary() - - -def _get_default_graph_uid_map(): - graph = ops.get_default_graph() - name_uid_map = PER_GRAPH_LAYER_NAME_UIDS.get(graph, None) - if name_uid_map is None: - name_uid_map = collections.defaultdict(int) - PER_GRAPH_LAYER_NAME_UIDS[graph] = name_uid_map - return name_uid_map - - -def _unique_layer_name(name, name_uid_map=None, avoid_names=None, namespace='', - zero_based=False): - """Makes a layer name (or arbitrary string) unique within a TensorFlow graph. - - Arguments: - name: String name to make unique. - name_uid_map: An optional defaultdict(int) to use when creating unique - names. If None (default), uses a per-Graph dictionary. - avoid_names: An optional set or dict with names which should not be used. If - None (default) does not avoid any names. - namespace: Gets a name which is unique within the (graph, namespace). Layers - which are not Networks use a blank namespace and so get graph-global - names. - zero_based: If True, name sequences start with no suffix (e.g. "dense", - "dense_1"). If False, naming is one-based ("dense_1", "dense_2"). - - Returns: - Unique string name. - - Example: - - ```python - _unique_layer_name('dense') # dense_1 - _unique_layer_name('dense') # dense_2 - ``` - """ - if name_uid_map is None: - name_uid_map = _get_default_graph_uid_map() - if avoid_names is None: - avoid_names = set() - proposed_name = None - while proposed_name is None or proposed_name in avoid_names: - name_key = (namespace, name) - if zero_based: - number = name_uid_map[name_key] - if number: - proposed_name = name + '_' + str(number) - else: - proposed_name = name - name_uid_map[name_key] += 1 - else: - name_uid_map[name_key] += 1 - proposed_name = name + '_' + str(name_uid_map[name_key]) - return proposed_name diff --git a/tensorflow/python/layers/base_test.py b/tensorflow/python/layers/base_test.py index 9ed4afeaba931c47d2a1e65f08489773f0b9eb1b..f08b552840f5ff9144edae1cb0f90a1bc3db0f8c 100644 --- a/tensorflow/python/layers/base_test.py +++ b/tensorflow/python/layers/base_test.py @@ -52,6 +52,12 @@ class BaseLayerTest(test.TestCase): layer = base_layers.Layer(name='my_layer', trainable=False) self.assertEqual(layer.trainable, False) + @test_util.run_in_graph_and_eager_modes() + def testInt64Layer(self): + layer = base_layers.Layer(name='my_layer', dtype='int64') + layer.add_variable('my_var', [2, 2]) + self.assertEqual(layer.name, 'my_layer') + @test_util.run_in_graph_and_eager_modes() def testAddWeight(self): layer = base_layers.Layer(name='my_layer') @@ -94,61 +100,6 @@ class BaseLayerTest(test.TestCase): with self.assertRaisesRegexp(ValueError, 'activity_regularizer'): core_layers.Dense(1, activity_regularizer=lambda *args, **kwargs: 0.) - def testGetVariable(self): - with self.test_session(): - - class MyLayer(base_layers.Layer): - - def build(self, input_shape): - self.my_var = self.add_variable( - 'my_var', [2, 2], initializer=init_ops.zeros_initializer()) - - def call(self, inputs): - return inputs * 2 - - layer = MyLayer(name='my_layer') - inputs = random_ops.random_uniform((5,), seed=1) - layer.apply(inputs) - layer.apply(inputs) - self.assertEqual([v.name for v in layer.variables], - ['my_layer/my_var:0']) - - # Creating a layer with no scope leads to lazy construction of - # the scope at apply() time. It uses scope "/base_name" - lazy_layer = MyLayer(_reuse=True) - with variable_scope.variable_scope('new_scope'): - with variable_scope.variable_scope('my_layer'): - variable_scope.get_variable('my_var', [2, 2]) - - # Smoke test: it runs. - lazy_layer.apply(inputs) - # The variables were created outside of the Layer, and - # reuse=True, so the Layer does not own them and they are not - # stored in its collection. - self.assertEqual(lazy_layer.variables, []) - self.assertEqual(lazy_layer._scope.name, 'new_scope/my_layer') - - # Creating a layer with no scope leads to lazy construction of - # the scope at apply() time. If 'scope' argument is passed to - # apply(), it uses that scope when accessing variables. - lazy_layer = MyLayer(_reuse=True) - with variable_scope.variable_scope('new_scope') as new_scope: - variable_scope.get_variable('my_var', [2, 2]) - - # Smoke test: it runs. - lazy_layer.apply(inputs, scope=new_scope) - # The variables were created outside of the Layer, and - # reuse=True, so the Layer does not own them and they are not - # stored in its collection. - self.assertEqual(lazy_layer.variables, []) - self.assertEqual(lazy_layer._scope.name, 'new_scope') - - # Checking for graph equality is only done in GRAPH mode. - with ops.Graph().as_default(): - inputs_ng = random_ops.random_uniform((5,), seed=1) - with self.assertRaisesRegexp(ValueError, r'graph are not the same'): - layer.apply(inputs_ng) - @test_util.run_in_graph_and_eager_modes() def testCall(self): @@ -165,38 +116,6 @@ class BaseLayerTest(test.TestCase): # op is only supported in GRAPH mode self.assertEqual(outputs.op.name, 'my_layer/Square') - def testFirstCallCanCreateVariablesButSecondCanNotWhenBuildEmpty(self): - # Note that this test is only run in Graph mode since with EAGER mode we can - # still create a new variable on second call. - - class MyLayer(base_layers.Layer): - - def build(self, _): - # Do not mark the layer as built. - pass - - def call(self, inputs): - self.my_var = self.add_variable('my_var', [2, 2]) - if self.built: - # Skip creating on the first call; try to create after it's - # built. This is expected to fail. - self.add_variable('this_will_break_on_second_call', [2, 2]) - return inputs + math_ops.square(self.my_var) - - layer = MyLayer(name='my_layer') - inputs = random_ops.random_uniform((2,), seed=1) - outputs = layer.apply(inputs) - self.assertEqual(layer.built, True) - self.assertEqual(outputs.op.name, 'my_layer/add') - self.assertEqual([v.name - for v in layer.variables], ['my_layer/my_var:0']) - with self.assertRaisesRegexp(ValueError, - 'my_layer/this_will_break_on_second_call'): - layer.apply(inputs) - # The list of variables hasn't changed. - self.assertEqual([v.name - for v in layer.variables], ['my_layer/my_var:0']) - @test_util.run_in_graph_and_eager_modes() def testDeepCopy(self): @@ -645,13 +564,14 @@ class BaseLayerTest(test.TestCase): def testLayerGraphSetInFirstApply(self): with ops.Graph().as_default(): - layer = core_layers.Dense(1) # Graph at construction time is ignored + # Graph at construction time is ignored + layer = core_layers.Dense(1) with ops.Graph().as_default(): - layer.apply(constant_op.constant([[1]])) + 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]])) + layer.apply(constant_op.constant([[1.]])) if __name__ == '__main__': diff --git a/tensorflow/python/layers/convolutional.py b/tensorflow/python/layers/convolutional.py index 2d99b1688f1b2736c0660ba2ac914018b21bf9ed..34a1487e748e41eebae8b87b17c34d0deda8597f 100644 --- a/tensorflow/python/layers/convolutional.py +++ b/tensorflow/python/layers/convolutional.py @@ -23,6 +23,7 @@ from __future__ import print_function from tensorflow.python.eager import context from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape +from tensorflow.python.keras._impl.keras import layers as keras_layers from tensorflow.python.layers import base from tensorflow.python.layers import utils from tensorflow.python.ops import array_ops @@ -32,201 +33,8 @@ from tensorflow.python.ops import nn_ops from tensorflow.python.util.tf_export import tf_export -class _Conv(base.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. If `use_bias` is True (and a `bias_initializer` is provided), - a bias vector is created and added to the outputs. Finally, if - `activation` is not `None`, it is applied to the outputs as well. - - 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. - use_bias: Boolean, whether the layer uses a bias. - kernel_initializer: An initializer for the convolution kernel. - bias_initializer: An initializer for the bias vector. If None, the default - initializer will be used. - kernel_regularizer: Optional regularizer for the convolution kernel. - bias_regularizer: Optional regularizer for the bias vector. - activity_regularizer: Optional regularizer function for the output. - kernel_constraint: Optional projection function to be applied to the - kernel after being updated by an `Optimizer` (e.g. used to implement - norm constraints or value constraints for layer weights). 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. - bias_constraint: Optional projection function to be applied to the - bias after being updated by an `Optimizer`. - trainable: Boolean, if `True` also add variables to the graph collection - `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`). - name: A string, the name of the layer. - """ - - def __init__(self, rank, - filters, - kernel_size, - strides=1, - padding='valid', - data_format='channels_last', - dilation_rate=1, - activation=None, - use_bias=True, - kernel_initializer=None, - bias_initializer=init_ops.zeros_initializer(), - kernel_regularizer=None, - bias_regularizer=None, - activity_regularizer=None, - kernel_constraint=None, - bias_constraint=None, - trainable=True, - name=None, - **kwargs): - super(_Conv, 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.use_bias = use_bias - self.kernel_initializer = kernel_initializer - self.bias_initializer = bias_initializer - self.kernel_regularizer = kernel_regularizer - self.bias_regularizer = bias_regularizer - self.kernel_constraint = kernel_constraint - self.bias_constraint = bias_constraint - self.input_spec = base.InputSpec(ndim=self.rank + 2) - - 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) - - self.kernel = self.add_variable(name='kernel', - shape=kernel_shape, - initializer=self.kernel_initializer, - regularizer=self.kernel_regularizer, - constraint=self.kernel_constraint, - trainable=True, - dtype=self.dtype) - if self.use_bias: - self.bias = self.add_variable(name='bias', - shape=(self.filters,), - initializer=self.bias_initializer, - regularizer=self.bias_regularizer, - constraint=self.bias_constraint, - trainable=True, - dtype=self.dtype) - else: - self.bias = None - self.input_spec = base.InputSpec(ndim=self.rank + 2, - axes={channel_axis: input_dim}) - self._convolution_op = nn_ops.Convolution( - input_shape, - filter_shape=self.kernel.get_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): - outputs = self._convolution_op(inputs, self.kernel) - - if self.use_bias: - 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, (1, self.filters, 1)) - outputs += bias - if self.rank == 2: - outputs = nn.bias_add(outputs, self.bias, 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() - if outputs_shape[0] is None: - outputs_shape[0] = -1 - 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, data_format='NCHW') - outputs = array_ops.reshape(outputs_4d, outputs_shape) - else: - outputs = nn.bias_add(outputs, self.bias, data_format='NHWC') - - if self.activation is not None: - return self.activation(outputs) - return outputs - - 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) - - @tf_export('layers.Conv1D') -class Conv1D(_Conv): +class Conv1D(keras_layers.Conv1D, base.Layer): """1D convolution layer (e.g. temporal convolution). This layer creates a convolution kernel that is convolved @@ -294,8 +102,7 @@ class Conv1D(_Conv): trainable=True, name=None, **kwargs): - super(Convolution1D, self).__init__( - rank=1, + super(Conv1D, self).__init__( filters=filters, kernel_size=kernel_size, strides=strides, @@ -417,7 +224,7 @@ def conv1d(inputs, @tf_export('layers.Conv2D') -class Conv2D(_Conv): +class Conv2D(keras_layers.Conv2D, base.Layer): """2D convolution layer (e.g. spatial convolution over images). This layer creates a convolution kernel that is convolved @@ -493,7 +300,6 @@ class Conv2D(_Conv): name=None, **kwargs): super(Conv2D, self).__init__( - rank=2, filters=filters, kernel_size=kernel_size, strides=strides, @@ -622,7 +428,7 @@ def conv2d(inputs, @tf_export('layers.Conv3D') -class Conv3D(_Conv): +class Conv3D(keras_layers.Conv3D, base.Layer): """3D convolution layer (e.g. spatial convolution over volumes). This layer creates a convolution kernel that is convolved @@ -699,7 +505,6 @@ class Conv3D(_Conv): name=None, **kwargs): super(Conv3D, self).__init__( - rank=3, filters=filters, kernel_size=kernel_size, strides=strides, @@ -828,169 +633,8 @@ def conv3d(inputs, return layer.apply(inputs) -class _SeparableConv(_Conv): - """Abstract base layer for separable nD convolution. - - This layer performs a depthwise convolution that acts separately on - channels, followed by a pointwise convolution that mixes channels. - If `use_bias` is True and a bias initializer is provided, - it adds a bias vector to the output. - It then optionally applies an activation function to produce the final output. - - 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: A tuple or list of integers specifying the spatial - dimensions of the filters. Can be a single integer to specify the same - value for all spatial dimensions. - strides: A tuple or list of integers specifying the strides - of the convolution. 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, ..., channels)` while `channels_first` corresponds to - inputs with shape `(batch, channels, ...)`. - 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. - depth_multiplier: The number of depthwise convolution output channels for - each input channel. The total number of depthwise convolution output - channels will be equal to `num_filters_in * depth_multiplier`. - activation: Activation function. Set it to None to maintain a - linear activation. - use_bias: Boolean, whether the layer uses a bias. - depthwise_initializer: An initializer for the depthwise convolution kernel. - pointwise_initializer: An initializer for the pointwise convolution kernel. - bias_initializer: An initializer for the bias vector. If None, the default - initializer will be used. - depthwise_regularizer: Optional regularizer for the depthwise - convolution kernel. - pointwise_regularizer: Optional regularizer for the pointwise - convolution kernel. - bias_regularizer: Optional regularizer for the bias vector. - activity_regularizer: Optional regularizer function for the output. - depthwise_constraint: Optional projection function to be applied to the - depthwise kernel after being updated by an `Optimizer` (e.g. used for - norm constraints or value constraints for layer weights). 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. - pointwise_constraint: Optional projection function to be applied to the - pointwise kernel after being updated by an `Optimizer`. - bias_constraint: Optional projection function to be applied to the - bias after being updated by an `Optimizer`. - trainable: Boolean, if `True` also add variables to the graph collection - `GraphKeys.TRAINABLE_VARIABLES` (see `tf.Variable`). - name: A string, the name of the layer. - """ - - def __init__(self, - rank, - filters, - kernel_size, - strides=1, - padding='valid', - data_format='channels_last', - dilation_rate=1, - depth_multiplier=1, - activation=None, - use_bias=True, - depthwise_initializer=None, - pointwise_initializer=None, - bias_initializer=init_ops.zeros_initializer(), - depthwise_regularizer=None, - pointwise_regularizer=None, - bias_regularizer=None, - activity_regularizer=None, - depthwise_constraint=None, - pointwise_constraint=None, - bias_constraint=None, - trainable=True, - name=None, - **kwargs): - super(_SeparableConv, self).__init__( - rank=rank, - filters=filters, - kernel_size=kernel_size, - strides=strides, - padding=padding, - data_format=data_format, - dilation_rate=dilation_rate, - activation=activation, - use_bias=use_bias, - bias_regularizer=bias_regularizer, - activity_regularizer=activity_regularizer, - bias_constraint=bias_constraint, - trainable=trainable, - name=name, - **kwargs) - self.depth_multiplier = depth_multiplier - self.depthwise_initializer = depthwise_initializer - self.pointwise_initializer = pointwise_initializer - self.depthwise_regularizer = depthwise_regularizer - self.pointwise_regularizer = pointwise_regularizer - self.depthwise_constraint = depthwise_constraint - self.pointwise_constraint = pointwise_constraint - - 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 - self.input_spec = base.InputSpec(ndim=self.rank + 2, - axes={channel_axis: input_dim}) - depthwise_kernel_shape = self.kernel_size + (input_dim, - self.depth_multiplier) - pointwise_kernel_shape = ( - 1,) * self.rank + (self.depth_multiplier * input_dim, self.filters) - - self.depthwise_kernel = self.add_variable( - name='depthwise_kernel', - shape=depthwise_kernel_shape, - initializer=self.depthwise_initializer, - regularizer=self.depthwise_regularizer, - constraint=self.depthwise_constraint, - trainable=True, - dtype=self.dtype) - self.pointwise_kernel = self.add_variable( - name='pointwise_kernel', - shape=pointwise_kernel_shape, - initializer=self.pointwise_initializer, - regularizer=self.pointwise_regularizer, - constraint=self.pointwise_constraint, - trainable=True, - dtype=self.dtype) - if self.use_bias: - self.bias = self.add_variable(name='bias', - shape=(self.filters,), - initializer=self.bias_initializer, - regularizer=self.bias_regularizer, - constraint=self.bias_constraint, - trainable=True, - dtype=self.dtype) - else: - self.bias = None - self.built = True - - def call(self, inputs): - raise NotImplementedError - - @tf_export('layers.SeparableConv1D') -class SeparableConv1D(_SeparableConv): +class SeparableConv1D(keras_layers.SeparableConv1D, base.Layer): """Depthwise separable 1D convolution. This layer performs a depthwise convolution that acts separately on @@ -1072,7 +716,6 @@ class SeparableConv1D(_SeparableConv): name=None, **kwargs): super(SeparableConv1D, self).__init__( - rank=1, filters=filters, kernel_size=kernel_size, strides=strides, @@ -1096,45 +739,9 @@ class SeparableConv1D(_SeparableConv): name=name, **kwargs) - def call(self, inputs): - if self.data_format == 'channels_last': - strides = (1,) + self.strides * 2 + (1,) - spatial_start_dim = 1 - else: - strides = (1, 1) + self.strides * 2 - spatial_start_dim = 2 - - # Explicitly broadcast inputs and kernels to 4D. - # TODO(fchollet): refactor when a native separable_conv1d op is available. - inputs = array_ops.expand_dims(inputs, spatial_start_dim) - depthwise_kernel = array_ops.expand_dims(self.depthwise_kernel, 0) - pointwise_kernel = array_ops.expand_dims(self.pointwise_kernel, 0) - dilation_rate = (1,) + self.dilation_rate - - outputs = nn.separable_conv2d( - inputs, - depthwise_kernel, - pointwise_kernel, - strides=strides, - padding=self.padding.upper(), - rate=dilation_rate, - data_format=utils.convert_data_format(self.data_format, ndim=4)) - - if self.use_bias: - outputs = nn.bias_add( - outputs, - self.bias, - data_format=utils.convert_data_format(self.data_format, ndim=4)) - - outputs = array_ops.squeeze(outputs, [spatial_start_dim]) - - if self.activation is not None: - return self.activation(outputs) - return outputs - @tf_export('layers.SeparableConv2D') -class SeparableConv2D(_SeparableConv): +class SeparableConv2D(keras_layers.SeparableConv2D, base.Layer): """Depthwise separable 2D convolution. This layer performs a depthwise convolution that acts separately on @@ -1221,7 +828,6 @@ class SeparableConv2D(_SeparableConv): name=None, **kwargs): super(SeparableConv2D, self).__init__( - rank=2, filters=filters, kernel_size=kernel_size, strides=strides, @@ -1245,31 +851,6 @@ class SeparableConv2D(_SeparableConv): name=name, **kwargs) - def call(self, inputs): - # Apply the actual ops. - if self.data_format == 'channels_last': - strides = (1,) + self.strides + (1,) - else: - strides = (1, 1) + self.strides - outputs = nn.separable_conv2d( - inputs, - self.depthwise_kernel, - self.pointwise_kernel, - strides=strides, - padding=self.padding.upper(), - rate=self.dilation_rate, - data_format=utils.convert_data_format(self.data_format, ndim=4)) - - if self.use_bias: - outputs = nn.bias_add( - outputs, - self.bias, - data_format=utils.convert_data_format(self.data_format, ndim=4)) - - if self.activation is not None: - return self.activation(outputs) - return outputs - @tf_export('layers.separable_conv1d') def separable_conv1d(inputs, @@ -1511,7 +1092,7 @@ def separable_conv2d(inputs, @tf_export('layers.Conv2DTranspose') -class Conv2DTranspose(Conv2D): +class Conv2DTranspose(keras_layers.Conv2DTranspose, base.Layer): """Transposed 2D convolution layer (sometimes called 2D Deconvolution). The need for transposed convolutions generally arises @@ -1576,8 +1157,8 @@ class Conv2DTranspose(Conv2D): name=None, **kwargs): super(Conv2DTranspose, self).__init__( - filters, - kernel_size, + filters=filters, + kernel_size=kernel_size, strides=strides, padding=padding, data_format=data_format, @@ -1593,120 +1174,6 @@ class Conv2DTranspose(Conv2D): trainable=trainable, name=name, **kwargs) - self.input_spec = base.InputSpec(ndim=4) - - def build(self, input_shape): - if len(input_shape) != 4: - raise ValueError('Inputs should have rank 4. Received input shape: ' + - str(input_shape)) - if self.data_format == 'channels_first': - channel_axis = 1 - else: - channel_axis = -1 - if input_shape[channel_axis] is None: - raise ValueError('The channel dimension of the inputs ' - 'should be defined. Found `None`.') - input_dim = input_shape[channel_axis] - self.input_spec = base.InputSpec(ndim=4, axes={channel_axis: input_dim}) - kernel_shape = self.kernel_size + (self.filters, input_dim) - - self.kernel = self.add_variable(name='kernel', - shape=kernel_shape, - initializer=self.kernel_initializer, - regularizer=self.kernel_regularizer, - constraint=self.kernel_constraint, - trainable=True, - dtype=self.dtype) - if self.use_bias: - self.bias = self.add_variable(name='bias', - shape=(self.filters,), - initializer=self.bias_initializer, - regularizer=self.bias_regularizer, - constraint=self.bias_constraint, - trainable=True, - dtype=self.dtype) - else: - self.bias = None - self.built = True - - def call(self, inputs): - inputs_shape = array_ops.shape(inputs) - batch_size = inputs_shape[0] - if self.data_format == 'channels_first': - c_axis, h_axis, w_axis = 1, 2, 3 - else: - c_axis, h_axis, w_axis = 3, 1, 2 - - height, width = inputs_shape[h_axis], inputs_shape[w_axis] - kernel_h, kernel_w = self.kernel_size - stride_h, stride_w = self.strides - - # Infer the dynamic output shape: - out_height = utils.deconv_output_length(height, - kernel_h, - self.padding, - stride_h) - out_width = utils.deconv_output_length(width, - kernel_w, - self.padding, - stride_w) - if self.data_format == 'channels_first': - output_shape = (batch_size, self.filters, out_height, out_width) - strides = (1, 1, stride_h, stride_w) - else: - output_shape = (batch_size, out_height, out_width, self.filters) - strides = (1, stride_h, stride_w, 1) - - output_shape_tensor = array_ops.stack(output_shape) - outputs = nn.conv2d_transpose( - inputs, - self.kernel, - output_shape_tensor, - strides, - padding=self.padding.upper(), - data_format=utils.convert_data_format(self.data_format, ndim=4)) - - if not context.executing_eagerly(): - # Infer the static output shape: - out_shape = inputs.get_shape().as_list() - out_shape[c_axis] = self.filters - out_shape[h_axis] = utils.deconv_output_length(out_shape[h_axis], - kernel_h, - self.padding, - stride_h) - out_shape[w_axis] = utils.deconv_output_length(out_shape[w_axis], - kernel_w, - self.padding, - stride_w) - outputs.set_shape(out_shape) - - if self.use_bias: - outputs = nn.bias_add( - outputs, - self.bias, - data_format=utils.convert_data_format(self.data_format, ndim=4)) - - if self.activation is not None: - return self.activation(outputs) - return outputs - - def compute_output_shape(self, input_shape): - input_shape = tensor_shape.TensorShape(input_shape).as_list() - output_shape = list(input_shape) - if self.data_format == 'channels_first': - c_axis, h_axis, w_axis = 1, 2, 3 - else: - c_axis, h_axis, w_axis = 3, 1, 2 - - kernel_h, kernel_w = self.kernel_size - stride_h, stride_w = self.strides - - output_shape[c_axis] = self.filters - output_shape[h_axis] = utils.deconv_output_length( - output_shape[h_axis], kernel_h, self.padding, stride_h) - output_shape[w_axis] = utils.deconv_output_length( - output_shape[w_axis], kernel_w, self.padding, stride_w) - return tensor_shape.TensorShape(output_shape) @tf_export('layers.conv2d_transpose') @@ -1806,7 +1273,7 @@ def conv2d_transpose(inputs, @tf_export('layers.Conv3DTranspose') -class Conv3DTranspose(Conv3D): +class Conv3DTranspose(keras_layers.Conv3DTranspose, base.Layer): """Transposed 3D convolution layer (sometimes called 3D Deconvolution). Arguments: @@ -1885,153 +1352,6 @@ class Conv3DTranspose(Conv3D): trainable=trainable, name=name, **kwargs) - self.input_spec = base.InputSpec(ndim=5) - - def build(self, input_shape): - if len(input_shape) != 5: - raise ValueError('Inputs should have rank 5, received input shape:', - str(input_shape)) - if self.data_format == 'channels_first': - channel_axis = 1 - else: - channel_axis = -1 - if input_shape[channel_axis] is None: - raise ValueError('The channel dimension of the inputs ' - 'should be defined, found None: ' + str(input_shape)) - input_dim = input_shape[channel_axis] - kernel_shape = self.kernel_size + (self.filters, input_dim) - - self.kernel = self.add_variable( - 'kernel', - shape=kernel_shape, - initializer=self.kernel_initializer, - regularizer=self.kernel_regularizer, - constraint=self.kernel_constraint, - trainable=True, - dtype=self.dtype) - if self.use_bias: - self.bias = self.add_variable( - 'bias', - shape=(self.filters,), - initializer=self.bias_initializer, - regularizer=self.bias_regularizer, - constraint=self.bias_constraint, - trainable=True, - dtype=self.dtype) - else: - self.bias = None - self.built = True - - def call(self, inputs): - inputs_shape = array_ops.shape(inputs) - batch_size = inputs_shape[0] - if self.data_format == 'channels_first': - c_axis, d_axis, h_axis, w_axis = 1, 2, 3, 4 - else: - c_axis, d_axis, h_axis, w_axis = 4, 1, 2, 3 - - self.input_spec = base.InputSpec(ndim=5, - axes={c_axis: inputs_shape[c_axis]}) - - depth = inputs_shape[d_axis] - height = inputs_shape[h_axis] - width = inputs_shape[w_axis] - - kernel_d, kernel_h, kernel_w = self.kernel_size - stride_d, stride_h, stride_w = self.strides - - # Infer the dynamic output shape: - out_depth = utils.deconv_output_length(depth, - kernel_d, - self.padding, - stride_d) - out_height = utils.deconv_output_length(height, - kernel_h, - self.padding, - stride_h) - out_width = utils.deconv_output_length(width, - kernel_w, - self.padding, - stride_w) - if self.data_format == 'channels_first': - output_shape = (batch_size, self.filters, out_depth, out_height, - out_width) - strides = (1, 1, stride_d, stride_h, stride_w) - else: - output_shape = (batch_size, out_depth, out_height, out_width, - self.filters) - strides = (1, stride_d, stride_h, stride_w, 1) - - output_shape_tensor = array_ops.stack(output_shape) - outputs = nn.conv3d_transpose( - inputs, - self.kernel, - output_shape_tensor, - strides, - data_format=utils.convert_data_format(self.data_format, ndim=5), - padding=self.padding.upper()) - - if not context.executing_eagerly(): - # Infer the static output shape: - out_shape = inputs.get_shape().as_list() - out_shape[c_axis] = self.filters - out_shape[d_axis] = utils.deconv_output_length(out_shape[d_axis], - kernel_d, - self.padding, - stride_d) - out_shape[h_axis] = utils.deconv_output_length(out_shape[h_axis], - kernel_h, - self.padding, - stride_h) - out_shape[w_axis] = utils.deconv_output_length(out_shape[w_axis], - kernel_w, - self.padding, - stride_w) - outputs.set_shape(out_shape) - - if self.use_bias: - outputs_shape = outputs.shape.as_list() - if outputs_shape[0] is None: - outputs_shape[0] = -1 - if self.data_format == 'channels_first': - outputs_4d = array_ops.reshape(outputs, [ - outputs_shape[0], outputs_shape[1], - outputs_shape[2] * outputs_shape[3], outputs_shape[4] - ]) - else: - 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, - data_format=utils.convert_data_format(self.data_format, ndim=4)) - outputs = array_ops.reshape(outputs_4d, outputs_shape) - - if self.activation is not None: - return self.activation(outputs) - return outputs - - def compute_output_shape(self, input_shape): - input_shape = tensor_shape.TensorShape(input_shape).as_list() - output_shape = list(input_shape) - if self.data_format == 'channels_first': - c_axis, d_axis, h_axis, w_axis = 1, 2, 3, 4 - else: - c_axis, d_axis, h_axis, w_axis = 4, 1, 2, 3 - - kernel_d, kernel_h, kernel_w = self.kernel_size - stride_d, stride_h, stride_w = self.strides - - output_shape[c_axis] = self.filters - output_shape[d_axis] = utils.deconv_output_length( - output_shape[d_axis], kernel_d, self.padding, stride_d) - output_shape[h_axis] = utils.deconv_output_length( - output_shape[h_axis], kernel_h, self.padding, stride_h) - output_shape[w_axis] = utils.deconv_output_length( - output_shape[w_axis], kernel_w, self.padding, stride_w) - return tensor_shape.TensorShape(output_shape) @tf_export('layers.conv3d_transpose') diff --git a/tensorflow/python/layers/core.py b/tensorflow/python/layers/core.py index e598d9f83ab21f2dd5fabb3dd37fa0bfb5f003a4..6d8e9eac878bb2eb65bfa29e872a0576a39af662 100644 --- a/tensorflow/python/layers/core.py +++ b/tensorflow/python/layers/core.py @@ -27,23 +27,14 @@ import six from six.moves import xrange # pylint: disable=redefined-builtin import numpy as np -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 layers as keras_layers from tensorflow.python.layers import base -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): +class Dense(keras_layers.Dense, base.Layer): """Densely-connected layer class. This layer implements the operation: @@ -108,73 +99,19 @@ class Dense(base.Layer): trainable=True, name=None, **kwargs): - super(Dense, self).__init__(trainable=trainable, name=name, + super(Dense, self).__init__(units=units, + activation=activation, + use_bias=use_bias, + kernel_initializer=kernel_initializer, + bias_initializer=bias_initializer, + kernel_regularizer=kernel_regularizer, + bias_regularizer=bias_regularizer, activity_regularizer=activity_regularizer, + kernel_constraint=kernel_constraint, + bias_constraint=bias_constraint, + trainable=trainable, + name=name, **kwargs) - self.units = units - self.activation = activation - self.use_bias = use_bias - self.kernel_initializer = kernel_initializer - self.bias_initializer = bias_initializer - self.kernel_regularizer = kernel_regularizer - self.bias_regularizer = bias_regularizer - self.kernel_constraint = kernel_constraint - self.bias_constraint = bias_constraint - self.input_spec = base.InputSpec(min_ndim=2) - - def build(self, input_shape): - input_shape = tensor_shape.TensorShape(input_shape) - if input_shape[-1].value is None: - raise ValueError('The last dimension of the inputs to `Dense` ' - 'should be defined. Found `None`.') - self.input_spec = base.InputSpec(min_ndim=2, - axes={-1: input_shape[-1].value}) - self.kernel = self.add_variable('kernel', - shape=[input_shape[-1].value, self.units], - initializer=self.kernel_initializer, - regularizer=self.kernel_regularizer, - constraint=self.kernel_constraint, - dtype=self.dtype, - trainable=True) - if self.use_bias: - self.bias = self.add_variable('bias', - shape=[self.units,], - initializer=self.bias_initializer, - regularizer=self.bias_regularizer, - constraint=self.bias_constraint, - dtype=self.dtype, - trainable=True) - else: - self.bias = None - self.built = True - - def call(self, inputs): - inputs = ops.convert_to_tensor(inputs, dtype=self.dtype) - shape = inputs.get_shape().as_list() - if len(shape) > 2: - # Broadcasting is required for the inputs. - outputs = standard_ops.tensordot(inputs, self.kernel, [[len(shape) - 1], - [0]]) - # Reshape the output back to the original ndim of the input. - if not context.executing_eagerly(): - output_shape = shape[:-1] + [self.units] - outputs.set_shape(output_shape) - else: - 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: - return self.activation(outputs) # pylint: disable=not-callable - return outputs - - def compute_output_shape(self, input_shape): - input_shape = tensor_shape.TensorShape(input_shape) - input_shape = input_shape.with_rank_at_least(2) - if input_shape[-1].value is None: - raise ValueError( - 'The innermost dimension of input_shape must be defined, but saw: %s' - % input_shape) - return input_shape[:-1].concatenate(self.units) @tf_export('layers.dense') @@ -254,7 +191,7 @@ def dense( @tf_export('layers.Dropout') -class Dropout(base.Layer): +class Dropout(keras_layers.Dropout, base.Layer): """Applies Dropout to the input. Dropout consists in randomly setting a fraction `rate` of input units to 0 @@ -282,31 +219,14 @@ class Dropout(base.Layer): seed=None, name=None, **kwargs): - super(Dropout, self).__init__(name=name, **kwargs) - self.rate = rate - self.noise_shape = noise_shape - self.seed = seed - - def _get_noise_shape(self, inputs): - # Subclasses of `Dropout` may implement `_get_noise_shape(self, inputs)`, - # which will override `self.noise_shape`, and allows for custom noise - # shapes with dynamically sized inputs. - if self.noise_shape is None: - return self.noise_shape - return nn_ops._get_noise_shape(inputs, self.noise_shape) + super(Dropout, self).__init__(rate=rate, + noise_shape=noise_shape, + seed=seed, + name=name, + **kwargs) def call(self, inputs, training=False): - - def dropped_inputs(): - return nn.dropout(inputs, 1 - self.rate, - noise_shape=self._get_noise_shape(inputs), - seed=self.seed) - return utils.smart_cond(training, - dropped_inputs, - lambda: array_ops.identity(inputs)) - - def compute_output_shape(self, input_shape): - return input_shape + return super(Dropout, self).call(inputs, training=training) @tf_export('layers.dropout') @@ -352,7 +272,7 @@ def dropout(inputs, @tf_export('layers.Flatten') -class Flatten(base.Layer): +class Flatten(keras_layers.Flatten, base.Layer): """Flattens an input tensor while preserving the batch axis (axis 0). Examples: @@ -367,25 +287,7 @@ class Flatten(base.Layer): # now `y` has shape `(None, None)` ``` """ - - def __init__(self, **kwargs): - super(Flatten, self).__init__(**kwargs) - self.input_spec = base.InputSpec(min_ndim=2) - - def call(self, inputs): - outputs = array_ops.reshape(inputs, (array_ops.shape(inputs)[0], -1)) - if not context.executing_eagerly(): - outputs.set_shape(self.compute_output_shape(inputs.get_shape())) - return outputs - - def compute_output_shape(self, input_shape): - input_shape = tensor_shape.TensorShape(input_shape).as_list() - output_shape = [input_shape[0]] - if all(input_shape[1:]): - output_shape += [np.prod(input_shape[1:])] - else: - output_shape += [None] - return tensor_shape.TensorShape(output_shape) + pass @tf_export('layers.flatten') diff --git a/tensorflow/python/layers/normalization.py b/tensorflow/python/layers/normalization.py index 83b201e6420d48cfab38048d6638a9f9185d7d6c..33284b0d695272db5a4e0d757d6f24b1930068de 100644 --- a/tensorflow/python/layers/normalization.py +++ b/tensorflow/python/layers/normalization.py @@ -24,26 +24,14 @@ import six from six.moves import xrange # pylint: disable=redefined-builtin 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 ops -from tensorflow.python.framework import tensor_shape +from tensorflow.python.keras._impl.keras import layers as keras_layers from tensorflow.python.layers import base -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 nn -from tensorflow.python.ops import resource_variable_ops -from tensorflow.python.ops import state_ops -from tensorflow.python.training import distribute as distribute_lib -from tensorflow.python.training import moving_averages from tensorflow.python.util.tf_export import tf_export @tf_export('layers.BatchNormalization') -class BatchNormalization(base.Layer): +class BatchNormalization(keras_layers.BatchNormalization, base.Layer): """Batch Normalization layer from http://arxiv.org/abs/1502.03167. "Batch Normalization: Accelerating Deep Network Training by Reducing @@ -143,485 +131,31 @@ class BatchNormalization(base.Layer): name=None, **kwargs): super(BatchNormalization, self).__init__( - name=name, trainable=trainable, **kwargs) - if isinstance(axis, list): - self.axis = axis[:] - else: - self.axis = axis - self.momentum = momentum - self.epsilon = epsilon - self.center = center - self.scale = scale - self.beta_initializer = beta_initializer - self.gamma_initializer = gamma_initializer - self.moving_mean_initializer = moving_mean_initializer - self.moving_variance_initializer = moving_variance_initializer - self.beta_regularizer = beta_regularizer - self.gamma_regularizer = gamma_regularizer - self.beta_constraint = beta_constraint - self.gamma_constraint = gamma_constraint - self.renorm = renorm - self.virtual_batch_size = virtual_batch_size - self.adjustment = adjustment - if fused is None: - fused = True - - self.fused = fused - self._bessels_correction_test_only = True - - if renorm: - renorm_clipping = renorm_clipping or {} - keys = ['rmax', 'rmin', 'dmax'] - if set(renorm_clipping) - set(keys): - raise ValueError('renorm_clipping %s contains keys not in %s' % - (renorm_clipping, keys)) - self.renorm_clipping = renorm_clipping - self.renorm_momentum = renorm_momentum - - def _add_tower_local_variable(self, *args, **kwargs): - tower_context = distribute_lib.get_tower_context() - with tower_context.tower_local_var_scope('mean'): - return self.add_variable(*args, **kwargs) - - def build(self, input_shape): - input_shape = tensor_shape.TensorShape(input_shape) - if not input_shape.ndims: - raise ValueError('Input has undefined rank:', input_shape) - ndims = len(input_shape) - - # Convert axis to list and resolve negatives - if isinstance(self.axis, int): - self.axis = [self.axis] - - if not isinstance(self.axis, list): - raise TypeError('axis must be int or list, type given: %s' - % type(self.axis)) - - for idx, x in enumerate(self.axis): - if x < 0: - self.axis[idx] = ndims + x - - # Validate axes - for x in self.axis: - if x < 0 or x >= ndims: - raise ValueError('Invalid axis: %d' % x) - if len(self.axis) != len(set(self.axis)): - raise ValueError('Duplicate axis: %s' % self.axis) - - if self.virtual_batch_size is not None: - if self.virtual_batch_size <= 0: - raise ValueError('virtual_batch_size must be a positive integer that ' - 'divides the true batch size of the input Tensor') - # If using virtual batches, the first dimension must be the batch - # dimension and cannot be the batch norm axis - if 0 in self.axis: - raise ValueError('When using virtual_batch_size, the batch dimension ' - 'must be 0 and thus axis cannot include 0') - if self.adjustment is not None: - raise ValueError('When using virtual_batch_size, adjustment cannot ' - 'be specified') - - if self.fused: - # Currently fused batch norm doesn't support renorm. It also only supports - # an input tensor of rank 4 and a channel dimension on axis 1 or 3. - # TODO(yaozhang): if input is not 4D, reshape it to 4D and reshape the - # output back to its original shape accordingly. - self.fused = (not self.renorm and - ndims == 4 and - self.axis in [[1], [3]] and - self.virtual_batch_size is None and - self.adjustment is None) - # TODO(chrisying): fused batch norm is currently not supported for - # multi-axis batch norm and by extension virtual batches. In some cases, - # it might be possible to use fused batch norm but would require reshaping - # the Tensor to 4D with the axis in 1 or 3 (preferred 1) which is - # particularly tricky. A compromise might be to just support the most - # common use case (turning 5D w/ virtual batch to NCHW) - - if self.fused: - if self.axis == [1]: - self._data_format = 'NCHW' - elif self.axis == [3]: - self._data_format = 'NHWC' - else: - raise ValueError('Unsupported axis, fused batch norm only supports ' - 'axis == [1] or axis == [3]') - - # Raise parameters of fp16 batch norm to fp32 - if self.dtype == dtypes.float16 or self.dtype == dtypes.bfloat16: - param_dtype = dtypes.float32 - else: - param_dtype = self.dtype or dtypes.float32 - - axis_to_dim = {x: input_shape[x].value for x in self.axis} - for x in axis_to_dim: - if axis_to_dim[x] is None: - raise ValueError('Input has undefined `axis` dimension. Input shape: ', - input_shape) - self.input_spec = base.InputSpec(ndim=ndims, axes=axis_to_dim) - - if len(axis_to_dim) == 1 and self.virtual_batch_size is None: - # Single axis batch norm (most common/default use-case) - param_shape = (list(axis_to_dim.values())[0],) - else: - # Parameter shape is the original shape but with 1 in all non-axis dims - param_shape = [axis_to_dim[i] if i in axis_to_dim - else 1 for i in range(ndims)] - if self.virtual_batch_size is not None: - # When using virtual batches, add an extra dim at index 1 - param_shape.insert(1, 1) - for idx, x in enumerate(self.axis): - self.axis[idx] = x + 1 # Account for added dimension - - if self.scale: - self.gamma = self.add_variable( - name='gamma', - shape=param_shape, - dtype=param_dtype, - initializer=self.gamma_initializer, - regularizer=self.gamma_regularizer, - constraint=self.gamma_constraint, - trainable=True) - else: - self.gamma = None - if self.fused: - self._gamma_const = array_ops.constant( - 1.0, dtype=param_dtype, shape=param_shape) - - if self.center: - self.beta = self.add_variable( - name='beta', - shape=param_shape, - dtype=param_dtype, - initializer=self.beta_initializer, - regularizer=self.beta_regularizer, - constraint=self.beta_constraint, - trainable=True) - else: - self.beta = None - if self.fused: - self._beta_const = array_ops.constant( - 0.0, dtype=param_dtype, shape=param_shape) - - # Disable variable partitioning when creating the moving mean and variance - try: - if self._scope: - partitioner = self._scope.partitioner - self._scope.set_partitioner(None) - else: - partitioner = None - self.moving_mean = self._add_tower_local_variable( - name='moving_mean', - shape=param_shape, - dtype=param_dtype, - initializer=self.moving_mean_initializer, - trainable=False) - - self.moving_variance = self._add_tower_local_variable( - name='moving_variance', - shape=param_shape, - dtype=param_dtype, - initializer=self.moving_variance_initializer, - trainable=False) - - if self.renorm: - # Create variables to maintain the moving mean and standard deviation. - # These are used in training and thus are different from the moving - # averages above. The renorm variables are colocated with moving_mean - # and moving_variance. - # NOTE: below, the outer `with device` block causes the current device - # stack to be cleared. The nested ones use a `lambda` to set the desired - # device and ignore any devices that may be set by the custom getter. - def _renorm_variable(name, shape): - var = self._add_tower_local_variable( - name=name, - shape=shape, - dtype=param_dtype, - initializer=init_ops.zeros_initializer(), - trainable=False) - return var - - with distribute_lib.get_distribution_strategy().colocate_vars_with( - self.moving_mean): - self.renorm_mean = _renorm_variable('renorm_mean', param_shape) - self.renorm_mean_weight = _renorm_variable('renorm_mean_weight', ()) - # We initialize renorm_stddev to 0, and maintain the (0-initialized) - # 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. - with distribute_lib.get_distribution_strategy().colocate_vars_with( - self.moving_variance): - self.renorm_stddev = _renorm_variable('renorm_stddev', param_shape) - self.renorm_stddev_weight = _renorm_variable('renorm_stddev_weight', - ()) - finally: - if partitioner: - self._scope.set_partitioner(partitioner) - self.built = True - - def _assign_moving_average(self, variable, value, momentum): - with ops.name_scope(None, 'AssignMovingAvg', - [variable, value, momentum]) as scope: - decay = ops.convert_to_tensor(1.0 - momentum, name='decay') - if decay.dtype != variable.dtype.base_dtype: - decay = math_ops.cast(decay, variable.dtype.base_dtype) - update_delta = (variable - value) * decay - return state_ops.assign_sub(variable, update_delta, name=scope) - - def _fused_batch_norm(self, inputs, training): - """Returns the output of fused batch norm.""" - beta = self.beta if self.center else self._beta_const - gamma = self.gamma if self.scale else self._gamma_const - - def _fused_batch_norm_training(): - return nn.fused_batch_norm( - inputs, - gamma, - beta, - epsilon=self.epsilon, - data_format=self._data_format) - - def _fused_batch_norm_inference(): - return nn.fused_batch_norm( - inputs, - gamma, - beta, - mean=self.moving_mean, - variance=self.moving_variance, - epsilon=self.epsilon, - is_training=False, - data_format=self._data_format) - - output, mean, variance = utils.smart_cond( - training, _fused_batch_norm_training, _fused_batch_norm_inference) - if not self._bessels_correction_test_only: - # Remove Bessel's correction to be consistent with non-fused batch norm. - # Note that the variance computed by fused batch norm is - # with Bessel's correction. - sample_size = math_ops.cast( - array_ops.size(inputs) / array_ops.size(variance), variance.dtype) - factor = (sample_size - math_ops.cast(1.0, variance.dtype)) / sample_size - variance *= factor - - training_value = utils.constant_value(training) - if training_value is None: - momentum = utils.smart_cond(training, lambda: self.momentum, lambda: 1.0) - else: - momentum = ops.convert_to_tensor(self.momentum) - if training_value or training_value is None: - mean_update = self._assign_moving_average(self.moving_mean, mean, - momentum) - variance_update = self._assign_moving_average(self.moving_variance, - variance, momentum) - self.add_update(mean_update, inputs=inputs) - self.add_update(variance_update, inputs=inputs) - - return output - - def _renorm_correction_and_moments(self, mean, variance, training): - """Returns the correction and update values for renorm.""" - stddev = math_ops.sqrt(variance + self.epsilon) - # Compute the average mean and standard deviation, as if they were - # initialized with this batch's moments. - mixed_renorm_mean = (self.renorm_mean + - (1. - self.renorm_mean_weight) * mean) - mixed_renorm_stddev = (self.renorm_stddev + - (1. - self.renorm_stddev_weight) * stddev) - # Compute the corrections for batch renorm. - r = stddev / mixed_renorm_stddev - d = (mean - mixed_renorm_mean) / mixed_renorm_stddev - # Ensure the corrections use pre-update moving averages. - with ops.control_dependencies([r, d]): - mean = array_ops.identity(mean) - stddev = array_ops.identity(stddev) - rmin, rmax, dmax = [self.renorm_clipping.get(key) - for key in ['rmin', 'rmax', 'dmax']] - if rmin is not None: - r = math_ops.maximum(r, rmin) - if rmax is not None: - r = math_ops.minimum(r, rmax) - if dmax is not None: - d = math_ops.maximum(d, -dmax) - d = math_ops.minimum(d, dmax) - # When not training, use r=1, d=0. - r = utils.smart_cond(training, lambda: r, lambda: array_ops.ones_like(r)) - d = utils.smart_cond(training, lambda: d, lambda: array_ops.zeros_like(d)) - - def _update_renorm_variable(var, weight, value): - """Updates a moving average and weight, returns the unbiased value.""" - value = array_ops.identity(value) - def _do_update(): - """Updates the var and weight, returns their updated ratio.""" - # Update the variables without zero debiasing. The debiasing will be - # accomplished by dividing the exponential moving average by the weight. - # For example, after a single update, the moving average would be - # (1-decay) * value. and the weight will be 1-decay, with their ratio - # giving the value. - # Make sure the weight is not updated until before r and d computation. - with ops.control_dependencies([value]): - weight_value = array_ops.constant(1., dtype=weight.dtype) - new_var = self._assign_moving_average(var, value, self.renorm_momentum) - new_weight = self._assign_moving_average(weight, weight_value, - self.renorm_momentum) - # TODO(yuefengz): the updates to var and weighted can not be batched - # together if we fetch their updated values here. Consider calculating - # new values and delaying the updates. - return new_var / new_weight - - def _fake_update(): - return array_ops.identity(var) - return utils.smart_cond(training, _do_update, _fake_update) - - # TODO(yuefengz): colocate the operations - new_mean = _update_renorm_variable(self.renorm_mean, - self.renorm_mean_weight, mean) - new_stddev = _update_renorm_variable(self.renorm_stddev, - self.renorm_stddev_weight, stddev) - # Make sqrt(moving_variance + epsilon) = new_stddev. - new_variance = math_ops.square(new_stddev) - self.epsilon - - return (r, d, new_mean, new_variance) + axis=axis, + momentum=momentum, + epsilon=epsilon, + center=center, + scale=scale, + beta_initializer=beta_initializer, + gamma_initializer=gamma_initializer, + moving_mean_initializer=moving_mean_initializer, + moving_variance_initializer=moving_variance_initializer, + beta_regularizer=beta_regularizer, + gamma_regularizer=gamma_regularizer, + beta_constraint=beta_constraint, + gamma_constraint=gamma_constraint, + renorm=renorm, + renorm_clipping=renorm_clipping, + renorm_momentum=renorm_momentum, + fused=fused, + trainable=trainable, + virtual_batch_size=virtual_batch_size, + adjustment=adjustment, + name=name, + **kwargs) 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 - original_shape = [-1] + inputs.shape.as_list()[1:] - expanded_shape = [self.virtual_batch_size, -1] + original_shape[1:] - - # Will cause errors if virtual_batch_size does not divide the batch size - inputs = array_ops.reshape(inputs, expanded_shape) - - def undo_virtual_batching(outputs): - outputs = array_ops.reshape(outputs, original_shape) - return outputs - - if self.fused: - outputs = self._fused_batch_norm(inputs, training=training) - if self.virtual_batch_size is not None: - # Currently never reaches here since fused_batch_norm does not support - # virtual batching - return undo_virtual_batching(outputs) - return outputs - - # Compute the axes along which to reduce the mean / variance - input_shape = inputs.get_shape() - ndims = len(input_shape) - reduction_axes = [i for i in range(ndims) if i not in self.axis] - if self.virtual_batch_size is not None: - del reduction_axes[1] # Do not reduce along virtual batch dim - - # Broadcasting only necessary for single-axis batch norm where the axis is - # not the last dimension - broadcast_shape = [1] * ndims - broadcast_shape[self.axis[0]] = input_shape[self.axis[0]].value - 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 - - scale, offset = _broadcast(self.gamma), _broadcast(self.beta) - - def _compose_transforms(scale, offset, then_scale, then_offset): - if then_scale is not None: - scale *= then_scale - offset *= then_scale - if then_offset is not None: - offset += then_offset - return (scale, offset) - - # Determine a boolean value for `training`: could be True, False, or None. - training_value = utils.constant_value(training) - if training_value is not False: - if self.adjustment: - adj_scale, adj_bias = self.adjustment(array_ops.shape(inputs)) - # Adjust only during training. - adj_scale = utils.smart_cond(training, - lambda: adj_scale, - lambda: array_ops.ones_like(adj_scale)) - adj_bias = utils.smart_cond(training, - lambda: adj_bias, - lambda: array_ops.zeros_like(adj_bias)) - scale, offset = _compose_transforms(adj_scale, adj_bias, scale, offset) - - # Some of the computations here are not necessary when training==False - # but not a constant. However, this makes the code simpler. - keep_dims = self.virtual_batch_size is not None or len(self.axis) > 1 - mean, variance = nn.moments(inputs, reduction_axes, keep_dims=keep_dims) - - moving_mean = self.moving_mean - moving_variance = self.moving_variance - - mean = utils.smart_cond(training, - lambda: mean, - lambda: moving_mean) - variance = utils.smart_cond(training, - lambda: variance, - lambda: moving_variance) - - if self.renorm: - r, d, new_mean, new_variance = self._renorm_correction_and_moments( - mean, variance, training) - # When training, the normalized values (say, x) will be transformed as - # x * gamma + beta without renorm, and (x * r + d) * gamma + beta - # = x * (r * gamma) + (d * gamma + beta) with renorm. - r = _broadcast(array_ops.stop_gradient(r, name='renorm_r')) - d = _broadcast(array_ops.stop_gradient(d, name='renorm_d')) - scale, offset = _compose_transforms(r, d, scale, offset) - else: - new_mean, new_variance = mean, variance - - if self.virtual_batch_size is not None: - # This isn't strictly correct since in ghost batch norm, you are - # supposed to sequentially update the moving_mean and moving_variance - # with each sub-batch. However, since the moving statistics are only - # used during evaluation, it is more efficient to just update in one - # step and should not make a significant difference in the result. - new_mean = math_ops.reduce_mean(new_mean, - axis=1, keep_dims=True) - new_variance = math_ops.reduce_mean(new_variance, - axis=1, keep_dims=True) - - def _do_update(var, value): - if in_eager_mode and not self.trainable: - return - - return self._assign_moving_average(var, value, self.momentum) - - mean_update = utils.smart_cond( - training, - lambda: _do_update(self.moving_mean, new_mean), - lambda: self.moving_mean) - variance_update = utils.smart_cond( - training, - lambda: _do_update(self.moving_variance, new_variance), - lambda: self.moving_variance) - if not context.executing_eagerly(): - self.add_update(mean_update, inputs=inputs) - self.add_update(variance_update, inputs=inputs) - - else: - mean, variance = self.moving_mean, self.moving_variance - - outputs = nn.batch_normalization(inputs, - _broadcast(mean), - _broadcast(variance), - offset, - scale, - self.epsilon) - # If some components of the shape got lost due to adjustments, fix that. - outputs.set_shape(input_shape) - - if self.virtual_batch_size is not None: - return undo_virtual_batching(outputs) - - return outputs - - def compute_output_shape(self, input_shape): - return input_shape + return super(BatchNormalization, self).call(inputs, training=training) @tf_export('layers.batch_normalization') diff --git a/tensorflow/python/layers/pooling.py b/tensorflow/python/layers/pooling.py index 50503ce093fbc251b11c4d5cbccb2a2683d92e7a..75abe56f51f2a206ea3e5a5dad032446c150293a 100644 --- a/tensorflow/python/layers/pooling.py +++ b/tensorflow/python/layers/pooling.py @@ -13,92 +13,19 @@ # limitations under the License. # ============================================================================= -# pylint: disable=unused-import,g-bad-import-order """Contains the pooling layer classes and their functional aliases. """ 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 layers as keras_layers 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): - """Pooling layer for arbitrary pooling functions, for 1D inputs. - - This class only exists for code reuse. It will never be an exposed API. - - Arguments: - pool_function: The pooling function to apply, e.g. `tf.nn.max_pool`. - pool_size: An integer or tuple/list of a single integer, - representing the size of the pooling window. - strides: An integer or tuple/list of a single integer, specifying the - strides of the pooling operation. - padding: A string. The padding method, either '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)`. - name: A string, the name of the layer. - """ - - def __init__(self, pool_function, pool_size, strides, - padding='valid', data_format='channels_last', - name=None, **kwargs): - super(_Pooling1D, self).__init__(name=name, **kwargs) - self.pool_function = pool_function - self.pool_size = utils.normalize_tuple(pool_size, 1, 'pool_size') - self.strides = utils.normalize_tuple(strides, 1, 'strides') - self.padding = utils.normalize_padding(padding) - self.data_format = utils.normalize_data_format(data_format) - self.input_spec = base.InputSpec(ndim=3) - - def call(self, inputs): - # There is no TF op for 1D pooling, hence we make the inputs 4D. - if self.data_format == 'channels_last': - # input is NWC, make it NHWC - inputs = array_ops.expand_dims(inputs, 1) - # pool on the W dim - pool_shape = (1, 1) + self.pool_size + (1,) - strides = (1, 1) + self.strides + (1,) - data_format = 'NHWC' - else: - # input is NCW, make it NCHW - inputs = array_ops.expand_dims(inputs, 2) - # pool on the W dim - pool_shape = (1, 1, 1) + self.pool_size - strides = (1, 1, 1) + self.strides - data_format = 'NCHW' - - outputs = self.pool_function( - inputs, - ksize=pool_shape, - strides=strides, - padding=self.padding.upper(), - data_format=data_format) - - if self.data_format == 'channels_last': - return array_ops.squeeze(outputs, 1) - else: - return array_ops.squeeze(outputs, 2) - - def compute_output_shape(self, input_shape): - input_shape = tensor_shape.TensorShape(input_shape).as_list() - length = utils.conv_output_length(input_shape[1], self.pool_size[0], - self.padding, self.strides[0]) - return tensor_shape.TensorShape([input_shape[0], length, input_shape[2]]) - - @tf_export('layers.AveragePooling1D') -class AveragePooling1D(_Pooling1D): +class AveragePooling1D(keras_layers.AveragePooling1D, base.Layer): """Average Pooling layer for 1D inputs. Arguments: @@ -119,8 +46,9 @@ class AveragePooling1D(_Pooling1D): def __init__(self, pool_size, strides, padding='valid', data_format='channels_last', name=None, **kwargs): + if strides is None: + raise ValueError('Argument `strides` must not be None.') super(AveragePooling1D, self).__init__( - nn.avg_pool, pool_size=pool_size, strides=strides, padding=padding, @@ -165,7 +93,7 @@ def average_pooling1d(inputs, pool_size, strides, @tf_export('layers.MaxPooling1D') -class MaxPooling1D(_Pooling1D): +class MaxPooling1D(keras_layers.MaxPooling1D, base.Layer): """Max Pooling layer for 1D inputs. Arguments: @@ -186,8 +114,9 @@ class MaxPooling1D(_Pooling1D): def __init__(self, pool_size, strides, padding='valid', data_format='channels_last', name=None, **kwargs): + if strides is None: + raise ValueError('Argument `strides` must not be None.') super(MaxPooling1D, self).__init__( - nn.max_pool, pool_size=pool_size, strides=strides, padding=padding, @@ -231,79 +160,8 @@ def max_pooling1d(inputs, pool_size, strides, return layer.apply(inputs) -class _Pooling2D(base.Layer): - """Pooling layer for arbitrary pooling functions, for 2D inputs (e.g. images). - - This class only exists for code reuse. It will never be an exposed API. - - Arguments: - pool_function: The pooling function to apply, e.g. `tf.nn.max_pool`. - pool_size: An integer or tuple/list of 2 integers: (pool_height, pool_width) - specifying the size of the pooling 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 pooling operation. - Can be a single integer to specify the same value for - all spatial dimensions. - padding: A string. The padding method, either '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)`. - name: A string, the name of the layer. - """ - - def __init__(self, pool_function, pool_size, strides, - padding='valid', data_format='channels_last', - name=None, **kwargs): - super(_Pooling2D, self).__init__(name=name, **kwargs) - self.pool_function = pool_function - self.pool_size = utils.normalize_tuple(pool_size, 2, 'pool_size') - self.strides = utils.normalize_tuple(strides, 2, 'strides') - self.padding = utils.normalize_padding(padding) - self.data_format = utils.normalize_data_format(data_format) - self.input_spec = base.InputSpec(ndim=4) - - def call(self, inputs): - if self.data_format == 'channels_last': - pool_shape = (1,) + self.pool_size + (1,) - strides = (1,) + self.strides + (1,) - else: - pool_shape = (1, 1) + self.pool_size - strides = (1, 1) + self.strides - outputs = self.pool_function( - inputs, - ksize=pool_shape, - strides=strides, - padding=self.padding.upper(), - data_format=utils.convert_data_format(self.data_format, 4)) - return outputs - - def compute_output_shape(self, input_shape): - input_shape = tensor_shape.TensorShape(input_shape).as_list() - if self.data_format == 'channels_first': - rows = input_shape[2] - cols = input_shape[3] - else: - rows = input_shape[1] - cols = input_shape[2] - rows = utils.conv_output_length(rows, self.pool_size[0], self.padding, - self.strides[0]) - cols = utils.conv_output_length(cols, self.pool_size[1], self.padding, - self.strides[1]) - if self.data_format == 'channels_first': - return tensor_shape.TensorShape( - [input_shape[0], input_shape[1], rows, cols]) - else: - return tensor_shape.TensorShape( - [input_shape[0], rows, cols, input_shape[3]]) - - @tf_export('layers.AveragePooling2D') -class AveragePooling2D(_Pooling2D): +class AveragePooling2D(keras_layers.AveragePooling2D, base.Layer): """Average pooling layer for 2D inputs (e.g. images). Arguments: @@ -328,8 +186,9 @@ class AveragePooling2D(_Pooling2D): def __init__(self, pool_size, strides, padding='valid', data_format='channels_last', name=None, **kwargs): + if strides is None: + raise ValueError('Argument `strides` must not be None.') super(AveragePooling2D, self).__init__( - nn.avg_pool, pool_size=pool_size, strides=strides, padding=padding, data_format=data_format, name=name, **kwargs) @@ -373,7 +232,7 @@ def average_pooling2d(inputs, @tf_export('layers.MaxPooling2D') -class MaxPooling2D(_Pooling2D): +class MaxPooling2D(keras_layers.MaxPooling2D, base.Layer): """Max pooling layer for 2D inputs (e.g. images). Arguments: @@ -398,8 +257,9 @@ class MaxPooling2D(_Pooling2D): def __init__(self, pool_size, strides, padding='valid', data_format='channels_last', name=None, **kwargs): + if strides is None: + raise ValueError('Argument `strides` must not be None.') super(MaxPooling2D, self).__init__( - nn.max_pool, pool_size=pool_size, strides=strides, padding=padding, data_format=data_format, name=name, **kwargs) @@ -442,90 +302,8 @@ def max_pooling2d(inputs, return layer.apply(inputs) -class _Pooling3D(base.Layer): - """Pooling layer for arbitrary pooling functions, for 3D inputs. - - This class only exists for code reuse. It will never be an exposed API. - - Arguments: - pool_function: The pooling function to apply, e.g. `tf.nn.max_pool`. - pool_size: An integer or tuple/list of 3 integers: - (pool_depth, pool_height, pool_width) - specifying the size of the pooling 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 pooling operation. - Can be a single integer to specify the same value for - all spatial dimensions. - padding: A string. The padding method, either '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)`. - name: A string, the name of the layer. - """ - - def __init__(self, pool_function, pool_size, strides, - padding='valid', data_format='channels_last', - name=None, **kwargs): - super(_Pooling3D, self).__init__(name=name, **kwargs) - self.pool_function = pool_function - self.pool_size = utils.normalize_tuple(pool_size, 3, 'pool_size') - self.strides = utils.normalize_tuple(strides, 3, 'strides') - self.padding = utils.normalize_padding(padding) - self.data_format = utils.normalize_data_format(data_format) - self.input_spec = base.InputSpec(ndim=5) - - def call(self, inputs): - pool_shape = (1,) + self.pool_size + (1,) - strides = (1,) + self.strides + (1,) - - if self.data_format == 'channels_first': - # TF does not support `channels_first` with 3D pooling operations, - # so we must handle this case manually. - # TODO(fchollet): remove this when TF pooling is feature-complete. - inputs = array_ops.transpose(inputs, (0, 2, 3, 4, 1)) - - outputs = self.pool_function( - inputs, - ksize=pool_shape, - strides=strides, - padding=self.padding.upper()) - - if self.data_format == 'channels_first': - outputs = array_ops.transpose(outputs, (0, 4, 1, 2, 3)) - return outputs - - def compute_output_shape(self, input_shape): - input_shape = tensor_shape.TensorShape(input_shape).as_list() - if self.data_format == 'channels_first': - len_dim1 = input_shape[2] - len_dim2 = input_shape[3] - len_dim3 = input_shape[4] - else: - len_dim1 = input_shape[1] - len_dim2 = input_shape[2] - len_dim3 = input_shape[3] - len_dim1 = utils.conv_output_length(len_dim1, self.pool_size[0], - self.padding, self.strides[0]) - len_dim2 = utils.conv_output_length(len_dim2, self.pool_size[1], - self.padding, self.strides[1]) - len_dim3 = utils.conv_output_length(len_dim3, self.pool_size[2], - self.padding, self.strides[2]) - if self.data_format == 'channels_first': - return tensor_shape.TensorShape( - [input_shape[0], input_shape[1], len_dim1, len_dim2, len_dim3]) - else: - return tensor_shape.TensorShape( - [input_shape[0], len_dim1, len_dim2, len_dim3, input_shape[4]]) - - @tf_export('layers.AveragePooling3D') -class AveragePooling3D(_Pooling3D): +class AveragePooling3D(keras_layers.AveragePooling3D, base.Layer): """Average pooling layer for 3D inputs (e.g. volumes). Arguments: @@ -552,8 +330,9 @@ class AveragePooling3D(_Pooling3D): def __init__(self, pool_size, strides, padding='valid', data_format='channels_last', name=None, **kwargs): + if strides is None: + raise ValueError('Argument `strides` must not be None.') super(AveragePooling3D, self).__init__( - nn.avg_pool3d, pool_size=pool_size, strides=strides, padding=padding, data_format=data_format, name=name, **kwargs) @@ -599,7 +378,7 @@ def average_pooling3d(inputs, @tf_export('layers.MaxPooling3D') -class MaxPooling3D(_Pooling3D): +class MaxPooling3D(keras_layers.MaxPooling3D, base.Layer): """Max pooling layer for 3D inputs (e.g. volumes). Arguments: @@ -626,8 +405,9 @@ class MaxPooling3D(_Pooling3D): def __init__(self, pool_size, strides, padding='valid', data_format='channels_last', name=None, **kwargs): + if strides is None: + raise ValueError('Argument `strides` must not be None.') super(MaxPooling3D, self).__init__( - nn.max_pool3d, pool_size=pool_size, strides=strides, padding=padding, data_format=data_format, name=name, **kwargs) diff --git a/tensorflow/python/layers/utils_test.py b/tensorflow/python/layers/utils_test.py index c941aad7bc63dbb891fbe78cd2a47dd6805bf231..7e94dda648166780af002ce6b979a751a0ced846 100644 --- a/tensorflow/python/layers/utils_test.py +++ b/tensorflow/python/layers/utils_test.py @@ -19,7 +19,6 @@ 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 @@ -89,33 +88,5 @@ class ConvUtilsTest(test.TestCase): 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/bfloat16.cc b/tensorflow/python/lib/core/bfloat16.cc index 7f07deebef3d8e8f24f73a42f29f4ade4cae568d..77fa2c1f66d2214dbb08e4d0ad3437fa4fe02822 100644 --- a/tensorflow/python/lib/core/bfloat16.cc +++ b/tensorflow/python/lib/core/bfloat16.cc @@ -616,8 +616,8 @@ bool Initialize() { }; // Comparisons - const std::array compare_types = {npy_bfloat16_, npy_bfloat16_, - NPY_BOOL}; + const std::array compare_types = { + {npy_bfloat16_, npy_bfloat16_, NPY_BOOL}}; if (!register_ufunc("equal", CompareUFunc, compare_types)) { diff --git a/tensorflow/python/lib/core/py_func.cc b/tensorflow/python/lib/core/py_func.cc index 22317a348c9d5472486ad118d865341ffb6ad829..8c6bb7955a4e29daddd92860e41d7105192eb24b 100644 --- a/tensorflow/python/lib/core/py_func.cc +++ b/tensorflow/python/lib/core/py_func.cc @@ -126,6 +126,9 @@ Status NumericNpDTypeToTfDType(const int np, DataType* tf) { case NPY_INT8: *tf = DT_INT8; break; + case NPY_UINT16: + *tf = DT_UINT16; + break; case NPY_INT16: *tf = DT_INT16; break; diff --git a/tensorflow/python/lib/io/py_record_reader.cc b/tensorflow/python/lib/io/py_record_reader.cc index 5fcb51b3b25e49192f651c648f48a1760ed23084..9500fc6a7c49d13ce0d6954f237f43ffba0869a0 100644 --- a/tensorflow/python/lib/io/py_record_reader.cc +++ b/tensorflow/python/lib/io/py_record_reader.cc @@ -43,9 +43,10 @@ PyRecordReader* PyRecordReader::New(const string& filename, uint64 start_offset, reader->offset_ = start_offset; reader->file_ = file.release(); + static const uint64 kReaderBufferSize = 16 * 1024 * 1024; RecordReaderOptions options = RecordReaderOptions::CreateRecordReaderOptions(compression_type_string); - + options.buffer_size = kReaderBufferSize; reader->reader_ = new RecordReader(reader->file_, options); return reader; } diff --git a/tensorflow/python/ops/array_grad.py b/tensorflow/python/ops/array_grad.py index 57d2657838760a6f0041dac850913035540dc939..3678bd4c1f6a4500622b6d9e8334cb1ebae46578 100644 --- a/tensorflow/python/ops/array_grad.py +++ b/tensorflow/python/ops/array_grad.py @@ -196,7 +196,7 @@ def _ConcatGradHelper(op, grad, start_value_index, end_value_index, dim_index): array_ops.where( math_ops.logical_and(grad.indices >= start, grad.indices < end)), - squeeze_dims=[1]) + axis=[1]) new_indices = array_ops.gather(grad.indices, indices_to_select) - start new_values = array_ops.gather(grad.values, indices_to_select) out_grads.append(ops.IndexedSlices(new_values, new_indices, size)) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index fa26e07c8535d537e366e06902ce28e67d6e6dfe..bbffff04831829e3ec44337416819b967b7f55ba 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -144,6 +144,7 @@ def identity(input, name=None): # pylint: disable=redefined-builtin # pylint: disable=redefined-builtin,protected-access @tf_export("expand_dims") +@deprecation.deprecated_args(None, "Use the `axis` argument instead", "dim") def expand_dims(input, axis=None, name=None, dim=None): """Inserts a dimension of 1 into a tensor's shape. @@ -193,11 +194,7 @@ def expand_dims(input, axis=None, name=None, dim=None): Raises: ValueError: if both `dim` and `axis` are specified. """ - # TODO(aselle): Remove argument dim - if dim is not None: - if axis is not None: - raise ValueError("can't specify both 'dim' and 'axis'") - axis = dim + axis = deprecation.deprecated_argument_lookup("axis", axis, "dim", dim) return gen_array_ops.expand_dims(input, axis, name) @@ -1060,9 +1057,7 @@ def unstack(value, num=None, axis=0, name="unstack"): `value[:, i, :, :]` and each tensor in `output` will have shape `(A, C, D)`. Etc. - This is the opposite of stack. The numpy equivalent is - - tf.unstack(x, n) = np.unstack(x) + This is the opposite of stack. Args: value: A rank `R > 0` `Tensor` to be unstacked. @@ -1235,7 +1230,7 @@ def boolean_mask(tensor, mask, name="boolean_mask", axis=None): def _apply_mask_1d(reshaped_tensor, mask, axis=None): """Mask tensor along dimension 0 with a 1-D mask.""" - indices = squeeze(where(mask), squeeze_dims=[1]) + indices = squeeze(where(mask), axis=[1]) return gather(reshaped_tensor, indices, axis=axis) with ops.name_scope(name, values=[tensor, mask]): @@ -2581,6 +2576,8 @@ def sequence_mask(lengths, maxlen=None, dtype=dtypes.bool, name=None): @tf_export("squeeze") +@deprecation.deprecated_args(None, "Use the `axis` argument instead", + "squeeze_dims") def squeeze(input, axis=None, name=None, squeeze_dims=None): # pylint: disable=redefined-builtin """Removes dimensions of size 1 from the shape of a tensor. @@ -2621,10 +2618,8 @@ def squeeze(input, axis=None, name=None, squeeze_dims=None): Raises: ValueError: When both `squeeze_dims` and `axis` are specified. """ - if squeeze_dims is not None: - if axis is not None: - raise ValueError("Cannot specify both 'squeeze_dims' and 'axis'") - axis = squeeze_dims + axis = deprecation.deprecated_argument_lookup( + "axis", axis, "squeeze_dims", squeeze_dims) if np.isscalar(axis): axis = [axis] return gen_array_ops.squeeze(input, axis, name) diff --git a/tensorflow/python/ops/clip_ops.py b/tensorflow/python/ops/clip_ops.py index 0829aa67ed5236a7c2af89fc104f1d203c8a0f23..75c459a9cf10a90f6043d304b302e0a0806bf045 100644 --- a/tensorflow/python/ops/clip_ops.py +++ b/tensorflow/python/ops/clip_ops.py @@ -27,7 +27,6 @@ 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_array_ops -from tensorflow.python.ops import gen_math_ops from tensorflow.python.ops import gen_nn_ops from tensorflow.python.ops import math_ops from tensorflow.python.util.tf_export import tf_export @@ -60,13 +59,26 @@ def clip_by_value(t, clip_value_min, clip_value_max, """ with ops.name_scope(name, "clip_by_value", [t, clip_value_min, clip_value_max]) as name: - return gen_math_ops.clip_by_value(t, - clip_value_min, - clip_value_max, - name=name) + t = ops.convert_to_tensor(t, name="t") + + # Go through list of tensors, for each value in each tensor clip + t_min = math_ops.minimum(t, clip_value_max) + # Assert that the shape is compatible with the initial shape, + # to prevent unintentional broadcasting. + _ = t.shape.merge_with(t_min.shape) + + t_max = math_ops.maximum(t_min, clip_value_min, name=name) + _ = t.shape.merge_with(t_max.shape) + + return t_max + # TODO(scottzhu): switch to use new implmentation in 2 weeks. + # return gen_math_ops.clip_by_value( + # t, clip_value_min, clip_value_max, name=name) + -@ops.RegisterGradient("ClipByValue") -def _ClipByValueGrad(op, grad): +# TODO(scottzhu): switch to use new implmentation in 2 weeks. +# @ops.RegisterGradient("ClipByValue") +def _clip_by_value_grad(op, grad): """Returns grad of clip_by_value.""" x = op.inputs[0] y = op.inputs[1] diff --git a/tensorflow/python/ops/control_flow_grad.py b/tensorflow/python/ops/control_flow_grad.py index 45955554cab130597e106660ff1fb4cdf7e9aeb1..6a551deb5ba55871b3a3fb144a6ecd2a3cbfcbd8 100644 --- a/tensorflow/python/ops/control_flow_grad.py +++ b/tensorflow/python/ops/control_flow_grad.py @@ -20,6 +20,7 @@ from __future__ import print_function from six.moves import xrange # pylint: disable=redefined-builtin +from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor from tensorflow.python.ops import control_flow_ops @@ -74,6 +75,11 @@ def _SwitchGrad(op, *grad): # At this point, we have created zero_grad guarded by the right switch. # Unfortunately, we may still get None here for not trainable data types. if zero_grad is None: + # For resource variables we get None always on the other branch, so bypass + # this. + if op.inputs[0].dtype == dtypes.resource: + return merge( + [grad[op_ctxt.branch]] * 2, name="cond_resource_grad")[0], None return None, None return merge(grad, name="cond_grad")[0], None else: diff --git a/tensorflow/python/ops/control_flow_ops.py b/tensorflow/python/ops/control_flow_ops.py index fb53d9ffea174ae4af6f664efb7e36f63e409124..a1bfe450c8986f448edba12b7431ccdd59c6f246 100644 --- a/tensorflow/python/ops/control_flow_ops.py +++ b/tensorflow/python/ops/control_flow_ops.py @@ -609,13 +609,13 @@ def _EnforceShapeInvariant(merge_var, next_var): """Check if the shapes of the loops variables are invariants. Args: - merge_vars: The list of tensors representing the initial values of the + merge_var: The list of tensors representing the initial values of the loop variables. - next_vars: The list of tensors representing the values of the loop + next_var: The list of tensors representing the values of the loop variables after one loop iteration. Raises: - ValueError: If any tensor in `merge_vars` has a more specific shape than + ValueError: If any tensor in `merge_var` has a more specific shape than its correspnding tensor in `next_var`. """ if isinstance(merge_var, ops.Tensor): @@ -833,7 +833,7 @@ class GradLoopState(object): if outer_grad_state: outer_forward_ctxt = outer_grad_state.forward_context else: - if not hasattr(forward_ctxt, 'outer_context'): + if not hasattr(forward_ctxt, "outer_context"): raise ValueError("Failed to call gradients on a while loop without" "properly serializing graph via MetaGraphDef") outer_forward_ctxt = forward_ctxt.outer_context @@ -2379,7 +2379,15 @@ class WhileContext(ControlFlowContext): def AddValue(self, val): """Add `val` to the current context and its outer context recursively.""" result = val - if val.name not in self._values: + new_value = val.name not in self._values + # Don't treat ops in this context as new values. Usually all known values + # are in self._values, except when we're importing a while loop inside this + # WhileContext. Since there's a cycle in this case, `val` may be part of the + # imported while loop but not yet processed by this context and added to + # self._values in _AddOpInternal. We only want to process external input + # tensors to the while loop here. + new_value &= val.op._control_flow_context is not self # pylint: disable=protected-access + if new_value: self._values.add(val.name) # If we are in a grad context and val is from its forward context, @@ -2965,7 +2973,7 @@ class WhileContext(ControlFlowContext): packed_exit_vars = nest.pack_sequence_as( structure=original_body_result, flat_sequence=exit_vars_with_tensor_arrays) - return (packed_exit_vars[0] if len(exit_vars) == 1 else packed_exit_vars) + return packed_exit_vars[0] if len(exit_vars) == 1 else packed_exit_vars def _FixControlInputsAndContext(self, enters): graph = ops.get_default_graph() diff --git a/tensorflow/python/ops/cudnn_rnn_grad.py b/tensorflow/python/ops/cudnn_rnn_grad.py new file mode 100644 index 0000000000000000000000000000000000000000..97331bb5b5cdecee7d65fe8d712e22fd80f25bae --- /dev/null +++ b/tensorflow/python/ops/cudnn_rnn_grad.py @@ -0,0 +1,47 @@ +# 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. +# ============================================================================== +"""Gradients for CuudnnRNN operators.""" +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 gen_cudnn_rnn_ops + + +@ops.RegisterGradient("CudnnRNN") +def _cudnn_rnn_backward(op, *grads): + """Gradients for the CudnnRNN op.""" + if not op.get_attr("is_training"): + raise ValueError( + "CudnnRNN must set is_training to True to be used in gradients") + return gen_cudnn_rnn_ops.cudnn_rnn_backprop( + input=op.inputs[0], + input_h=op.inputs[1], + input_c=op.inputs[2], + params=op.inputs[3], + output=op.outputs[0], + output_h=op.outputs[1], + output_c=op.outputs[2], + output_backprop=grads[0], + output_h_backprop=grads[1], + output_c_backprop=grads[2], + reserve_space=op.outputs[3], + dropout=op.get_attr("dropout"), + seed=op.get_attr("seed"), + seed2=op.get_attr("seed2"), + rnn_mode=op.get_attr("rnn_mode"), + input_mode=op.get_attr("input_mode"), + direction=op.get_attr("direction")) diff --git a/tensorflow/python/ops/data_flow_ops.py b/tensorflow/python/ops/data_flow_ops.py index cb725199a8501d2a894f18a9b57c23de4e349374..62c5adc385a2e87d27298c72f8dd2f67303119df 100644 --- a/tensorflow/python/ops/data_flow_ops.py +++ b/tensorflow/python/ops/data_flow_ops.py @@ -571,7 +571,7 @@ class QueueBase(object): name=name) def is_closed(self, name=None): - """ Returns true if queue is closed. + """Returns true if queue is closed. This operation returns true if the queue is closed and false if the queue is open. @@ -1563,7 +1563,7 @@ class BaseStagingArea(object): of the staging area. Args: - vals: A tensor, a list or tuple of tensors, or a dictionary.. + vals: A tensor, a list or tuple of tensors, or a dictionary. Returns: A (tensors, indices) tuple where `tensors` is a list of `Tensor` objects @@ -1582,7 +1582,7 @@ class BaseStagingArea(object): (sorted(vals.keys()), sorted(self._names))) # The order of values in `self._names` indicates the order in which the # tensors in the dictionary `vals` must be listed. - vals, indices, n = zip(*[(vals[k], i, k) + vals, indices, _ = zip(*[(vals[k], i, k) for i, k in enumerate(self._names) if k in vals]) else: @@ -1612,7 +1612,7 @@ class BaseStagingArea(object): for val, i in zip(vals, indices): dtype, shape = self._dtypes[i], self._shapes[i] # Check dtype - if not val.dtype == dtype: + if val.dtype != dtype: raise ValueError("Datatypes do not match. '%s' != '%s'" % (str(val.dtype), str(dtype))) @@ -1626,7 +1626,7 @@ class BaseStagingArea(object): def _create_device_transfers(self, tensors): """Encode inter-device transfers if the current device - is not the same as the Staging Area's device + is not the same as the Staging Area's device. """ if not isinstance(tensors, (tuple, list)): @@ -1739,11 +1739,6 @@ class StagingArea(BaseStagingArea): Args: dtypes: A list of types. The length of dtypes must equal the number of tensors in each element. - capacity: (Optional.) Maximum number of elements. - An integer. If zero, the Staging Area is unbounded - memory_limit: (Optional.) Maximum number of bytes of all tensors - in the Staging Area. - An integer. If zero, the Staging Area is unbounded shapes: (Optional.) Constraints on the shapes of tensors in an element. A list of shape tuples or None. This list is the same length as dtypes. If the shape of any tensors in the element are constrained, @@ -1754,6 +1749,11 @@ class StagingArea(BaseStagingArea): shared_name: (Optional.) A name to be used for the shared object. By passing the same name to two different python objects they will share the underlying staging area. Must be a string. + capacity: (Optional.) Maximum number of elements. + An integer. If zero, the Staging Area is unbounded + memory_limit: (Optional.) Maximum number of bytes of all tensors + in the Staging Area. + An integer. If zero, the Staging Area is unbounded Raises: ValueError: If one of the arguments is invalid. @@ -1782,7 +1782,7 @@ class StagingArea(BaseStagingArea): """ with ops.name_scope(name, "%s_put" % self._name, self._scope_vals(values)) as scope: - + if not isinstance(values, (list, tuple, dict)): values = [values] @@ -1911,7 +1911,8 @@ class StagingArea(BaseStagingArea): class MapStagingArea(BaseStagingArea): - """A `MapStagingArea` is a TensorFlow data structure that stores tensors across multiple steps, and exposes operations that can put and get tensors. + """A `MapStagingArea` is a TensorFlow data structure that stores tensors + across multiple steps, and exposes operations that can put and get tensors. Each `MapStagingArea` element is a (key, value) pair. Only int64 keys are supported, other types should be @@ -2375,7 +2376,7 @@ class RecordInput(object): return records else: with ops.name_scope(self._name): - batch_list = [[] for i in six.moves.range(self._batches)] + batch_list = [[] for _ in six.moves.range(self._batches)] records = array_ops.split(records, self._batch_size, 0) records = [array_ops.reshape(record, []) for record in records] for index, protobuf in zip(six.moves.range(len(records)), records): diff --git a/tensorflow/python/ops/distributions/BUILD b/tensorflow/python/ops/distributions/BUILD index 9d9ede7ad75f4eafa91ad051458afbcb6dc8f7b5..e7ad028376b841fc485f7c9cded2a3cdd9dcc153 100644 --- a/tensorflow/python/ops/distributions/BUILD +++ b/tensorflow/python/ops/distributions/BUILD @@ -8,9 +8,13 @@ licenses(["notice"]) # Apache 2.0 py_library( name = "distributions", - srcs = glob(["*.py"]), + srcs = glob( + ["*.py"], + exclude = ["util.py"], + ), srcs_version = "PY2AND3", deps = [ + ":util", "//tensorflow/python:array_ops", "//tensorflow/python:check_ops", "//tensorflow/python:control_flow_ops", @@ -26,3 +30,23 @@ py_library( "@six_archive//:six", ], ) + +py_library( + name = "util", + srcs = ["util.py"], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/python:array_ops", + "//tensorflow/python:check_ops", + "//tensorflow/python:control_flow_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:math_ops", + "//tensorflow/python:nn", + "//tensorflow/python:nn_ops", + "//tensorflow/python:random_ops", + "//tensorflow/python:special_math_ops", + "//tensorflow/python:tensor_util", + "//third_party/py/numpy", + "@six_archive//:six", + ], +) diff --git a/tensorflow/python/ops/distributions/bernoulli.py b/tensorflow/python/ops/distributions/bernoulli.py index 68aaf3815e7e2b21c9550562aa49195569c8ea43..2c9f0e9a32dd3f2caa81befebc06dcc740c832cb 100644 --- a/tensorflow/python/ops/distributions/bernoulli.py +++ b/tensorflow/python/ops/distributions/bernoulli.py @@ -72,7 +72,7 @@ class Bernoulli(distribution.Distribution): ValueError: If p and logits are passed, or if neither are passed. """ parameters = locals() - with ops.name_scope(name): + with ops.name_scope(name) as name: self._logits, self._probs = distribution_util.get_logits_and_probs( logits=logits, probs=probs, diff --git a/tensorflow/python/ops/distributions/beta.py b/tensorflow/python/ops/distributions/beta.py index 469bcadb8ea3a0ec2a85d3a72c0ca5ba08796856..8beab99bf868cdb16dc842eae113cf55dd565131 100644 --- a/tensorflow/python/ops/distributions/beta.py +++ b/tensorflow/python/ops/distributions/beta.py @@ -151,7 +151,7 @@ class Beta(distribution.Distribution): name: Python `str` name prefixed to Ops created by this class. """ parameters = locals() - with ops.name_scope(name, values=[concentration1, concentration0]): + with ops.name_scope(name, values=[concentration1, concentration0]) as name: self._concentration1 = self._maybe_assert_valid_concentration( ops.convert_to_tensor(concentration1, name="concentration1"), validate_args) @@ -323,7 +323,7 @@ class BetaWithSoftplusConcentration(Beta): name="BetaWithSoftplusConcentration"): parameters = locals() with ops.name_scope(name, values=[concentration1, - concentration0]) as ns: + concentration0]) as name: super(BetaWithSoftplusConcentration, self).__init__( concentration1=nn.softplus(concentration1, name="softplus_concentration1"), @@ -331,7 +331,7 @@ class BetaWithSoftplusConcentration(Beta): name="softplus_concentration0"), validate_args=validate_args, allow_nan_stats=allow_nan_stats, - name=ns) + name=name) self._parameters = parameters diff --git a/tensorflow/python/ops/distributions/bijector_impl.py b/tensorflow/python/ops/distributions/bijector_impl.py index ed435557fde7a2e8a0a4f7eef4e240daef0565e7..4ebc600d034603a80d4fae93b1339e1a1feea038 100644 --- a/tensorflow/python/ops/distributions/bijector_impl.py +++ b/tensorflow/python/ops/distributions/bijector_impl.py @@ -23,7 +23,6 @@ import collections import contextlib import re -import numpy as np import six from tensorflow.python.framework import dtypes @@ -31,8 +30,8 @@ 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 math_ops -from tensorflow.python.util.tf_export import tf_export __all__ = [ @@ -41,23 +40,24 @@ __all__ = [ class _Mapping(collections.namedtuple( - "_Mapping", ["x", "y", "ildj", "kwargs"])): + "_Mapping", ["x", "y", "ildj_map", "kwargs"])): """Helper class to make it easier to manage caching in `Bijector`.""" - def __new__(cls, x=None, y=None, ildj=None, kwargs=None): + def __new__(cls, x=None, y=None, ildj_map=None, kwargs=None): """Custom __new__ so namedtuple items have defaults. Args: x: `Tensor`. Forward. y: `Tensor`. Inverse. - ildj: `Tensor`. Inverse log det Jacobian. + ildj_map: `Dictionary`. This is a mapping from event_ndims to a `Tensor` + representing the inverse log det jacobian. kwargs: Python dictionary. Extra args supplied to forward/inverse/etc functions. Returns: mapping: New instance of _Mapping. """ - return super(_Mapping, cls).__new__(cls, x, y, ildj, kwargs) + return super(_Mapping, cls).__new__(cls, x, y, ildj_map, kwargs) @property def x_key(self): @@ -69,13 +69,14 @@ class _Mapping(collections.namedtuple( """Returns key used for caching X=g^{-1}(Y).""" return (self.y,) + self._deep_tuple(tuple(sorted(self.kwargs.items()))) - def merge(self, x=None, y=None, ildj=None, kwargs=None, mapping=None): + def merge(self, x=None, y=None, ildj_map=None, kwargs=None, mapping=None): """Returns new _Mapping with args merged with self. Args: x: `Tensor`. Forward. y: `Tensor`. Inverse. - ildj: `Tensor`. Inverse log det Jacobian. + ildj_map: `Dictionary`. This is a mapping from event_ndims to a `Tensor` + representing the inverse log det jacobian. kwargs: Python dictionary. Extra args supplied to forward/inverse/etc functions. mapping: Instance of _Mapping to merge. Can only be specified if no other @@ -88,15 +89,30 @@ class _Mapping(collections.namedtuple( ValueError: if mapping and any other arg is not `None`. """ if mapping is None: - mapping = _Mapping(x=x, y=y, ildj=ildj, kwargs=kwargs) - elif not all(arg is None for arg in [x, y, ildj, kwargs]): - raise ValueError("Cannot specify mapping and individual args.") + mapping = _Mapping(x=x, y=y, ildj_map=ildj_map, kwargs=kwargs) + elif any(arg is not None for arg in [x, y, ildj_map, kwargs]): + raise ValueError("Cannot simultaneously specify mapping and individual " + "arguments.") + return _Mapping( x=self._merge(self.x, mapping.x), y=self._merge(self.y, mapping.y), - ildj=self._merge(self.ildj, mapping.ildj), + ildj_map=self._merge_dicts(self.ildj_map, mapping.ildj_map), kwargs=self._merge(self.kwargs, mapping.kwargs)) + def _merge_dicts(self, old=None, new=None): + """Helper to merge two dictionaries.""" + old = dict() if old is None else old + new = dict() if new is None else new + for k, v in six.iteritems(new): + val = old.get(k, None) + if val is not None and val != v: + raise ValueError("Found different value for existing key " + "(key:{} old_value:{} new_value:{}".format( + k, old[k], v)) + old[k] = v + return old + def _merge(self, old, new): """Helper to merge which handles merging one value.""" if old is None: @@ -112,7 +128,6 @@ class _Mapping(collections.namedtuple( @six.add_metaclass(abc.ABCMeta) -@tf_export("distributions.bijectors.Bijector") class Bijector(object): r"""Interface for transformations of a `Distribution` sample. @@ -137,11 +152,11 @@ class Bijector(object): 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_det_jacobian(x)`\ "The log of the determinant of the matrix of all first-order partial 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 + of another. Geometrically, the Jacobian determinant is the volume of the transformation and is used to scale the probability. By convention, transformations of random variables are named in terms of the @@ -164,7 +179,7 @@ class Bijector(object): ```python def transformed_log_prob(bijector, log_prob, x): - return (bijector.inverse_log_det_jacobian(x) + + return (bijector.inverse_log_det_jacobian(x, event_ndims=0) + log_prob(bijector.inverse(x))) ``` @@ -199,9 +214,11 @@ class Bijector(object): ```python class Exp(Bijector): - def __init__(self, event_ndims=0, validate_args=False, name="exp"): + def __init__(self, validate_args=False, name="exp"): super(Exp, self).__init__( - event_ndims=event_ndims, validate_args=validate_args, name=name) + validate_args=validate_args, + forward_min_event_ndims=0, + name=name) def _forward(self, x): return math_ops.exp(x) @@ -213,10 +230,11 @@ class Bijector(object): return -self._forward_log_det_jacobian(self._inverse(y)) def _forward_log_det_jacobian(self, x): - if self.event_ndims is None: - raise ValueError("Jacobian requires known event_ndims.") - event_dims = array_ops.shape(x)[-self.event_ndims:] - return math_ops.reduce_sum(x, axis=event_dims) + # Notice that we needn't do any reducing, even when`event_ndims > 0`. + # The base Bijector class will handle reducing for us; it knows how + # to do so because we called `super` `__init__` with + # `forward_min_event_ndims = 0`. + return x ``` - "Affine" @@ -237,18 +255,50 @@ class Bijector(object): MultivariateNormal(inv(sqrtSigma) * (y - mu); 0, I_d) ``` - #### Jacobian + #### Min_event_ndims and Naming + + Bijectors are named for the dimensionality of data they act on (i.e. without + broadcasting). We can think of bijectors having an intrinsic `min_event_ndims` + , which is the minimum number of dimensions for the bijector act on. For + instance, a Cholesky decomposition requires a matrix, and hence + `min_event_ndims=2`. + + Some examples: + + `AffineScalar: min_event_ndims=0` + `Affine: min_event_ndims=1` + `Cholesky: min_event_ndims=2` + `Exp: min_event_ndims=0` + `Sigmoid: min_event_ndims=0` + `SoftmaxCentered: min_event_ndims=1` + + Note the difference between `Affine` and `AffineScalar`. `AffineScalar` + operates on scalar events, whereas `Affine` operates on vector-valued events. - The Jacobian is a reduction over event dims. To see this, consider the `Exp` - `Bijector` applied to a `Tensor` which has sample, batch, and event (S, B, E) - shape semantics. Suppose the `Tensor`'s partitioned-shape is `(S=[4], B=[2], - E=[3, 3])`. The shape of the `Tensor` returned by `forward` and `inverse` is - unchanged, i.e., `[4, 2, 3, 3]`. However the shape returned by - `inverse_log_det_jacobian` is `[4, 2]` because the Jacobian is a reduction - over the event dimensions. + More generally, there is a `forward_min_event_ndims` and an + `inverse_min_event_ndims`. In most cases, these will be the same. + However, for some shape changing bijectors, these will be different + (e.g. a bijector which pads an extra dimension at the end, might have + `forward_min_event_ndims=0` and `inverse_min_event_ndims=1`. - It is sometimes useful to implement the inverse Jacobian as the negative - forward Jacobian. For example, + + #### Jacobian Determinant + + The Jacobian determinant is a reduction over `event_ndims - min_event_ndims` + (`forward_min_event_ndims` for `forward_log_det_jacobian` and + `inverse_min_event_ndims` for `inverse_log_det_jacobian`). + To see this, consider the `Exp` `Bijector` applied to a `Tensor` which has + sample, batch, and event (S, B, E) shape semantics. Suppose the `Tensor`'s + partitioned-shape is `(S=[4], B=[2], E=[3, 3])`. The shape of the `Tensor` + returned by `forward` and `inverse` is unchanged, i.e., `[4, 2, 3, 3]`. + However the shape returned by `inverse_log_det_jacobian` is `[4, 2]` because + the Jacobian determinant is a reduction over the event dimensions. + + Another example is the `Affine` `Bijector`. Because `min_event_ndims = 1`, the + Jacobian determinant reduction is over `event_ndims - 1`. + + It is sometimes useful to implement the inverse Jacobian determinant as the + negative forward Jacobian determinant. For example, ```python def _inverse_log_det_jacobian(self, y): @@ -279,9 +329,54 @@ class Bijector(object): The claim follows from [properties of determinant]( https://en.wikipedia.org/wiki/Determinant#Multiplicativity_and_matrix_groups). - Generally its preferable to directly implement the inverse Jacobian. This - should have superior numerical stability and will often share subgraphs with - the `_inverse` implementation. + Generally its preferable to directly implement the inverse Jacobian + determinant. This should have superior numerical stability and will often + share subgraphs with the `_inverse` implementation. + + #### Is_constant_jacobian + + Certain bijectors will have constant jacobian matrices. For instance, the + `Affine` bijector encodes multiplication by a matrix plus a shift, with + jacobian matrix, the same aforementioned matrix. + + `is_constant_jacobian` encodes the fact that the jacobian matrix is constant. + The semantics of this argument are the following: + + * Repeated calls to "log_det_jacobian" functions with the same + `event_ndims` (but not necessarily same input), will return the first + computed jacobian (because the matrix is constant, and hence is input + independent). + * `log_det_jacobian` implementations are merely broadcastable to the true + `log_det_jacobian` (because, again, the jacobian matrix is input + independent). Specifically, `log_det_jacobian` is implemented as the + log jacobian determinant for a single input. + + ```python + class Identity(Bijector): + + def __init__(self, validate_args=False, name="identity"): + super(Identity, self).__init__( + is_constant_jacobian=True, + validate_args=validate_args, + forward_min_event_ndims=0, + name=name) + + def _forward(self, x): + return x + + def _inverse(self, y): + return y + + def _inverse_log_det_jacobian(self, y): + return -self._forward_log_det_jacobian(self._inverse(y)) + + def _forward_log_det_jacobian(self, x): + # The full log jacobian determinant would be array_ops.zero_like(x). + # However, we circumvent materializing that, since the jacobian + # calculation is input independent, and we specify it for one input. + return constant_op.constant(0., x.dtype.base_dtype) + + ``` #### Subclass Requirements @@ -364,14 +459,14 @@ class Bijector(object): ==> (-1., 1.) # The |dX/dY| is constant, == 1. So Log|dX/dY| == 0. - abs.inverse_log_det_jacobian(1.) + abs.inverse_log_det_jacobian(1., event_ndims=0) ==> (0., 0.) # Special case handling of 0. abs.inverse(0.) ==> (0., 0.) - abs.inverse_log_det_jacobian(0.) + abs.inverse_log_det_jacobian(0., event_ndims=0) ==> (0., 0.) ``` @@ -379,11 +474,12 @@ class Bijector(object): @abc.abstractmethod def __init__(self, - event_ndims=None, graph_parents=None, is_constant_jacobian=False, validate_args=False, dtype=None, + forward_min_event_ndims=None, + inverse_min_event_ndims=None, name=None): """Constructs Bijector. @@ -392,42 +488,61 @@ class Bijector(object): Examples: ```python - # Create the Y = g(X) = X transform which operates on vector events. - identity = Identity(event_ndims=1) + # Create the Y = g(X) = X transform. + identity = Identity() - # Create the Y = g(X) = exp(X) transform which operates on matrices. - exp = Exp(event_ndims=2) + # Create the Y = g(X) = exp(X) transform. + exp = Exp() ``` See `Bijector` subclass docstring for more details and specific examples. Args: - event_ndims: number of dimensions associated with event coordinates. graph_parents: Python list of graph prerequisites of this `Bijector`. - is_constant_jacobian: Python `bool` indicating that the Jacobian is not a - function of the input. + is_constant_jacobian: Python `bool` indicating that the Jacobian matrix is + not a function of the input. validate_args: Python `bool`, default `False`. Whether to validate input with asserts. If `validate_args` is `False`, and the inputs are invalid, correct behavior is not guaranteed. dtype: `tf.dtype` supported by this `Bijector`. `None` means dtype is not enforced. + forward_min_event_ndims: Python `integer` indicating the minimum number of + dimensions `forward` operates on. + inverse_min_event_ndims: Python `integer` indicating the minimum number of + dimensions `inverse` operates on. Will be set to + `forward_min_event_ndims` by default, if no value is provided. name: The name to give Ops created by the initializer. Raises: + ValueError: If neither `forward_min_event_ndims` and + `inverse_min_event_ndims` are specified, or if either of them is + negative. ValueError: If a member of `graph_parents` is not a `Tensor`. """ - self._event_ndims = ( - ops.convert_to_tensor(event_ndims, dtype=dtypes.int32) - if event_ndims is not None else None) self._graph_parents = graph_parents or [] + + if forward_min_event_ndims is None and inverse_min_event_ndims is None: + raise ValueError("Must specify at least one of `forward_min_event_ndims` " + "and `inverse_min_event_ndims`.") + elif inverse_min_event_ndims is None: + inverse_min_event_ndims = forward_min_event_ndims + elif forward_min_event_ndims is None: + forward_min_event_ndims = inverse_min_event_ndims + + if forward_min_event_ndims < 0: + raise ValueError("forward_min_event_ndims must be a non-negative " + "integer.") + if inverse_min_event_ndims < 0: + raise ValueError("inverse_min_event_ndims must be a non-negative " + "integer.") + self._forward_min_event_ndims = forward_min_event_ndims + self._inverse_min_event_ndims = inverse_min_event_ndims self._is_constant_jacobian = is_constant_jacobian + self._constant_ildj_map = {} self._validate_args = validate_args self._dtype = dtype self._from_y = {} self._from_x = {} - # Using abbreviation ildj for "inverse log det Jacobian." - # This variable is not `None` iff is_constant_jacobian is `True`. - self._constant_ildj = None if name: self._name = name else: @@ -442,21 +557,27 @@ class Bijector(object): if t is None or not tensor_util.is_tensor(t): raise ValueError("Graph parent item %d is not a Tensor; %s." % (i, t)) - @property - def event_ndims(self): - """Returns then number of event dimensions this bijector operates on.""" - return self._event_ndims - @property def graph_parents(self): """Returns this `Bijector`'s graph_parents as a Python list.""" return self._graph_parents + @property + def forward_min_event_ndims(self): + """Returns the minimal number of dimensions bijector.forward operates on.""" + return self._forward_min_event_ndims + + @property + def inverse_min_event_ndims(self): + """Returns the minimal number of dimensions bijector.inverse operates on.""" + return self._inverse_min_event_ndims + @property def is_constant_jacobian(self): - """Returns true iff the Jacobian is not a function of x. + """Returns true iff the Jacobian matrix is not a function of x. - Note: Jacobian is either constant for both forward and inverse or neither. + Note: Jacobian matrix is either constant for both forward and inverse or + neither. Returns: is_constant_jacobian: Python `bool`. @@ -653,36 +774,57 @@ class Bijector(object): return self._call_inverse(y, name) def _inverse_log_det_jacobian(self, y): - """Subclass implementation of `inverse_log_det_jacobian` public function.""" + """Subclass implementation of `inverse_log_det_jacobian` public function. + + In particular, this method differs from the public function, in that it + does not take `event_ndims`. Thus, this implements the minimal Jacobian + determinant calculation (i.e. over `inverse_min_event_ndims`). + + Args: + y: `Tensor`. The input to the "inverse_log_det_jacobian" evaluation. + Returns: + inverse_log_det_jacobian: `Tensor`, if this bijector is injective. + If not injective, returns the k-tuple containing jacobians for the + unique `k` points `(x1, ..., xk)` such that `g(xi) = y`. + """ raise NotImplementedError("inverse_log_det_jacobian not implemented.") - def _call_inverse_log_det_jacobian(self, y, name, **kwargs): + def _call_inverse_log_det_jacobian(self, y, event_ndims, name, **kwargs): with self._name_scope(name, [y]): - if self._constant_ildj is not None: - return self._constant_ildj + if event_ndims in self._constant_ildj_map: + return self._constant_ildj_map[event_ndims] y = ops.convert_to_tensor(y, name="y") self._maybe_assert_dtype(y) if not self._is_injective: # No caching for non-injective - return self._inverse_log_det_jacobian(y, **kwargs) + ildjs = self._inverse_log_det_jacobian(y, **kwargs) + return tuple(self._reduce_jacobian_det_over_event( + y, ildj, self.inverse_min_event_ndims, event_ndims) + for ildj in ildjs) mapping = self._lookup(y=y, kwargs=kwargs) - if mapping.ildj is not None: - return mapping.ildj + if mapping.ildj_map is not None and event_ndims in mapping.ildj_map: + return mapping.ildj_map[event_ndims] try: x = None # Not needed; leave cache as is. ildj = self._inverse_log_det_jacobian(y, **kwargs) + ildj = self._reduce_jacobian_det_over_event( + y, ildj, self.inverse_min_event_ndims, event_ndims) except NotImplementedError as original_exception: try: x = mapping.x if mapping.x is not None else self._inverse(y, **kwargs) ildj = -self._forward_log_det_jacobian(x, **kwargs) + ildj = self._reduce_jacobian_det_over_event( + x, ildj, self.forward_min_event_ndims, event_ndims) except NotImplementedError: raise original_exception - mapping = mapping.merge(x=x, ildj=ildj) + + mapping = mapping.merge(x=x, ildj_map={event_ndims: ildj}) self._cache(mapping) if self.is_constant_jacobian: - self._constant_ildj = mapping.ildj - return mapping.ildj + self._constant_ildj_map[event_ndims] = ildj + return ildj - def inverse_log_det_jacobian(self, y, name="inverse_log_det_jacobian"): + def inverse_log_det_jacobian( + self, y, event_ndims, name="inverse_log_det_jacobian"): """Returns the (log o det o Jacobian o inverse)(y). Mathematically, returns: `log(det(dX/dY))(Y)`. (Recall that: `X=g^{-1}(Y)`.) @@ -691,7 +833,12 @@ class Bijector(object): evaluated at `g^{-1}(y)`. Args: - y: `Tensor`. The input to the "inverse" Jacobian evaluation. + y: `Tensor`. The input to the "inverse" Jacobian determinant evaluation. + event_ndims: Number of dimensions in the probabilistic events being + transformed. Must be greater than or equal to + `self.inverse_min_event_ndims`. The result is summed over the final + dimensions to produce a scalar Jacobian determinant for each event, + i.e. it has shape `y.shape.ndims - event_ndims` dimensions. name: The name to give this op. Returns: @@ -705,45 +852,74 @@ class Bijector(object): `self.dtype`. NotImplementedError: if `_inverse_log_det_jacobian` is not implemented. """ - return self._call_inverse_log_det_jacobian(y, name) + with ops.control_dependencies(self._check_valid_event_ndims( + min_event_ndims=self.inverse_min_event_ndims, event_ndims=event_ndims)): + return self._call_inverse_log_det_jacobian(y, event_ndims, name) def _forward_log_det_jacobian(self, x): - """Subclass implementation of `forward_log_det_jacobian`.""" + """Subclass implementation of `forward_log_det_jacobian` public function. + + In particular, this method differs from the public function, in that it + does not take `event_ndims`. Thus, this implements the minimal Jacobian + determinant calculation (i.e. over `forward_min_event_ndims`). + + Args: + x: `Tensor`. The input to the "forward_log_det_jacobian" evaluation. + + Returns: + forward_log_det_jacobian: `Tensor`, if this bijector is injective. + If not injective, returns the k-tuple containing jacobians for the + unique `k` points `(x1, ..., xk)` such that `g(xi) = y`. + """ + raise NotImplementedError( "forward_log_det_jacobian not implemented.") - def _call_forward_log_det_jacobian(self, x, name, **kwargs): + def _call_forward_log_det_jacobian(self, x, event_ndims, name, **kwargs): with self._name_scope(name, [x]): - if self._constant_ildj is not None: + if event_ndims in self._constant_ildj_map: # Need "-1. *" to avoid invalid-unary-operand-type linter warning. - return -1. * self._constant_ildj + return -1. * self._constant_ildj_map[event_ndims] x = ops.convert_to_tensor(x, name="x") self._maybe_assert_dtype(x) if not self._is_injective: - return self._forward_log_det_jacobian(x, **kwargs) # No caching. + fldjs = self._forward_log_det_jacobian(x, **kwargs) # No caching. + return tuple(self._reduce_jacobian_det_over_event( + x, fldj, self.forward_min_event_ndims, event_ndims) + for fldj in fldjs) mapping = self._lookup(x=x, kwargs=kwargs) - if mapping.ildj is not None: - return -mapping.ildj + if mapping.ildj_map is not None and event_ndims in mapping.ildj_map: + return -mapping.ildj_map[event_ndims] try: y = None # Not needed; leave cache as is. ildj = -self._forward_log_det_jacobian(x, **kwargs) + ildj = self._reduce_jacobian_det_over_event( + x, ildj, self.forward_min_event_ndims, event_ndims) except NotImplementedError as original_exception: try: y = mapping.y if mapping.y is not None else self._forward(x, **kwargs) ildj = self._inverse_log_det_jacobian(y, **kwargs) + ildj = self._reduce_jacobian_det_over_event( + y, ildj, self.inverse_min_event_ndims, event_ndims) except NotImplementedError: raise original_exception - mapping = mapping.merge(y=y, ildj=ildj) + mapping = mapping.merge(y=y, ildj_map={event_ndims: ildj}) self._cache(mapping) if self.is_constant_jacobian: - self._constant_ildj = mapping.ildj - return -mapping.ildj + self._constant_ildj_map[event_ndims] = ildj + return -ildj - def forward_log_det_jacobian(self, x, name="forward_log_det_jacobian"): + def forward_log_det_jacobian( + self, x, event_ndims, name="forward_log_det_jacobian"): """Returns both the forward_log_det_jacobian. Args: - x: `Tensor`. The input to the "forward" Jacobian evaluation. + x: `Tensor`. The input to the "forward" Jacobian determinant evaluation. + event_ndims: Number of dimensions in the probabilistic events being + transformed. Must be greater than or equal to + `self.forward_min_event_ndims`. The result is summed over the final + dimensions to produce a scalar Jacobian determinant for each event, + i.e. it has shape `x.shape.ndims - event_ndims` dimensions. name: The name to give this op. Returns: @@ -761,7 +937,9 @@ class Bijector(object): raise NotImplementedError( "forward_log_det_jacobian cannot be implemented for non-injective " "transforms.") - return self._call_forward_log_det_jacobian(x, name) + with ops.control_dependencies(self._check_valid_event_ndims( + min_event_ndims=self.forward_min_event_ndims, event_ndims=event_ndims)): + return self._call_forward_log_det_jacobian(x, event_ndims, name) @contextlib.contextmanager def _name_scope(self, name=None, values=None): @@ -779,9 +957,6 @@ class Bijector(object): def _cache(self, mapping): """Helper which stores mapping info in forward/inverse dicts.""" - if self._constant_ildj is not None: - # Fold in ildj if known constant Jacobian. - mapping = mapping.merge(ildj=self._constant_ildj) # Merging from lookup is an added check that we're not overwriting anything # which is not None. mapping = mapping.merge(mapping=self._lookup( @@ -803,22 +978,66 @@ class Bijector(object): return self._from_y.get(mapping.y_key, mapping) return mapping - def _event_dims_tensor(self, sample): - """Return a 1D `int32` tensor: `range(rank(sample))[-event_ndims:]`.""" - if self.event_ndims is None: - raise ValueError("Jacobian cannot be computed with unknown event_ndims") - static_event_ndims = tensor_util.constant_value(self.event_ndims) - static_rank = sample.get_shape().ndims - if static_event_ndims is not None and static_rank is not None: - return ops.convert_to_tensor( - static_rank + np.arange(-static_event_ndims, 0).astype(np.int32)) - - if static_event_ndims is not None: - event_range = np.arange(-static_event_ndims, 0).astype(np.int32) - else: - event_range = math_ops.range(-self.event_ndims, 0, dtype=dtypes.int32) - - if static_rank is not None: - return event_range + static_rank + def _reduce_jacobian_det_over_event( + self, y, ildj, min_event_ndims, event_ndims): + """Reduce jacobian over event_ndims - min_event_ndims.""" + if not self.is_constant_jacobian: + return math_ops.reduce_sum( + ildj, + self._get_event_reduce_dims(min_event_ndims, event_ndims)) + + # In this case, we need to tile the jacobian over the event and reduce. + y_rank = array_ops.rank(y) + y_shape = array_ops.shape(y)[ + y_rank - event_ndims : y_rank - min_event_ndims] + + ones = array_ops.ones(y_shape, ildj.dtype) + reduced_ildj = math_ops.reduce_sum( + ones * ildj, + axis=self._get_event_reduce_dims(min_event_ndims, event_ndims)) + # The multiplication by ones can change the inferred static shape so we try + # to recover as much as possible. + if (isinstance(event_ndims, int) and + y.get_shape().ndims and ildj.get_shape().ndims): + y_shape = y.get_shape() + y_shape = y_shape[y_shape.ndims - event_ndims : + y_shape.ndims - min_event_ndims] + ildj_shape = ildj.get_shape() + broadcast_shape = array_ops.broadcast_static_shape( + ildj_shape, y_shape) + reduced_ildj.set_shape( + broadcast_shape[: broadcast_shape.ndims - ( + event_ndims - min_event_ndims)]) + + return reduced_ildj + + def _get_event_reduce_dims(self, min_event_ndims, event_ndims): + """Compute the reduction dimensions given event_ndims.""" + min_event_ndims_ = (min_event_ndims if isinstance(min_event_ndims, int) + else tensor_util.constant_value(min_event_ndims)) + event_ndims_ = (event_ndims if isinstance(event_ndims, int) + else tensor_util.constant_value(event_ndims)) + + if min_event_ndims_ is not None and event_ndims_ is not None: + return [-index for index in range(1, event_ndims_ - min_event_ndims_ + 1)] else: - return event_range + array_ops.rank(sample) + reduce_ndims = event_ndims - min_event_ndims + return math_ops.range(-reduce_ndims, 0) + + def _check_valid_event_ndims(self, min_event_ndims, event_ndims): + """Check whether event_ndims is atleast min_event_ndims.""" + min_event_ndims_ = (min_event_ndims if isinstance(min_event_ndims, int) + else tensor_util.constant_value(min_event_ndims)) + event_ndims_ = (event_ndims if isinstance(event_ndims, int) + else tensor_util.constant_value(event_ndims)) + + if min_event_ndims_ is not None and event_ndims_ is not None: + if min_event_ndims_ > event_ndims_: + raise ValueError("event_ndims ({}) must be larger than " + "min_event_ndims ({})".format( + event_ndims_, min_event_ndims_)) + return [] + + if self.validate_args: + return [check_ops.assert_greater_equal(event_ndims, min_event_ndims)] + return [] diff --git a/tensorflow/python/ops/distributions/bijector_test_util.py b/tensorflow/python/ops/distributions/bijector_test_util.py index ff3535c62642d98bdd9b18808f45deae27d6d88d..784bfd58352f4035cd1bd4caa91eba6e6dc8d30d 100644 --- a/tensorflow/python/ops/distributions/bijector_test_util.py +++ b/tensorflow/python/ops/distributions/bijector_test_util.py @@ -79,9 +79,7 @@ def assert_scalar_congruency(bijector, Raises: AssertionError: If tests fail. """ - # Checks and defaults. - assert bijector.event_ndims.eval() == 0 if sess is None: sess = ops.get_default_session() @@ -111,7 +109,10 @@ def assert_scalar_congruency(bijector, # (b - a) = \int_a^b dx = \int_{y(a)}^{y(b)} |dx/dy| dy # "change_measure_dy_dx" below is a Monte Carlo approximation to the right # hand side, which should then be close to the left, which is (b - a). - dy_dx = math_ops.exp(bijector.inverse_log_det_jacobian(uniform_y_samps)) + # We assume event_ndims=0 because we assume scalar -> scalar. The log_det + # methods will handle whether they expect event_ndims > 0. + dy_dx = math_ops.exp(bijector.inverse_log_det_jacobian( + uniform_y_samps, event_ndims=0)) # E[|dx/dy|] under Uniform[lower_y, upper_y] # = \int_{y(a)}^{y(b)} |dx/dy| dP(u), where dP(u) is the uniform measure expectation_of_dy_dx_under_uniform = math_ops.reduce_mean(dy_dx) @@ -121,7 +122,8 @@ def assert_scalar_congruency(bijector, # We'll also check that dy_dx = 1 / dx_dy. dx_dy = math_ops.exp( - bijector.forward_log_det_jacobian(bijector.inverse(uniform_y_samps))) + bijector.forward_log_det_jacobian( + bijector.inverse(uniform_y_samps), event_ndims=0)) [ forward_on_10_pts_v, @@ -158,7 +160,8 @@ def assert_scalar_congruency(bijector, dy_dx_v, np.divide(1., dx_dy_v), atol=1e-5, rtol=1e-3) -def assert_bijective_and_finite(bijector, x, y, atol=0, rtol=1e-5, sess=None): +def assert_bijective_and_finite( + bijector, x, y, event_ndims, atol=0, rtol=1e-5, sess=None): """Assert that forward/inverse (along with jacobians) are inverses and finite. It is recommended to use x and y values that are very very close to the edge @@ -168,6 +171,8 @@ def assert_bijective_and_finite(bijector, x, y, atol=0, rtol=1e-5, sess=None): bijector: A Bijector instance. x: np.array of values in the domain of bijector.forward. y: np.array of values in the domain of bijector.inverse. + event_ndims: Integer describing the number of event dimensions this bijector + operates on. atol: Absolute tolerance. rtol: Relative tolerance. sess: TensorFlow session. Defaults to the default session. @@ -197,10 +202,10 @@ def assert_bijective_and_finite(bijector, x, y, atol=0, rtol=1e-5, sess=None): ] = sess.run([ bijector.inverse(f_x), bijector.forward(g_y), - bijector.inverse_log_det_jacobian(f_x), - bijector.forward_log_det_jacobian(x), - bijector.inverse_log_det_jacobian(y), - bijector.forward_log_det_jacobian(g_y), + bijector.inverse_log_det_jacobian(f_x, event_ndims=event_ndims), + bijector.forward_log_det_jacobian(x, event_ndims=event_ndims), + bijector.inverse_log_det_jacobian(y, event_ndims=event_ndims), + bijector.forward_log_det_jacobian(g_y, event_ndims=event_ndims), f_x, g_y, ]) diff --git a/tensorflow/python/ops/distributions/categorical.py b/tensorflow/python/ops/distributions/categorical.py index 9161e3fa9f5f7f844e7f4926992c954acae246d6..8f25b1149c3c8b220fe2bea95c33854ee2f6275a 100644 --- a/tensorflow/python/ops/distributions/categorical.py +++ b/tensorflow/python/ops/distributions/categorical.py @@ -183,7 +183,7 @@ class Categorical(distribution.Distribution): name: Python `str` name prefixed to Ops created by this class. """ parameters = locals() - with ops.name_scope(name, values=[logits, probs]): + with ops.name_scope(name, values=[logits, probs]) as name: self._logits, self._probs = distribution_util.get_logits_and_probs( logits=logits, probs=probs, @@ -311,7 +311,7 @@ class Categorical(distribution.Distribution): nn_ops.log_softmax(self.logits) * self.probs, axis=-1) def _mode(self): - ret = math_ops.argmax(self.logits, dimension=self._batch_rank) + ret = math_ops.argmax(self.logits, axis=self._batch_rank) ret = math_ops.cast(ret, self.dtype) ret.set_shape(self.batch_shape) return ret diff --git a/tensorflow/python/ops/distributions/dirichlet.py b/tensorflow/python/ops/distributions/dirichlet.py index 25afeec936069b9cbf926cdc3bbb79226a79aa30..eafcd5c78f7752dc3f9f5be8f8dbe7ac368fd3be 100644 --- a/tensorflow/python/ops/distributions/dirichlet.py +++ b/tensorflow/python/ops/distributions/dirichlet.py @@ -155,7 +155,7 @@ class Dirichlet(distribution.Distribution): name: Python `str` name prefixed to Ops created by this class. """ parameters = locals() - with ops.name_scope(name, values=[concentration]): + with ops.name_scope(name, values=[concentration]) as name: self._concentration = self._maybe_assert_valid_concentration( ops.convert_to_tensor(concentration, name="concentration"), validate_args) diff --git a/tensorflow/python/ops/distributions/dirichlet_multinomial.py b/tensorflow/python/ops/distributions/dirichlet_multinomial.py index 03a98c56ba509ea1f70f12a74ba67b903013cf70..fe0ed7e07d59654b8a4c7ac989212913697c6832 100644 --- a/tensorflow/python/ops/distributions/dirichlet_multinomial.py +++ b/tensorflow/python/ops/distributions/dirichlet_multinomial.py @@ -192,7 +192,7 @@ class DirichletMultinomial(distribution.Distribution): name: Python `str` name prefixed to Ops created by this class. """ parameters = locals() - with ops.name_scope(name, values=[total_count, concentration]): + with ops.name_scope(name, values=[total_count, concentration]) as name: # Broadcasting works because: # * The broadcasting convention is to prepend dimensions of size [1], and # we use the last dimension for the distribution, whereas diff --git a/tensorflow/python/ops/distributions/distribution.py b/tensorflow/python/ops/distributions/distribution.py index 7c43bf54fc783815127f03cc287ab0fc4349beb5..3815abf72de1e2b5278706315a10899dccca4182 100644 --- a/tensorflow/python/ops/distributions/distribution.py +++ b/tensorflow/python/ops/distributions/distribution.py @@ -434,13 +434,17 @@ class Distribution(_BaseDistribution): for i, t in enumerate(graph_parents): if t is None or not tensor_util.is_tensor(t): raise ValueError("Graph parent item %d is not a Tensor; %s." % (i, t)) + if not name or name[-1] != "/": # `name` is not a name scope + non_unique_name = name or type(self).__name__ + with ops.name_scope(non_unique_name) as name: + pass self._dtype = dtype self._reparameterization_type = reparameterization_type self._allow_nan_stats = allow_nan_stats self._validate_args = validate_args self._parameters = parameters or {} self._graph_parents = graph_parents - self._name = name or type(self).__name__ + self._name = name @classmethod def param_shapes(cls, sample_shape, name="DistributionParamShapes"): diff --git a/tensorflow/python/ops/distributions/distributions.py b/tensorflow/python/ops/distributions/distributions.py index 9df7d148a583e533475276e090bcb02cb705290f..7c4b8697d818a88000fc2e84aec461d2918386a6 100644 --- a/tensorflow/python/ops/distributions/distributions.py +++ b/tensorflow/python/ops/distributions/distributions.py @@ -19,7 +19,6 @@ from __future__ import print_function # pylint: disable=wildcard-import,unused-import -from tensorflow.python.ops.distributions import bijectors from tensorflow.python.ops.distributions.bernoulli import Bernoulli from tensorflow.python.ops.distributions.beta import Beta from tensorflow.python.ops.distributions.categorical import Categorical @@ -40,7 +39,6 @@ from tensorflow.python.util.all_util import remove_undocumented _allowed_symbols = [ - "bijectors", "Bernoulli", "Beta", "Categorical", diff --git a/tensorflow/python/ops/distributions/exponential.py b/tensorflow/python/ops/distributions/exponential.py index 6345a76d485c64659aa01fa1611cd27426d8c8a5..cf0e729e1a189d94381d56aedb5a753862a20962 100644 --- a/tensorflow/python/ops/distributions/exponential.py +++ b/tensorflow/python/ops/distributions/exponential.py @@ -95,7 +95,7 @@ class Exponential(gamma.Gamma): # true in the parent class "Gamma." Therefore, passing # allow_nan_stats=True # through to the parent class results in unnecessary asserts. - with ops.name_scope(name, values=[rate]): + with ops.name_scope(name, values=[rate]) as name: self._rate = ops.convert_to_tensor(rate, name="rate") super(Exponential, self).__init__( concentration=array_ops.ones([], dtype=self._rate.dtype), @@ -144,7 +144,7 @@ class ExponentialWithSoftplusRate(Exponential): allow_nan_stats=True, name="ExponentialWithSoftplusRate"): parameters = locals() - with ops.name_scope(name, values=[rate]): + with ops.name_scope(name, values=[rate]) as name: super(ExponentialWithSoftplusRate, self).__init__( rate=nn.softplus(rate, name="softplus_rate"), validate_args=validate_args, diff --git a/tensorflow/python/ops/distributions/gamma.py b/tensorflow/python/ops/distributions/gamma.py index adb1f4f9a879e44cf8cb4cafd22b92554f487712..d39f7c56d39ae1681c32bacd9771f4f2372d3f98 100644 --- a/tensorflow/python/ops/distributions/gamma.py +++ b/tensorflow/python/ops/distributions/gamma.py @@ -127,7 +127,7 @@ class Gamma(distribution.Distribution): TypeError: if `concentration` and `rate` are different dtypes. """ parameters = locals() - with ops.name_scope(name, values=[concentration, rate]): + with ops.name_scope(name, values=[concentration, rate]) as name: with ops.control_dependencies([ check_ops.assert_positive(concentration), check_ops.assert_positive(rate), @@ -262,7 +262,7 @@ class GammaWithSoftplusConcentrationRate(Gamma): allow_nan_stats=True, name="GammaWithSoftplusConcentrationRate"): parameters = locals() - with ops.name_scope(name, values=[concentration, rate]): + with ops.name_scope(name, values=[concentration, rate]) as name: super(GammaWithSoftplusConcentrationRate, self).__init__( concentration=nn.softplus(concentration, name="softplus_concentration"), diff --git a/tensorflow/python/ops/distributions/identity_bijector.py b/tensorflow/python/ops/distributions/identity_bijector.py index 2972c3554b3639a1ae30a4167f73613b1ff8add2..8628e68f967337fb81187bae9576a168e1cd5a36 100644 --- a/tensorflow/python/ops/distributions/identity_bijector.py +++ b/tensorflow/python/ops/distributions/identity_bijector.py @@ -20,7 +20,6 @@ from __future__ import print_function from tensorflow.python.framework import constant_op from tensorflow.python.ops.distributions import bijector -from tensorflow.python.util.tf_export import tf_export __all__ = [ @@ -28,7 +27,6 @@ __all__ = [ ] -@tf_export("distributions.bijectors.Identity") class Identity(bijector.Bijector): """Compute Y = g(X) = X. @@ -37,7 +35,7 @@ class Identity(bijector.Bijector): ```python # Create the Y=g(X)=X transform which is intended for Tensors with 1 batch # ndim and 1 event ndim (i.e., vector of vectors). - identity = Identity(event_ndims=1) + identity = Identity() x = [[1., 2], [3, 4]] x == identity.forward(x) == identity.inverse(x) @@ -45,10 +43,10 @@ class Identity(bijector.Bijector): """ - def __init__(self, validate_args=False, event_ndims=0, name="identity"): + def __init__(self, validate_args=False, name="identity"): super(Identity, self).__init__( + forward_min_event_ndims=0, is_constant_jacobian=True, - event_ndims=event_ndims, validate_args=validate_args, name=name) diff --git a/tensorflow/python/ops/distributions/laplace.py b/tensorflow/python/ops/distributions/laplace.py index e98ac855c58efa1ef3ccef2de24f329d839bac26..3ccfc618d1157722bbc4ce2d5f88fbcd4fb1ba10 100644 --- a/tensorflow/python/ops/distributions/laplace.py +++ b/tensorflow/python/ops/distributions/laplace.py @@ -101,7 +101,7 @@ class Laplace(distribution.Distribution): TypeError: if `loc` and `scale` are of different dtype. """ parameters = locals() - with ops.name_scope(name, values=[loc, scale]): + with ops.name_scope(name, values=[loc, scale]) as name: with ops.control_dependencies([check_ops.assert_positive(scale)] if validate_args else []): self._loc = array_ops.identity(loc, name="loc") @@ -218,7 +218,7 @@ class LaplaceWithSoftplusScale(Laplace): allow_nan_stats=True, name="LaplaceWithSoftplusScale"): parameters = locals() - with ops.name_scope(name, values=[loc, scale]): + with ops.name_scope(name, values=[loc, scale]) as name: super(LaplaceWithSoftplusScale, self).__init__( loc=loc, scale=nn.softplus(scale, name="softplus_scale"), diff --git a/tensorflow/python/ops/distributions/multinomial.py b/tensorflow/python/ops/distributions/multinomial.py index 4ae67a009b0a4052f6e23e2e42262bb7c42f1c14..ab77f5c1f815f36edb989c5b315ef90fb198b001 100644 --- a/tensorflow/python/ops/distributions/multinomial.py +++ b/tensorflow/python/ops/distributions/multinomial.py @@ -183,7 +183,7 @@ class Multinomial(distribution.Distribution): name: Python `str` name prefixed to Ops created by this class. """ parameters = locals() - with ops.name_scope(name, values=[total_count, logits, probs]): + with ops.name_scope(name, values=[total_count, logits, probs]) as name: self._total_count = ops.convert_to_tensor(total_count, name="total_count") if validate_args: self._total_count = ( diff --git a/tensorflow/python/ops/distributions/normal.py b/tensorflow/python/ops/distributions/normal.py index 32e8a49c81bc4b23d8897639998dd33942b41a80..20d4420e91886cf84855bd585cfc4869fc674178 100644 --- a/tensorflow/python/ops/distributions/normal.py +++ b/tensorflow/python/ops/distributions/normal.py @@ -132,7 +132,7 @@ class Normal(distribution.Distribution): TypeError: if `loc` and `scale` have different `dtype`. """ parameters = locals() - with ops.name_scope(name, values=[loc, scale]): + with ops.name_scope(name, values=[loc, scale]) as name: with ops.control_dependencies([check_ops.assert_positive(scale)] if validate_args else []): self._loc = array_ops.identity(loc, name="loc") @@ -244,7 +244,7 @@ class NormalWithSoftplusScale(Normal): allow_nan_stats=True, name="NormalWithSoftplusScale"): parameters = locals() - with ops.name_scope(name, values=[scale]): + with ops.name_scope(name, values=[scale]) as name: super(NormalWithSoftplusScale, self).__init__( loc=loc, scale=nn.softplus(scale, name="softplus_scale"), diff --git a/tensorflow/python/ops/distributions/student_t.py b/tensorflow/python/ops/distributions/student_t.py index 9d9e65b4e8d6d2e40bf9c263339f899439c842c3..961b07a7bdac34b76ca1a594fa3df5e97951c76b 100644 --- a/tensorflow/python/ops/distributions/student_t.py +++ b/tensorflow/python/ops/distributions/student_t.py @@ -158,7 +158,7 @@ class StudentT(distribution.Distribution): TypeError: if loc and scale are different dtypes. """ parameters = locals() - with ops.name_scope(name, values=[df, loc, scale]): + with ops.name_scope(name, values=[df, loc, scale]) as name: with ops.control_dependencies([check_ops.assert_positive(df)] if validate_args else []): self._df = array_ops.identity(df, name="df") @@ -350,7 +350,7 @@ class StudentTWithAbsDfSoftplusScale(StudentT): allow_nan_stats=True, name="StudentTWithAbsDfSoftplusScale"): parameters = locals() - with ops.name_scope(name, values=[df, scale]): + with ops.name_scope(name, values=[df, scale]) as name: super(StudentTWithAbsDfSoftplusScale, self).__init__( df=math_ops.floor(math_ops.abs(df)), loc=loc, diff --git a/tensorflow/python/ops/distributions/transformed_distribution.py b/tensorflow/python/ops/distributions/transformed_distribution.py index 1efcf9d32e9ea9924bb080459efb7015e33ccd54..6aa6ec40d9b2ac4d156ba132e75e1846c870b772 100644 --- a/tensorflow/python/ops/distributions/transformed_distribution.py +++ b/tensorflow/python/ops/distributions/transformed_distribution.py @@ -197,8 +197,7 @@ class TransformedDistribution(distribution_lib.Distribution): distribution=ds.Normal(loc=0., scale=1.), bijector=ds.bijectors.Affine( shift=-1., - scale_identity_multiplier=2., - event_ndims=0), + scale_identity_multiplier=2.) name="NormalTransformedDistribution") ``` @@ -258,7 +257,7 @@ class TransformedDistribution(distribution_lib.Distribution): parameters = locals() name = name or (("" if bijector is None else bijector.name) + distribution.name) - with ops.name_scope(name, values=[event_shape, batch_shape]): + with ops.name_scope(name, values=[event_shape, batch_shape]) as name: # For convenience we define some handy constants. self._zero = constant_op.constant(0, dtype=dtypes.int32, name="zero") self._empty = constant_op.constant([], dtype=dtypes.int32, name="empty") @@ -419,48 +418,51 @@ class TransformedDistribution(distribution_lib.Distribution): # For caching to work, it is imperative that the bijector is the first to # modify the input. x = self.bijector.inverse(y) - ildj = self.bijector.inverse_log_det_jacobian(y) + event_ndims = self._maybe_get_event_ndims_statically() + + ildj = self.bijector.inverse_log_det_jacobian(y, event_ndims=event_ndims) if self.bijector._is_injective: # pylint: disable=protected-access - return self._finish_log_prob_for_one_fiber(y, x, ildj) + return self._finish_log_prob_for_one_fiber(y, x, ildj, event_ndims) lp_on_fibers = [ - self._finish_log_prob_for_one_fiber(y, x_i, ildj_i) + self._finish_log_prob_for_one_fiber(y, x_i, ildj_i, event_ndims) for x_i, ildj_i in zip(x, ildj)] return math_ops.reduce_logsumexp(array_ops.stack(lp_on_fibers), axis=0) - def _finish_log_prob_for_one_fiber(self, y, x, ildj): + def _finish_log_prob_for_one_fiber(self, y, x, ildj, event_ndims): """Finish computation of log_prob on one element of the inverse image.""" x = self._maybe_rotate_dims(x, rotate_right=True) log_prob = self.distribution.log_prob(x) if self._is_maybe_event_override: log_prob = math_ops.reduce_sum(log_prob, self._reduce_event_indices) log_prob += math_ops.cast(ildj, log_prob.dtype) - if self._is_maybe_event_override: + if self._is_maybe_event_override and isinstance(event_ndims, int): log_prob.set_shape(array_ops.broadcast_static_shape( - y.get_shape().with_rank_at_least(1)[:-1], self.batch_shape)) + x.get_shape().with_rank_at_least(1)[:-event_ndims], self.batch_shape)) return log_prob def _prob(self, y): x = self.bijector.inverse(y) - ildj = self.bijector.inverse_log_det_jacobian(y) + event_ndims = self._maybe_get_event_ndims_statically() + ildj = self.bijector.inverse_log_det_jacobian(y, event_ndims=event_ndims) if self.bijector._is_injective: # pylint: disable=protected-access - return self._finish_prob_for_one_fiber(y, x, ildj) + return self._finish_prob_for_one_fiber(y, x, ildj, event_ndims) prob_on_fibers = [ - self._finish_prob_for_one_fiber(y, x_i, ildj_i) + self._finish_prob_for_one_fiber(y, x_i, ildj_i, event_ndims) for x_i, ildj_i in zip(x, ildj)] return sum(prob_on_fibers) - def _finish_prob_for_one_fiber(self, y, x, ildj): + def _finish_prob_for_one_fiber(self, y, x, ildj, event_ndims): """Finish computation of prob on one element of the inverse image.""" x = self._maybe_rotate_dims(x, rotate_right=True) prob = self.distribution.prob(x) if self._is_maybe_event_override: prob = math_ops.reduce_prod(prob, self._reduce_event_indices) prob *= math_ops.exp(math_ops.cast(ildj, prob.dtype)) - if self._is_maybe_event_override: + if self._is_maybe_event_override and isinstance(event_ndims, int): prob.set_shape(array_ops.broadcast_static_shape( - y.get_shape().with_rank_at_least(1)[:-1], self.batch_shape)) + y.get_shape().with_rank_at_least(1)[:-event_ndims], self.batch_shape)) return prob def _log_cdf(self, y): @@ -545,10 +547,17 @@ class TransformedDistribution(distribution_lib.Distribution): _ones_like(self.distribution.batch_shape_tensor()) ], 0) entropy = array_ops.tile(entropy, multiples) - dummy = array_ops.zeros([], self.dtype) - entropy -= math_ops.cast( - self.bijector.inverse_log_det_jacobian(dummy), - entropy.dtype) + dummy = array_ops.zeros( + shape=array_ops.concat( + [self.batch_shape_tensor(), self.event_shape_tensor()], + 0), + dtype=self.dtype) + event_ndims = (self.event_shape.ndims if self.event_shape.ndims is not None + else array_ops.size(self.event_shape_tensor())) + ildj = self.bijector.inverse_log_det_jacobian( + dummy, event_ndims=event_ndims) + + entropy -= math_ops.cast(ildj, entropy.dtype) entropy.set_shape(self.batch_shape) return entropy @@ -610,3 +619,16 @@ class TransformedDistribution(distribution_lib.Distribution): n = (ndims - self._rotate_ndims) if rotate_right else self._rotate_ndims return array_ops.transpose( x, _concat_vectors(math_ops.range(n, ndims), math_ops.range(0, n))) + + def _maybe_get_event_ndims_statically(self): + if self.event_shape.ndims is not None: + return self.event_shape.ndims + + event_ndims = array_ops.size(self.event_shape_tensor()) + + static_event_ndims = tensor_util.constant_value(event_ndims) + + if static_event_ndims is not None: + return static_event_ndims + + return event_ndims diff --git a/tensorflow/python/ops/distributions/uniform.py b/tensorflow/python/ops/distributions/uniform.py index 0891bffdd554828586c5b23919f955f685632694..087797c653bb3ee294b747dd7da8ff2896e79fc6 100644 --- a/tensorflow/python/ops/distributions/uniform.py +++ b/tensorflow/python/ops/distributions/uniform.py @@ -103,7 +103,7 @@ class Uniform(distribution.Distribution): InvalidArgumentError: if `low >= high` and `validate_args=False`. """ parameters = locals() - with ops.name_scope(name, values=[low, high]): + with ops.name_scope(name, values=[low, high]) as name: with ops.control_dependencies([ check_ops.assert_less( low, high, message="uniform not defined when low >= high.") diff --git a/tensorflow/python/ops/distributions/util.py b/tensorflow/python/ops/distributions/util.py index 0fe6aa30f945dc7682a53fa6495823288cf111b7..2e067eab459050e30d220bdb7ff0d65cb9c552f7 100644 --- a/tensorflow/python/ops/distributions/util.py +++ b/tensorflow/python/ops/distributions/util.py @@ -58,8 +58,7 @@ def assert_close( if data is None: data = [ message, - "Condition x ~= y did not hold element-wise: x = ", x.name, x, "y = ", - y.name, y + "Condition x ~= y did not hold element-wise: x = ", x, "y = ", y ] if x.dtype.is_integer: @@ -95,7 +94,7 @@ def assert_integer_form( x = ops.convert_to_tensor(x, name="x") if x.dtype.is_integer: return control_flow_ops.no_op() - message = message or "{} has non-integer components".format(x.op.name) + message = message or "{} has non-integer components".format(x) if int_dtype is None: try: int_dtype = { @@ -123,13 +122,13 @@ def embed_check_nonnegative_integer_form( x = ops.convert_to_tensor(x, name="x") assertions = [ check_ops.assert_non_negative( - x, message="'{}' must be non-negative.".format(x.op.name)), + x, message="'{}' must be non-negative.".format(x)), ] if not x.dtype.is_integer: assertions += [ assert_integer_form( x, message="'{}' cannot contain fractional components.".format( - x.op.name)), + x)), ] return control_flow_ops.with_dependencies(assertions, x) @@ -434,7 +433,7 @@ def embed_check_integer_casting_closed( and not _is_integer_like_by_dtype(target_dtype)): raise TypeError("At least one of {}.dtype ({}) and target_dtype ({}) " "must be integer-type.".format( - x.op.name, x.dtype.name, target_dtype.name)) + x, x.dtype.name, target_dtype.name)) assertions = [] if assert_nonnegative: @@ -683,7 +682,7 @@ def pick_vector(cond, cond = ops.convert_to_tensor(cond, name="cond") if cond.dtype != dtypes.bool: raise TypeError("%s.dtype=%s which is not %s" % - (cond.name, cond.dtype, dtypes.bool)) + (cond, cond.dtype, dtypes.bool)) cond_value_static = tensor_util.constant_value(cond) if cond_value_static is not None: return true_vector if cond_value_static else false_vector @@ -692,8 +691,8 @@ def pick_vector(cond, if true_vector.dtype != false_vector.dtype: raise TypeError( "%s.dtype=%s does not match %s.dtype=%s" - % (true_vector.name, true_vector.dtype, - false_vector.name, false_vector.dtype)) + % (true_vector, true_vector.dtype, + false_vector, false_vector.dtype)) n = array_ops.shape(true_vector)[0] return array_ops.slice( array_ops.concat([true_vector, false_vector], 0), diff --git a/tensorflow/python/ops/embedding_ops.py b/tensorflow/python/ops/embedding_ops.py index f0120f2957db12caf6a513fde9aa8c756aff8bad..9e46739bc1b6899b9b1fea898254f6e61fb7b8ba 100644 --- a/tensorflow/python/ops/embedding_ops.py +++ b/tensorflow/python/ops/embedding_ops.py @@ -331,11 +331,11 @@ def embedding_lookup_sparse(params, representing sharded embedding tensors. Alternatively, a `PartitionedVariable`, created by partitioning along dimension 0. Each element must be appropriately sized for the given `partition_strategy`. - sp_ids: N x M SparseTensor of int64 ids (typically from FeatureValueToId), + sp_ids: N x M `SparseTensor` of int64 ids (typically from FeatureValueToId), where N is typically batch size and M is arbitrary. - sp_weights: either a SparseTensor of float / double weights, or None to - indicate all weights should be taken to be 1. If specified, sp_weights - must have exactly the same shape and indices as sp_ids. + sp_weights: either a `SparseTensor` of float / double weights, or `None` to + indicate all weights should be taken to be 1. If specified, `sp_weights` + must have exactly the same shape and indices as `sp_ids`. partition_strategy: A string specifying the partitioning strategy, relevant if `len(params) > 1`. Currently `"div"` and `"mod"` are supported. Default is `"mod"`. See `tf.nn.embedding_lookup` for more details. @@ -351,39 +351,43 @@ def embedding_lookup_sparse(params, Returns: A dense tensor representing the combined embeddings for the - sparse ids. For each row in the dense tensor represented by sp_ids, the op + sparse ids. For each row in the dense tensor represented by `sp_ids`, the op looks up the embeddings for all ids in that row, multiplies them by the corresponding weight, and combines these embeddings as specified. In other words, if - shape(combined params) = [p0, p1, ..., pm] + `shape(combined params) = [p0, p1, ..., pm]` and - shape(sp_ids) = shape(sp_weights) = [d0, d1, ..., dn] + `shape(sp_ids) = shape(sp_weights) = [d0, d1, ..., dn]` then - shape(output) = [d0, d1, ..., dn-1, p1, ..., pm]. + `shape(output) = [d0, d1, ..., dn-1, p1, ..., pm]`. For instance, if params is a 10x20 matrix, and sp_ids / sp_weights are + ```python [0, 0]: id 1, weight 2.0 [0, 1]: id 3, weight 0.5 [1, 0]: id 0, weight 1.0 [2, 3]: id 1, weight 3.0 + ``` with `combiner`="mean", then the output will be a 3x20 matrix where + ```python output[0, :] = (params[1, :] * 2.0 + params[3, :] * 0.5) / (2.0 + 0.5) 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 - None nor SparseTensor. - ValueError: If combiner is not one of {"mean", "sqrtn", "sum"}. + TypeError: If `sp_ids` is not a `SparseTensor`, or if `sp_weights` is + neither `None` nor `SparseTensor`. + ValueError: If `combiner` is not one of {"mean", "sqrtn", "sum"}. """ if combiner is None: logging.warn("The default value of combiner will change from \"mean\" " diff --git a/tensorflow/python/ops/gradients_test.py b/tensorflow/python/ops/gradients_test.py index c94f1396b28e2124c6e5123cf711ac86abf174ab..0603d3b6706b960a0fa9d9b33d383dd0c9063780 100644 --- a/tensorflow/python/ops/gradients_test.py +++ b/tensorflow/python/ops/gradients_test.py @@ -44,6 +44,7 @@ from tensorflow.python.ops import gradients_impl from tensorflow.python.ops import math_grad # pylint: disable=unused-import from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn_grad # pylint: disable=unused-import +from tensorflow.python.ops import resource_variable_ops from tensorflow.python.ops import state_grad # pylint: disable=unused-import from tensorflow.python.ops import tensor_array_grad # pylint: disable=unused-import from tensorflow.python.ops import tensor_array_ops @@ -810,5 +811,29 @@ class OnlyRealGradientsTest(test_util.TensorFlowTestCase): gradients.gradients(y, x) +class ResourceCondTest(test_util.TensorFlowTestCase): + + def testBasic(self): + gamma = resource_variable_ops.ResourceVariable( + np.random.random((3,)), + dtype="float32", name="gamma") + + inputs = array_ops.ones(shape=(3,), dtype="float32") + + def TestFn(): + output = inputs + gamma + return output + + training = array_ops.placeholder_with_default(True, shape=()) + output = control_flow_ops.cond( + training, TestFn, lambda: inputs) + + loss = output + + grads = gradients.gradients( + loss, [gamma]) + self.assertTrue(None not in grads) + + if __name__ == "__main__": googletest.main() diff --git a/tensorflow/python/ops/histogram_ops.py b/tensorflow/python/ops/histogram_ops.py index 4a1ef54fb50013881aa832f83674ac66ecccd9bc..ec38d89a0ec04469e14e2f1f99b98a51d0a9f513 100644 --- a/tensorflow/python/ops/histogram_ops.py +++ b/tensorflow/python/ops/histogram_ops.py @@ -32,7 +32,6 @@ 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') diff --git a/tensorflow/python/ops/image_ops_impl.py b/tensorflow/python/ops/image_ops_impl.py index 3369fe3c9b37ca05311c5548dbfa3228ba04ee80..bd5b2ae83b5dd16e15dcf87de9714c082fa10ef9 100644 --- a/tensorflow/python/ops/image_ops_impl.py +++ b/tensorflow/python/ops/image_ops_impl.py @@ -269,17 +269,7 @@ def random_flip_up_down(image, seed=None): Raises: ValueError: if the shape of `image` not supported. """ - with ops.name_scope(None, 'random_flip_up_down', [image]) as scope: - image = ops.convert_to_tensor(image, name='image') - image = _Assert3DImage(image) - uniform_random = random_ops.random_uniform([], 0, 1.0, seed=seed) - mirror_cond = math_ops.less(uniform_random, .5) - result = control_flow_ops.cond( - mirror_cond, - lambda: array_ops.reverse(image, [0]), - lambda: image, - name=scope) - return fix_image_flip_shape(image, result) + return _random_flip(image, 0, seed, 'random_flip_up_down') @tf_export('image.random_flip_left_right') @@ -301,14 +291,34 @@ def random_flip_left_right(image, seed=None): Raises: ValueError: if the shape of `image` not supported. """ - with ops.name_scope(None, 'random_flip_left_right', [image]) as scope: + return _random_flip(image, 1, seed, 'random_flip_left_right') + + +def _random_flip(image, flip_index, seed, scope_name): + """Randomly (50% chance) flip an image along axis `flip_index`. + Args: + image: A 3-D tensor of shape `[height, width, channels].` + flip_index: The dimension along which to flip the image. + Vertical: 0, Horizontal: 1 + seed: A Python integer. Used to create a random seed. See + @{tf.set_random_seed} + for behavior. + scope_name: Name of the scope in which the ops are added. + + Returns: + A 3-D tensor of the same type and shape as `image`. + + Raises: + ValueError: if the shape of `image` not supported. + """ + with ops.name_scope(None, scope_name, [image]) as scope: image = ops.convert_to_tensor(image, name='image') image = _Assert3DImage(image) uniform_random = random_ops.random_uniform([], 0, 1.0, seed=seed) mirror_cond = math_ops.less(uniform_random, .5) result = control_flow_ops.cond( mirror_cond, - lambda: array_ops.reverse(image, [1]), + lambda: array_ops.reverse(image, [flip_index]), lambda: image, name=scope) return fix_image_flip_shape(image, result) @@ -332,16 +342,7 @@ def flip_left_right(image): Raises: ValueError: if the shape of `image` not supported. """ - with ops.name_scope(None, 'flip_left_right', [image]): - image = ops.convert_to_tensor(image, name='image') - 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.') + return _flip(image, 1, 'flip_left_right') @tf_export('image.flip_up_down') @@ -362,14 +363,35 @@ def flip_up_down(image): Raises: ValueError: if the shape of `image` not supported. """ - with ops.name_scope(None, 'flip_up_down', [image]): + return _flip(image, 0, 'flip_up_down') + + +def _flip(image, flip_index, scope_name): + """Flip an image either horizontally or vertically. + + Outputs the contents of `image` flipped along the dimension `flip_index`. + + See also `reverse()`. + + Args: + image: 4-D Tensor of shape `[batch, height, width, channels]` or + 3-D Tensor of shape `[height, width, channels]`. + flip_index: 0 For vertical, 1 for horizontal. + + Returns: + A tensor of the same type and shape as `image`. + + Raises: + ValueError: if the shape of `image` not supported. + """ + with ops.name_scope(None, scope_name, [image]): image = ops.convert_to_tensor(image, name='image') 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])) + return fix_image_flip_shape(image, array_ops.reverse(image, [flip_index])) elif shape.ndims == 4: - return array_ops.reverse(image, [1]) + return array_ops.reverse(image, [flip_index+1]) else: raise ValueError('\'image\' must have either 3 or 4 dimensions.') @@ -630,7 +652,7 @@ def pad_to_bounding_box(image, offset_height, offset_width, target_height, padded.set_shape(padded_shape) if not is_batch: - padded = array_ops.squeeze(padded, squeeze_dims=[0]) + padded = array_ops.squeeze(padded, axis=[0]) return padded @@ -710,7 +732,7 @@ def crop_to_bounding_box(image, offset_height, offset_width, target_height, cropped.set_shape(cropped_shape) if not is_batch: - cropped = array_ops.squeeze(cropped, squeeze_dims=[0]) + cropped = array_ops.squeeze(cropped, axis=[0]) return cropped @@ -827,7 +849,7 @@ def resize_image_with_crop_or_pad(image, target_height, target_width): resized = control_flow_ops.with_dependencies(assert_ops, resized) if not is_batch: - resized = array_ops.squeeze(resized, squeeze_dims=[0]) + resized = array_ops.squeeze(resized, axis=[0]) return resized @@ -920,7 +942,7 @@ def resize_images(images, for x in [new_width_const, width, new_height_const, height]) and ( width == new_width_const and height == new_height_const): if not is_batch: - images = array_ops.squeeze(images, squeeze_dims=[0]) + images = array_ops.squeeze(images, axis=[0]) return images if method == ResizeMethod.BILINEAR: @@ -943,7 +965,7 @@ def resize_images(images, images.set_shape([None, new_height_const, new_width_const, None]) if not is_batch: - images = array_ops.squeeze(images, squeeze_dims=[0]) + images = array_ops.squeeze(images, axis=[0]) return images diff --git a/tensorflow/python/ops/init_ops.py b/tensorflow/python/ops/init_ops.py index 9dfe5ffbf42bcfc9657739b6fc6ad1f3c4823a7d..f93bf0a17f31b4451cad9c5bc525d9f237f97bef 100644 --- a/tensorflow/python/ops/init_ops.py +++ b/tensorflow/python/ops/init_ops.py @@ -39,10 +39,10 @@ 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 linalg_ops +from tensorflow.python.ops import linalg_ops_impl +from tensorflow.python.ops import gen_linalg_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops -from tensorflow.python.ops import random_ops from tensorflow.python.util.deprecation import deprecated from tensorflow.python.util.tf_export import tf_export @@ -499,10 +499,10 @@ class Orthogonal(Initializer): Args: gain: multiplicative factor to apply to the orthogonal matrix - dtype: The type of the output. seed: A Python integer. Used to create random seeds. See @{tf.set_random_seed} for behavior. + dtype: The data type. """ def __init__(self, gain=1.0, seed=None, dtype=dtypes.float32): @@ -529,7 +529,7 @@ class Orthogonal(Initializer): # Generate a random matrix a = random_ops.random_normal(flat_shape, dtype=dtype, seed=self.seed) # Compute the qr factorization - q, r = linalg_ops.qr(a, full_matrices=False) + q, r = gen_linalg_ops.qr(a, full_matrices=False) # Make Q uniform d = array_ops.diag_part(r) q *= math_ops.sign(d) @@ -549,13 +549,12 @@ class ConvolutionDeltaOrthogonal(Initializer): 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. + 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. + @{tf.set_random_seed} for behavior. + dtype: The data type. """ def __init__(self, gain=1.0, seed=None, dtype=dtypes.float32): @@ -578,10 +577,9 @@ class ConvolutionDeltaOrthogonal(Initializer): a = random_ops.random_normal([shape[-1], shape[-1]], dtype=dtype, seed=self.seed) # Compute the qr factorization - q, r = linalg_ops.qr(a, full_matrices=False) + q, r = gen_linalg_ops.qr(a, full_matrices=False) # Make Q uniform d = array_ops.diag_part(r) - # ph = d / math_ops.abs(d) q *= math_ops.sign(d) q = q[:shape[-2], :] q *= math_ops.sqrt(math_ops.cast(self.gain, dtype=dtype)) @@ -601,6 +599,469 @@ class ConvolutionDeltaOrthogonal(Initializer): return {"gain": self.gain, "seed": self.seed, "dtype": self.dtype.name} +class ConvolutionOrthogonal(Initializer): + """Initializer that generates orthogonal kernel for ConvNets. + + Base class used to construct 1D, 2D and 3D orthogonal kernels for convolution. + + 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. + seed: A Python integer. Used to create random seeds. See + @{tf.set_random_seed} for behavior. + dtype: The data type. + """ + + 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): + raise NotImplementedError + + def get_config(self): + return {"gain": self.gain, "seed": self.seed, "dtype": self.dtype.name} + + # Helper functions. + def _orthogonal_matrix(self, n): + """Construct an n x n orthogonal matrix. + + Args: + n: Dimension. + Returns: + A n x n orthogonal matrix. + """ + a = random_ops.random_normal([n, n], dtype=self.dtype, seed=self.seed) + if self.seed: + self.seed += 1 + q, r = gen_linalg_ops.qr(a) + d = array_ops.diag_part(r) + # make q uniform + q *= math_ops.sign(d) + return q + + def _symmetric_projection(self, n): + """Compute a n x n symmetric projection matrix. + + Args: + n: Dimension. + Returns: + A n x n symmetric projection matrix, i.e. a matrix P s.t. P=P*P, P=P^T. + """ + q = self._orthogonal_matrix(n) + # randomly zeroing out some columns + mask = math_ops.cast(random_ops.random_normal([n], seed=self.seed) > 0, + self.dtype) + if self.seed: + self.seed += 1 + c = math_ops.multiply(q, mask) + return math_ops.matmul(c, array_ops.matrix_transpose(c)) + + +class ConvolutionOrthogonal2D(ConvolutionOrthogonal): + """Initializer that generates a 2D orthogonal kernel for ConvNets. + + The shape of the tensor must have length 4. The number of input + filters must not exceed the number of output filters. + The orthogonality(==isometry) is exact when the inputs are circular padded. + There are finite-width effects with non-circular padding (e.g. zero padding). + + Args: + gain: Multiplicative factor to apply to the orthogonal matrix. Default is 1. + This has the effect of scaling the output 2-norm by a factor of + `sqrt(gain)`. + seed: A Python integer. Used to create random seeds. See + @{tf.set_random_seed} for behavior. + dtype: The data type. + """ + + def __call__(self, shape, dtype=None, partition_info=None): + if dtype is None: + dtype = self.dtype + if len(shape) != 4: + raise ValueError("The tensor to initialize must be four-dimensional") + + if shape[-2] > shape[-1]: + raise ValueError("In_filters cannot be greater than out_filters.") + + if shape[0] != shape[1]: + raise ValueError("Kernel sizes must be equal.") + + kernel = self._orthogonal_kernel(shape[0], shape[2], shape[3]) + kernel *= math_ops.sqrt(math_ops.cast(self.gain, dtype=dtype)) + return kernel + + def _dict_to_tensor(self, x, k1, k2): + """Convert a dictionary to a tensor. + + Args: + x: A k1 * k2 dictionary. + k1: First dimension of x. + k2: Second dimension of x. + Returns: + A k1 * k2 tensor. + """ + + return array_ops.stack([array_ops.stack([x[i, j] for j in range(k2)]) + for i in range(k1)]) + + def _block_orth(self, p1, p2): + """Construct a 2 x 2 kernel. Used to construct orthgonal kernel. + + Args: + p1: A symmetric projection matrix. + p2: A symmetric projection matrix. + Returns: + A 2 x 2 kernel [[p1p2, p1(1-p2)], + [(1-p1)p2, (1-p1)(1-p2)]]. + Raises: + ValueError: If the dimensions of p1 and p2 are different. + """ + if p1.shape.as_list() != p2.shape.as_list(): + raise ValueError("The dimension of the matrices must be the same.") + n = p1.shape.as_list()[0] + kernel2x2 = {} + eye = linalg_ops_impl.eye(n, dtype=self.dtype) + kernel2x2[0, 0] = math_ops.matmul(p1, p2) + kernel2x2[0, 1] = math_ops.matmul(p1, (eye - p2)) + kernel2x2[1, 0] = math_ops.matmul((eye - p1), p2) + kernel2x2[1, 1] = math_ops.matmul((eye - p1), (eye - p2)) + + return kernel2x2 + + def _matrix_conv(self, m1, m2): + """Matrix convolution. + + Args: + m1: A k x k dictionary, each element is a n x n matrix. + m2: A l x l dictionary, each element is a n x n matrix. + + Returns: + (k + l - 1) * (k + l - 1) dictionary each element is a n x n matrix. + Raises: + ValueError: if the entries of m1 and m2 are of different dimensions. + """ + + n = (m1[0, 0]).shape.as_list()[0] + if n != (m2[0, 0]).shape.as_list()[0]: + raise ValueError("The entries in matrices m1 and m2 " + "must have the same dimensions!") + k = int(np.sqrt(len(m1))) + l = int(np.sqrt(len(m2))) + result = {} + size = k + l - 1 + # Compute matrix convolution between m1 and m2. + for i in range(size): + for j in range(size): + result[i, j] = array_ops.zeros([n, n], self.dtype) + for index1 in range(min(k, i + 1)): + for index2 in range(min(k, j + 1)): + if (i - index1) < l and (j - index2) < l: + result[i, j] += math_ops.matmul(m1[index1, index2], + m2[i - index1, j - index2]) + return result + + def _orthogonal_kernel(self, ksize, cin, cout): + """Construct orthogonal kernel for convolution. + + Args: + ksize: Kernel size. + cin: Number of input channels. + cout: Number of output channels. + Returns: + An [ksize, ksize, cin, cout] orthogonal kernel. + Raises: + ValueError: If cin > cout. + """ + if cin > cout: + raise ValueError("The number of input channels cannot exceed " + "the number of output channels.") + orth = self._orthogonal_matrix(cout)[0:cin, :] + if ksize == 1: + return array_ops.expand_dims(array_ops.expand_dims(orth, 0), 0) + + p = self._block_orth(self._symmetric_projection(cout), + self._symmetric_projection(cout)) + for _ in range(ksize - 2): + temp = self._block_orth(self._symmetric_projection(cout), + self._symmetric_projection(cout)) + p = self._matrix_conv(p, temp) + for i in range(ksize): + for j in range(ksize): + p[i, j] = math_ops.matmul(orth, p[i, j]) + + return self._dict_to_tensor(p, ksize, ksize) + + +class ConvolutionOrthogonal1D(ConvolutionOrthogonal): + """Initializer that generates a 1D orthogonal kernel for ConvNets. + + The shape of the tensor must have length 3. The number of input + filters must not exceed the number of output filters. + The orthogonality(==isometry) is exact when the inputs are circular padded. + There are finite-width effects with non-circular padding (e.g. zero padding). + + 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. + seed: A Python integer. Used to create random seeds. See + @{tf.set_random_seed} + for behavior. + dtype: The data type. + """ + + def __call__(self, shape, dtype=None, partition_info=None): + if dtype is None: + dtype = self.dtype + if len(shape) != 3: + raise ValueError("The tensor to initialize must be three-dimensional") + + if shape[-2] > shape[-1]: + raise ValueError("In_filters cannot be greater than out_filters.") + + kernel = self._orthogonal_kernel(shape[0], shape[-2], shape[-1]) + kernel *= math_ops.sqrt(math_ops.cast(self.gain, dtype=dtype)) + return kernel + + def _dict_to_tensor(self, x, k): + """Convert a dictionary to a tensor. + + Args: + x: A dictionary of length k. + k: Dimension of x. + Returns: + A tensor with the same dimension. + """ + + return array_ops.stack([x[i] for i in range(k)]) + + def _block_orth(self, projection_matrix): + """Construct a kernel. Used to construct orthgonal kernel. + + Args: + projection_matrix: A symmetric projection matrix of size n x n. + Returns: + [projection_matrix, (1 - projection_matrix)]. + """ + n = projection_matrix.shape.as_list()[0] + kernel = {} + eye = linalg_ops_impl.eye(n, dtype=self.dtype) + kernel[0] = projection_matrix + kernel[1] = eye - projection_matrix + return kernel + + def _matrix_conv(self, m1, m2): + """Matrix convolution. + + Args: + m1: A dictionary of length k, each element is a n x n matrix. + m2: A dictionary of length l, each element is a n x n matrix. + + Returns: + (k + l - 1) dictionary each element is a n x n matrix. + Raises: + ValueError: Ff the entries of m1 and m2 are of different dimensions. + """ + + n = (m1[0]).shape.as_list()[0] + if n != (m2[0]).shape.as_list()[0]: + raise ValueError("The entries in matrices m1 and m2 " + "must have the same dimensions!") + k = len(m1) + l = len(m2) + result = {} + size = k + l - 1 + # Compute matrix convolution between m1 and m2. + for i in range(size): + result[i] = array_ops.zeros([n, n], self.dtype) + for index in range(min(k, i + 1)): + if (i - index) < l: + result[i] += math_ops.matmul(m1[index], m2[i - index]) + return result + + def _orthogonal_kernel(self, ksize, cin, cout): + """Construct orthogonal kernel for convolution. + + Args: + ksize: Kernel size. + cin: Number of input channels. + cout: Number of output channels. + Returns: + An [ksize, ksize, cin, cout] orthogonal kernel. + Raises: + ValueError: If cin > cout. + """ + if cin > cout: + raise ValueError("The number of input channels cannot exceed " + "the number of output channels.") + orth = self._orthogonal_matrix(cout)[0:cin, :] + if ksize == 1: + return array_ops.expand_dims(orth, 0) + + p = self._block_orth(self._symmetric_projection(cout)) + for _ in range(ksize - 2): + temp = self._block_orth(self._symmetric_projection(cout)) + p = self._matrix_conv(p, temp) + for i in range(ksize): + p[i] = math_ops.matmul(orth, p[i]) + + return self._dict_to_tensor(p, ksize) + + +class ConvolutionOrthogonal3D(ConvolutionOrthogonal): + """Initializer that generates a 3D orthogonal kernel for ConvNets. + + The shape of the tensor must have length 5. The number of input + filters must not exceed the number of output filters. + The orthogonality(==isometry) is exact when the inputs are circular padded. + There are finite-width effects with non-circular padding (e.g. zero padding). + + 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. + seed: A Python integer. Used to create random seeds. See + @{tf.set_random_seed} for behavior. + dtype: The data type. + """ + + def __call__(self, shape, dtype=None, partition_info=None): + if dtype is None: + dtype = self.dtype + if len(shape) != 5: + raise ValueError("The tensor to initialize must be five-dimensional") + + if shape[-2] > shape[-1]: + raise ValueError("In_filters cannot be greater than out_filters.") + + if shape[0] != shape[1] or shape[0] != shape[2]: + raise ValueError("Kernel sizes must be equal.") + + kernel = self._orthogonal_kernel(shape[0], shape[-2], shape[-1]) + kernel *= math_ops.sqrt(math_ops.cast(self.gain, dtype=dtype)) + return kernel + + def _dict_to_tensor(self, x, k1, k2, k3): + """Convert a dictionary to a tensor. + + Args: + x: A k1 * k2 dictionary. + k1: First dimension of x. + k2: Second dimension of x. + k3: Third dimension of x. + Returns: + A k1 * k2 * k3 tensor. + """ + + return array_ops.stack([array_ops.stack( + [array_ops.stack([x[i, j, k] for k in range(k3)]) + for j in range(k2)]) for i in range(k1)]) + + def _block_orth(self, p1, p2, p3): + """Construct a 3 x 3 kernel. Used to construct orthgonal kernel. + + Args: + p1: A symmetric projection matrix. + p2: A symmetric projection matrix. + p3: A symmetric projection matrix. + Returns: + A 2 x 2 x 2 kernel. + Raises: + ValueError: If the dimensions of p1, p2 and p3 are different. + """ + p1_shape = p1.shape.as_list() + if p1_shape != p2.shape.as_list() or p1_shape != p3.shape.as_list(): + raise ValueError("The dimension of the matrices must be the same.") + n = p1_shape[0] + eye = linalg_ops_impl.eye(n, dtype=self.dtype) + kernel2x2x2 = {} + def matmul(p1, p2, p3): + return math_ops.matmul(math_ops.matmul(p1, p2), p3) + def cast(i, p): + """Return p or (1-p).""" + return i * p + (1-i) * (eye - p) + for i in [0, 1]: + for j in [0, 1]: + for k in [0, 1]: + kernel2x2x2[i, j, k] = matmul(cast(i, p1), cast(j, p2), cast(k, p3)) + return kernel2x2x2 + + def _matrix_conv(self, m1, m2): + """Matrix convolution. + + Args: + m1: is a k x k x k dictionary, each element is a n x n matrix. + m2: is a l x l x l dictionary, each element is a n x n matrix. + + Returns: + (k + l - 1) x (k + l - 1) x (k + l - 1) dictionary each + element is a n x n matrix. + Raises: + ValueError: if the entries of m1 and m2 are of different dimensions. + """ + + n = (m1[0, 0, 0]).shape.as_list()[0] + if n != (m2[0, 0, 0]).shape.as_list()[0]: + raise ValueError("The entries in matrices m1 and m2 " + "must have the same dimensions!") + k = int(np.cbrt(len(m1))) + l = int(np.cbrt(len(m2))) + result = {} + size = k + l - 1 + # Compute matrix convolution between m1 and m2. + for i in range(size): + for j in range(size): + for r in range(size): + result[i, j, r] = array_ops.zeros([n, n], self.dtype) + for index1 in range(min(k, i + 1)): + for index2 in range(min(k, j + 1)): + for index3 in range(min(k, r + 1)): + if (i - index1) < l and (j - index2) < l and (r - index3) < l: + result[i, j, r] += math_ops.matmul(m1[index1, index2, index3], + m2[i - index1, j - index2, + r - index3]) + return result + + def _orthogonal_kernel(self, ksize, cin, cout): + """Construct orthogonal kernel for convolution. + + Args: + ksize: Kernel size. + cin: Number of input channels. + cout: Number of output channels. + Returns: + An [ksize, ksize, ksize, cin, cout] orthogonal kernel. + Raises: + ValueError: If cin > cout. + """ + if cin > cout: + raise ValueError("The number of input channels cannot exceed " + "the number of output channels.") + orth = self._orthogonal_matrix(cout)[0:cin, :] + if ksize == 1: + return array_ops.expand_dims( + array_ops.expand_dims( + array_ops.expand_dims(orth, 0), 0), 0) + + p = self._block_orth(self._symmetric_projection(cout), + self._symmetric_projection(cout), + self._symmetric_projection(cout)) + for _ in range(ksize - 2): + temp = self._block_orth(self._symmetric_projection(cout), + self._symmetric_projection(cout), + self._symmetric_projection(cout)) + p = self._matrix_conv(p, temp) + for i in range(ksize): + for j in range(ksize): + for k in range(ksize): + p[i, j, k] = math_ops.matmul(orth, p[i, j, k]) + + return self._dict_to_tensor(p, ksize, ksize, ksize) + + @tf_export("keras.initializers.Identity", "initializers.identity") class Identity(Initializer): """Initializer that generates the identity matrix. @@ -623,7 +1084,7 @@ class Identity(Initializer): "Identity matrix initializer can only be used for 2D matrices.") if dtype is None: dtype = self.dtype - initializer = linalg_ops.eye(*full_shape, dtype=dtype) + initializer = linalg_ops_impl.eye(*full_shape, dtype=dtype) if partition_info is not None: initializer = array_ops.slice(initializer, partition_info.var_offset, shape) @@ -646,6 +1107,9 @@ variance_scaling_initializer = VarianceScaling orthogonal_initializer = Orthogonal identity_initializer = Identity convolutional_delta_orthogonal = ConvolutionDeltaOrthogonal +convolutional_orthogonal_1d = ConvolutionOrthogonal1D +convolutional_orthogonal_2d = ConvolutionOrthogonal2D +convolutional_orthogonal_3d = ConvolutionOrthogonal3D # pylint: enable=invalid-name diff --git a/tensorflow/python/ops/linalg/linear_operator.py b/tensorflow/python/ops/linalg/linear_operator.py index 193c787baa2ac68feec7e5d8bb03b251fc78d781..8cfe964b1c0a572f43a14c66885e74ea105b0916 100644 --- a/tensorflow/python/ops/linalg/linear_operator.py +++ b/tensorflow/python/ops/linalg/linear_operator.py @@ -699,9 +699,10 @@ class LinearOperator(object): " Requires conversion to a dense matrix and O(N^3) operations.") rhs = linalg.adjoint(rhs) if adjoint_arg else rhs if self._can_use_cholesky(): - return linalg_ops.cholesky_solve( + return linear_operator_util.cholesky_solve_with_broadcast( linalg_ops.cholesky(self.to_dense()), rhs) - return linalg_ops.matrix_solve(self.to_dense(), rhs, adjoint=adjoint) + return linear_operator_util.matrix_solve_with_broadcast( + self.to_dense(), rhs, adjoint=adjoint) def solve(self, rhs, adjoint=False, adjoint_arg=False, name="solve"): """Solve (exact or approx) `R` (batch) systems of equations: `A X = rhs`. diff --git a/tensorflow/python/ops/linalg/linear_operator_full_matrix.py b/tensorflow/python/ops/linalg/linear_operator_full_matrix.py index 5ba3b090ae9decaba239b31226db84c2d7b254bd..746da8df1ce957e86bc2e730b5709a699adbf612 100644 --- a/tensorflow/python/ops/linalg/linear_operator_full_matrix.py +++ b/tensorflow/python/ops/linalg/linear_operator_full_matrix.py @@ -21,8 +21,8 @@ from __future__ import print_function from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops -from tensorflow.python.ops import math_ops 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__ = ["LinearOperatorFullMatrix"] @@ -176,7 +176,7 @@ class LinearOperatorFullMatrix(linear_operator.LinearOperator): return array_ops.shape(self._matrix) def _matmul(self, x, adjoint=False, adjoint_arg=False): - return math_ops.matmul( + return linear_operator_util.matmul_with_broadcast( self._matrix, x, adjoint_a=adjoint, adjoint_b=adjoint_arg) def _to_dense(self): 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 be911029095920d424ac90b406e7b85b73884b3b..08e5896e1034fb1782beacfb18fef16da083bded 100644 --- a/tensorflow/python/ops/linalg/linear_operator_low_rank_update.py +++ b/tensorflow/python/ops/linalg/linear_operator_low_rank_update.py @@ -27,6 +27,7 @@ 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.ops.linalg import linear_operator_util from tensorflow.python.util.tf_export import tf_export __all__ = [ @@ -365,14 +366,17 @@ class LinearOperatorLowRankUpdate(linear_operator.LinearOperator): leading_term = l.matmul(x, adjoint=adjoint, adjoint_arg=adjoint_arg) if adjoint: - uh_x = math_ops.matmul(u, x, adjoint_a=True, adjoint_b=adjoint_arg) + uh_x = linear_operator_util.matmul_with_broadcast( + u, x, adjoint_a=True, adjoint_b=adjoint_arg) d_uh_x = d.matmul(uh_x, adjoint=adjoint) - v_d_uh_x = math_ops.matmul(v, d_uh_x) + v_d_uh_x = linear_operator_util.matmul_with_broadcast( + v, d_uh_x) return leading_term + v_d_uh_x else: - vh_x = math_ops.matmul(v, x, adjoint_a=True, adjoint_b=adjoint_arg) + vh_x = linear_operator_util.matmul_with_broadcast( + v, x, adjoint_a=True, adjoint_b=adjoint_arg) d_vh_x = d.matmul(vh_x, adjoint=adjoint) - u_d_vh_x = math_ops.matmul(u, d_vh_x) + u_d_vh_x = linear_operator_util.matmul_with_broadcast(u, d_vh_x) return leading_term + u_d_vh_x def _determinant(self): @@ -431,16 +435,18 @@ class LinearOperatorLowRankUpdate(linear_operator.LinearOperator): # L^{-1} rhs linv_rhs = l.solve(rhs, adjoint=adjoint, adjoint_arg=adjoint_arg) # V^H L^{-1} rhs - vh_linv_rhs = math_ops.matmul(v, linv_rhs, adjoint_a=True) + vh_linv_rhs = linear_operator_util.matmul_with_broadcast( + v, linv_rhs, adjoint_a=True) # C^{-1} V^H L^{-1} rhs if self._use_cholesky: - capinv_vh_linv_rhs = linalg_ops.cholesky_solve( + capinv_vh_linv_rhs = linear_operator_util.cholesky_solve_with_broadcast( self._chol_capacitance, vh_linv_rhs) else: - capinv_vh_linv_rhs = linalg_ops.matrix_solve( + capinv_vh_linv_rhs = linear_operator_util.matrix_solve_with_broadcast( self._capacitance, vh_linv_rhs, adjoint=adjoint) # U C^{-1} V^H M^{-1} rhs - u_capinv_vh_linv_rhs = math_ops.matmul(u, capinv_vh_linv_rhs) + u_capinv_vh_linv_rhs = linear_operator_util.matmul_with_broadcast( + u, capinv_vh_linv_rhs) # L^{-1} U C^{-1} V^H L^{-1} rhs linv_u_capinv_vh_linv_rhs = l.solve(u_capinv_vh_linv_rhs, adjoint=adjoint) @@ -454,7 +460,8 @@ class LinearOperatorLowRankUpdate(linear_operator.LinearOperator): # L^{-1} U linv_u = self.base_operator.solve(self.u) # V^H L^{-1} U - vh_linv_u = math_ops.matmul(self.v, linv_u, adjoint_a=True) + vh_linv_u = linear_operator_util.matmul_with_broadcast( + self.v, linv_u, adjoint_a=True) # D^{-1} + V^H L^{-1} V capacitance = self._diag_inv_operator.add_to_tensor(vh_linv_u) diff --git a/tensorflow/python/ops/linalg/linear_operator_lower_triangular.py b/tensorflow/python/ops/linalg/linear_operator_lower_triangular.py index c4d386ccb4efc7dede8310243e517fe2f6b45bd9..fb1eb2fedba5b47ce38f9635527b91e18d894a8f 100644 --- a/tensorflow/python/ops/linalg/linear_operator_lower_triangular.py +++ b/tensorflow/python/ops/linalg/linear_operator_lower_triangular.py @@ -21,7 +21,6 @@ from __future__ import print_function from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops -from tensorflow.python.ops import linalg_ops 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 @@ -194,7 +193,7 @@ class LinearOperatorLowerTriangular(linear_operator.LinearOperator): message="Singular operator: Diagonal contained zero values.") def _matmul(self, x, adjoint=False, adjoint_arg=False): - return math_ops.matmul( + return linear_operator_util.matmul_with_broadcast( self._tril, x, adjoint_a=adjoint, adjoint_b=adjoint_arg) def _determinant(self): @@ -206,7 +205,7 @@ class LinearOperatorLowerTriangular(linear_operator.LinearOperator): def _solve(self, rhs, adjoint=False, adjoint_arg=False): rhs = linalg.adjoint(rhs) if adjoint_arg else rhs - return linalg_ops.matrix_triangular_solve( + return linear_operator_util.matrix_triangular_solve_with_broadcast( self._tril, rhs, lower=True, adjoint=adjoint) def _to_dense(self): diff --git a/tensorflow/python/ops/linalg/linear_operator_test_util.py b/tensorflow/python/ops/linalg/linear_operator_test_util.py index ce1a112ad584a14298be6e471578858ef31573d5..7e4fb6a6fc31960d570c78398ae211ba45a4a4a7 100644 --- a/tensorflow/python/ops/linalg/linear_operator_test_util.py +++ b/tensorflow/python/ops/linalg/linear_operator_test_util.py @@ -32,6 +32,7 @@ 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.linalg import linalg_impl as linalg +from tensorflow.python.ops.linalg import linear_operator_util from tensorflow.python.platform import test @@ -126,13 +127,16 @@ class LinearOperatorDerivedClassTest(test.TestCase): raise NotImplementedError("Not implemented yet.") @abc.abstractmethod - def _make_rhs(self, operator, adjoint): + def _make_rhs(self, operator, adjoint, with_batch=True): """Make a rhs appropriate for calling operator.solve(rhs). Args: operator: A `LinearOperator` adjoint: Python `bool`. If `True`, we are making a 'rhs' value for the adjoint operator. + with_batch: Python `bool`. If `True`, create `rhs` with the same batch + shape as operator, and otherwise create a matrix without any batch + shape. Returns: A `Tensor` @@ -140,13 +144,15 @@ class LinearOperatorDerivedClassTest(test.TestCase): raise NotImplementedError("_make_rhs is not defined.") @abc.abstractmethod - def _make_x(self, operator, adjoint): + def _make_x(self, operator, adjoint, with_batch=True): """Make an 'x' appropriate for calling operator.matmul(x). Args: operator: A `LinearOperator` adjoint: Python `bool`. If `True`, we are making an 'x' value for the adjoint operator. + with_batch: Python `bool`. If `True`, create `x` with the same batch shape + as operator, and otherwise create a matrix without any batch shape. Returns: A `Tensor` @@ -224,10 +230,15 @@ class LinearOperatorDerivedClassTest(test.TestCase): [op_log_abs_det, mat_log_abs_det], feed_dict=feed_dict) self.assertAC(op_log_abs_det_v, mat_log_abs_det_v) - def test_matmul(self): - self._skip_if_tests_to_skip_contains("matmul") + def _test_matmul(self, with_batch): for use_placeholder in self._use_placeholder_options: for build_info in self._operator_build_infos: + # If batch dimensions are omitted, but there are + # no batch dimensions for the linear operator, then + # skip the test case. This is already checked with + # with_batch=True. + if not with_batch and len(build_info.shape) <= 2: + continue for dtype in self._dtypes_to_test: for adjoint in self._adjoint_options: for adjoint_arg in self._adjoint_arg_options: @@ -235,7 +246,8 @@ class LinearOperatorDerivedClassTest(test.TestCase): sess.graph.seed = random_seed.DEFAULT_GRAPH_SEED operator, mat, feed_dict = self._operator_and_mat_and_feed_dict( build_info, dtype, use_placeholder=use_placeholder) - x = self._make_x(operator, adjoint=adjoint) + x = self._make_x( + operator, adjoint=adjoint, with_batch=with_batch) # If adjoint_arg, compute A X^H^H = A X. if adjoint_arg: op_matmul = operator.matmul( @@ -244,7 +256,8 @@ class LinearOperatorDerivedClassTest(test.TestCase): adjoint_arg=adjoint_arg) else: op_matmul = operator.matmul(x, adjoint=adjoint) - mat_matmul = math_ops.matmul(mat, x, adjoint_a=adjoint) + mat_matmul = linear_operator_util.matmul_with_broadcast( + mat, x, adjoint_a=adjoint) if not use_placeholder: self.assertAllEqual(op_matmul.get_shape(), mat_matmul.get_shape()) @@ -252,10 +265,23 @@ class LinearOperatorDerivedClassTest(test.TestCase): [op_matmul, mat_matmul], feed_dict=feed_dict) self.assertAC(op_matmul_v, mat_matmul_v) - def test_solve(self): - self._skip_if_tests_to_skip_contains("solve") + def test_matmul(self): + self._skip_if_tests_to_skip_contains("matmul") + self._test_matmul(with_batch=True) + + def test_matmul_with_broadcast(self): + self._skip_if_tests_to_skip_contains("matmul_with_broadcast") + self._test_matmul(with_batch=False) + + def _test_solve(self, with_batch): for use_placeholder in self._use_placeholder_options: for build_info in self._operator_build_infos: + # If batch dimensions are omitted, but there are + # no batch dimensions for the linear operator, then + # skip the test case. This is already checked with + # with_batch=True. + if not with_batch and len(build_info.shape) <= 2: + continue for dtype in self._dtypes_to_test: for adjoint in self._adjoint_options: for adjoint_arg in self._adjoint_arg_options: @@ -263,7 +289,8 @@ class LinearOperatorDerivedClassTest(test.TestCase): sess.graph.seed = random_seed.DEFAULT_GRAPH_SEED operator, mat, feed_dict = self._operator_and_mat_and_feed_dict( build_info, dtype, use_placeholder=use_placeholder) - rhs = self._make_rhs(operator, adjoint=adjoint) + rhs = self._make_rhs( + operator, adjoint=adjoint, with_batch=with_batch) # If adjoint_arg, solve A X = (rhs^H)^H = rhs. if adjoint_arg: op_solve = operator.solve( @@ -273,7 +300,8 @@ class LinearOperatorDerivedClassTest(test.TestCase): else: op_solve = operator.solve( rhs, adjoint=adjoint, adjoint_arg=adjoint_arg) - mat_solve = linalg_ops.matrix_solve(mat, rhs, adjoint=adjoint) + mat_solve = linear_operator_util.matrix_solve_with_broadcast( + mat, rhs, adjoint=adjoint) if not use_placeholder: self.assertAllEqual(op_solve.get_shape(), mat_solve.get_shape()) @@ -281,6 +309,14 @@ class LinearOperatorDerivedClassTest(test.TestCase): [op_solve, mat_solve], feed_dict=feed_dict) self.assertAC(op_solve_v, mat_solve_v) + def test_solve(self): + self._skip_if_tests_to_skip_contains("solve") + self._test_solve(with_batch=True) + + def test_solve_with_broadcast(self): + self._skip_if_tests_to_skip_contains("solve_with_broadcast") + self._test_solve(with_batch=False) + def test_trace(self): self._skip_if_tests_to_skip_contains("trace") for use_placeholder in self._use_placeholder_options: @@ -358,13 +394,13 @@ class SquareLinearOperatorDerivedClassTest(LinearOperatorDerivedClassTest): build_info((3, 4, 4)), build_info((2, 1, 4, 4))] - def _make_rhs(self, operator, adjoint): + def _make_rhs(self, operator, adjoint, with_batch=True): # This operator is square, so rhs and x will have same shape. # adjoint value makes no difference because the operator shape doesn't # change since it is square, but be pedantic. - return self._make_x(operator, adjoint=not adjoint) + return self._make_x(operator, adjoint=not adjoint, with_batch=with_batch) - def _make_x(self, operator, adjoint): + def _make_x(self, operator, adjoint, with_batch=True): # Value of adjoint makes no difference because the operator is square. # Return the number of systems to solve, R, equal to 1 or 2. r = self._get_num_systems(operator) @@ -373,11 +409,17 @@ class SquareLinearOperatorDerivedClassTest(LinearOperatorDerivedClassTest): if operator.shape.is_fully_defined(): batch_shape = operator.batch_shape.as_list() n = operator.domain_dimension.value - x_shape = batch_shape + [n, r] + if with_batch: + x_shape = batch_shape + [n, r] + else: + x_shape = [n, r] else: batch_shape = operator.batch_shape_tensor() n = operator.domain_dimension_tensor() - x_shape = array_ops.concat((batch_shape, [n, r]), 0) + if with_batch: + x_shape = array_ops.concat((batch_shape, [n, r]), 0) + else: + x_shape = [n, r] return random_normal(x_shape, dtype=operator.dtype) @@ -404,7 +446,7 @@ class NonSquareLinearOperatorDerivedClassTest(LinearOperatorDerivedClassTest): @property def _tests_to_skip(self): """List of test names to skip.""" - return ["solve", "det", "log_abs_det"] + return ["solve", "solve_with_broadcast", "det", "log_abs_det"] @property def _operator_build_infos(self): @@ -417,12 +459,12 @@ class NonSquareLinearOperatorDerivedClassTest(LinearOperatorDerivedClassTest): build_info((3, 3, 4)), build_info((2, 1, 2, 4))] - def _make_rhs(self, operator, adjoint): + def _make_rhs(self, operator, adjoint, with_batch=True): # TODO(langmore) Add once we're testing solve_ls. raise NotImplementedError( "_make_rhs not implemented because we don't test solve") - def _make_x(self, operator, adjoint): + def _make_x(self, operator, adjoint, with_batch=True): # Return the number of systems for the argument 'x' for .matmul(x) r = self._get_num_systems(operator) # If operator.shape = [B1,...,Bb, M, N] this returns a random matrix of @@ -433,14 +475,20 @@ class NonSquareLinearOperatorDerivedClassTest(LinearOperatorDerivedClassTest): n = operator.range_dimension.value else: n = operator.domain_dimension.value - x_shape = batch_shape + [n, r] + if with_batch: + x_shape = batch_shape + [n, r] + else: + x_shape = [n, r] else: batch_shape = operator.batch_shape_tensor() if adjoint: n = operator.range_dimension_tensor() else: n = operator.domain_dimension_tensor() - x_shape = array_ops.concat((batch_shape, [n, r]), 0) + if with_batch: + x_shape = array_ops.concat((batch_shape, [n, r]), 0) + else: + x_shape = [n, r] return random_normal(x_shape, dtype=operator.dtype) diff --git a/tensorflow/python/ops/linalg_ops.py b/tensorflow/python/ops/linalg_ops.py index 170861b43fd980ab0e107fc0b2e3d6f02339ed34..a0dfa543f9b3aee15f11b073dc683b1d2d14388f 100644 --- a/tensorflow/python/ops/linalg_ops.py +++ b/tensorflow/python/ops/linalg_ops.py @@ -24,12 +24,13 @@ 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 functional_ops from tensorflow.python.ops import gen_linalg_ops +from tensorflow.python.ops import linalg_ops_impl from tensorflow.python.ops import math_ops # pylint: disable=wildcard-import from tensorflow.python.ops.gen_linalg_ops import * # pylint: enable=wildcard-import -from tensorflow.python.util import compat from tensorflow.python.util import deprecation from tensorflow.python.util.tf_export import tf_export @@ -159,36 +160,11 @@ def eye(num_rows, Returns: A `Tensor` of shape `batch_shape + [num_rows, num_columns]` """ - with ops.name_scope( - name, default_name='eye', values=[num_rows, num_columns, batch_shape]): - is_square = num_columns is None - batch_shape = [] if batch_shape is None else batch_shape - num_columns = num_rows if num_columns is None else num_columns - if isinstance(num_rows, ops.Tensor) or isinstance( - num_columns, ops.Tensor) or isinstance(batch_shape, ops.Tensor): - batch_shape = ops.convert_to_tensor( - batch_shape, name='shape', dtype=dtypes.int32) - diag_size = math_ops.minimum(num_rows, num_columns) - diag_shape = array_ops.concat((batch_shape, [diag_size]), 0) - if not is_square: - shape = array_ops.concat((batch_shape, [num_rows, num_columns]), 0) - else: - if not isinstance(num_rows, compat.integral_types) or not isinstance( - num_columns, compat.integral_types): - raise TypeError( - 'num_rows and num_columns must be positive integer values.') - batch_shape = [dim for dim in batch_shape] - is_square = num_rows == num_columns - diag_shape = batch_shape + [np.minimum(num_rows, num_columns)] - if not is_square: - shape = batch_shape + [num_rows, num_columns] - - diag_ones = array_ops.ones(diag_shape, dtype=dtype) - if is_square: - return array_ops.matrix_diag(diag_ones) - else: - zero_matrix = array_ops.zeros(shape, dtype=dtype) - return array_ops.matrix_set_diag(zero_matrix, diag_ones) + return linalg_ops_impl.eye(num_rows, + num_columns=num_columns, + batch_shape=batch_shape, + dtype=dtype, + name=name) @tf_export('matrix_solve_ls', 'linalg.lstsq') @@ -454,7 +430,7 @@ def norm(tensor, This function can compute several different vector norms (the 1-norm, the Euclidean or 2-norm, the inf-norm, and in general the p-norm for p > 0) and - matrix norms (Frobenius, 1-norm, and inf-norm). + matrix norms (Frobenius, 1-norm, 2-norm and inf-norm). Args: tensor: `Tensor` of types `float32`, `float64`, `complex64`, `complex128` @@ -465,7 +441,7 @@ def norm(tensor, Some restrictions apply: a) The Frobenius norm `fro` is not defined for vectors, b) If axis is a 2-tuple (matrix norm), only 'euclidean', 'fro', `1`, - `np.inf` are supported. + `2`, `np.inf` are supported. See the description of `axis` on how to compute norms for a batch of vectors or matrices stored in a tensor. axis: If `axis` is `None` (the default), the input is considered a vector @@ -521,8 +497,7 @@ def norm(tensor, axis[0] == axis[1]): raise ValueError( "'axis' must be None, an integer, or a tuple of 2 unique integers") - # TODO(rmlarsen): Implement matrix 2-norm using tf.svd(). - supported_matrix_norms = ['euclidean', 'fro', 1, np.inf] + supported_matrix_norms = ['euclidean', 'fro', 1, 2, np.inf] if ord not in supported_matrix_norms: raise ValueError("'ord' must be a supported matrix norm in %s, got %s" % (supported_matrix_norms, ord)) @@ -539,12 +514,34 @@ def norm(tensor, with ops.name_scope(name, 'norm', [tensor]): tensor = ops.convert_to_tensor(tensor) + if ord in ['fro', 'euclidean', 2, 2.0]: - # TODO(rmlarsen): Move 2-norm to a separate clause once we support it for - # matrices. - result = math_ops.sqrt( - math_ops.reduce_sum( - tensor * math_ops.conj(tensor), axis, keepdims=True)) + if is_matrix_norm and ord in [2, 2.0]: + rank = array_ops.rank(tensor) + positive_axis = functional_ops.map_fn( + lambda i: control_flow_ops.cond(i >= 0, lambda: i, lambda: i + rank), + ops.convert_to_tensor(axis)) + axes = math_ops.range(rank) + perm_before = array_ops.concat( + [array_ops.setdiff1d(axes, positive_axis)[0], positive_axis], + axis=0) + perm_after = functional_ops.map_fn( + lambda i: math_ops.cast( + array_ops.squeeze( + array_ops.where(math_ops.equal(perm_before, i))), + dtype=dtypes.int32), axes) + permed = array_ops.transpose(tensor, perm=perm_before) + matrix_2_norm = array_ops.expand_dims( + math_ops.reduce_max( + math_ops.abs(gen_linalg_ops.svd(permed, compute_uv=False)[0]), + axis=-1, + keepdims=True), + axis=-1) + result = array_ops.transpose(matrix_2_norm, perm=perm_after) + else: + result = math_ops.sqrt( + math_ops.reduce_sum( + tensor * math_ops.conj(tensor), axis, keepdims=True)) else: result = math_ops.abs(tensor) if ord == 1: diff --git a/tensorflow/python/ops/linalg_ops_impl.py b/tensorflow/python/ops/linalg_ops_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..e7c89f6ae3e9c517920e6c9afce99a8b192be164 --- /dev/null +++ b/tensorflow/python/ops/linalg_ops_impl.py @@ -0,0 +1,73 @@ +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Operations for linear algebra.""" + +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 +from tensorflow.python.util import compat + +# Names below are lower_case. +# pylint: disable=invalid-name + + +def eye(num_rows, + num_columns=None, + batch_shape=None, + dtype=dtypes.float32, + name=None): + """Construct an identity matrix, or a batch of matrices. + + See `linalg_ops.eye`. + """ + with ops.name_scope( + name, default_name='eye', values=[num_rows, num_columns, batch_shape]): + is_square = num_columns is None + batch_shape = [] if batch_shape is None else batch_shape + num_columns = num_rows if num_columns is None else num_columns + if isinstance(num_rows, ops.Tensor) or isinstance( + num_columns, ops.Tensor) or isinstance(batch_shape, ops.Tensor): + batch_shape = ops.convert_to_tensor( + batch_shape, name='shape', dtype=dtypes.int32) + diag_size = math_ops.minimum(num_rows, num_columns) + diag_shape = array_ops.concat((batch_shape, [diag_size]), 0) + if not is_square: + shape = array_ops.concat((batch_shape, [num_rows, num_columns]), 0) + else: + if not isinstance(num_rows, compat.integral_types) or not isinstance( + num_columns, compat.integral_types): + raise TypeError( + 'num_rows and num_columns must be positive integer values.') + batch_shape = [dim for dim in batch_shape] + is_square = num_rows == num_columns + diag_shape = batch_shape + [np.minimum(num_rows, num_columns)] + if not is_square: + shape = batch_shape + [num_rows, num_columns] + + diag_ones = array_ops.ones(diag_shape, dtype=dtype) + if is_square: + return array_ops.matrix_diag(diag_ones) + else: + zero_matrix = array_ops.zeros(shape, dtype=dtype) + return array_ops.matrix_set_diag(zero_matrix, diag_ones) + +# pylint: enable=invalid-name,redefined-builtin diff --git a/tensorflow/python/ops/list_ops.py b/tensorflow/python/ops/list_ops.py index bdf0774bbf834ec10f68423e89e3b8b9b96ad9a1..d9ede875301c52219cc1e3f05a892ee887a70e67 100644 --- a/tensorflow/python/ops/list_ops.py +++ b/tensorflow/python/ops/list_ops.py @@ -29,6 +29,10 @@ from tensorflow.python.ops.gen_list_ops import * # pylint: enable=wildcard-import +ops.NotDifferentiable("TensorListConcat") +ops.NotDifferentiable("TensorListPushBackBatch") + + @ops.RegisterGradient("TensorListPushBack") def _PushBackGrad(op, dresult): return gen_list_ops.tensor_list_pop_back( diff --git a/tensorflow/python/ops/losses/losses_impl.py b/tensorflow/python/ops/losses/losses_impl.py index 34ca1adc3e13dc67560fb21d70c16cd42dc40552..9fc545c9678e7eb33a7ad35e2a84f890885e09af 100644 --- a/tensorflow/python/ops/losses/losses_impl.py +++ b/tensorflow/python/ops/losses/losses_impl.py @@ -29,6 +29,7 @@ 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.deprecation import deprecated_argument_lookup from tensorflow.python.util.tf_export import tf_export @@ -306,11 +307,8 @@ def cosine_distance( ValueError: If `predictions` shape doesn't match `labels` shape, or `axis`, `labels`, `predictions` or `weights` is `None`. """ - if dim is not None: - if axis is not None: - raise ValueError("Cannot specify both 'axis' and 'dim'") - axis = dim - if axis is None and dim is None: + axis = deprecated_argument_lookup("axis", axis, "dim", dim) + if axis is None: raise ValueError("You must specify 'axis'.") if labels is None: raise ValueError("labels must not be None.") @@ -696,7 +694,7 @@ def softmax_cross_entropy( onehot_labels, logits, weights=1.0, label_smoothing=0, scope=None, loss_collection=ops.GraphKeys.LOSSES, reduction=Reduction.SUM_BY_NONZERO_WEIGHTS): - """Creates a cross-entropy loss using tf.nn.softmax_cross_entropy_with_logits. + """Creates a cross-entropy loss using tf.nn.softmax_cross_entropy_with_logits_v2. `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 @@ -707,11 +705,16 @@ def softmax_cross_entropy( new_onehot_labels = onehot_labels * (1 - label_smoothing) + label_smoothing / num_classes + Note that `onehot_labels` and `logits` must have the same shape, + e.g. `[batch_size, num_classes]`. The shape of `weights` must be + broadcastable to loss, whose shape is decided by the shape of `logits`. + In case the shape of `logits` is `[batch_size, num_classes]`, loss is + a `Tensor` of shape `[batch_size]`. + Args: - onehot_labels: `[batch_size, num_classes]` target one-hot-encoded labels. - logits: `[batch_size, num_classes]` logits outputs of the network . - weights: Optional `Tensor` whose rank is either 0, or rank 1 and is - broadcastable to the loss which is a `Tensor` of shape `[batch_size]`. + onehot_labels: One-hot-encoded labels. + logits: Logits outputs of the network. + weights: Optional `Tensor` that is broadcastable to loss. label_smoothing: If greater than 0 then smooth the labels. scope: the scope for the operations performed in computing the loss. loss_collection: collection to which the loss will be added. diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py index 01d670ea2d93d54f558a9be6aabf7b6b237f65ef..30ac001c2510dd3c8cc8afd9a8c29b1af22f6c3a 100644 --- a/tensorflow/python/ops/math_ops.py +++ b/tensorflow/python/ops/math_ops.py @@ -211,11 +211,9 @@ def argmax(input, name=None, dimension=None, output_type=dtypes.int64): - if dimension is not None: - if axis is not None: - raise ValueError("Cannot specify both 'axis' and 'dimension'") - axis = dimension - elif axis is None: + axis = deprecation.deprecated_argument_lookup( + "axis", axis, "dimension", dimension) + if axis is None: axis = 0 return gen_math_ops.arg_max(input, axis, name=name, output_type=output_type) @@ -231,11 +229,9 @@ def argmin(input, name=None, dimension=None, output_type=dtypes.int64): - if dimension is not None: - if axis is not None: - raise ValueError("Cannot specify both 'axis' and 'dimension'") - axis = dimension - elif axis is None: + axis = deprecation.deprecated_argument_lookup( + "axis", axis, "dimension", dimension) + if axis is None: axis = 0 return gen_math_ops.arg_min(input, axis, name=name, output_type=output_type) @@ -761,13 +757,25 @@ def cast(x, dtype, name=None): tf.cast(x, tf.int32) # [1, 2], dtype=tf.int32 ``` + The operation supports data types (for `x` and `dtype`) of + `uint8`, `int8`, `uint16`, `int16`, `int32`, `int64`, `float16`, `float32`, + `float64`, `complex64`, `complex128`, `bfloat16`. In case of casting from + complex types (`complex64`, `complex128`) to real types, only the real part + of `x` is returned. In case of casting from real types to complex types + (`complex64`, `complex128`), the imaginary part of the returned value is set + to `0`. The handling of complex types here matches the behavior of numpy. + Args: - x: A `Tensor` or `SparseTensor`. - dtype: The destination type. + x: A `Tensor` or `SparseTensor` of numeric type. It could be + `uint8`, `int8`, `uint16`, `int16`, `int32`, `int64`, + `float16`, `float32`, `float64`, `complex64`, `complex128`, `bfloat16`. + dtype: The destination type. The list of supported dtypes is the same + as `x`. name: A name for the operation (optional). Returns: - A `Tensor` or `SparseTensor` with same shape as `x`. + A `Tensor` or `SparseTensor` with same shape as `x` and + same type as `dtype`. Raises: TypeError: If `x` cannot be cast to the `dtype`. @@ -1458,8 +1466,18 @@ def count_nonzero(input_tensor, tf.count_nonzero(x, [0, 1]) # 3 ``` + **NOTE** Strings are compared against zero-length empty string `""`. Any + string with a size greater than zero is already considered as nonzero. + + For example: + ```python + x = tf.constant(["", "a", " ", "b", ""]) + tf.count_nonzero(x) # 3, with "a", " ", and "b" as nonzero strings. + ``` + Args: - input_tensor: The tensor to reduce. Should be of numeric type, or `bool`. + input_tensor: The tensor to reduce. Should be of numeric type, `bool`, + or `string`. axis: The dimensions to reduce. If `None` (the default), reduces all dimensions. Must be in the range `[-rank(input_tensor), rank(input_tensor))`. @@ -1479,7 +1497,8 @@ def count_nonzero(input_tensor, with ops.name_scope(name, "count_nonzero", [input_tensor]): input_tensor = ops.convert_to_tensor(input_tensor, name="input_tensor") - zero = input_tensor.dtype.as_numpy_dtype() + # A scalar of 'zero' is enough as `not_equal` will broadcast. + zero = array_ops.zeros([], dtype=input_tensor.dtype) return cast( reduce_sum( # int64 reduction happens on GPU @@ -1632,7 +1651,7 @@ def reduce_min(input_tensor, tensor with a single element is returned. Args: - input_tensor: The tensor to reduce. Should have numeric type. + input_tensor: The tensor to reduce. Should have real numeric type. axis: The dimensions to reduce. If `None` (the default), reduces all dimensions. Must be in the range `[-rank(input_tensor), rank(input_tensor))`. @@ -1681,7 +1700,7 @@ def reduce_max(input_tensor, tensor with a single element is returned. Args: - input_tensor: The tensor to reduce. Should have numeric type. + input_tensor: The tensor to reduce. Should have real numeric type. axis: The dimensions to reduce. If `None` (the default), reduces all dimensions. Must be in the range `[-rank(input_tensor), rank(input_tensor))`. diff --git a/tensorflow/python/ops/nn.py b/tensorflow/python/ops/nn.py index ee1a00623a734e18d4aebe6c84f77ba53ee1050c..1d0d9a52a125b32c8fba65f8b2c9907331296836 100644 --- a/tensorflow/python/ops/nn.py +++ b/tensorflow/python/ops/nn.py @@ -98,6 +98,7 @@ See the @{$python/nn} guide. @@fixed_unigram_candidate_sampler @@compute_accidental_hits @@quantized_conv2d +@@quantized_relu @@quantized_relu_x @@quantized_max_pool @@quantized_avg_pool @@ -126,8 +127,6 @@ from tensorflow.python.ops.nn_impl import * from tensorflow.python.ops.nn_ops import * from tensorflow.python.ops.candidate_sampling_ops import * from tensorflow.python.ops.embedding_ops import * -from tensorflow.python.ops.rnn import * -from tensorflow.python.ops import rnn_cell # pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/python/ops/nn_grad.py b/tensorflow/python/ops/nn_grad.py index 4af5bd26dd80b984b1c898411c2a23827bed1b4b..3a41391340edbe25bd97cfadc58587d91bef9de2 100644 --- a/tensorflow/python/ops/nn_grad.py +++ b/tensorflow/python/ops/nn_grad.py @@ -94,6 +94,7 @@ def _Conv3DGrad(op, grad): array_ops.shape(op.inputs[0]), op.inputs[1], grad, + dilations=op.get_attr("dilations"), strides=op.get_attr("strides"), padding=op.get_attr("padding"), data_format=data_format), @@ -101,6 +102,7 @@ def _Conv3DGrad(op, grad): op.inputs[0], array_ops.shape(op.inputs[1]), grad, + dilations=op.get_attr("dilations"), strides=op.get_attr("strides"), padding=op.get_attr("padding"), data_format=data_format) @@ -116,12 +118,14 @@ def _Conv3DBackpropInputGrad(op, grad): grad, array_ops.shape(op.inputs[1]), op.inputs[2], + dilations=op.get_attr("dilations"), strides=op.get_attr("strides"), padding=op.get_attr("padding"), data_format=data_format), nn_ops.conv3d( grad, op.inputs[1], + dilations=op.get_attr("dilations"), strides=op.get_attr("strides"), padding=op.get_attr("padding"), data_format=data_format) @@ -136,12 +140,14 @@ def _Conv3DBackpropFilterGrad(op, grad): array_ops.shape(op.inputs[0]), grad, op.inputs[2], + dilations=op.get_attr("dilations"), strides=op.get_attr("strides"), padding=op.get_attr("padding"), data_format=data_format), None, nn_ops.conv3d( op.inputs[0], grad, + dilations=op.get_attr("dilations"), strides=op.get_attr("strides"), padding=op.get_attr("padding"), data_format=data_format) diff --git a/tensorflow/python/ops/nn_impl.py b/tensorflow/python/ops/nn_impl.py index 47cc4da7f2abd1f5b00e193a76c8391be94ca27d..576627e78ed10d832f82477487140e470e70cfb2 100644 --- a/tensorflow/python/ops/nn_impl.py +++ b/tensorflow/python/ops/nn_impl.py @@ -765,9 +765,9 @@ def weighted_moments(x, axes, frequency_weights, name=None, keep_dims=False): weighted_variance = math_ops.multiply(weighted_distsq, divisor) if not keep_dims: - weighted_mean = array_ops.squeeze(weighted_mean, squeeze_dims=axes) + weighted_mean = array_ops.squeeze(weighted_mean, axis=axes) weighted_variance = array_ops.squeeze( - weighted_variance, squeeze_dims=axes) + weighted_variance, axis=axes) if needs_cast: weighted_mean = math_ops.cast(weighted_mean, dtypes.float16) @@ -987,7 +987,7 @@ def _compute_sampled_logits(weights, class biases. labels: A `Tensor` of type `int64` and shape `[batch_size, num_true]`. The target classes. Note that this format differs from - the `labels` argument of `nn.softmax_cross_entropy_with_logits`. + the `labels` argument of `nn.softmax_cross_entropy_with_logits_v2`. inputs: A `Tensor` of shape `[batch_size, dim]`. The forward activations of the input network. num_sampled: An `int`. The number of classes to randomly sample per batch. @@ -1012,7 +1012,7 @@ def _compute_sampled_logits(weights, out_logits: `Tensor` object with shape `[batch_size, num_true + num_sampled]`, for passing to either `nn.sigmoid_cross_entropy_with_logits` (NCE) or - `nn.softmax_cross_entropy_with_logits` (sampled softmax). + `nn.softmax_cross_entropy_with_logits_v2` (sampled softmax). out_labels: A Tensor object with the same shape as `out_logits`. """ @@ -1285,7 +1285,7 @@ def sampled_softmax_loss(weights, logits = tf.matmul(inputs, tf.transpose(weights)) logits = tf.nn.bias_add(logits, biases) labels_one_hot = tf.one_hot(labels, n_classes) - loss = tf.nn.softmax_cross_entropy_with_logits( + loss = tf.nn.softmax_cross_entropy_with_logits_v2( labels=labels_one_hot, logits=logits) ``` @@ -1303,7 +1303,7 @@ def sampled_softmax_loss(weights, biases: A `Tensor` of shape `[num_classes]`. The class biases. labels: A `Tensor` of type `int64` and shape `[batch_size, num_true]`. The target classes. Note that this format differs from - the `labels` argument of `nn.softmax_cross_entropy_with_logits`. + the `labels` argument of `nn.softmax_cross_entropy_with_logits_v2`. inputs: A `Tensor` of shape `[batch_size, dim]`. The forward activations of the input network. num_sampled: An `int`. The number of classes to randomly sample per batch. @@ -1340,7 +1340,8 @@ def sampled_softmax_loss(weights, partition_strategy=partition_strategy, name=name, seed=seed) - sampled_losses = nn_ops.softmax_cross_entropy_with_logits( + labels = array_ops.stop_gradient(labels, name="labels_stop_gradient") + sampled_losses = nn_ops.softmax_cross_entropy_with_logits_v2( labels=labels, logits=logits) # sampled_losses is a [batch_size] tensor. return sampled_losses diff --git a/tensorflow/python/ops/nn_ops.py b/tensorflow/python/ops/nn_ops.py index 07ca32953f796466964d4555e45052fcf3c53ce0..cd07550d2ee31e9c830afa8e9c5e17deb35d1135 100644 --- a/tensorflow/python/ops/nn_ops.py +++ b/tensorflow/python/ops/nn_ops.py @@ -1155,7 +1155,7 @@ def atrous_conv2d(value, filters, rate, padding, name=None): Returns: A `Tensor` with the same type as `value`. - Output shape with `'VALID`` padding is: + Output shape with `'VALID'` padding is: [batch, height - 2 * (filter_width - 1), width - 2 * (filter_height - 1), out_channels]. @@ -1458,10 +1458,10 @@ def conv3d_transpose( if isinstance(output_shape, (list, np.ndarray)): # output_shape's shape should be == [5] if reached this point. - if not filter.get_shape()[3].is_compatible_with(output_shape[4]): + if not filter.get_shape()[3].is_compatible_with(output_shape[axis]): raise ValueError( "output_shape does not match filter's output channels, " - "{} != {}".format(output_shape[4], + "{} != {}".format(output_shape[axis], filter.get_shape()[3])) if padding != "VALID" and padding != "SAME": @@ -1803,8 +1803,11 @@ def softmax_cross_entropy_with_logits_v2( on `logits` internally for efficiency. Do not call this op with the output of `softmax`, as it will produce incorrect results. - `logits` and `labels` must have the same shape, e.g. - `[batch_size, num_classes]` and the same dtype (either `float16`, `float32`, + A common use case is to have logits and labels of shape + `[batch_size, num_classes]`, but higher dimensions are supported, with + the `dim` argument specifying the class dimension. + + `logits` and `labels` must have the same dtype (either `float16`, `float32`, or `float64`). Backpropagation will happen into both `logits` and `labels`. To disallow @@ -1816,14 +1819,17 @@ def softmax_cross_entropy_with_logits_v2( Args: _sentinel: Used to prevent positional parameters. Internal, do not use. - labels: Each row `labels[i]` must be a valid probability distribution. + labels: Each vector along the class dimension should hold a valid + probability distribution e.g. for the case in which labels are of shape + `[batch_size, num_classes]`, each row of `labels[i]` must be a valid + probability distribution. logits: Unscaled log probabilities. dim: The class dimension. Defaulted to -1 which is the last dimension. name: A name for the operation (optional). Returns: - A 1-D `Tensor` of length `batch_size` of the same type as `logits` with the - softmax cross entropy loss. + A `Tensor` of the same shape as `labels` and of the same type as `logits` + with the softmax cross entropy loss. """ _ensure_xent_args("softmax_cross_entropy_with_logits", _sentinel, labels, logits) @@ -1926,9 +1932,9 @@ def softmax_cross_entropy_with_logits( on `logits` internally for efficiency. Do not call this op with the output of `softmax`, as it will produce incorrect results. - `logits` and `labels` must have the same shape, e.g. - `[batch_size, num_classes]` and the same dtype (either `float16`, `float32`, - or `float64`). + A common use case is to have logits and labels of shape + `[batch_size, num_classes]`, but higher dimensions are supported, with + the `dim` argument specifying the class dimension. Backpropagation will happen only into `logits`. To calculate a cross entropy loss that allows backpropagation into both `logits` and `labels`, see @@ -1939,14 +1945,17 @@ def softmax_cross_entropy_with_logits( Args: _sentinel: Used to prevent positional parameters. Internal, do not use. - labels: Each row `labels[i]` must be a valid probability distribution. + labels: Each vector along the class dimension should hold a valid + probability distribution e.g. for the case in which labels are of shape + `[batch_size, num_classes]`, each row of `labels[i]` must be a valid + probability distribution. logits: Unscaled log probabilities. dim: The class dimension. Defaulted to -1 which is the last dimension. name: A name for the operation (optional). Returns: - A 1-D `Tensor` of length `batch_size` of the same type as `logits` with the - softmax cross entropy loss. + A `Tensor` of the same shape as `labels` and of the same type as `logits` + with the softmax cross entropy loss. """ _ensure_xent_args("softmax_cross_entropy_with_logits", _sentinel, labels, logits) @@ -1977,14 +1986,17 @@ def sparse_softmax_cross_entropy_with_logits( must provide a single specific index for the true class for each row of `logits` (each minibatch entry). For soft softmax classification with a probability distribution for each entry, see - `softmax_cross_entropy_with_logits`. + `softmax_cross_entropy_with_logits_v2`. **WARNING:** This op expects unscaled logits, since it performs a `softmax` on `logits` internally for efficiency. Do not call this op with the output of `softmax`, as it will produce incorrect results. - A common use case is to have logits of shape `[batch_size, num_classes]` and - labels of shape `[batch_size]`. But higher dimensions are supported. + A common use case is to have logits and labels of shape + `[batch_size, num_classes]`, but higher dimensions are supported, in which + case the `dim`-th dimension is assumed to be of size `num_classes`. + `logits` and `labels` must have the same dtype (either `float16`, `float32`, + or `float64`). **Note that to avoid confusion, it is required to pass only named arguments to this function.** diff --git a/tensorflow/python/ops/parsing_ops.py b/tensorflow/python/ops/parsing_ops.py index 075b38d743d13329e646c0b268e938b5c5704e47..d8d9af545f17fe3e0133b51b1eab82f7732dc299 100644 --- a/tensorflow/python/ops/parsing_ops.py +++ b/tensorflow/python/ops/parsing_ops.py @@ -1176,8 +1176,13 @@ def _parse_single_sequence_example_raw(serialized, # Swap `name` and `na_value` for backward compatibility. @tf_export("decode_csv") -def decode_csv(records, record_defaults, field_delim=",", - use_quote_delim=True, name=None, na_value=""): +def decode_csv(records, + record_defaults, + field_delim=",", + use_quote_delim=True, + name=None, + na_value="", + select_cols=None): """Convert CSV records to tensors. Each column maps to one tensor. RFC 4180 format is expected for the CSV records. @@ -1200,19 +1205,32 @@ def decode_csv(records, record_defaults, field_delim=",", Bullet 5). name: A name for the operation (optional). na_value: Additional string to recognize as NA/NaN. + select_cols: Optional sorted list of column indices to select. If specified, + only this subset of columns will be parsed and returned. Returns: A list of `Tensor` objects. Has the same type as `record_defaults`. Each tensor will have the same shape as records. + + Raises: + ValueError: If any of the arguments is malformed. """ - # TODO(martinwicke), remove the wrapper when new Python API generator is done. + if select_cols is not None and any(select_cols[i] >= select_cols[i + 1] + for i in range(len(select_cols) - 1)): + raise ValueError("select_cols is not strictly increasing.") + if select_cols is not None and select_cols[0] < 0: + raise ValueError("select_cols contains negative values.") + if select_cols is not None and len(select_cols) != len(record_defaults): + raise ValueError("Length of select_cols and record_defaults do not match.") 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) + name=name, + select_cols=select_cols, + ) # TODO(b/70890287): Combine the implementation of this op and diff --git a/tensorflow/python/ops/resource_variable_ops.py b/tensorflow/python/ops/resource_variable_ops.py index 508ba9bfeeb4dcb627288357c1c7b6ab4ef14c5c..4d26b2f46e3f7982305b1163c7f3762660ccef2f 100644 --- a/tensorflow/python/ops/resource_variable_ops.py +++ b/tensorflow/python/ops/resource_variable_ops.py @@ -24,6 +24,8 @@ 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 c_api_util +from tensorflow.python.framework import cpp_shape_inference_pb2 from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape @@ -41,6 +43,19 @@ from tensorflow.python.training import checkpointable from tensorflow.python.util import compat +def get_resource_handle_data(graph_op): + assert ops._USE_C_SHAPES # pylint: disable=protected-access + assert type(graph_op) == ops.Tensor # pylint: disable=unidiomatic-typecheck + + with c_api_util.tf_buffer() as buf: + pywrap_tensorflow.TFE_GetResourceHandleShapeAndType( + graph_op.graph._c_graph, graph_op._as_tf_output(), buf) # pylint: disable=protected-access + data = pywrap_tensorflow.TF_GetBuffer(buf) + + return cpp_shape_inference_pb2.CppShapeInferenceResult.HandleData.FromString( + compat.as_bytes(data)) + + def _eager_safe_variable_handle(shape, dtype, shared_name, name, graph_mode): """Creates a variable handle with information to do shape inference.""" container = ops.get_default_graph()._container # pylint: disable=protected-access @@ -73,9 +88,12 @@ def _eager_safe_variable_handle(shape, dtype, shared_name, name, graph_mode): # shape inference doesn't run in eager mode we copy this data here for when # the handle is captured by an eager mode function. # pylint: disable=protected-access - if h._handle_data is None: - ops.set_shape_and_handle_data_for_outputs(h.op) - handle._handle_data = h._handle_data + if ops._USE_C_SHAPES: + handle._handle_data = get_resource_handle_data(h) + else: + if h._handle_data is None: + ops.set_shape_and_handle_data_for_outputs(h.op) + handle._handle_data = h._handle_data # pylint: enable=protected-access # Clean up our reference cycles to avoid making the garbage collector run. @@ -522,11 +540,19 @@ class ResourceVariable(variables.Variable): else: self._initial_value = None if variable_def.snapshot_name: - self._cached_value = g.as_graph_element( + snapshot = g.as_graph_element( ops.prepend_name_scope( variable_def.snapshot_name, import_scope=import_scope)) + self._cached_value = snapshot + while snapshot.op.type != "ReadVariableOp": + snapshot = snapshot.op.inputs[0] + self._graph_element = snapshot else: self._cached_value = None + # Legacy case for protos without the snapshot name; assume it's the + # following. + self._graph_element = g.get_tensor_by_name( + self._handle.op.name + "/Read/ReadVariableOp:0") if variable_def.HasField("save_slice_info_def"): self._save_slice_info = variables.Variable.SaveSliceInfo( save_slice_info_def=variable_def.save_slice_info_def, @@ -535,8 +561,6 @@ 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 = g.get_tensor_by_name( - self._handle.op.name + "/Read/ReadVariableOp:0") self._constraint = None self._cached_shape_as_list = None @@ -745,6 +769,10 @@ class ResourceVariable(variables.Variable): if self._cached_value is not None: var_def.snapshot_name = ops.strip_name_scope(self._cached_value.name, export_scope) + else: + # Store the graph_element here + var_def.snapshot_name = ops.strip_name_scope(self._graph_element.name, + export_scope) var_def.is_resource = True if self._save_slice_info: var_def.save_slice_info_def.MergeFrom( @@ -910,7 +938,6 @@ class ResourceVariable(variables.Variable): def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False): del name if dtype is not None and dtype != self.dtype: - print("trying to switch the dtype to ", dtype, " from ", self.dtype) return NotImplemented if as_ref: return self.read_value().op.inputs[0] diff --git a/tensorflow/python/ops/rnn_cell_impl.py b/tensorflow/python/ops/rnn_cell_impl.py index cbc2dcf419b95fe3ea16602db0919ccb782335de..86dc053c0fb8d0d91eda87c40f41076391b3d9b8 100644 --- a/tensorflow/python/ops/rnn_cell_impl.py +++ b/tensorflow/python/ops/rnn_cell_impl.py @@ -352,10 +352,17 @@ class BasicRNNCell(LayerRNNCell): name: String, the name of the layer. Layers with the same name will share weights, but to avoid mistakes we require reuse=True in such cases. + dtype: Default dtype of the layer (default of `None` means use the type + of the first input). Required when `build` is called before `call`. """ - def __init__(self, num_units, activation=None, reuse=None, name=None): - super(BasicRNNCell, self).__init__(_reuse=reuse, name=name) + def __init__(self, + num_units, + activation=None, + reuse=None, + name=None, + dtype=None): + super(BasicRNNCell, self).__init__(_reuse=reuse, name=name, dtype=dtype) # Inputs must be 2-dimensional. self.input_spec = base_layer.InputSpec(ndim=2) @@ -413,6 +420,8 @@ class GRUCell(LayerRNNCell): name: String, the name of the layer. Layers with the same name will share weights, but to avoid mistakes we require reuse=True in such cases. + dtype: Default dtype of the layer (default of `None` means use the type + of the first input). Required when `build` is called before `call`. """ def __init__(self, @@ -421,8 +430,9 @@ class GRUCell(LayerRNNCell): reuse=None, kernel_initializer=None, bias_initializer=None, - name=None): - super(GRUCell, self).__init__(_reuse=reuse, name=name) + name=None, + dtype=None): + super(GRUCell, self).__init__(_reuse=reuse, name=name, dtype=dtype) # Inputs must be 2-dimensional. self.input_spec = base_layer.InputSpec(ndim=2) @@ -531,8 +541,14 @@ class BasicLSTMCell(LayerRNNCell): that follows. """ - def __init__(self, num_units, forget_bias=1.0, - state_is_tuple=True, activation=None, reuse=None, name=None): + def __init__(self, + num_units, + forget_bias=1.0, + state_is_tuple=True, + activation=None, + reuse=None, + name=None, + dtype=None): """Initialize the basic LSTM cell. Args: @@ -550,11 +566,13 @@ class BasicLSTMCell(LayerRNNCell): name: String, the name of the layer. Layers with the same name will share weights, but to avoid mistakes we require reuse=True in such cases. + dtype: Default dtype of the layer (default of `None` means use the type + of the first input). Required when `build` is called before `call`. When restoring from CudnnLSTM-trained checkpoints, must use `CudnnCompatibleLSTMCell` instead. """ - super(BasicLSTMCell, self).__init__(_reuse=reuse, name=name) + super(BasicLSTMCell, self).__init__(_reuse=reuse, name=name, dtype=dtype) if not state_is_tuple: logging.warn("%s: Using a concatenated state is slower and will soon be " "deprecated. Use state_is_tuple=True.", self) @@ -599,9 +617,9 @@ class BasicLSTMCell(LayerRNNCell): Args: inputs: `2-D` tensor with shape `[batch_size, input_size]`. state: An `LSTMStateTuple` of state tensors, each shaped - `[batch_size, self.state_size]`, if `state_is_tuple` has been set to + `[batch_size, num_units]`, if `state_is_tuple` has been set to `True`. Otherwise, a `Tensor` shaped - `[batch_size, 2 * self.state_size]`. + `[batch_size, 2 * num_units]`. Returns: A pair containing the new hidden state, and the new state (either a @@ -668,7 +686,7 @@ class LSTMCell(LayerRNNCell): initializer=None, num_proj=None, proj_clip=None, num_unit_shards=None, num_proj_shards=None, forget_bias=1.0, state_is_tuple=True, - activation=None, reuse=None, name=None): + activation=None, reuse=None, name=None, dtype=None): """Initialize the parameters for an LSTM cell. Args: @@ -701,11 +719,13 @@ class LSTMCell(LayerRNNCell): name: String, the name of the layer. Layers with the same name will share weights, but to avoid mistakes we require reuse=True in such cases. + dtype: Default dtype of the layer (default of `None` means use the type + of the first input). Required when `build` is called before `call`. When restoring from CudnnLSTM-trained checkpoints, use `CudnnCompatibleLSTMCell` instead. """ - super(LSTMCell, self).__init__(_reuse=reuse, name=name) + super(LSTMCell, self).__init__(_reuse=reuse, name=name, dtype=dtype) if not state_is_tuple: logging.warn("%s: Using a concatenated state is slower and will soon be " "deprecated. Use state_is_tuple=True.", self) diff --git a/tensorflow/python/ops/standard_ops.py b/tensorflow/python/ops/standard_ops.py index e90ff0746a8e86b4b462b71028fd677632c9075d..f71f98aa12c458b74228cfac53a00bd4f6d9a013 100644 --- a/tensorflow/python/ops/standard_ops.py +++ b/tensorflow/python/ops/standard_ops.py @@ -22,12 +22,13 @@ from __future__ import print_function import sys as _sys +# pylint: disable=g-bad-import-order # Imports the following modules so that @RegisterGradient get executed. from tensorflow.python.ops import array_grad +from tensorflow.python.ops import cudnn_rnn_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 from tensorflow.python.ops import spectral_grad from tensorflow.python.ops import state_grad @@ -96,6 +97,7 @@ from tensorflow.python.ops.tensor_array_ops import * from tensorflow.python.ops.variable_scope import * from tensorflow.python.ops.variables import * # pylint: enable=wildcard-import +# pylint: enable=g-bad-import-order #### For use in remove_undocumented below: from tensorflow.python.framework import constant_op as _constant_op diff --git a/tensorflow/contrib/summary/summary_ops.py b/tensorflow/python/ops/summary_ops_v2.py similarity index 90% rename from tensorflow/contrib/summary/summary_ops.py rename to tensorflow/python/ops/summary_ops_v2.py index bc763fe655edc455e2538e536d6efab314c8228c..12f361c513fcebf8ce4b9c367d101b11ab10260b 100644 --- a/tensorflow/contrib/summary/summary_ops.py +++ b/tensorflow/python/ops/summary_ops_v2.py @@ -31,7 +31,7 @@ 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.layers import utils +from tensorflow.python.framework import smart_cond from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import gen_summary_ops @@ -108,8 +108,10 @@ class SummaryWriter(object): - @{tf.contrib.summary.create_db_writer} """ - def __init__(self, resource): + def __init__(self, resource, init_op_fn): self._resource = resource + # TODO(nickfelt): cache constructed ops in graph mode + self._init_op_fn = init_op_fn if context.executing_eagerly() and self._resource is not None: self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._resource, handle_device="cpu:0") @@ -129,10 +131,32 @@ class SummaryWriter(object): yield self # Flushes the summary writer in eager mode or in graph functions, but not # in legacy graph mode (you're on your own there). - with ops.device("cpu:0"): - gen_summary_ops.flush_summary_writer(self._resource) + self.flush() context.context().summary_writer_resource = old + def init(self): + """Operation to initialize the summary writer resource.""" + if self._resource is not None: + return self._init_op_fn() + + def _flush(self): + return _flush_fn(writer=self) + + def flush(self): + """Operation to force the summary writer to flush any buffered data.""" + if self._resource is not None: + return self._flush() + + def _close(self): + with ops.control_dependencies([self.flush()]): + with ops.device("cpu:0"): + return gen_summary_ops.close_summary_writer(self._resource) + + def close(self): + """Operation to flush and close the summary writer resource.""" + if self._resource is not None: + return self._close() + def initialize( graph=None, # pylint: disable=redefined-outer-name @@ -178,7 +202,7 @@ def create_file_writer(logdir, flush_millis=None, filename_suffix=None, name=None): - """Creates a summary file writer in the current context. + """Creates a summary file writer in the current context under the given name. Args: logdir: a string, or None. If a string, creates a summary file writer @@ -186,18 +210,20 @@ def create_file_writer(logdir, a mock object which acts like a summary writer but does nothing, useful to use as a context manager. max_queue: the largest number of summaries to keep in a queue; will - flush once the queue gets bigger than this. - flush_millis: the largest interval between flushes. - filename_suffix: optional suffix for the event file name. + flush once the queue gets bigger than this. Defaults to 10. + flush_millis: the largest interval between flushes. Defaults to 120,000. + filename_suffix: optional suffix for the event file name. Defaults to `.v2`. name: Shared name for this SummaryWriter resource stored to default - Graph. + Graph. Defaults to the provided logdir prefixed with `logdir:`. Note: if a + summary writer resource with this shared name already exists, the returned + SummaryWriter wraps that resource and the other arguments have no effect. Returns: Either a summary writer or an empty object which can be used as a summary writer. """ if logdir is None: - return SummaryWriter(None) + return SummaryWriter(None, None) with ops.device("cpu:0"): if max_queue is None: max_queue = constant_op.constant(10) @@ -205,6 +231,8 @@ def create_file_writer(logdir, flush_millis = constant_op.constant(2 * 60 * 1000) if filename_suffix is None: filename_suffix = constant_op.constant(".v2") + if name is None: + name = "logdir:" + logdir return _make_summary_writer( name, gen_summary_ops.create_summary_file_writer, @@ -267,13 +295,12 @@ def create_db_writer(db_uri, def _make_summary_writer(name, factory, **kwargs): resource = gen_summary_ops.summary_writer(shared_name=name) + init_op_fn = lambda: factory(resource, **kwargs) # TODO(apassos): Consider doing this instead. - # node = factory(resource, **kwargs) # if not context.executing_eagerly(): - # ops.get_default_session().run(node) - ops.add_to_collection(_SUMMARY_WRITER_INIT_COLLECTION_NAME, - factory(resource, **kwargs)) - return SummaryWriter(resource) + # ops.get_default_session().run(init_op) + ops.add_to_collection(_SUMMARY_WRITER_INIT_COLLECTION_NAME, init_op_fn()) + return SummaryWriter(resource, init_op_fn) def _cleanse_string(name, pattern, value): @@ -341,7 +368,7 @@ def summary_writer_function(name, tensor, function, family=None): if context.context().summary_writer_resource is None: return control_flow_ops.no_op() with ops.device("cpu:0"): - op = utils.smart_cond( + op = smart_cond.smart_cond( should_record_summaries(), record, _nothing, name="") ops.add_to_collection(ops.GraphKeys._SUMMARY_COLLECTION, op) # pylint: disable=protected-access return op @@ -538,7 +565,14 @@ def flush(writer=None, name=None): writer = context.context().summary_writer_resource if writer is None: return control_flow_ops.no_op() - return gen_summary_ops.flush_summary_writer(writer, name=name) + else: + if isinstance(writer, SummaryWriter): + writer = writer._resource # pylint: disable=protected-access + with ops.device("cpu:0"): + return gen_summary_ops.flush_summary_writer(writer, name=name) + + +_flush_fn = flush # for within SummaryWriter.flush() def eval_dir(model_dir, name=None): diff --git a/tensorflow/python/ops/template.py b/tensorflow/python/ops/template.py index 0294ecee548d1e7f507a5e4195e4ee320a0b9918..9b6b8c508fcd7e73e6259d14b466892544ec65d7 100644 --- a/tensorflow/python/ops/template.py +++ b/tensorflow/python/ops/template.py @@ -452,8 +452,7 @@ class Template(checkpointable.CheckpointableBase): # Only reuse variables if they were already created. with variable_scope.variable_scope( self._variable_scope, reuse=self._variables_created): - result = self._call_func(args, kwargs) - return result + return self._call_func(args, kwargs) else: # The scope was not created at construction time, so create it here. # Subsequent calls should reuse variables. @@ -461,8 +460,7 @@ class Template(checkpointable.CheckpointableBase): self._unique_name, self._name, custom_getter=self._custom_getter) as vs: self._variable_scope = vs - result = self._call_func(args, kwargs) - return result + return self._call_func(args, kwargs) @property def name(self): @@ -730,8 +728,7 @@ class EagerTemplate(Template): 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 + return self._call_func(args, kwargs) else: # The scope was not created at construction time, so create it here. # Subsequent calls should reuse variables. @@ -743,8 +740,7 @@ class EagerTemplate(Template): # store's variable scope name is unset; set it here. self._template_store.set_variable_scope_name(vs.name) with self._template_store.as_default(): - result = self._call_func(args, kwargs) - return result + return self._call_func(args, kwargs) @property def name(self): diff --git a/tensorflow/python/ops/variable_scope.py b/tensorflow/python/ops/variable_scope.py index e33085ba626a7645be64941dd4da8e6943292e7e..ba213ef884165f7f72094d27932913e39c9a5901 100644 --- a/tensorflow/python/ops/variable_scope.py +++ b/tensorflow/python/ops/variable_scope.py @@ -307,6 +307,17 @@ class _VariableStore(object): raise ValueError( "Passed a custom_getter which is not callable: %s" % custom_getter) + with ops.init_scope(): + if context.executing_eagerly(): + # Variable creation and initialization takes place in `init_scope`s; + # as such, if an `init_scope` lifts us into the eager context, then we + # need to use `ResourceVariable`s. + use_resource = True + + # Note that it's fine to reuse eager variables whose initialization was + # lifted from a function-building graph into the eager context (that's why + # the following clause is not wrapped in an `init_scope`); lifted variables + # are tracked by the graph's `VariableStore`. if context.executing_eagerly(): if not self._store_eager_variables and reuse: raise RuntimeError( @@ -315,7 +326,6 @@ class _VariableStore(object): " EagerVariableStore for example usage.") if self._store_eager_variables: reuse = AUTO_REUSE - use_resource = True # If a *_ref type is passed in an error would be triggered further down the # stack. We prevent this using base_dtype to get a non-ref version of the diff --git a/tensorflow/python/profiler/tfprof_logger_test.py b/tensorflow/python/profiler/tfprof_logger_test.py index 141144f98776f3aa7c95b9ef743022aeca5084e1..caf3869f56d94a272218a0e318f2fdf5cb238795 100644 --- a/tensorflow/python/profiler/tfprof_logger_test.py +++ b/tensorflow/python/profiler/tfprof_logger_test.py @@ -38,7 +38,7 @@ class TFProfLoggerTest(test.TestCase): return math_ops.matmul(a, b) # pylint: disable=pointless-string-statement - """# TODO(xpan): This this out of core so it doesn't depend on contrib. + """# TODO(xpan): This out of core so it doesn't depend on contrib. def testFillMissingShape(self): a, b, y = self._BuildSmallPlaceholderlModel() run_options = config_pb2.RunOptions( diff --git a/tensorflow/python/pywrap_tfe.i b/tensorflow/python/pywrap_tfe.i index 5ee55301df986998b22b8b57b5f01b1f6b4918ac..0982a67deebb2fbf086df8352069e498b2972b72 100644 --- a/tensorflow/python/pywrap_tfe.i +++ b/tensorflow/python/pywrap_tfe.i @@ -59,6 +59,8 @@ limitations under the License. %rename("%s") TFE_ContextOptionsSetAsync; %rename("%s") TFE_DeleteContextOptions; %rename("%s") TFE_Py_TensorShapeSlice; +%rename("%s") TFE_GetResourceHandleShapeAndType; +%rename("%s") TFE_SetResourceHandleShapeAndType; %{ #include "tensorflow/python/eager/pywrap_tfe.h" diff --git a/tensorflow/python/summary/writer/event_file_writer_v2.py b/tensorflow/python/summary/writer/event_file_writer_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..5c66c0f7a8fd4c65d3539d257b4e4fa89f839a98 --- /dev/null +++ b/tensorflow/python/summary/writer/event_file_writer_v2.py @@ -0,0 +1,140 @@ +# 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. +# ============================================================================== +"""Writes events to disk in a logdir.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import summary_ops_v2 +from tensorflow.python.platform import gfile + + +class EventFileWriterV2(object): + """Writes `Event` protocol buffers to an event file via the graph. + + The `EventFileWriterV2` class is backed by the summary file writer in the v2 + summary API (currently in tf.contrib.summary), so it uses a shared summary + writer resource and graph ops to write events. + + As with the original EventFileWriter, this class will asynchronously write + Event protocol buffers to the backing file. The Event file is encoded using + the tfrecord format, which is similar to RecordIO. + """ + + def __init__(self, session, logdir, max_queue=10, flush_secs=120, + filename_suffix=''): + """Creates an `EventFileWriterV2` and an event file to write to. + + On construction, this calls `tf.contrib.summary.create_file_writer` within + the graph from `session.graph` to look up a shared summary writer resource + for `logdir` if one exists, and create one if not. Creating the summary + writer resource in turn creates a new event file in `logdir` to be filled + with `Event` protocol buffers passed to `add_event`. Graph ops to control + this writer resource are added to `session.graph` during this init call; + stateful methods on this class will call `session.run()` on these ops. + + Note that because the underlying resource is shared, it is possible that + other parts of the code using the same session may interact independently + with the resource, e.g. by flushing or even closing it. It is the caller's + responsibility to avoid any undesirable sharing in this regard. + + The remaining arguments to the constructor (`flush_secs`, `max_queue`, and + `filename_suffix`) control the construction of the shared writer resource + if one is created. If an existing resource is reused, these arguments have + no effect. See `tf.contrib.summary.create_file_writer` for details. + + Args: + session: A `tf.Session`. Session that will hold shared writer resource. + The writer ops will be added to session.graph during this init call. + logdir: A string. Directory where event file will be written. + max_queue: Integer. Size of the queue for pending events and summaries. + flush_secs: Number. How often, in seconds, to flush the + pending events and summaries to disk. + filename_suffix: A string. Every event file's name is suffixed with + `filename_suffix`. + """ + self._session = session + self._logdir = logdir + self._closed = False + if not gfile.IsDirectory(self._logdir): + gfile.MakeDirs(self._logdir) + + with self._session.graph.as_default(): + with ops.name_scope('filewriter'): + file_writer = summary_ops_v2.create_file_writer( + logdir=self._logdir, + max_queue=max_queue, + flush_millis=flush_secs * 1000, + filename_suffix=filename_suffix) + with summary_ops_v2.always_record_summaries(), file_writer.as_default(): + self._event_placeholder = array_ops.placeholder_with_default( + constant_op.constant('unused', dtypes.string), + shape=[]) + self._add_event_op = summary_ops_v2.import_event( + self._event_placeholder) + self._init_op = file_writer.init() + self._flush_op = file_writer.flush() + self._close_op = file_writer.close() + self._session.run(self._init_op) + + def get_logdir(self): + """Returns the directory where event file will be written.""" + return self._logdir + + def reopen(self): + """Reopens the EventFileWriter. + + Can be called after `close()` to add more events in the same directory. + The events will go into a new events file. + + Does nothing if the EventFileWriter was not closed. + """ + if self._closed: + self._closed = False + self._session.run(self._init_op) + + def add_event(self, event): + """Adds an event to the event file. + + Args: + event: An `Event` protocol buffer. + """ + if not self._closed: + event_pb = event.SerializeToString() + self._session.run( + self._add_event_op, feed_dict={self._event_placeholder: event_pb}) + + def flush(self): + """Flushes the event file to disk. + + Call this method to make sure that all pending events have been written to + disk. + """ + self._session.run(self._flush_op) + + def close(self): + """Flushes the event file to disk and close the file. + + Call this method when you do not need the summary writer anymore. + """ + if not self._closed: + self.flush() + self._session.run(self._close_op) + self._closed = True diff --git a/tensorflow/python/summary/writer/writer.py b/tensorflow/python/summary/writer/writer.py index 57f78c156b1334a5486b29f2ddec957e49156e73..aca084fc9168e710316e4c988594cff69e54ebab 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.summary.writer.event_file_writer_v2 import EventFileWriterV2 from tensorflow.python.util.tf_export import tf_export _PLUGINS_DIR = "plugins" @@ -286,6 +287,11 @@ class FileWriter(SummaryToEventTransformer): file contents asynchronously. This allows a training program to call methods to add data to the file directly from the training loop, without slowing down training. + + When constructed with a `tf.Session` parameter, a `FileWriter` instead forms + a compatibility layer over new graph-based summaries (`tf.contrib.summary`) + to facilitate the use of new summary writing with pre-existing code that + expects a `FileWriter` instance. """ def __init__(self, @@ -294,10 +300,11 @@ class FileWriter(SummaryToEventTransformer): max_queue=10, flush_secs=120, graph_def=None, - filename_suffix=None): - """Creates a `FileWriter` and an event file. + filename_suffix=None, + session=None): + """Creates a `FileWriter`, optionally shared within the given session. - On construction the summary writer creates a new event file in `logdir`. + Typically, constructing a file writer creates a new event file in `logdir`. This event file will contain `Event` protocol buffers constructed when you call one of the following functions: `add_summary()`, `add_session_log()`, `add_event()`, or `add_graph()`. @@ -317,13 +324,16 @@ class FileWriter(SummaryToEventTransformer): writer = tf.summary.FileWriter(, sess.graph) ``` - The other arguments to the constructor control the asynchronous writes to - the event file: - - * `flush_secs`: How often, in seconds, to flush the added summaries - and events to disk. - * `max_queue`: Maximum number of summaries or events pending to be - written to disk before one of the 'add' calls block. + The `session` argument to the constructor makes the returned `FileWriter` a + a compatibility layer over new graph-based summaries (`tf.contrib.summary`). + Crucially, this means the underlying writer resource and events file will + be shared with any other `FileWriter` using the same `session` and `logdir`, + and with any `tf.contrib.summary.SummaryWriter` in this session using the + the same shared resource name (which by default scoped to the logdir). If + no such resource exists, one will be created using the remaining arguments + to this constructor, but if one already exists those arguments are ignored. + In either case, ops will be added to `session.graph` to control the + underlying file writer resource. See `tf.contrib.summary` for more details. Args: logdir: A string. Directory where event file will be written. @@ -334,6 +344,7 @@ class FileWriter(SummaryToEventTransformer): graph_def: DEPRECATED: Use the `graph` argument instead. filename_suffix: A string. Every event file's name is suffixed with `suffix`. + session: A `tf.Session` object. See details above. Raises: RuntimeError: If called with eager execution enabled. @@ -347,9 +358,12 @@ class FileWriter(SummaryToEventTransformer): raise RuntimeError( "tf.summary.FileWriter is not compatible with eager execution. " "Use tf.contrib.summary instead.") - - event_writer = EventFileWriter(logdir, max_queue, flush_secs, - filename_suffix) + if session is not None: + event_writer = EventFileWriterV2( + session, logdir, max_queue, flush_secs, filename_suffix) + else: + event_writer = EventFileWriter(logdir, max_queue, flush_secs, + filename_suffix) super(FileWriter, self).__init__(event_writer, graph, graph_def) def __enter__(self): diff --git a/tensorflow/python/summary/writer/writer_test.py b/tensorflow/python/summary/writer/writer_test.py index 88ade0aac33f1cd8f9d8cb30344aabca76a13511..dc990c2602427049ecdb7588ff217207a69cbcd2 100644 --- a/tensorflow/python/summary/writer/writer_test.py +++ b/tensorflow/python/summary/writer/writer_test.py @@ -29,10 +29,12 @@ from tensorflow.core.protobuf import config_pb2 from tensorflow.core.protobuf import meta_graph_pb2 from tensorflow.core.util import event_pb2 from tensorflow.core.util.event_pb2 import SessionLog +from tensorflow.python.client import session 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 +from tensorflow.python.ops import summary_ops_v2 from tensorflow.python.platform import gfile from tensorflow.python.platform import test from tensorflow.python.summary import plugin_asset @@ -42,7 +44,10 @@ from tensorflow.python.summary.writer import writer_cache from tensorflow.python.util import compat -class SummaryWriterTestCase(test.TestCase): +class FileWriterTestCase(test.TestCase): + + def _FileWriter(self, *args, **kwargs): + return writer.FileWriter(*args, **kwargs) def _TestDir(self, test_name): test_dir = os.path.join(self.get_temp_dir(), test_name) @@ -96,7 +101,7 @@ class SummaryWriterTestCase(test.TestCase): def testAddingSummaryGraphAndRunMetadata(self): test_dir = self._CleanTestDir("basics") - sw = writer.FileWriter(test_dir) + sw = self._FileWriter(test_dir) sw.add_session_log(event_pb2.SessionLog(status=SessionLog.START), 1) sw.add_summary( @@ -171,7 +176,7 @@ class SummaryWriterTestCase(test.TestCase): test_dir = self._CleanTestDir("basics_named_graph") with ops.Graph().as_default() as g: constant_op.constant([12], name="douze") - sw = writer.FileWriter(test_dir, graph=g) + sw = self._FileWriter(test_dir, graph=g) sw.close() self._assertEventsWithGraph(test_dir, g, True) @@ -179,7 +184,7 @@ class SummaryWriterTestCase(test.TestCase): test_dir = self._CleanTestDir("basics_positional_graph") with ops.Graph().as_default() as g: constant_op.constant([12], name="douze") - sw = writer.FileWriter(test_dir, g) + sw = self._FileWriter(test_dir, g) sw.close() self._assertEventsWithGraph(test_dir, g, True) @@ -188,7 +193,7 @@ class SummaryWriterTestCase(test.TestCase): with ops.Graph().as_default() as g: constant_op.constant([12], name="douze") gd = g.as_graph_def() - sw = writer.FileWriter(test_dir, graph_def=gd) + sw = self._FileWriter(test_dir, graph_def=gd) sw.close() self._assertEventsWithGraph(test_dir, g, False) @@ -197,7 +202,7 @@ class SummaryWriterTestCase(test.TestCase): with ops.Graph().as_default() as g: constant_op.constant([12], name="douze") gd = g.as_graph_def() - sw = writer.FileWriter(test_dir, gd) + sw = self._FileWriter(test_dir, gd) sw.close() self._assertEventsWithGraph(test_dir, g, False) @@ -207,18 +212,18 @@ class SummaryWriterTestCase(test.TestCase): with ops.Graph().as_default() as g: constant_op.constant([12], name="douze") gd = g.as_graph_def() - sw = writer.FileWriter(test_dir, graph=g, graph_def=gd) + sw = self._FileWriter(test_dir, graph=g, graph_def=gd) sw.close() def testNeitherGraphNorGraphDef(self): with self.assertRaises(TypeError): test_dir = self._CleanTestDir("basics_string_instead_of_graph") - sw = writer.FileWriter(test_dir, "string instead of graph object") + sw = self._FileWriter(test_dir, "string instead of graph object") sw.close() def testCloseAndReopen(self): test_dir = self._CleanTestDir("close_and_reopen") - sw = writer.FileWriter(test_dir) + sw = self._FileWriter(test_dir) sw.add_session_log(event_pb2.SessionLog(status=SessionLog.START), 1) sw.close() # Sleep at least one second to make sure we get a new event file name. @@ -261,7 +266,7 @@ class SummaryWriterTestCase(test.TestCase): def testNonBlockingClose(self): test_dir = self._CleanTestDir("non_blocking_close") - sw = writer.FileWriter(test_dir) + sw = self._FileWriter(test_dir) # Sleep 1.2 seconds to make sure event queue is empty. time.sleep(1.2) time_before_close = time.time() @@ -270,7 +275,7 @@ class SummaryWriterTestCase(test.TestCase): def testWithStatement(self): test_dir = self._CleanTestDir("with_statement") - with writer.FileWriter(test_dir) as sw: + with self._FileWriter(test_dir) as sw: sw.add_session_log(event_pb2.SessionLog(status=SessionLog.START), 1) event_paths = sorted(glob.glob(os.path.join(test_dir, "event*"))) self.assertEquals(1, len(event_paths)) @@ -280,7 +285,7 @@ class SummaryWriterTestCase(test.TestCase): # protocol buffers correctly. def testAddingSummariesFromSessionRunCalls(self): test_dir = self._CleanTestDir("global_step") - sw = writer.FileWriter(test_dir) + sw = self._FileWriter(test_dir) with self.test_session(): i = constant_op.constant(1, dtype=dtypes.int32, shape=[]) l = constant_op.constant(2, dtype=dtypes.int64, shape=[]) @@ -327,7 +332,7 @@ class SummaryWriterTestCase(test.TestCase): def testPluginMetadataStrippedFromSubsequentEvents(self): test_dir = self._CleanTestDir("basics") - sw = writer.FileWriter(test_dir) + sw = self._FileWriter(test_dir) sw.add_session_log(event_pb2.SessionLog(status=SessionLog.START), 1) @@ -386,7 +391,7 @@ class SummaryWriterTestCase(test.TestCase): def testFileWriterWithSuffix(self): test_dir = self._CleanTestDir("test_suffix") - sw = writer.FileWriter(test_dir, filename_suffix="_test_suffix") + sw = self._FileWriter(test_dir, filename_suffix="_test_suffix") for _ in range(10): sw.add_summary( summary_pb2.Summary(value=[ @@ -400,9 +405,178 @@ class SummaryWriterTestCase(test.TestCase): for filename in event_filenames: self.assertTrue(filename.endswith("_test_suffix")) + def testPluginAssetSerialized(self): + class ExamplePluginAsset(plugin_asset.PluginAsset): + plugin_name = "example" + + def assets(self): + return {"foo.txt": "foo!", "bar.txt": "bar!"} + + with ops.Graph().as_default() as g: + plugin_asset.get_plugin_asset(ExamplePluginAsset) + + logdir = self.get_temp_dir() + fw = self._FileWriter(logdir) + fw.add_graph(g) + plugin_dir = os.path.join(logdir, writer._PLUGINS_DIR, "example") + + with gfile.Open(os.path.join(plugin_dir, "foo.txt"), "r") as f: + content = f.read() + self.assertEqual(content, "foo!") + + with gfile.Open(os.path.join(plugin_dir, "bar.txt"), "r") as f: + content = f.read() + self.assertEqual(content, "bar!") -class SummaryWriterCacheTest(test.TestCase): - """SummaryWriterCache tests.""" + +class SessionBasedFileWriterTestCase(FileWriterTestCase): + """Tests for FileWriter behavior when passed a Session argument.""" + + def _FileWriter(self, *args, **kwargs): + if "session" not in kwargs: + # Pass in test_session() as the session. It will be cached during this + # test method invocation so that any other use of test_session() with no + # graph should result in re-using the same underlying Session. + with self.test_session() as sess: + kwargs["session"] = sess + return writer.FileWriter(*args, **kwargs) + return writer.FileWriter(*args, **kwargs) + + def _createTaggedSummary(self, tag): + summary = summary_pb2.Summary() + summary.value.add(tag=tag) + return summary + + def testSharing_withOtherSessionBasedFileWriters(self): + logdir = self.get_temp_dir() + with session.Session() as sess: + # Initial file writer + writer1 = writer.FileWriter(session=sess, logdir=logdir) + writer1.add_summary(self._createTaggedSummary("one"), 1) + writer1.flush() + + # File writer, should share file with writer1 + writer2 = writer.FileWriter(session=sess, logdir=logdir) + writer2.add_summary(self._createTaggedSummary("two"), 2) + writer2.flush() + + # File writer with different logdir (shouldn't be in this logdir at all) + writer3 = writer.FileWriter(session=sess, logdir=logdir + "-other") + writer3.add_summary(self._createTaggedSummary("three"), 3) + writer3.flush() + + # File writer in a different session (should be in separate file) + time.sleep(1.1) # Ensure filename has a different timestamp + with session.Session() as other_sess: + writer4 = writer.FileWriter(session=other_sess, logdir=logdir) + writer4.add_summary(self._createTaggedSummary("four"), 4) + writer4.flush() + + # One more file writer, should share file with writer1 + writer5 = writer.FileWriter(session=sess, logdir=logdir) + writer5.add_summary(self._createTaggedSummary("five"), 5) + writer5.flush() + + event_paths = iter(sorted(glob.glob(os.path.join(logdir, "event*")))) + + # First file should have tags "one", "two", and "five" + events = summary_iterator.summary_iterator(next(event_paths)) + self.assertEqual("brain.Event:2", next(events).file_version) + self.assertEqual("one", next(events).summary.value[0].tag) + self.assertEqual("two", next(events).summary.value[0].tag) + self.assertEqual("five", next(events).summary.value[0].tag) + self.assertRaises(StopIteration, lambda: next(events)) + + # Second file should have just "four" + events = summary_iterator.summary_iterator(next(event_paths)) + self.assertEqual("brain.Event:2", next(events).file_version) + self.assertEqual("four", next(events).summary.value[0].tag) + self.assertRaises(StopIteration, lambda: next(events)) + + # No more files + self.assertRaises(StopIteration, lambda: next(event_paths)) + + # Just check that the other logdir file exists to be sure we wrote it + self.assertTrue(glob.glob(os.path.join(logdir + "-other", "event*"))) + + def testSharing_withExplicitSummaryFileWriters(self): + logdir = self.get_temp_dir() + with session.Session() as sess: + # Initial file writer via FileWriter(session=?) + writer1 = writer.FileWriter(session=sess, logdir=logdir) + writer1.add_summary(self._createTaggedSummary("one"), 1) + writer1.flush() + + # Next one via create_file_writer(), should use same file + writer2 = summary_ops_v2.create_file_writer(logdir=logdir) + with summary_ops_v2.always_record_summaries(), writer2.as_default(): + summary2 = summary_ops_v2.scalar("two", 2.0, step=2) + sess.run(writer2.init()) + sess.run(summary2) + sess.run(writer2.flush()) + + # Next has different shared name, should be in separate file + time.sleep(1.1) # Ensure filename has a different timestamp + writer3 = summary_ops_v2.create_file_writer(logdir=logdir, name="other") + with summary_ops_v2.always_record_summaries(), writer3.as_default(): + summary3 = summary_ops_v2.scalar("three", 3.0, step=3) + sess.run(writer3.init()) + sess.run(summary3) + sess.run(writer3.flush()) + + # Next uses a second session, should be in separate file + time.sleep(1.1) # Ensure filename has a different timestamp + with session.Session() as other_sess: + writer4 = summary_ops_v2.create_file_writer(logdir=logdir) + with summary_ops_v2.always_record_summaries(), writer4.as_default(): + summary4 = summary_ops_v2.scalar("four", 4.0, step=4) + other_sess.run(writer4.init()) + other_sess.run(summary4) + other_sess.run(writer4.flush()) + + # Next via FileWriter(session=?) uses same second session, should be in + # same separate file. (This checks sharing in the other direction) + writer5 = writer.FileWriter(session=other_sess, logdir=logdir) + writer5.add_summary(self._createTaggedSummary("five"), 5) + writer5.flush() + + # One more via create_file_writer(), should use same file + writer6 = summary_ops_v2.create_file_writer(logdir=logdir) + with summary_ops_v2.always_record_summaries(), writer6.as_default(): + summary6 = summary_ops_v2.scalar("six", 6.0, step=6) + sess.run(writer6.init()) + sess.run(summary6) + sess.run(writer6.flush()) + + event_paths = iter(sorted(glob.glob(os.path.join(logdir, "event*")))) + + # First file should have tags "one", "two", and "six" + events = summary_iterator.summary_iterator(next(event_paths)) + self.assertEqual("brain.Event:2", next(events).file_version) + self.assertEqual("one", next(events).summary.value[0].tag) + self.assertEqual("two", next(events).summary.value[0].tag) + self.assertEqual("six", next(events).summary.value[0].tag) + self.assertRaises(StopIteration, lambda: next(events)) + + # Second file should have just "three" + events = summary_iterator.summary_iterator(next(event_paths)) + self.assertEqual("brain.Event:2", next(events).file_version) + self.assertEqual("three", next(events).summary.value[0].tag) + self.assertRaises(StopIteration, lambda: next(events)) + + # Third file should have "four" and "five" + events = summary_iterator.summary_iterator(next(event_paths)) + self.assertEqual("brain.Event:2", next(events).file_version) + self.assertEqual("four", next(events).summary.value[0].tag) + self.assertEqual("five", next(events).summary.value[0].tag) + self.assertRaises(StopIteration, lambda: next(events)) + + # No more files + self.assertRaises(StopIteration, lambda: next(event_paths)) + + +class FileWriterCacheTest(test.TestCase): + """FileWriterCache tests.""" def _test_dir(self, test_name): """Create an empty dir to use for tests. @@ -448,32 +622,5 @@ class SummaryWriterCacheTest(test.TestCase): self.assertFalse(sw1 == sw2) -class ExamplePluginAsset(plugin_asset.PluginAsset): - plugin_name = "example" - - def assets(self): - return {"foo.txt": "foo!", "bar.txt": "bar!"} - - -class PluginAssetsTest(test.TestCase): - - def testPluginAssetSerialized(self): - with ops.Graph().as_default() as g: - plugin_asset.get_plugin_asset(ExamplePluginAsset) - - logdir = self.get_temp_dir() - fw = writer.FileWriter(logdir) - fw.add_graph(g) - plugin_dir = os.path.join(logdir, writer._PLUGINS_DIR, "example") - - with gfile.Open(os.path.join(plugin_dir, "foo.txt"), "r") as f: - content = f.read() - self.assertEqual(content, "foo!") - - with gfile.Open(os.path.join(plugin_dir, "bar.txt"), "r") as f: - content = f.read() - self.assertEqual(content, "bar!") - - if __name__ == "__main__": test.main() diff --git a/tensorflow/python/tools/BUILD b/tensorflow/python/tools/BUILD index 6e39ce8c808a1716ff9263982e99a14592472c76..84d20f8e3628ed28fc408e9100b2999973911a33 100644 --- a/tensorflow/python/tools/BUILD +++ b/tensorflow/python/tools/BUILD @@ -28,7 +28,7 @@ py_library( name = "saved_model_utils", srcs = ["saved_model_utils.py"], srcs_version = "PY2AND3", - deps = ["//tensorflow:tensorflow_py"], + deps = ["//tensorflow/contrib/saved_model:reader"], ) py_library( @@ -41,8 +41,10 @@ py_library( "//tensorflow/python", # TODO(b/34059704): remove when fixed "//tensorflow/python:client", "//tensorflow/python:framework", + "//tensorflow/python:parsing_ops", "//tensorflow/python:platform", "//tensorflow/python:training", + "//tensorflow/python/saved_model:loader", "@six_archive//:six", ], ) @@ -52,14 +54,7 @@ py_binary( srcs = ["freeze_graph.py"], srcs_version = "PY2AND3", deps = [ - ":saved_model_utils", - "//tensorflow/core:protos_all_py", - "//tensorflow/python", # TODO(b/34059704): remove when fixed - "//tensorflow/python:client", - "//tensorflow/python:framework", - "//tensorflow/python:platform", - "//tensorflow/python:training", - "@six_archive//:six", + ":freeze_graph_lib", ], ) diff --git a/tensorflow/python/tools/saved_model_cli.py b/tensorflow/python/tools/saved_model_cli.py index b88be4ae04d5dc7a7641fb8dbd7e56e61035869f..73ea85ab0c4c5896efe6106cc12c62043e738d11 100644 --- a/tensorflow/python/tools/saved_model_cli.py +++ b/tensorflow/python/tools/saved_model_cli.py @@ -41,6 +41,7 @@ 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 # pylint: disable=unused-import +from tensorflow.python.lib.io import file_io from tensorflow.python.saved_model import loader from tensorflow.python.tools import saved_model_utils @@ -543,7 +544,7 @@ def load_inputs_from_input_arg_string(inputs_str, input_exprs_str, input_examples = preprocess_input_examples_arg_string(input_examples_str) for input_tensor_key, (filename, variable_name) in inputs.items(): - data = np.load(filename) + data = np.load(file_io.FileIO(filename, mode='r')) # When a variable_name key is specified for the input file if variable_name: diff --git a/tensorflow/python/training/adam.py b/tensorflow/python/training/adam.py index 006e360389b404a8edd97c9a8bf4b8876c828004..6fa3ff66583ce07a6ee7b0d8158c851ea578637c 100644 --- a/tensorflow/python/training/adam.py +++ b/tensorflow/python/training/adam.py @@ -43,23 +43,19 @@ class AdamOptimizer(optimizer.Optimizer): Initialization: - ``` - m_0 <- 0 (Initialize initial 1st moment vector) - v_0 <- 0 (Initialize initial 2nd moment vector) - t <- 0 (Initialize timestep) - ``` + $$m_0 := 0 (Initialize initial 1st moment vector)$$ + $$v_0 := 0 (Initialize initial 2nd moment vector)$$ + $$t := 0 (Initialize timestep)$$ The update rule for `variable` with gradient `g` uses an optimization described at the end of section2 of the paper: - ``` - t <- t + 1 - lr_t <- learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t) + $$t := t + 1$$ + $$lr_t := \text{learning_rate} * \sqrt{(1 - beta_2^t) / (1 - beta_1^t)}$$ - m_t <- beta1 * m_{t-1} + (1 - beta1) * g - v_t <- beta2 * v_{t-1} + (1 - beta2) * g * g - variable <- variable - lr_t * m_t / (sqrt(v_t) + epsilon) - ``` + $$m_t := beta_1 * m_{t-1} + (1 - beta_1) * g$$ + $$v_t := beta_2 * v_{t-1} + (1 - beta_2) * g * g$$ + $$variable := variable - lr_t * m_t / (\sqrt{v_t} + \epsilon)$$ The default value of 1e-8 for epsilon might not be a good default in general. For example, when training an Inception network on ImageNet a diff --git a/tensorflow/python/training/basic_session_run_hooks.py b/tensorflow/python/training/basic_session_run_hooks.py index 094a9e886ba87b639c143983d77949d664c8255a..3651291bdfcc77d6a6c288e1f53a78ef696a99d1 100644 --- a/tensorflow/python/training/basic_session_run_hooks.py +++ b/tensorflow/python/training/basic_session_run_hooks.py @@ -391,7 +391,8 @@ class CheckpointSaverHook(session_run_hook.SessionRunHook): saver=None, checkpoint_basename="model.ckpt", scaffold=None, - listeners=None): + listeners=None, + steps_per_run=1): """Initializes a `CheckpointSaverHook`. Args: @@ -404,6 +405,9 @@ class CheckpointSaverHook(session_run_hook.SessionRunHook): listeners: List of `CheckpointSaverListener` subclass instances. Used for callbacks that run immediately before or after this hook saves the checkpoint. + steps_per_run: `int`, number of steps that occur between each invocation + of the hook. Primarily used for TPU workloads which run multiple steps + in a while loop in a single Session.run. Raises: ValueError: One of `save_steps` or `save_secs` should be set. @@ -419,6 +423,7 @@ class CheckpointSaverHook(session_run_hook.SessionRunHook): self._timer = SecondOrStepTimer(every_secs=save_secs, every_steps=save_steps) self._listeners = listeners or [] + self._steps_per_run = steps_per_run def begin(self): self._summary_writer = SummaryWriterCache.get(self._checkpoint_dir) @@ -450,7 +455,8 @@ class CheckpointSaverHook(session_run_hook.SessionRunHook): def after_run(self, run_context, run_values): stale_global_step = run_values.results - if self._timer.should_trigger_for_step(stale_global_step+1): + if self._timer.should_trigger_for_step( + stale_global_step + self._steps_per_run): # get the real value after train op. global_step = run_context.session.run(self._global_step_tensor) if self._timer.should_trigger_for_step(global_step): diff --git a/tensorflow/python/training/basic_session_run_hooks_test.py b/tensorflow/python/training/basic_session_run_hooks_test.py index f39a5261a93c3dc2df3a2364010448df116461e9..25962f6bf7abf5f11785ccb6ba0e31c5c85f4ebd 100644 --- a/tensorflow/python/training/basic_session_run_hooks_test.py +++ b/tensorflow/python/training/basic_session_run_hooks_test.py @@ -719,6 +719,99 @@ class CheckpointSaverHookTest(test.TestCase): fake_summary_writer.FakeSummaryWriter.uninstall() +class CheckpointSaverHookMultiStepTest(test.TestCase): + + def setUp(self): + self.model_dir = tempfile.mkdtemp() + self.graph = ops.Graph() + self.steps_per_run = 5 + with self.graph.as_default(): + self.scaffold = monitored_session.Scaffold() + self.global_step = variables.get_or_create_global_step() + self.train_op = training_util._increment_global_step(self.steps_per_run) + + def tearDown(self): + shutil.rmtree(self.model_dir, ignore_errors=True) + + def test_save_steps_saves_in_first_step(self): + with self.graph.as_default(): + hook = basic_session_run_hooks.CheckpointSaverHook( + self.model_dir, + save_steps=2*self.steps_per_run, + scaffold=self.scaffold, + steps_per_run=self.steps_per_run) + hook.begin() + self.scaffold.finalize() + with session_lib.Session() as sess: + sess.run(self.scaffold.init_op) + mon_sess = monitored_session._HookedSession(sess, [hook]) + mon_sess.run(self.train_op) + self.assertEqual(5, + checkpoint_utils.load_variable(self.model_dir, + self.global_step.name)) + + def test_save_steps_saves_periodically(self): + with self.graph.as_default(): + hook = basic_session_run_hooks.CheckpointSaverHook( + self.model_dir, + save_steps=2*self.steps_per_run, + scaffold=self.scaffold, + steps_per_run=self.steps_per_run) + hook.begin() + self.scaffold.finalize() + with session_lib.Session() as sess: + sess.run(self.scaffold.init_op) + mon_sess = monitored_session._HookedSession(sess, [hook]) + mon_sess.run(self.train_op) + # Saved (step=5) + self.assertEqual(5, + checkpoint_utils.load_variable(self.model_dir, + self.global_step.name)) + + mon_sess.run(self.train_op) + # Not saved (step=10) + self.assertEqual(5, + checkpoint_utils.load_variable(self.model_dir, + self.global_step.name)) + + mon_sess.run(self.train_op) + # Saved (step=15) + self.assertEqual(15, + checkpoint_utils.load_variable(self.model_dir, + self.global_step.name)) + + mon_sess.run(self.train_op) + # Not saved (step=20) + self.assertEqual(15, + checkpoint_utils.load_variable(self.model_dir, + self.global_step.name)) + + mon_sess.run(self.train_op) + # Saved (step=25) + self.assertEqual(25, + checkpoint_utils.load_variable(self.model_dir, + self.global_step.name)) + + def test_save_steps_saves_at_end(self): + with self.graph.as_default(): + hook = basic_session_run_hooks.CheckpointSaverHook( + self.model_dir, + save_steps=2*self.steps_per_run, + scaffold=self.scaffold, + steps_per_run=self.steps_per_run) + hook.begin() + self.scaffold.finalize() + with session_lib.Session() as sess: + sess.run(self.scaffold.init_op) + mon_sess = monitored_session._HookedSession(sess, [hook]) + mon_sess.run(self.train_op) + mon_sess.run(self.train_op) + hook.end(sess) + self.assertEqual(10, + checkpoint_utils.load_variable(self.model_dir, + self.global_step.name)) + + class ResourceCheckpointSaverHookTest(test.TestCase): def setUp(self): diff --git a/tensorflow/python/training/checkpointable.py b/tensorflow/python/training/checkpointable.py index bbbe1e8ac5b985b11f2481ddcadedc06ed70a4fb..0b8473742c1b46c49d66a0bdbbd329aff58201ab 100644 --- a/tensorflow/python/training/checkpointable.py +++ b/tensorflow/python/training/checkpointable.py @@ -26,6 +26,11 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import gen_io_ops as io_ops from tensorflow.python.util import nest + +# Key where the object graph proto is saved in a TensorBundle +OBJECT_GRAPH_PROTO_KEY = "_CHECKPOINTABLE_OBJECT_GRAPH" + + # 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 @@ -94,12 +99,13 @@ class _CheckpointPosition(object): 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) + with ops.init_scope(): + 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. @@ -409,28 +415,29 @@ class CheckpointableBase(object): "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 - + with ops.init_scope(): + 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) diff --git a/tensorflow/python/training/checkpointable_utils.py b/tensorflow/python/training/checkpointable_utils.py index 32123f87ef2d12497077ab0e2f7d4d4cad1ec5dd..4769e15120c32707c2048c4fdcdce0ba3672b9d0 100644 --- a/tensorflow/python/training/checkpointable_utils.py +++ b/tensorflow/python/training/checkpointable_utils.py @@ -17,14 +17,47 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import abc +import collections import weakref +from tensorflow.core.protobuf 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.training import checkpointable +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 resource_variable_ops +from tensorflow.python.ops import variable_scope +from tensorflow.python.training import checkpointable as checkpointable_lib +from tensorflow.python.training import optimizer as optimizer_lib from tensorflow.python.training import saver as saver_lib +from tensorflow.python.util import deprecation +from tensorflow.python.util.tf_export import tf_export -class _Checkpoint(object): +_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" + + +class _CheckpointRestoreCoordinator(object): """Holds the status of an object-based checkpoint load.""" def __init__(self, object_graph_proto, save_path, dtype_map=None): @@ -72,7 +105,973 @@ class _Checkpoint(object): # `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 + checkpointable_lib._SlotVariableRestoration( # pylint: disable=protected-access optimizer_id=node_index, slot_variable_id=slot_reference.slot_variable_node_id, slot_name=slot_reference.slot_name)) + + +# 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 + .CheckpointableObject.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) + + +def _copy_saver_with_new_var_list(old_saver, new_var_list): + """Copy a `tf.train.Saver`'s state to a new Saver with different variables.""" + new_saver = saver_lib.Saver(var_list=new_var_list) + # TODO(allenl): Move to copying functionality to Saver? + # pylint: disable=protected-access + new_saver._last_checkpoints = old_saver._last_checkpoints + new_saver._checkpoints_to_be_deleted = old_saver._checkpoints_to_be_deleted + new_saver._next_checkpoint_time = old_saver._next_checkpoint_time + # pylint: enable=protected-access + return new_saver + + +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._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 checkpointable_lib.OBJECT_GRAPH_PROTO_KEY not in named_variables + named_variables[checkpointable_lib.OBJECT_GRAPH_PROTO_KEY] = ( + _NoRestoreSaveable( + tensor=object_graph_tensor, + name=checkpointable_lib.OBJECT_GRAPH_PROTO_KEY)) + if (self._last_save_object_graph != graph_proto + # When executing eagerly, we need to re-create SaveableObjects each time + # save() is called so they pick up new Tensors passed to their + # constructors. That means the Saver needs to be copied with a new + # var_list. + or context.executing_eagerly()): + if self._last_save_object_graph is not None: + self._last_save_saver = _copy_saver_with_new_var_list( + old_saver=self._last_save_saver, new_var_list=named_variables) + else: + self._last_save_saver = saver_lib.Saver(var_list=named_variables) + self._last_save_object_graph = graph_proto + with ops.device("/cpu:0"): + save_path = self._last_save_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): + """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. + + 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. + + 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: + file_prefix_tensor = self._file_prefix_placeholder + file_prefix_feed_dict = {self._file_prefix_placeholder: save_path} + else: + with ops.device("/cpu:0"): + file_prefix_tensor = constant_op.constant(save_path) + file_prefix_feed_dict = None + reader = pywrap_tensorflow.NewCheckpointReader(save_path) + try: + object_graph_string = reader.get_tensor( + checkpointable_lib.OBJECT_GRAPH_PROTO_KEY) + 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: + dtype_map = reader.get_variable_to_dtype_map() + checkpoint = _CheckpointRestoreCoordinator( + 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 + checkpointable_lib._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 + + +@tf_export("train.Checkpoint") +class Checkpoint(checkpointable_lib.Checkpointable): + """Groups checkpointable objects, saving and restoring them. + + `Checkpoint`'s constructor accepts keyword arguments whose values are types + that contain checkpointable state, such as `tf.train.Optimizer` + implementations, `tf.Variable`, `tf.keras.Layer` implementations, or + `tf.keras.Model` implementations. It saves these values with a checkpoint, and + maintains a `save_counter` for numbering checkpoints. + + Example usage when graph building: + + ```python + import tensorflow as tf + import os + + checkpoint_directory = "/tmp/training_checkpoints" + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + + checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model) + status = checkpoint.restore(tf.train.latest_checkpoint(checkpoint_directory)) + train_op = optimizer.minimize( ... ) + status.assert_consumed() # Optional sanity checks. + with tf.Session() as session: + # Use the Session to restore variables, or initialize them if + # tf.train.latest_checkpoint returned None. + status.initialize_or_restore(session) + for _ in range(num_training_steps): + session.run(train_op) + checkpoint.save(file_prefix=checkpoint_prefix) + ``` + + Example usage with eager execution enabled: + + ```python + import tensorflow as tf + import os + + tf.enable_eager_execution() + + checkpoint_directory = "/tmp/training_checkpoints" + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + + checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model) + status = checkpoint.restore(tf.train.latest_checkpoint(checkpoint_directory)) + for _ in range(num_training_steps): + optimizer.minimize( ... ) # Variables will be restored on creation. + status.assert_consumed() # Optional sanity checks. + checkpoint.save(file_prefix=checkpoint_prefix) + ``` + + `Checkpoint.save` and `Checkpoint.restore` write and read object-based + checkpoints, in contrast to `tf.train.Saver` which writes and reads + `variable.name` based checkpoints. Object-based checkpointing saves a graph of + dependencies between Python objects (`Layer`s, `Optimizer`s, `Variable`s, + etc.) with named edges, and this graph is used to match variables when + restoring a checkpoint. It can be more robust to changes in the Python + program, and helps to support restore-on-create for variables when executing + eagerly. Prefer `tf.train.Checkpoint` over `tf.train.Saver` for new code. + + `Checkpoint` objects have dependencies on the objects passed as keyword + arguments to their constructors, and each dependency is given a name that is + identical to the name of the keyword argument for which it was created. + TensorFlow classes like `Layer`s and `Optimizer`s will automatically add + dependencies on their variables (e.g. "kernel" and "bias" for + `tf.keras.layers.Dense`). Inheriting from `tf.keras.Model` makes managing + dependencies easy in user-defined classes, since `Model` hooks into attribute + assignment. For example: + + ```python + class Regress(tf.keras.Model): + + def __init__(self): + super(Regress, self).__init__() + self.input_transform = tf.keras.layers.Dense(10) + # ... + + def call(self, inputs): + x = self.input_transform(inputs) + # ... + ``` + + This `Model` has a dependency named "input_transform" on its `Dense` layer, + which in turn depends on its variables. As a result, saving an instance of + `Regress` using `tf.train.Checkpoint` will also save all the variables created + by the `Dense` layer. + + 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. Values must be checkpointable objects. + 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, checkpointable_lib.CheckpointableBase): + raise ValueError( + ("`Checkpoint` was expecting a checkpointable object (an object " + "derived from `CheckpointableBase`), got %s. If you believe this " + "object should be checkpointable (i.e. it is part of the " + "TensorFlow Python API and manages state), please open an issue.") + % (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 training checkpoint. + + The saved checkpoint includes variables created by this object and any + checkpointable objects it depends on at the time `Checkpoint.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.save_counter`. + 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. + """ + 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 training checkpoint. + + Restores this `Checkpoint` and any objects it depends on. + + When executing eagerly, either assigns values immediately if variables to + restore have been created already, or defers restoration until the variables + are created. Dependencies added after this call will be matched if they have + a corresponding object in the checkpoint (the restore request will queue in + any checkpointable object waiting for the expected dependency to be added). + + When graph building, restoration ops are added to the graph but not run + immediately. + + To ensure that loading is complete and no more assignments will take place, + use the `assert_consumed()` method of the status object returned by + `restore`: + + ```python + checkpoint = tf.train.Checkpoint( ... ) + checkpoint.restore(path).assert_consumed() + ``` + + An exception will be raised if any Python objects in the dependency graph + were not found in the checkpoint, or if any checkpointed values do not have + a matching Python object. + + When graph building, `assert_consumed()` indicates that all of the restore + ops that will be created for this checkpoint have been created. They can be + run via the `run_restore_ops()` method of the status object: + + ```python + checkpoint.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 + `tf.train.Checkpoint.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. + + Returns: + A load status object, which can be used to make assertions about the + status of a checkpoint restoration and run initialization/restore ops. + + The returned status object has the following methods: + - `assert_consumed()`: + Raises an exception if any variables/objects are unmatched: either + checkpointed values which don't have a matching Python object or + Python objects in the dependency graph with no values in the + checkpoint. This method returns the status object, and so may be + chained with `initialize_or_restore` or `run_restore_ops`. + - `initialize_or_restore(session=None)`: + When graph building, runs variable initializers if `save_path` is + `None`, but otherwise runs restore operations. If no `session` is + explicitly specified, the default session is used. No effect for + object-based checkpoints when executing eagerly (variables are + initialized or restored eagerly). + - `run_restore_ops(session=None)`: + When graph building, runs restore operations. If no `session` is + explicitly specified, the default session is used. No effect for + object-based checkpoints when executing eagerly (restore operations + are run eagerly). May only be called when `save_path` is not `None`. + """ + 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/python/training/checkpointable_utils_test.py similarity index 93% rename from tensorflow/contrib/eager/python/checkpointable_utils_test.py rename to tensorflow/python/training/checkpointable_utils_test.py index 1dd0f21a077470c69c58c4aea987788c0b290f6d..29fcdb70b418b08f543d75ff643b82d2a456754d 100644 --- a/tensorflow/contrib/eager/python/checkpointable_utils_test.py +++ b/tensorflow/python/training/checkpointable_utils_test.py @@ -21,10 +21,10 @@ 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 function from tensorflow.python.eager import test from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes @@ -32,8 +32,7 @@ from tensorflow.python.framework import ops from tensorflow.python.framework import test_util from tensorflow.python.keras._impl.keras.engine import sequential from tensorflow.python.keras._impl.keras.engine import training -from tensorflow.python.layers import core -from tensorflow.python.ops import array_ops +from tensorflow.python.keras._impl.keras.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 @@ -42,7 +41,8 @@ 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 checkpointable_utils +from tensorflow.python.training import saver as saver_lib from tensorflow.python.training import training_util @@ -70,87 +70,6 @@ class MyModel(training.Model): return ret -def _split_variable_closure(variable): - def _fill_save_buffer_fn(save_buffer): - save_buffer["first_half"] = variable[:2] - save_buffer["second_half"] = variable[2:] - return _fill_save_buffer_fn - - -def _combine_variable_closure(variable): - def _consume_restore_buffer_fn(restore_buffer): - return variable.assign( - array_ops.concat([restore_buffer["first_half"], - restore_buffer["second_half"]], - axis=0)) - return _consume_restore_buffer_fn - - -class SaveTensorSlicesAsDeps(checkpointable.CheckpointableBase): - - def __init__(self): - self.combined = resource_variable_ops.ResourceVariable([0., 0., 0., 0.]) - split_dependencies = checkpointable_utils.split_dependency( - component_names=("first_half", "second_half"), - component_dtypes=(self.combined.dtype,) * 2, - fill_save_buffer_fn=_split_variable_closure( - self.combined), - consume_restore_buffer_fn=_combine_variable_closure( - self.combined)) - for name, dep in split_dependencies.items(): - self._track_checkpointable(dep, name=name) - - -class HasRegularDeps(checkpointable.Checkpointable): - - def __init__(self): - self.first_half = resource_variable_ops.ResourceVariable([0., 0.]) - self.second_half = resource_variable_ops.ResourceVariable([0., 0.]) - - -class OnlyOneDep(checkpointable.Checkpointable): - - def __init__(self): - self.first_half = resource_variable_ops.ResourceVariable([0., 0.]) - - -class SplitTests(test.TestCase): - - @test_util.run_in_graph_and_eager_modes() - def testSaveRestoreSplitDep(self): - save_checkpoint = checkpointable_utils.Checkpoint( - dep=SaveTensorSlicesAsDeps()) - self.evaluate(save_checkpoint.dep.combined.assign([1., 2., 3., 4.])) - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - save_path = save_checkpoint.save(checkpoint_prefix) - - regular_deps = HasRegularDeps() - regular_restore_checkpoint = checkpointable_utils.Checkpoint( - dep=regular_deps) - regular_restore_checkpoint.restore( - save_path).assert_consumed().run_restore_ops() - self.assertAllEqual([1., 2.], self.evaluate(regular_deps.first_half)) - self.assertAllEqual([3., 4.], self.evaluate(regular_deps.second_half)) - - one_dep = OnlyOneDep() - one_dep_restore_checkpoint = checkpointable_utils.Checkpoint(dep=one_dep) - status = one_dep_restore_checkpoint.restore(save_path) - with self.assertRaises(AssertionError): - # Missing the second dependency. - status.assert_consumed() - status.run_restore_ops() - self.assertAllEqual([1., 2.], self.evaluate(one_dep.first_half)) - - restore_checkpoint = checkpointable_utils.Checkpoint() - status = restore_checkpoint.restore(save_path) - restore_checkpoint.dep = SaveTensorSlicesAsDeps() - status.assert_consumed().run_restore_ops() - self.assertAllEqual( - [1., 2., 3., 4.], - self.evaluate(restore_checkpoint.dep.combined)) - - class InterfaceTests(test.TestCase): @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) @@ -237,13 +156,13 @@ class InterfaceTests(test.TestCase): self.assertAllEqual([1., 1., 1.], self.evaluate(v2)) -class _MirroringSaveable(core_saver.BaseSaverBuilder.SaveableObject): +class _MirroringSaveable(saver_lib.BaseSaverBuilder.SaveableObject): def __init__(self, primary_variable, mirrored_variable, name): self._primary_variable = primary_variable self._mirrored_variable = mirrored_variable tensor = self._primary_variable.read_value() - spec = core_saver.BaseSaverBuilder.SaveSpec( + spec = saver_lib.BaseSaverBuilder.SaveSpec( tensor=tensor, slice_spec="", name=name) @@ -414,7 +333,7 @@ class CheckpointingTests(test.TestCase): def testMoreComplexSaveableReturnedWithGlobalName(self): # The same object can also be saved using the name-based saver. v = _OwnsMirroredVariables() - saver = core_saver.Saver(var_list=[v]) + saver = saver_lib.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.)) @@ -501,7 +420,7 @@ class CheckpointingTests(test.TestCase): 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)) + root.restore(saver_lib.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.]]) @@ -529,7 +448,7 @@ class CheckpointingTests(test.TestCase): train_op = optimizer.minimize( model(input_value), global_step=root.global_step) - checkpoint_path = core_saver.latest_checkpoint(checkpoint_directory) + checkpoint_path = saver_lib.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) @@ -562,7 +481,7 @@ class CheckpointingTests(test.TestCase): 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) + checkpoint_path = saver_lib.latest_checkpoint(checkpoint_directory) status = root.restore(save_path=checkpoint_path) input_value = constant_op.constant([[3.]]) train_fn = functools.partial( @@ -580,6 +499,50 @@ class CheckpointingTests(test.TestCase): self.assertEqual(training_continuation + 1, self.evaluate(root.save_counter)) + # pylint: disable=cell-var-from-loop + @test_util.run_in_graph_and_eager_modes() + def testWithDefun(self): + num_training_steps = 2 + 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() + # Don't actually train so we can test variable values + optimizer = adam.AdamOptimizer(0.) + root = checkpointable_utils.Checkpoint( + optimizer=optimizer, model=model, + global_step=training_util.get_or_create_global_step()) + checkpoint_path = saver_lib.latest_checkpoint(checkpoint_directory) + status = root.restore(save_path=checkpoint_path) + def train_fn(): + @function.defun + def _call_model(x): + return model(x) + with backprop.GradientTape() as tape: + loss = _call_model(constant_op.constant([[3.]])) + gradients = tape.gradient(loss, model.variables) + return optimizer.apply_gradients(zip(gradients, model.variables), + 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() + if training_continuation > 0: + status.assert_consumed() + self.assertAllClose([[42.]], self.evaluate(model.variables[0])) + else: + self.evaluate(model.variables[0].assign([[42.]])) + 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)) + # pylint: enable=cell-var-from-loop + def _get_checkpoint_name(self, name): root = checkpointable.Checkpointable() checkpointable_utils.add_variable( @@ -1281,7 +1244,7 @@ class CheckpointCompatibilityTests(test.TestCase): with save_graph.as_default(), self.test_session( graph=save_graph) as session: root = self._initialized_model() - name_saver = core_saver.Saver() + name_saver = saver_lib.Saver() return name_saver.save( sess=session, save_path=checkpoint_prefix, global_step=root.optimizer_step) @@ -1305,9 +1268,6 @@ class CheckpointCompatibilityTests(test.TestCase): 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") diff --git a/tensorflow/python/training/distribute.py b/tensorflow/python/training/distribute.py index c6b2dcdf98bd8436357a26c2ff69072b5009bed3..21ec5292adb5bb6463b5d46d934750fac7f8fcef 100644 --- a/tensorflow/python/training/distribute.py +++ b/tensorflow/python/training/distribute.py @@ -20,6 +20,7 @@ from __future__ import print_function import threading +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops @@ -391,7 +392,8 @@ class DistributionStrategy(object): ``` with my_distribution.scope(): - iterator = my_distribution.distribute_dataset(dataset) + iterator = my_distribution.distribute_dataset( + dataset).make_one_shot_iterator() tower_train_ops = my_distribution.call_for_each_tower( tower_fn, iterator.get_next()) train_op = tf.group(my_distribution.unwrap(tower_train_ops)) @@ -404,8 +406,14 @@ class DistributionStrategy(object): `tower_fn` can use the `get_tower_context()` API to get enhanced behavior in this case. - Note that in the future we will add support for initializable - Dataset iterators, at which point this example code will change. + You can also create an initializable iterator instead of one shot iterator. + In that case, you will need to ensure that you initialize the iterator + before calling get_next. + ``` + iterator = my_distribution.distribute_dataset( + dataset).make_initializable_iterator()) + session.run(iterator.initializer) + ``` * If you want to write a distributed algorithm, you may use any of the `DistributionStrategy` APIs inside a @@ -486,8 +494,8 @@ class DistributionStrategy(object): a variable (which by definition will have locality V(`v`), though will match another locality if inside a `colocate_vars_with` scope). - * `d.distribute_dataset(dataset)`: in cross-tower context, produces an - iterator with locality T + * `d.distribute_dataset(dataset).make_one_shot_iterator()`: in cross-tower + context, produces an iterator with locality T * `d.broadcast(t)`: in cross-tower context, produces a value with locality M * `d.broadcast(t, v)`: in cross-tower context, produces a value with locality V(`v`) @@ -510,7 +518,7 @@ class DistributionStrategy(object): The standard pattern for updating variables is to: - 1. Wrap your input dataset in `d.distribute_dataset()`. + 1. Wrap your input dataset in `d.distribute_dataset()` and create an iterator. 2. Define each tower `d.call_for_each_tower()` up to the point of getting a list of gradient, variable pairs. 3. Call `d.reduce("sum", t, v)` or `d.batch_reduce()` to sum the @@ -665,25 +673,38 @@ class DistributionStrategy(object): _require_distribution_strategy_scope(self) return variable_scope.variable_creator_scope(create_colocated_variable) - # TODO(josh11b): Currently this returns an iterator, but should return - # something implementing (a subset of) the Dataset API. - def distribute_dataset(self, dataset): - """Return an iterator into `dataset` split across all towers. + def _call_dataset_fn(self, dataset_fn): + result = dataset_fn() + if not isinstance(result, dataset_ops.Dataset): + raise ValueError( + "dataset_fn() must return a tf.data.Dataset when using a " + "DistributionStrategy.") + return result + + # TODO(josh11b): `PerDeviceDataset` currently only implements a few methods of + # Dataset API such as make_one_shot_iterator and make_initializable_iterator. + # Extend to implement more functionality of datasets. + def distribute_dataset(self, dataset_fn): + """Return a `dataset` split across all towers. - Suitable for providing input to for `call_for_each_tower()`, as in: + Suitable for providing input to for `call_for_each_tower()` by creating an + iterator: ``` + def dataset_fn(): + return tf.data.Dataset.from_tensors([[1.]]).repeat() with distribution_strategy.scope(): - iterator = distribution_strategy.distribute_dataset(dataset) + distributed_dataset = distribution_strategy.distribute_dataset(dataset_fn) + iterator = distributed_dataset.make_one_shot_iterator() tower_results = distribution_strategy.call_for_each_tower( tower_fn, iterator.get_next()) ``` Args: - dataset: A `tf.data.Dataset`. + dataset_fn: A function that returns a `tf.data.Dataset`. Returns: - A Dataset iterator that will produce separate splits for each tower. + A `PerDeviceDataset` that will produce data for each tower. """ raise NotImplementedError("must be implemented in descendants") @@ -1125,10 +1146,8 @@ class _DefaultDistributionStrategy(DistributionStrategy): _require_distribution_strategy_scope(self) return ops.colocate_with(colocate_with_variable) - def distribute_dataset(self, dataset): - # TODO(josh11b): Support for this when executing eagerly is currently only - # in contrib. - return dataset.make_one_shot_iterator() + def distribute_dataset(self, dataset_fn): + return self._call_dataset_fn(dataset_fn) def _broadcast(self, tensor, destinations): if destinations is None: diff --git a/tensorflow/python/training/momentum_test.py b/tensorflow/python/training/momentum_test.py index 297a8bbde5447cff9465be36c0bb71f2490c60fc..7bd57ad3d854534e196fa7b72bebbd7195e6bca8 100644 --- a/tensorflow/python/training/momentum_test.py +++ b/tensorflow/python/training/momentum_test.py @@ -237,7 +237,17 @@ class MomentumOptimizerTest(test.TestCase): @test_util.run_in_graph_and_eager_modes(reset_test=True) def testMinimizeSparseResourceVariable(self): for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype) + # This test invokes the ResourceSparseApplyMomentum operation, which + # did not have a registered GPU kernel as of April 2018. With graph + # execution, the placement algorithm notices this and automatically + # places the variable in CPU (host) memory. With eager execution, + # the variable would be placed in GPU memory if available, which + # would then conflict with the future invocation of the + # ResourceSparseApplyMomentum operation. + # To work around this discrepancy, for now we force the variable + # to be placed on CPU. + with ops.device("/cpu:0"): + var0 = resource_variable_ops.ResourceVariable([[1.0, 2.0]], dtype=dtype) # pylint: disable=cell-var-from-loop def loss(): @@ -256,7 +266,17 @@ class MomentumOptimizerTest(test.TestCase): @test_util.run_in_graph_and_eager_modes(reset_test=True) def testMinimizeWith2DIndiciesForEmbeddingLookup(self): - var0 = resource_variable_ops.ResourceVariable(array_ops.ones([2, 2])) + # This test invokes the ResourceSparseApplyMomentum operation, which + # did not have a registered GPU kernel as of April 2018. With graph + # execution, the placement algorithm notices this and automatically + # places the variable in CPU (host) memory. With eager execution, + # the variable would be placed in GPU memory if available, which + # would then conflict with the future invocation of the + # ResourceSparseApplyMomentum operation. + # To work around this discrepancy, for now we force the variable + # to be placed on CPU. + with ops.device("/cpu:0"): + var0 = resource_variable_ops.ResourceVariable(array_ops.ones([2, 2])) def loss(): return math_ops.reduce_sum(embedding_ops.embedding_lookup(var0, [[1]])) diff --git a/tensorflow/python/training/monitored_session.py b/tensorflow/python/training/monitored_session.py index 4ce6f6d00267410626f7d7a9e2251d3f40b6bb6e..f584a009d946a193f1ab76b3030db4f8a4954d27 100644 --- a/tensorflow/python/training/monitored_session.py +++ b/tensorflow/python/training/monitored_session.py @@ -202,7 +202,7 @@ class Scaffold(object): if self._local_init_op is None: self._local_init_op = Scaffold.get_or_default( 'local_init_op', ops.GraphKeys.LOCAL_INIT_OP, - Scaffold._default_local_init_op) + Scaffold.default_local_init_op) if self._summary_op is None: self._summary_op = Scaffold.get_or_default('summary_op', ops.GraphKeys.SUMMARY_OP, @@ -267,7 +267,17 @@ class Scaffold(object): return op @staticmethod - def _default_local_init_op(): + def default_local_init_op(): + """Returns an op that groups the default local init ops. + + This op is used during session initialization when a Scaffold is + initialized without specifying the local_init_op arg. It includes + `tf.local_variables_initializer`, `tf.tables_initializer`, and also + initializes local session resources. + + Returns: + The default Scaffold local init op. + """ return control_flow_ops.group( variables.local_variables_initializer(), lookup_ops.tables_initializer(), diff --git a/tensorflow/python/training/optimizer.py b/tensorflow/python/training/optimizer.py index 75665fc2840797dd53dd863721d2744cb1b08af5..f126d3847b6b0b43495c63b31ca915c107ede969 100644 --- a/tensorflow/python/training/optimizer.py +++ b/tensorflow/python/training/optimizer.py @@ -689,9 +689,7 @@ class Optimizer( # device_policy is set because non-mirrored tensors will be read in # `update_op`. `_resource_apply_dense`, `lr_t`, `beta1_t` and `beta2_t` # is an example. - with ops.name_scope( - "update_" + scope_name), context.context().device_policy( - context.DEVICE_PLACEMENT_SILENT): + with ops.name_scope("update_" + scope_name): return p.update_op(self, g) with ops.name_scope(name, self._name) as name: @@ -707,11 +705,8 @@ class Optimizer( return self._finish(update_ops, "update") non_slot_devices = distribution.non_slot_devices(var_list) - # Device policy is needed because hyperparameter tensors (such as - # AdamOptimizer's beta1_t) need to be copied across devices in Eager. - with context.context().device_policy(context.DEVICE_PLACEMENT_SILENT): - finish_updates = distribution.update_non_slot( - non_slot_devices, finish, self, update_ops) + finish_updates = distribution.update_non_slot( + non_slot_devices, finish, self, update_ops) if global_step is None: apply_updates = distribution.group(finish_updates, name=name) else: @@ -823,13 +818,13 @@ class Optimizer( 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) + # 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 diff --git a/tensorflow/python/training/saver.py b/tensorflow/python/training/saver.py index e40b8d22ed2ab0f4c9ff65e953f0f1cf681c8068..a74d629a8f81bacd019417b7a359d51f1080b16c 100644 --- a/tensorflow/python/training/saver.py +++ b/tensorflow/python/training/saver.py @@ -22,6 +22,7 @@ from __future__ import print_function import collections import os.path import re +import sys import time import uuid @@ -30,8 +31,10 @@ import six from google.protobuf import text_format +from tensorflow.core.protobuf import checkpointable_object_graph_pb2 from tensorflow.core.protobuf import meta_graph_pb2 from tensorflow.core.protobuf import saver_pb2 +from tensorflow.python import pywrap_tensorflow from tensorflow.python.client import session from tensorflow.python.eager import context from tensorflow.python.framework import constant_op @@ -1340,6 +1343,9 @@ class Saver(object): self._check_saver_def() self._write_version = self.saver_def.version self._save_relative_paths = save_relative_paths + # For compatibility with object-based checkpoints, we may build a second + # Saver to read the renamed keys. + self._object_restore_saver = None def build(self): if context.executing_eagerly(): @@ -1795,11 +1801,63 @@ class Saver(object): if save_path is None: raise ValueError("Can't load save_path when it is None.") logging.info("Restoring parameters from %s", save_path) - 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}) + try: + 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}) + except errors.NotFoundError: + exception_type, exception_value, exception_traceback = sys.exc_info() + # The checkpoint would not be loaded successfully as is. Try to parse it + # as an object-based checkpoint. + try: + reader = pywrap_tensorflow.NewCheckpointReader(save_path) + object_graph_string = reader.get_tensor( + checkpointable.OBJECT_GRAPH_PROTO_KEY) + except errors.NotFoundError: + # This is not an object-based checkpoint, or the checkpoint doesn't + # exist. Re-raise the original exception. + six.reraise(exception_type, exception_value, exception_traceback) + del exception_traceback # avoid reference cycles + + # This is an object-based checkpoint. We'll print a warning and then do + # the restore. + logging.warning( + "Restoring an object-based checkpoint using a name-based saver. This " + "may be somewhat fragile, and will re-build the Saver. Instead, " + "consider loading object-based checkpoints using " + "tf.train.Checkpoint().") + self._restore_from_object_based_checkpoint( + sess=sess, save_path=save_path, + object_graph_string=object_graph_string) + + def _restore_from_object_based_checkpoint(self, sess, save_path, + object_graph_string): + """A compatibility mode for reading object-based checkpoints.""" + object_graph_proto = ( + checkpointable_object_graph_pb2.CheckpointableObjectGraph()) + object_graph_proto.ParseFromString(object_graph_string) + names_to_keys = {} + for node in object_graph_proto.nodes: + for attribute in node.attributes: + names_to_keys[attribute.full_name] = attribute.checkpoint_key + saveables = self._builder._ValidateAndSliceInputs(self._var_list) # pylint: disable=protected-access + for saveable in saveables: + for spec in saveable.specs: + if spec.name not in names_to_keys: + raise errors.NotFoundError( + None, None, + message=("Attempting to load an object-based checkpoint using " + "variable names, but could not find %s in the " + "checkpoint.") % spec.name) + spec.name = names_to_keys[spec.name] + if self._object_restore_saver is None: + # Cache the Saver so multiple restore() calls don't pollute the graph when + # graph building. This assumes keys are consistent (i.e. this is the same + # type of object-based checkpoint we saw previously). + self._object_restore_saver = Saver(saveables) + self._object_restore_saver.restore(sess=sess, save_path=save_path) @staticmethod def _add_collection_def(meta_graph_def, key, export_scope=None): diff --git a/tensorflow/python/training/saver_test.py b/tensorflow/python/training/saver_test.py index 14dda7997948ead7b12dee953a0b2ee3b2ee8fc9..70495291bc57629252423d86e35987c5c0d2b1ee 100644 --- a/tensorflow/python/training/saver_test.py +++ b/tensorflow/python/training/saver_test.py @@ -19,6 +19,7 @@ from __future__ import division from __future__ import print_function import contextlib +import functools import math import os import random @@ -50,6 +51,8 @@ from tensorflow.python.framework import graph_io from tensorflow.python.framework import meta_graph from tensorflow.python.framework import ops as ops_lib from tensorflow.python.framework import test_util +from tensorflow.python.keras._impl.keras.engine import training +from tensorflow.python.keras._impl.keras.layers import core from tensorflow.python.lib.io import file_io from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops @@ -69,10 +72,12 @@ 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 checkpointable_utils from tensorflow.python.training import gradient_descent from tensorflow.python.training import queue_runner_impl from tensorflow.python.training import saver as saver_module from tensorflow.python.training import saver_test_utils +from tensorflow.python.training import training_util from tensorflow.python.training.checkpoint_state_pb2 import CheckpointState from tensorflow.python.util import compat @@ -2726,7 +2731,7 @@ class ScopedGraphTest(test.TestCase): # The rest of the variables. rest_variables = list( set(variables.global_variables()) - set(var_list.keys())) - init_rest_op = variables.initialize_variables(rest_variables) + init_rest_op = variables.variables_initializer(rest_variables) with self.test_session(graph=graph) as sess: saver = saver_module.Saver(var_list=var_list, max_to_keep=1) @@ -2948,6 +2953,29 @@ class _OwnsMirroredVariables(checkpointable.CheckpointableBase): return self.non_dep_variable.name +class NonLayerCheckpointable(checkpointable.Checkpointable): + + def __init__(self): + super(NonLayerCheckpointable, self).__init__() + self.a_variable = checkpointable_utils.add_variable( + self, name="a_variable", shape=[]) + + +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 + + @test_util.with_c_api class CheckpointableCompatibilityTests(test.TestCase): @@ -3011,6 +3039,128 @@ class CheckpointableCompatibilityTests(test.TestCase): saver.restore(sess, save_path) self.assertEqual(1, v.eval_count) + 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 testVariableNotFoundErrorRaised(self): + # Restore does some tricky exception handling to figure out if it should + # load an object-based checkpoint. Tests that the exception handling isn't + # too broad. + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + + a = resource_variable_ops.ResourceVariable(1., name="a") + b = resource_variable_ops.ResourceVariable(1., name="b") + a_saver = saver_module.Saver([a]) + b_saver = saver_module.Saver([b]) + with self.test_session() as sess: + sess.run(a.initializer) + save_path = a_saver.save(sess=sess, save_path=checkpoint_prefix) + with self.assertRaisesRegexp( + errors.NotFoundError, "Key b not found in checkpoint"): + b_saver.restore(sess=sess, save_path=save_path) + + def testCheckpointNotFoundErrorRaised(self): + # Restore does some tricky exception handling to figure out if it should + # load an object-based checkpoint. Tests that the exception handling isn't + # too broad. + a = resource_variable_ops.ResourceVariable(1., name="a") + saver = saver_module.Saver([a]) + with self.test_session() as sess: + with self.assertRaisesRegexp( + errors.NotFoundError, + "Failed to find any matching files for path_which_does_not_exist"): + saver.restore(sess=sess, save_path="path_which_does_not_exist") + + def testLoadFromObjectBasedGraph(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + + save_graph = ops_lib.Graph() + with save_graph.as_default(), self.test_session(graph=save_graph) as sess: + root = self._initialized_model() + object_saver = checkpointable_utils.CheckpointableSaver(root) + save_path = object_saver.save(file_prefix=checkpoint_prefix) + + # An incompatible object-based checkpoint to check error messages + var = resource_variable_ops.ResourceVariable(1., name="a") + self.evaluate(var.initializer) + second_saver = checkpointable_utils.CheckpointableSaver(var) + second_path = second_saver.save(file_prefix=os.path.join( + checkpoint_directory, "second")) + + restore_graph = ops_lib.Graph() + with restore_graph.as_default(), self.test_session( + graph=restore_graph) as sess: + root = self._initialized_model() + self._set_sentinels(root) + saver = saver_module.Saver() + saver.restore(sess=sess, save_path=save_path) + self._check_sentinels(root) + before_second_restore_ops = restore_graph.get_operations() + # Test that multiple restores do not pollute the graph + saver.restore(sess=sess, save_path=save_path) + self.assertEqual(before_second_restore_ops, + restore_graph.get_operations()) + with self.assertRaisesRegexp(errors.NotFoundError, + "could not find a_variable"): + saver.restore(sess=sess, save_path=second_path) + + def testLoadFromObjectBasedEager(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + + save_graph = ops_lib.Graph() + with save_graph.as_default(), self.test_session(graph=save_graph): + root = self._initialized_model() + object_saver = checkpointable_utils.CheckpointableSaver(root) + save_path = object_saver.save(file_prefix=checkpoint_prefix) + + with context.eager_mode(): + root = self._initialized_model() + self._set_sentinels(root) + saver = saver_module.Saver( + root.model.variables + root.optimizer.variables()) + saver.restore(sess=None, save_path=save_path) + self._check_sentinels(root) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/training/session_manager.py b/tensorflow/python/training/session_manager.py index a00ceb90211e371c3b2f2b32f2042d1556158595..3cb3877cc2fb846477640058b5bcdfea97ee6828 100644 --- a/tensorflow/python/training/session_manager.py +++ b/tensorflow/python/training/session_manager.py @@ -263,8 +263,6 @@ class SessionManager(object): Raises: RuntimeError: If the model cannot be initialized or recovered. - - Raises: ValueError: If both checkpoint_dir and checkpoint_filename_with_path are set. """ diff --git a/tensorflow/python/training/training.py b/tensorflow/python/training/training.py index b759b156d78cf8d869b49375058cc7ed42e82b34..d7e5078be7be6704597599efa2c9ccdad17833e1 100644 --- a/tensorflow/python/training/training.py +++ b/tensorflow/python/training/training.py @@ -156,6 +156,7 @@ from tensorflow.python.training.basic_session_run_hooks import FinalOpsHook from tensorflow.python.training.basic_session_run_hooks import FeedFnHook from tensorflow.python.training.basic_session_run_hooks import ProfilerHook from tensorflow.python.training.basic_loops import basic_train_loop +from tensorflow.python.training.checkpointable_utils import Checkpoint from tensorflow.python.training.checkpoint_utils import init_from_checkpoint from tensorflow.python.training.checkpoint_utils import list_variables from tensorflow.python.training.checkpoint_utils import load_checkpoint diff --git a/tensorflow/python/training/warm_starting_util_test.py b/tensorflow/python/training/warm_starting_util_test.py index 6e445d8bd14cc13010541c1ab0f737f96a4b1e03..7e8cbd6baeea160075b61d1191c8f1da5fe2163c 100644 --- a/tensorflow/python/training/warm_starting_util_test.py +++ b/tensorflow/python/training/warm_starting_util_test.py @@ -946,18 +946,20 @@ class WarmStartingUtilTest(test.TestCase): # emb_vocab should be correctly warm-started after vocab remapping. # Missing values are filled in with the EmbeddingColumn's initializer. self._assert_cols_to_vars( - cols_to_vars, { + cols_to_vars, + { emb_vocab: [ - # embedding_weights part 0. - np.array([[3., 3.3], [2., 2.2], [1., 1.1]]), - # embedding_weights part 1. - np.array([[0.5, 0.4], [0.42, 0.42], [0.42, 0.42]]), # linear weights part 0. np.array([[0.69]]), # linear weights part 1. - np.array([[0.71]]) + np.array([[0.71]]), + # embedding_weights part 0. + np.array([[3., 3.3], [2., 2.2], [1., 1.1]]), + # embedding_weights part 1. + np.array([[0.5, 0.4], [0.42, 0.42], [0.42, 0.42]]) ] - }, sess) + }, + sess) def testErrorConditions(self): x = variable_scope.get_variable( diff --git a/tensorflow/python/util/compat.py b/tensorflow/python/util/compat.py index 738479c946d5c3f9e106d4222ffb15dea86c6c62..3358ffe5264b16b53ae38b2108ad0fe016ea8972 100644 --- a/tensorflow/python/util/compat.py +++ b/tensorflow/python/util/compat.py @@ -42,7 +42,6 @@ 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 def as_bytes(bytes_or_text, encoding='utf-8'): diff --git a/tensorflow/python/util/tf_inspect.py b/tensorflow/python/util/tf_inspect.py index 286028b8bbe6bc89b9d7e479a8f1b5c06d7ba5cb..663036de8a01c3937ae93f8d2bfa27f7add48e39 100644 --- a/tensorflow/python/util/tf_inspect.py +++ b/tensorflow/python/util/tf_inspect.py @@ -17,21 +17,21 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import inspect as _inspect -import six from collections import namedtuple +import inspect as _inspect from tensorflow.python.util import tf_decorator ArgSpec = _inspect.ArgSpec -if six.PY3: - FullArgSpec = _inspect.FullArgSpec +if hasattr(_inspect, 'FullArgSpec'): + FullArgSpec = _inspect.FullArgSpec # pylint: disable=invalid-name else: - FullArgSpec = namedtuple( - 'FullArgSpec', ['args', 'varargs', 'varkw', 'defaults', - 'kwonlyargs', 'kwonlydefaults', 'annotations']) + FullArgSpec = namedtuple('FullArgSpec', [ + 'args', 'varargs', 'varkw', 'defaults', 'kwonlyargs', 'kwonlydefaults', + 'annotations' + ]) def currentframe(): @@ -70,8 +70,20 @@ def getfullargspec(obj): # pylint: disable=redefined-builtin callable is not decorated, `inspect.getfullargspec()` will be called directly on the callable. """ - if six.PY2: + if hasattr(_inspect, 'getfullargspec'): + spec_fn = _inspect.getfullargspec + else: def spec_fn(target): + """Spec function that adding default value from FullArgSpec. + + It is used when getfullargspec is not available (eg in PY2). + + Args: + target: the target object to inspect. + Returns: + The full argument specs with empty kwonlyargs, kwonlydefaults and + annotations. + """ argspecs = _inspect.getargspec(target) fullargspecs = FullArgSpec( args=argspecs.args, @@ -82,8 +94,6 @@ def getfullargspec(obj): # pylint: disable=redefined-builtin kwonlydefaults=None, annotations={}) return fullargspecs - else: - spec_fn = _inspect.getfullargspec decorators, target = tf_decorator.unwrap(obj) return next((d.decorator_argspec for d in decorators diff --git a/tensorflow/stream_executor/BUILD b/tensorflow/stream_executor/BUILD index 1913fc20ee0212b3d9588828fe4da4ba7ebca030..80fc9ff2926c53b557a7ba9e242d597a89acf79f 100644 --- a/tensorflow/stream_executor/BUILD +++ b/tensorflow/stream_executor/BUILD @@ -33,6 +33,7 @@ cc_library( }), visibility = ["//visibility:public"], deps = [ + "//tensorflow/compiler/xla:statusor", "//tensorflow/core:lib", "@local_config_cuda//cuda:cuda_headers", ], @@ -45,6 +46,7 @@ cc_library( visibility = ["//visibility:public"], deps = [ "//tensorflow/core:lib", + "//tensorflow/compiler/xla:statusor", "@local_config_cuda//cuda:cuda_headers", ] + if_static([":stream_executor_impl"]), ) diff --git a/tensorflow/stream_executor/blas.cc b/tensorflow/stream_executor/blas.cc index 31724cf6c9b97e45975b9e053459f7b8f5918dfa..906d6fb7020ce35adb1438d394b34983c332f182 100644 --- a/tensorflow/stream_executor/blas.cc +++ b/tensorflow/stream_executor/blas.cc @@ -17,8 +17,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/strcat.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace blas { string TransposeString(Transpose t) { @@ -95,5 +94,4 @@ std::ostream& operator<<(std::ostream& os, ComputationType ty) { } } // namespace blas -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/blas.h b/tensorflow/stream_executor/blas.h index c5f778a5c74519c0f35cea5d59aac3d0d4564c56..6e62b85728acd8db551a7e27459a6804a9f3d4b2 100644 --- a/tensorflow/stream_executor/blas.h +++ b/tensorflow/stream_executor/blas.h @@ -49,8 +49,7 @@ namespace Eigen { struct half; } // namespace Eigen -namespace perftools { -namespace gputools { +namespace stream_executor { class Stream; class ScratchAllocator; @@ -2100,7 +2099,6 @@ class BlasSupport { DeviceMemory> *b, int ldb) override; } // namespace blas -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_BLAS_H_ diff --git a/tensorflow/stream_executor/cuda/cuda_activation.cc b/tensorflow/stream_executor/cuda/cuda_activation.cc index 5f4cf9dbd781e865345143a964fc16ca4fdb18d4..cf6b9e2c6e4b327c06ecce318f8e8809308f6f02 100644 --- a/tensorflow/stream_executor/cuda/cuda_activation.cc +++ b/tensorflow/stream_executor/cuda/cuda_activation.cc @@ -19,8 +19,7 @@ limitations under the License. #include "tensorflow/stream_executor/stream_executor.h" #include "tensorflow/stream_executor/stream_executor_internal.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { CudaContext* ExtractCudaContext(CUDAExecutor *cuda_exec); @@ -40,5 +39,4 @@ ScopedActivateExecutorContext::~ScopedActivateExecutorContext() { } } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/cuda/cuda_activation.h b/tensorflow/stream_executor/cuda/cuda_activation.h index c9d43a9766e049ae002df3e4ce4d8625e71f99db..04ffaef3646bb3df03407f8e74e620455bb9e5cb 100644 --- a/tensorflow/stream_executor/cuda/cuda_activation.h +++ b/tensorflow/stream_executor/cuda/cuda_activation.h @@ -25,8 +25,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { class StreamExecutor; @@ -56,7 +55,6 @@ class ScopedActivateExecutorContext { }; } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_ACTIVATION_H_ diff --git a/tensorflow/stream_executor/cuda/cuda_blas.cc b/tensorflow/stream_executor/cuda/cuda_blas.cc index 1c550dbb13657d39aecb3f7373efd12bae80aabb..007c0f1c86c77b035a31b4fee433ec71c26bdf9b 100644 --- a/tensorflow/stream_executor/cuda/cuda_blas.cc +++ b/tensorflow/stream_executor/cuda/cuda_blas.cc @@ -75,15 +75,14 @@ limitations under the License. #include "tensorflow/stream_executor/scratch_allocator.h" #include "tensorflow/stream_executor/stream_executor.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { PLUGIN_REGISTRY_DEFINE_PLUGIN_ID(kCuBlasPlugin); namespace wrap { -#define PERFTOOLS_GPUTOOLS_CUBLAS_WRAP(__name) \ +#define STREAM_EXECUTOR_CUBLAS_WRAP(__name) \ struct WrapperShim__##__name { \ static const char *kName; \ template \ @@ -94,8 +93,8 @@ namespace wrap { } __name; \ const char *WrapperShim__##__name::kName = #__name; -#define PERFTOOLS_GPUTOOLS_CUBLAS_V2_WRAP(__name) \ - PERFTOOLS_GPUTOOLS_CUBLAS_WRAP(__name) +#define STREAM_EXECUTOR_CUBLAS_V2_WRAP(__name) \ + STREAM_EXECUTOR_CUBLAS_WRAP(__name) #define CUBLAS_BLAS_ROUTINE_EACH(__macro) \ __macro(cublasSnrm2) \ @@ -269,28 +268,28 @@ namespace wrap { __macro(cublasCdgmm) \ __macro(cublasZdgmm) -PERFTOOLS_GPUTOOLS_CUBLAS_V2_WRAP(cublasCreate) -PERFTOOLS_GPUTOOLS_CUBLAS_V2_WRAP(cublasDestroy) -PERFTOOLS_GPUTOOLS_CUBLAS_V2_WRAP(cublasSetStream) -PERFTOOLS_GPUTOOLS_CUBLAS_V2_WRAP(cublasSetPointerMode) -PERFTOOLS_GPUTOOLS_CUBLAS_V2_WRAP(cublasGetPointerMode) -PERFTOOLS_GPUTOOLS_CUBLAS_WRAP(cublasSgemmBatched) -PERFTOOLS_GPUTOOLS_CUBLAS_WRAP(cublasDgemmBatched) -PERFTOOLS_GPUTOOLS_CUBLAS_WRAP(cublasCgemmBatched) -PERFTOOLS_GPUTOOLS_CUBLAS_WRAP(cublasZgemmBatched) -CUBLAS_BLAS_ROUTINE_EACH(PERFTOOLS_GPUTOOLS_CUBLAS_V2_WRAP) +STREAM_EXECUTOR_CUBLAS_V2_WRAP(cublasCreate) +STREAM_EXECUTOR_CUBLAS_V2_WRAP(cublasDestroy) +STREAM_EXECUTOR_CUBLAS_V2_WRAP(cublasSetStream) +STREAM_EXECUTOR_CUBLAS_V2_WRAP(cublasSetPointerMode) +STREAM_EXECUTOR_CUBLAS_V2_WRAP(cublasGetPointerMode) +STREAM_EXECUTOR_CUBLAS_WRAP(cublasSgemmBatched) +STREAM_EXECUTOR_CUBLAS_WRAP(cublasDgemmBatched) +STREAM_EXECUTOR_CUBLAS_WRAP(cublasCgemmBatched) +STREAM_EXECUTOR_CUBLAS_WRAP(cublasZgemmBatched) +CUBLAS_BLAS_ROUTINE_EACH(STREAM_EXECUTOR_CUBLAS_V2_WRAP) #if CUDA_VERSION >= 7050 -PERFTOOLS_GPUTOOLS_CUBLAS_WRAP(cublasSgemmEx) +STREAM_EXECUTOR_CUBLAS_WRAP(cublasSgemmEx) #endif #if CUDA_VERSION >= 8000 -PERFTOOLS_GPUTOOLS_CUBLAS_WRAP(cublasGemmEx) +STREAM_EXECUTOR_CUBLAS_WRAP(cublasGemmEx) #endif #if CUDA_VERSION >= 9000 -PERFTOOLS_GPUTOOLS_CUBLAS_WRAP(cublasGetMathMode) -PERFTOOLS_GPUTOOLS_CUBLAS_WRAP(cublasSetMathMode) +STREAM_EXECUTOR_CUBLAS_WRAP(cublasGetMathMode) +STREAM_EXECUTOR_CUBLAS_WRAP(cublasSetMathMode) #endif } // namespace wrap @@ -2803,46 +2802,39 @@ bool CUDABlas::DoBlasTrsm(Stream *stream, blas::Side side, } // namespace cuda -namespace gpu = ::perftools::gputools; - void initialize_cublas() { - gpu::port::Status status = - gpu::PluginRegistry::Instance() - ->RegisterFactory( - gpu::cuda::kCudaPlatformId, gpu::cuda::kCuBlasPlugin, "cuBLAS", - [](gpu::internal::StreamExecutorInterface - *parent) -> gpu::blas::BlasSupport * { - gpu::cuda::CUDAExecutor *cuda_executor = - dynamic_cast(parent); - if (cuda_executor == nullptr) { - LOG(ERROR) - << "Attempting to initialize an instance of the cuBLAS " - << "support library with a non-CUDA StreamExecutor"; - return nullptr; - } - - gpu::cuda::CUDABlas *blas = - new gpu::cuda::CUDABlas(cuda_executor); - if (!blas->Init()) { - // Note: Init() will log a more specific error. - delete blas; - return nullptr; - } - return blas; - }); + port::Status status = + PluginRegistry::Instance()->RegisterFactory( + cuda::kCudaPlatformId, cuda::kCuBlasPlugin, "cuBLAS", + [](internal::StreamExecutorInterface *parent) -> blas::BlasSupport * { + cuda::CUDAExecutor *cuda_executor = + dynamic_cast(parent); + if (cuda_executor == nullptr) { + LOG(ERROR) + << "Attempting to initialize an instance of the cuBLAS " + << "support library with a non-CUDA StreamExecutor"; + return nullptr; + } + + cuda::CUDABlas *blas = new cuda::CUDABlas(cuda_executor); + if (!blas->Init()) { + // Note: Init() will log a more specific error. + delete blas; + return nullptr; + } + return blas; + }); if (!status.ok()) { LOG(ERROR) << "Unable to register cuBLAS factory: " << status.error_message(); } - gpu::PluginRegistry::Instance()->SetDefaultFactory(gpu::cuda::kCudaPlatformId, - gpu::PluginKind::kBlas, - gpu::cuda::kCuBlasPlugin); + PluginRegistry::Instance()->SetDefaultFactory( + cuda::kCudaPlatformId, PluginKind::kBlas, cuda::kCuBlasPlugin); } -} // namespace gputools -} // namespace perftools +} // namespace stream_executor REGISTER_MODULE_INITIALIZER(register_cublas, - { perftools::gputools::initialize_cublas(); }); + { stream_executor::initialize_cublas(); }); diff --git a/tensorflow/stream_executor/cuda/cuda_blas.h b/tensorflow/stream_executor/cuda/cuda_blas.h index deb211c04bcaa9e98ee04c5e9066a2a13092cb06..55c414a1f921c40efa5fd7a3d1bf0564fc5fbb9c 100644 --- a/tensorflow/stream_executor/cuda/cuda_blas.h +++ b/tensorflow/stream_executor/cuda/cuda_blas.h @@ -29,8 +29,7 @@ limitations under the License. typedef struct cublasContext *cublasHandle_t; -namespace perftools { -namespace gputools { +namespace stream_executor { class Stream; @@ -162,7 +161,6 @@ class CUDABlas : public blas::BlasSupport { }; } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_BLAS_H_ diff --git a/tensorflow/stream_executor/cuda/cuda_diagnostics.cc b/tensorflow/stream_executor/cuda/cuda_diagnostics.cc index 933c103f524ef37f840c9e13b9e4024289e274c1..feb529297e8fff99030cb14fd4f8a9298bbc5935 100644 --- a/tensorflow/stream_executor/cuda/cuda_diagnostics.cc +++ b/tensorflow/stream_executor/cuda/cuda_diagnostics.cc @@ -51,8 +51,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/str_util.h" #include "tensorflow/stream_executor/lib/inlined_vector.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { #ifdef __APPLE__ @@ -384,5 +383,4 @@ port::StatusOr Diagnostician::FindKernelDriverVersion() { } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/cuda/cuda_diagnostics.h b/tensorflow/stream_executor/cuda/cuda_diagnostics.h index aa68321acc858c902d1a43600a14ac5d88edb0be..f2db2eb20a18c671e055b910809dfde940a5e3f8 100644 --- a/tensorflow/stream_executor/cuda/cuda_diagnostics.h +++ b/tensorflow/stream_executor/cuda/cuda_diagnostics.h @@ -22,8 +22,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/statusor.h" #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { // e.g. DriverVersion{346, 3, 4} @@ -93,7 +92,6 @@ class Diagnostician { }; } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_DIAGNOSTICS_H_ diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc index 1dc7f991b38c00d9df06322b869911d2866dd965..102419a264984aa2edb4fee756e9a53ef615d5bc 100644 --- a/tensorflow/stream_executor/cuda/cuda_dnn.cc +++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc @@ -59,8 +59,7 @@ NarrowT CheckedNarrowing(const WideT& wide) { } // namespace -namespace perftools { -namespace gputools { +namespace stream_executor { using dnn::BatchDescriptor; using dnn::FilterDescriptor; @@ -159,7 +158,7 @@ static port::ThreadPool* GetCudaThreadpool() { return cudnn_threadpool; } -#define PERFTOOLS_GPUTOOLS_CUDNN_WRAP(__name) \ +#define STREAM_EXECUTOR_CUDNN_WRAP(__name) \ struct WrapperShim__##__name { \ template \ cudnnStatus_t operator()(CUDAExecutor* parent, Args... args) { \ @@ -169,11 +168,34 @@ static port::ThreadPool* GetCudaThreadpool() { } \ } __name; +#define STREAM_EXECUTOR_CUDNN_WRAP_WITH_CHECKED_STREAM(__name) \ + struct WrapperShim__##__name { \ + template \ + cudnnStatus_t operator()(CudnnSupport* dnn, Stream* s, Args... args) \ + SHARED_LOCKS_REQUIRED(dnn->dnn_handle_mutex_) { \ + CHECK_NOTNULL(s); \ + CHECK_EQ(s, dnn->GetCurrentDnnStream()) \ + << "Stream is not set correctly!"; \ + cuda::ScopedActivateExecutorContext sac{dnn->GetParentExecutor()}; \ + cudnnStatus_t retval = ::__name(args...); \ + return retval; \ + } \ + } __name; + +// Handles cudnnSetStream differently in order to add debug information. +struct WrapperShim__cudnnSetStream { + cudnnStatus_t operator()(CudnnSupport* dnn, Stream* stream, + cudnnHandle_t handle) + EXCLUSIVE_LOCKS_REQUIRED(dnn->dnn_handle_mutex_) { + dnn->SetCurrentDnnStream(stream); + cuda::ScopedActivateExecutorContext sac{dnn->GetParentExecutor()}; + cudnnStatus_t retval = ::cudnnSetStream(handle, AsCUDAStreamValue(stream)); + return retval; + } +} cudnnSetStream; + // clang-format off #define CUDNN_DNN_ROUTINE_EACH(__macro) \ - __macro(cudnnBatchNormalizationBackward) \ - __macro(cudnnBatchNormalizationForwardInference) \ - __macro(cudnnBatchNormalizationForwardTraining) \ __macro(cudnnGetConvolutionNdForwardOutputDim) \ __macro(cudnnGetConvolutionForwardAlgorithm) \ __macro(cudnnCreateTensorDescriptor) \ @@ -190,16 +212,25 @@ static port::ThreadPool* GetCudaThreadpool() { __macro(cudnnDestroyConvolutionDescriptor) \ __macro(cudnnCreate) \ __macro(cudnnDestroy) \ - __macro(cudnnSetStream) \ - __macro(cudnnActivationForward) \ - __macro(cudnnConvolutionForward) \ - __macro(cudnnConvolutionBackwardBias) \ __macro(cudnnGetConvolutionForwardWorkspaceSize) \ - __macro(cudnnTransformTensor) \ __macro(cudnnSetConvolutionNdDescriptor) \ __macro(cudnnSetTensor4dDescriptor) \ __macro(cudnnSetTensorNdDescriptor) \ - __macro(cudnnSetFilterNdDescriptor) \ + __macro(cudnnSetFilterNdDescriptor) + +// clang-format on +CUDNN_DNN_ROUTINE_EACH(STREAM_EXECUTOR_CUDNN_WRAP) +#undef CUDNN_DNN_ROUTINE_EACH + +// clang-format off +#define CUDNN_DNN_ROUTINE_EACH_WITH_STREAM(__macro) \ + __macro(cudnnBatchNormalizationBackward) \ + __macro(cudnnBatchNormalizationForwardInference) \ + __macro(cudnnBatchNormalizationForwardTraining) \ + __macro(cudnnActivationForward) \ + __macro(cudnnConvolutionForward) \ + __macro(cudnnConvolutionBackwardBias) \ + __macro(cudnnTransformTensor) \ __macro(cudnnPoolingForward) \ __macro(cudnnPoolingBackward) \ __macro(cudnnLRNCrossChannelForward) \ @@ -207,9 +238,11 @@ static port::ThreadPool* GetCudaThreadpool() { __macro(cudnnAddTensor) \ __macro(cudnnConvolutionBackwardData) \ __macro(cudnnConvolutionBackwardFilter) -// clang-format on -CUDNN_DNN_ROUTINE_EACH(PERFTOOLS_GPUTOOLS_CUDNN_WRAP) +// clang-format on +CUDNN_DNN_ROUTINE_EACH_WITH_STREAM( + STREAM_EXECUTOR_CUDNN_WRAP_WITH_CHECKED_STREAM) +#undef CUDNN_DNN_ROUTINE_EACH_WITH_STREAM // APIs available after R3: #if CUDNN_VERSION >= 3000 @@ -218,21 +251,22 @@ CUDNN_DNN_ROUTINE_EACH(PERFTOOLS_GPUTOOLS_CUDNN_WRAP) __macro(cudnnGetConvolutionBackwardDataAlgorithm) \ __macro(cudnnGetConvolutionBackwardFilterAlgorithm) \ __macro(cudnnGetConvolutionBackwardDataWorkspaceSize) -CUDNN_DNN_ROUTINE_EACH_AFTER_R3(PERFTOOLS_GPUTOOLS_CUDNN_WRAP) +CUDNN_DNN_ROUTINE_EACH_AFTER_R3(STREAM_EXECUTOR_CUDNN_WRAP) #undef CUDNN_DNN_ROUTINE_EACH_AFTER_R3 #endif // APIs in R3 but not in R5 // clang-format off #if CUDNN_VERSION >= 3000 && CUDNN_VERSION < 5000 -#define CUDNN_DNN_ROUTINE_EACH_R3(__macro) \ +#define CUDNN_DNN_ROUTINE_EACH_R3_WITH_STREAM(__macro) \ __macro(cudnnAddTensor_v3) \ __macro(cudnnConvolutionBackwardData_v3) \ __macro(cudnnConvolutionBackwardFilter_v3) // clang-format on -CUDNN_DNN_ROUTINE_EACH_R3(PERFTOOLS_GPUTOOLS_CUDNN_WRAP) -#undef CUDNN_DNN_ROUTINE_EACH_R3 +CUDNN_DNN_ROUTINE_EACH_R3_WITH_STREAM( + STREAM_EXECUTOR_CUDNN_WRAP_WITH_CHECKED_STREAM) +#undef CUDNN_DNN_ROUTINE_EACH_R3_WITH_STREAM #endif // APIs in R5 @@ -254,29 +288,44 @@ CUDNN_DNN_ROUTINE_EACH_R3(PERFTOOLS_GPUTOOLS_CUDNN_WRAP) __macro(cudnnGetRNNTrainingReserveSize) \ __macro(cudnnGetRNNLinLayerMatrixParams) \ __macro(cudnnGetRNNLinLayerBiasParams) \ - __macro(cudnnRNNForwardInference) \ - __macro(cudnnRNNForwardTraining) \ - __macro(cudnnRNNBackwardData) \ - __macro(cudnnRNNBackwardWeights) \ __macro(cudnnSetRNNDescriptor) \ __macro(cudnnGetFilterNdDescriptor) // clang-format on - -CUDNN_DNN_ROUTINE_EACH_R5(PERFTOOLS_GPUTOOLS_CUDNN_WRAP) +CUDNN_DNN_ROUTINE_EACH_R5(STREAM_EXECUTOR_CUDNN_WRAP) #undef CUDNN_DNN_ROUTINE_EACH_R5 + +// clang-format off +#define CUDNN_DNN_ROUTINE_EACH_R5_WITH_STREAM(__macro) \ + __macro(cudnnRNNForwardInference) \ + __macro(cudnnRNNForwardTraining) \ + __macro(cudnnRNNBackwardData) \ + __macro(cudnnRNNBackwardWeights) + +// clang-format on +CUDNN_DNN_ROUTINE_EACH_R5_WITH_STREAM( + STREAM_EXECUTOR_CUDNN_WRAP_WITH_CHECKED_STREAM) +#undef CUDNN_DNN_ROUTINE_EACH_R5_WITH_STREAM #endif // APIs in R6 // clang-format off #if CUDNN_VERSION >= 6000 #define CUDNN_DNN_ROUTINE_EACH_R6(__macro) \ - __macro(cudnnConvolutionBiasActivationForward) \ __macro(cudnnSetRNNDescriptor_v6) // clang-format on -CUDNN_DNN_ROUTINE_EACH_R6(PERFTOOLS_GPUTOOLS_CUDNN_WRAP) +CUDNN_DNN_ROUTINE_EACH_R6(STREAM_EXECUTOR_CUDNN_WRAP) #undef CUDNN_DNN_ROUTINE_EACH_R6 + +// clang-format off +#define CUDNN_DNN_ROUTINE_EACH_R6_WITH_STREAM(__macro) \ + __macro(cudnnConvolutionBiasActivationForward) + +// clang-format on +CUDNN_DNN_ROUTINE_EACH_R6_WITH_STREAM( + STREAM_EXECUTOR_CUDNN_WRAP_WITH_CHECKED_STREAM) +#undef CUDNN_DNN_ROUTINE_EACH_R6_WITH_STREAM #endif // APIs in R7 @@ -287,12 +336,10 @@ CUDNN_DNN_ROUTINE_EACH_R6(PERFTOOLS_GPUTOOLS_CUDNN_WRAP) __macro(cudnnSetRNNMatrixMathType) // clang-format on -CUDNN_DNN_ROUTINE_EACH_R7(PERFTOOLS_GPUTOOLS_CUDNN_WRAP) +CUDNN_DNN_ROUTINE_EACH_R7(STREAM_EXECUTOR_CUDNN_WRAP) #undef CUDNN_DNN_ROUTINE_EACH_R7 #endif -#undef CUDNN_DNN_ROUTINE_EACH - } // namespace wrap namespace { @@ -419,7 +466,7 @@ port::Status GetLoadedCudnnVersion(CudnnVersion* version) { } // namespace CudnnSupport::CudnnSupport(CUDAExecutor* parent) - : parent_(parent), dnn_handle_(nullptr) {} + : parent_(parent), dnn_handle_(nullptr), current_dnn_stream_(nullptr) {} CudnnSupport::~CudnnSupport() { auto status = wrap::cudnnDestroy(parent_, ToHandle(dnn_handle_)); @@ -477,11 +524,12 @@ port::Status CudnnSupport::Init() { ToString(status))}; } -port::StatusOr> CudnnSupport::GetVersion() { +port::StatusOr +CudnnSupport::GetVersion() { CudnnVersion version; TF_RETURN_IF_ERROR(GetLoadedCudnnVersion(&version)); - return std::make_tuple(version.major_version, version.minor_version, - version.patch_level); + return perftools::gputools::dnn::VersionInfo( + version.major_version, version.minor_version, version.patch_level); } // Turns a BatchDescriptor structure into a cudnn tensor handle within a scope. @@ -1660,6 +1708,12 @@ bool CudnnSupport::DoRnnForwardImpl( // check params size mutex_lock lock{dnn_handle_mutex_}; + auto set_stream_status = + wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); + if (set_stream_status != CUDNN_STATUS_SUCCESS) { + LOG(FATAL) << "failed to set stream for cudnn handle: " + << ToString(set_stream_status); + } if (!CheckRNNParameterSize(parent_, ToHandle(dnn_handle_), rnn_desc, input_desc)) { @@ -1720,7 +1774,7 @@ bool CudnnSupport::DoRnnForwardImpl( cudnnStatus_t status; if (!is_training) { status = wrap::cudnnRNNForwardInference( - parent_, ToHandle(dnn_handle_) /*handle*/, + this, stream, ToHandle(dnn_handle_) /*handle*/, rnn_desc.handle() /*rnnDesc*/, model_dims.seq_length /*seqLength*/, input_desc.handles() /*xDesc*/, input_data.opaque() /*x*/, input_h_desc.handle() /*hxDesc*/, input_h_data.opaque() /*hx*/, @@ -1733,7 +1787,7 @@ bool CudnnSupport::DoRnnForwardImpl( workspace.size() /*workSpaceSizeInBytes*/); } else { status = wrap::cudnnRNNForwardTraining( - parent_, ToHandle(dnn_handle_) /*handle*/, + this, stream, ToHandle(dnn_handle_) /*handle*/, rnn_desc.handle() /*rnnDesc*/, model_dims.seq_length /*seqLength*/, input_desc.handles() /*xDesc*/, input_data.opaque() /*x*/, input_h_desc.handle() /*hxDesc*/, input_h_data.opaque() /*hx*/, @@ -1810,6 +1864,12 @@ bool CudnnSupport::DoRnnBackwardImpl( // check params size mutex_lock lock{dnn_handle_mutex_}; + auto set_stream_status = + wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); + if (set_stream_status != CUDNN_STATUS_SUCCESS) { + LOG(FATAL) << "failed to set stream for cudnn handle: " + << ToString(set_stream_status); + } if (!CheckRNNParameterSize(parent_, ToHandle(dnn_handle_), rnn_desc, input_desc)) { @@ -1841,10 +1901,11 @@ bool CudnnSupport::DoRnnBackwardImpl( } // make the backward data call cudnnStatus_t status = wrap::cudnnRNNBackwardData( - parent_, ToHandle(dnn_handle_) /*handle*/, rnn_desc.handle() /*rnnDesc*/, - model_dims.seq_length /*seqLength*/, output_desc.handles() /*yDesc*/, - output_data.opaque() /*y*/, output_desc.handles() /*dyDesc*/, - output_backprop_data.opaque() /*dy*/, output_h_desc.handle() /*dhyDesc*/, + this, stream, ToHandle(dnn_handle_) /*handle*/, + rnn_desc.handle() /*rnnDesc*/, model_dims.seq_length /*seqLength*/, + output_desc.handles() /*yDesc*/, output_data.opaque() /*y*/, + output_desc.handles() /*dyDesc*/, output_backprop_data.opaque() /*dy*/, + output_h_desc.handle() /*dhyDesc*/, output_h_backprop_data.opaque() /*dhy*/, output_c_desc.handle() /*dcyDesc*/, output_c_backprop_data.opaque() /*dcy*/, @@ -1873,7 +1934,7 @@ bool CudnnSupport::DoRnnBackwardImpl( stream->ThenMemZero(params_backprop_data, params_backprop_data->size()); // make the backward weight call status = wrap::cudnnRNNBackwardWeights( - parent_, ToHandle(dnn_handle_) /*handle*/, + this, stream, ToHandle(dnn_handle_) /*handle*/, rnn_desc.handle() /*rnnDesc*/, model_dims.seq_length /*seqLength*/, input_desc.handles() /*xDesc*/, input_data.opaque() /*x*/, input_h_desc.handle() /*hxDesc*/, input_h_data.opaque() /*hx*/, @@ -2469,12 +2530,20 @@ cudnnDataType_t GetConvComputeType() { } // A helper struct to decide whether to use FP32 as the internal compute type -// for rnn when the input data type is FP16. By default it is turned on, -// users can explicitly disable them (choose to use FP16 as the internal compute -// type) through an env-var "TF_FP16_RNN_USE_FP32_COMPUTE=0". +// for rnn when the input data type is FP16. At present it is turned off, +// users can explicitly control them through an env-var +// TF_FP16_RNN_USE_FP32_COMPUTE. +// After the TODO below is fixed, users should almost always use fp32 compute +// type for training. Using fp16 might suffer suboptimal accuracy due to loss +// in precision. struct RnnDoFP32ComputationFP16Input { static constexpr const char* kName = "TF_FP16_RNN_USE_FP32_COMPUTE"; - static constexpr bool kDefaultFlag = true; + // TODO(jamesqin): b/78182362 flip to true when cudnn 7.1.4 fixes the bug. + // Before cudnn 7.1.4 RNN are always done in fp32, no matter what math + // precision is set. + // Set it temporary to false s.t. no error is raised when using fp16 inputs, + // fp32 math precision. + static constexpr bool kDefaultFlag = false; }; // A helper function to return the internal compute type for @@ -2517,8 +2586,7 @@ bool CudnnSupport::DoConvolveImpl( GetConvComputeType()}; mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(FATAL) << "failed to set stream for cudnn handle: " << ToString(status); } @@ -2668,7 +2736,7 @@ bool CudnnSupport::DoConvolveImpl( } } status = wrap::cudnnConvolutionForward( - parent_, ToHandle(dnn_handle_), + this, stream, ToHandle(dnn_handle_), /*alpha=*/alpha, /*srcDesc=*/input_nd.handle(), /*srcData=*/input_data.opaque(), /*filterDesc=*/filter.handle(), /*filterData=*/filter_data.opaque(), /*convDesc=*/conv.handle(), @@ -2737,8 +2805,7 @@ bool CudnnSupport::DoFusedConvolveImpl( static_cast(cudnn_compute_type)}; mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); CHECK(status == CUDNN_STATUS_SUCCESS) << "failed to set stream for cudnn handle: " << ToString(status); @@ -2804,7 +2871,7 @@ bool CudnnSupport::DoFusedConvolveImpl( << "\noutput_data->opaque() = " << output_data->opaque(); status = wrap::cudnnConvolutionBiasActivationForward( - parent_, ToHandle(dnn_handle_), /*alpha1=*/&conv_input_scale, + this, stream, ToHandle(dnn_handle_), /*alpha1=*/&conv_input_scale, /*srcDesc=*/conv_input_nd.handle(), /*srcData=*/conv_input_data.opaque(), /*filterDesc=*/filter.handle(), /*filterData=*/filter_data.opaque(), /*convDesc=*/conv.handle(), algo, /*workSpace=*/scratch.opaque(), @@ -3009,8 +3076,7 @@ bool CudnnSupport::DoBatchNormalizationForwardImpl( bool is_training, std::function&()> var_to_inv_var, std::function inv_var_to_var) { mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status); return false; @@ -3046,7 +3112,7 @@ bool CudnnSupport::DoBatchNormalizationForwardImpl( } status = wrap::cudnnBatchNormalizationForwardTraining( - parent_, ToHandle(dnn_handle_), mode, &one, &zero, + this, stream, ToHandle(dnn_handle_), mode, &one, &zero, x_descriptor.handle(), x.opaque(), x_descriptor.handle(), y->opaque(), scale_offset_descriptor.handle(), scale.opaque(), offset.opaque(), 1.0, batch_mean_opaque, batch_var_opaque, epsilon, saved_mean->opaque(), @@ -3063,7 +3129,7 @@ bool CudnnSupport::DoBatchNormalizationForwardImpl( const void* maybe_inv_var = estimated_variance.opaque(); #endif status = wrap::cudnnBatchNormalizationForwardInference( - parent_, ToHandle(dnn_handle_), mode, &one, &zero, + this, stream, ToHandle(dnn_handle_), mode, &one, &zero, x_descriptor.handle(), x.opaque(), x_descriptor.handle(), y->opaque(), scale_offset_descriptor.handle(), scale.opaque(), offset.opaque(), estimated_mean.opaque(), maybe_inv_var, epsilon); @@ -3114,8 +3180,7 @@ bool CudnnSupport::DoBatchNormalizationBackwardImpl( DeviceMemory* x_backprop, DeviceMemory* scale_backprop, DeviceMemory* offset_backprop) { mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status); return false; @@ -3136,7 +3201,7 @@ bool CudnnSupport::DoBatchNormalizationBackwardImpl( float zero = 0.0; status = wrap::cudnnBatchNormalizationBackward( - parent_, ToHandle(dnn_handle_), mode, &one, &zero, &one, &zero, + this, stream, ToHandle(dnn_handle_), mode, &one, &zero, &one, &zero, x_descriptor.handle(), x.opaque(), x_descriptor.handle(), y_backprop.opaque(), x_descriptor.handle(), x_backprop->opaque(), scale_offset_descriptor.handle(), scale.opaque(), @@ -3326,7 +3391,7 @@ DeviceMemory CudnnSupport::MaybeTransformLayout( float alpha = 1.0f; float beta = 0.0f; auto status = wrap::cudnnTransformTensor( - parent_, ToHandle(dnn_handle_), &alpha, orig_out_back_nd.handle(), + this, stream, ToHandle(dnn_handle_), &alpha, orig_out_back_nd.handle(), backward_output_data.opaque(), &beta, transformed_out_back_nd.handle(), (*transform_scratch)->mutable_device_memory()->opaque()); @@ -3345,8 +3410,7 @@ bool CudnnSupport::DoTransformTensor(Stream* stream, dnn::DataType output_type, float scale, DeviceMemoryBase* output_data) { mutex_lock lock{dnn_handle_mutex_}; - cudnnStatus_t status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(FATAL) << "failed to set stream for cudnn handle: " << ToString(status); } @@ -3357,7 +3421,7 @@ bool CudnnSupport::DoTransformTensor(Stream* stream, ScopedTensorDescriptor output_tensor_desc( parent_, output_desc, ToCudnnDataType(output_type, output_desc.layout())); status = wrap::cudnnTransformTensor( - parent_, ToHandle(dnn_handle_), &scale, input_tensor_desc.handle(), + this, stream, ToHandle(dnn_handle_), &scale, input_tensor_desc.handle(), input_data.opaque(), &beta, output_tensor_desc.handle(), output_data->opaque()); if (status != CUDNN_STATUS_SUCCESS) { @@ -3384,8 +3448,7 @@ bool CudnnSupport::DoConvolveBackwardDataImpl( const dnn::AlgorithmConfig& algorithm_config, dnn::ProfileResult* output_profile_result) { mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(FATAL) << "failed to set stream for cudnn handle: " << ToString(status); } @@ -3554,7 +3617,7 @@ bool CudnnSupport::DoConvolveBackwardDataImpl( #else status = wrap::cudnnConvolutionBackwardData_v3( #endif - parent_, ToHandle(dnn_handle_), + this, stream, ToHandle(dnn_handle_), /*alpha=*/alpha, /*filterDesc=*/filter.handle(), /*filterData=*/filter_data.opaque(), @@ -3655,8 +3718,7 @@ bool CudnnSupport::DoConvolveBackwardFilterImpl( const dnn::AlgorithmConfig& algorithm_config, dnn::ProfileResult* output_profile_result) { mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(FATAL) << "failed to set stream for cudnn handle: " << ToString(status); } @@ -3826,7 +3888,7 @@ bool CudnnSupport::DoConvolveBackwardFilterImpl( #else status = wrap::cudnnConvolutionBackwardFilter_v3( #endif - parent_, ToHandle(dnn_handle_), /*alpha=*/alpha, + this, stream, ToHandle(dnn_handle_), /*alpha=*/alpha, /*srcDesc=*/input_nd.handle(), /*srcData=*/input_data.opaque(), /*diffDesc=*/out_back_nd.handle(), @@ -3922,8 +3984,7 @@ bool CudnnSupport::DoConvolveBackwardBiasImpl( const dnn::BatchDescriptor& bias_descriptor, DeviceMemory* backward_bias_data) { mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(FATAL) << "failed to set stream for cudnn handle: " << ToString(status); } @@ -3938,7 +3999,7 @@ bool CudnnSupport::DoConvolveBackwardBiasImpl( float beta = 0.0; status = wrap::cudnnConvolutionBackwardBias( - parent_, ToHandle(dnn_handle_), &alpha, input_nd.handle(), + this, stream, ToHandle(dnn_handle_), &alpha, input_nd.handle(), input_data.opaque(), &beta, bias_nd.handle(), backward_bias_data->opaque()); if (status != CUDNN_STATUS_SUCCESS) { @@ -4143,8 +4204,7 @@ bool CudnnSupport::DoBiasAdd(Stream* stream, } mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status); return false; @@ -4158,7 +4218,7 @@ bool CudnnSupport::DoBiasAdd(Stream* stream, #else status = wrap::cudnnAddTensor_v3( #endif - parent_, ToHandle(dnn_handle_), &alpha, bias_descriptor.handle(), + this, stream, ToHandle(dnn_handle_), &alpha, bias_descriptor.handle(), biases.opaque(), &beta, input_descriptor.handle(), output_data->opaque()); if (status != CUDNN_STATUS_SUCCESS) { @@ -4176,8 +4236,7 @@ bool CudnnSupport::DoActivate(Stream* stream, DeviceMemory* output_data, uint64 options) { mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status); return false; @@ -4221,7 +4280,7 @@ bool CudnnSupport::DoActivate(Stream* stream, // Beta is the output scaling factor. float beta = 0.0; status = wrap::cudnnActivationForward( - parent_, ToHandle(dnn_handle_), + this, stream, ToHandle(dnn_handle_), #if CUDNN_VERSION >= 5000 activation_desc.handle(), #else @@ -4245,8 +4304,7 @@ bool CudnnSupport::DoPoolForward( const dnn::BatchDescriptor& output_dimensions, DeviceMemory* output_data) { mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status); return false; @@ -4262,7 +4320,7 @@ bool CudnnSupport::DoPoolForward( CUDNN_DATA_DOUBLE}; ScopedPoolingDescriptor pooling_desc{parent_, pooling_dimensions}; status = wrap::cudnnPoolingForward( - parent_, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha, + this, stream, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha, src_desc.handle(), input_data.opaque(), &beta, dest_desc.handle(), output_data->opaque()); if (status != CUDNN_STATUS_SUCCESS) { @@ -4280,8 +4338,7 @@ bool CudnnSupport::DoPoolForward( const dnn::BatchDescriptor& output_dimensions, DeviceMemory* output_data) { mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status); return false; @@ -4297,7 +4354,7 @@ bool CudnnSupport::DoPoolForward( CUDNN_DATA_FLOAT}; ScopedPoolingDescriptor pooling_desc{parent_, pooling_dimensions}; status = wrap::cudnnPoolingForward( - parent_, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha, + this, stream, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha, src_desc.handle(), input_data.opaque(), &beta, dest_desc.handle(), output_data->opaque()); if (status != CUDNN_STATUS_SUCCESS) { @@ -4315,8 +4372,7 @@ bool CudnnSupport::DoPoolForward( const dnn::BatchDescriptor& output_dimensions, DeviceMemory* output_data) { mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status); return false; @@ -4331,7 +4387,7 @@ bool CudnnSupport::DoPoolForward( ScopedTensorDescriptor dest_desc{parent_, output_dimensions, CUDNN_DATA_HALF}; ScopedPoolingDescriptor pooling_desc{parent_, pooling_dimensions}; status = wrap::cudnnPoolingForward( - parent_, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha, + this, stream, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha, src_desc.handle(), input_data.opaque(), &beta, dest_desc.handle(), output_data->opaque()); if (status != CUDNN_STATUS_SUCCESS) { @@ -4351,8 +4407,7 @@ bool CudnnSupport::DoPoolBackward( const DeviceMemory& input_diff_data, DeviceMemory* output_diff_data) { mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status); return false; @@ -4368,7 +4423,7 @@ bool CudnnSupport::DoPoolBackward( CUDNN_DATA_DOUBLE}; ScopedPoolingDescriptor pooling_desc{parent_, pooling_dimensions}; status = wrap::cudnnPoolingBackward( - parent_, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha, + this, stream, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha, dest_desc.handle(), output_data.opaque(), dest_desc.handle(), input_diff_data.opaque(), src_desc.handle(), input_data.opaque(), &beta, src_desc.handle(), output_diff_data->opaque()); @@ -4389,8 +4444,7 @@ bool CudnnSupport::DoPoolBackward( const DeviceMemory& input_diff_data, DeviceMemory* output_diff_data) { mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status); return false; @@ -4406,7 +4460,7 @@ bool CudnnSupport::DoPoolBackward( CUDNN_DATA_FLOAT}; ScopedPoolingDescriptor pooling_desc{parent_, pooling_dimensions}; status = wrap::cudnnPoolingBackward( - parent_, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha, + this, stream, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha, dest_desc.handle(), output_data.opaque(), dest_desc.handle(), input_diff_data.opaque(), src_desc.handle(), input_data.opaque(), &beta, src_desc.handle(), output_diff_data->opaque()); @@ -4427,8 +4481,7 @@ bool CudnnSupport::DoPoolBackward( const DeviceMemory& input_diff_data, DeviceMemory* output_diff_data) { mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status); return false; @@ -4443,7 +4496,7 @@ bool CudnnSupport::DoPoolBackward( ScopedTensorDescriptor dest_desc{parent_, output_dimensions, CUDNN_DATA_HALF}; ScopedPoolingDescriptor pooling_desc{parent_, pooling_dimensions}; status = wrap::cudnnPoolingBackward( - parent_, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha, + this, stream, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha, dest_desc.handle(), output_data.opaque(), dest_desc.handle(), input_diff_data.opaque(), src_desc.handle(), input_data.opaque(), &beta, src_desc.handle(), output_diff_data->opaque()); @@ -4478,8 +4531,7 @@ bool CudnnSupport::DoNormalizeWithDimensions( // Launch the normalization. mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status); return false; @@ -4494,7 +4546,7 @@ bool CudnnSupport::DoNormalizeWithDimensions( float beta = 0.0f; status = wrap::cudnnLRNCrossChannelForward( - parent_, ToHandle(dnn_handle_), normalize.handle(), + this, stream, ToHandle(dnn_handle_), normalize.handle(), CUDNN_LRN_CROSS_CHANNEL_DIM1, &alpha, dims.handle(), input_data.opaque(), &beta, dims.handle(), output_data->opaque()); if (status != CUDNN_STATUS_SUCCESS) { @@ -4521,8 +4573,7 @@ bool CudnnSupport::DoNormalizeBackwardWithDimensions( } mutex_lock lock{dnn_handle_mutex_}; - auto status = wrap::cudnnSetStream(parent_, ToHandle(dnn_handle_), - AsCUDAStreamValue(stream)); + auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_)); if (status != CUDNN_STATUS_SUCCESS) { LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status); return false; @@ -4535,7 +4586,7 @@ bool CudnnSupport::DoNormalizeBackwardWithDimensions( float beta = 0.0f; status = wrap::cudnnLRNCrossChannelBackward( - parent_, ToHandle(dnn_handle_), normalize.handle(), + this, stream, ToHandle(dnn_handle_), normalize.handle(), CUDNN_LRN_CROSS_CHANNEL_DIM1, &alpha, dims.handle(), normalized_data.opaque(), dims.handle(), normalized_variable_gradient.opaque(), dims.handle(), raw_data.opaque(), @@ -4684,46 +4735,39 @@ bool CudnnSupport::DeriveOutputBatchDescriptor( } // namespace cuda -namespace gpu = ::perftools::gputools; - void initialize_cudnn() { - gpu::port::Status status = - gpu::PluginRegistry::Instance() - ->RegisterFactory( - gpu::cuda::kCudaPlatformId, gpu::cuda::kCuDnnPlugin, "cuDNN", - [](gpu::internal::StreamExecutorInterface* - parent) -> gpu::dnn::DnnSupport* { - gpu::cuda::CUDAExecutor* cuda_executor = - dynamic_cast(parent); - if (cuda_executor == nullptr) { - LOG(ERROR) - << "Attempting to initialize an instance of the cuBLAS " - << "support library with a non-CUDA StreamExecutor"; - return nullptr; - } - - gpu::cuda::CudnnSupport* dnn = - new gpu::cuda::CudnnSupport(cuda_executor); - if (!dnn->Init().ok()) { - // Note: Init() will log a more specific error. - delete dnn; - return nullptr; - } - return dnn; - }); + port::Status status = + PluginRegistry::Instance()->RegisterFactory( + cuda::kCudaPlatformId, cuda::kCuDnnPlugin, "cuDNN", + [](internal::StreamExecutorInterface* parent) -> dnn::DnnSupport* { + cuda::CUDAExecutor* cuda_executor = + dynamic_cast(parent); + if (cuda_executor == nullptr) { + LOG(ERROR) + << "Attempting to initialize an instance of the cuBLAS " + << "support library with a non-CUDA StreamExecutor"; + return nullptr; + } + + cuda::CudnnSupport* dnn = new cuda::CudnnSupport(cuda_executor); + if (!dnn->Init().ok()) { + // Note: Init() will log a more specific error. + delete dnn; + return nullptr; + } + return dnn; + }); if (!status.ok()) { LOG(ERROR) << "Unable to register cuDNN factory: " << status.error_message(); } - gpu::PluginRegistry::Instance()->SetDefaultFactory(gpu::cuda::kCudaPlatformId, - gpu::PluginKind::kDnn, - gpu::cuda::kCuDnnPlugin); + PluginRegistry::Instance()->SetDefaultFactory( + cuda::kCudaPlatformId, PluginKind::kDnn, cuda::kCuDnnPlugin); } -} // namespace gputools -} // namespace perftools +} // namespace stream_executor REGISTER_MODULE_INITIALIZER(register_cudnn, - { perftools::gputools::initialize_cudnn(); }); + { stream_executor::initialize_cudnn(); }); diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.h b/tensorflow/stream_executor/cuda/cuda_dnn.h index 0e5368aca82e7fa4d0281062dd3fbb59bebc0a3d..5ded7cf1543ee2cf3b874bcbd27b0fc25df2d844 100644 --- a/tensorflow/stream_executor/cuda/cuda_dnn.h +++ b/tensorflow/stream_executor/cuda/cuda_dnn.h @@ -26,8 +26,7 @@ limitations under the License. #include "tensorflow/stream_executor/plugin_registry.h" #include "tensorflow/stream_executor/temporary_device_memory.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { class CUDAExecutor; @@ -46,7 +45,7 @@ class CudnnSupport : public dnn::DnnSupport { ~CudnnSupport() override; port::Status Init() override; - port::StatusOr> GetVersion() override; + port::StatusOr GetVersion() override; port::StatusOr> createRnnDescriptor( int num_layers, int hidden_size, int input_size, @@ -625,10 +624,27 @@ class CudnnSupport : public dnn::DnnSupport { dnn::DataType output_type, float scale, DeviceMemoryBase* output_data) override; - private: - // Guards the enqueueing of DNN operations via the dnn_handle_ below. + const Stream* GetCurrentDnnStream() const + SHARED_LOCKS_REQUIRED(dnn_handle_mutex_) { + return current_dnn_stream_; + } + + void SetCurrentDnnStream(Stream* stream) + EXCLUSIVE_LOCKS_REQUIRED(dnn_handle_mutex_) { + current_dnn_stream_ = stream; + } + + CUDAExecutor* GetParentExecutor() { return parent_; } + + // Guards the enqueueing of DNN operations via the dnn_handle_ below, and + // access to current_dnn_stream_. + // + // This is a public member because we need to add thread safty annotations in + // the cudnn wrapper functions in the cc file, which need to access this + // mutex (the annotations require C++ permission checks). mutex dnn_handle_mutex_; + private: CUDAExecutor* parent_; // Parent executor object. Not owned. // cudnn library handle. cudnnHandle_t type is not present in this header to @@ -636,6 +652,9 @@ class CudnnSupport : public dnn::DnnSupport { // single cuda_dnn translation unit. void* dnn_handle_ GUARDED_BY(dnn_handle_mutex_); + // The current cudnn stream that is set by cudnnSetStream(). + Stream* current_dnn_stream_ GUARDED_BY(dnn_handle_mutex_); + // NOTE(keveman): Temporary data layout transformation until cuDNN supports // kBatchYXDepth for backward pass. This function allocates temporary memory, // lays out the source data into the temporary but in the kBatchDepthXY @@ -790,7 +809,6 @@ class CudnnSupport : public dnn::DnnSupport { }; } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_DNN_H_ diff --git a/tensorflow/stream_executor/cuda/cuda_driver.cc b/tensorflow/stream_executor/cuda/cuda_driver.cc index 58e1e58c593a3d938d97baff2356bce2c215a7a1..71cab145b9bb5a8320df8fb78241dd0d37d54239 100644 --- a/tensorflow/stream_executor/cuda/cuda_driver.cc +++ b/tensorflow/stream_executor/cuda/cuda_driver.cc @@ -37,14 +37,6 @@ limitations under the License. #include "tensorflow/stream_executor/platform/port.h" #include "tensorflow/stream_executor/lib/inlined_vector.h" -#if defined(PLATFORM_WINDOWS) -// TODO: in windows ARRAYSIZE is defined in winnt.h but including it -// here creates a conflict with cuda.h - for now define it here. -#define ARRAYSIZE(a) \ - ((sizeof(a) / sizeof(*(a))) / \ - static_cast(!(sizeof(a) % sizeof(*(a))))) -#endif - bool FLAGS_gpuexec_cuda_driver_inject_init_error = false; bool FLAGS_gpuexec_cuda_sync_around_driver_calls = false; bool FLAGS_gpuexec_cuda_device_0_only = false; @@ -53,8 +45,7 @@ bool FLAGS_gpuexec_cuda_device_0_only = false; // matches the expected one. constexpr bool kVerifyCudaContext = false; -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { namespace { @@ -720,15 +711,15 @@ CUDADriver::ContextGetSharedMemConfig(CudaContext* context) { port::bit_cast(uintptr_t(info_log_buffer_bytes)), port::bit_cast(info_log_buffer.data()), port::bit_cast(uintptr_t(log_verbose))}; - CHECK(ARRAYSIZE(options) == ARRAYSIZE(option_values)); + CHECK(TF_ARRAYSIZE(options) == TF_ARRAYSIZE(option_values)); CUresult res; { // TODO(leary) Need to see if NVIDIA can expunge the leakiness in their // module loading: see http://b/13248943 - res = cuModuleLoadDataEx(module, ptx_data, ARRAYSIZE(options), options, - option_values); + res = cuModuleLoadDataEx(module, ptx_data, TF_ARRAYSIZE(options), + options, option_values); } // The PTX JIT mutates the values in the option values array to reflect the @@ -1649,5 +1640,4 @@ static port::StatusOr GetSimpleAttribute(CUdevice device, } } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/cuda/cuda_driver.h b/tensorflow/stream_executor/cuda/cuda_driver.h index fa9172b3f008d3083309126bbfa4a1ab961030e1..a9969e247e181599f2b3707f6c65c6527dd4683d 100644 --- a/tensorflow/stream_executor/cuda/cuda_driver.h +++ b/tensorflow/stream_executor/cuda/cuda_driver.h @@ -27,8 +27,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/port.h" #include "cuda/include/cuda.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { // Identifies the memory space where an allocation resides. See @@ -506,7 +505,6 @@ class CudaContext { }; } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_DRIVER_H_ diff --git a/tensorflow/stream_executor/cuda/cuda_event.cc b/tensorflow/stream_executor/cuda/cuda_event.cc index 1b41502300d446b55708c2afb6f5538bdf6cf220..96dcf173566087db475e3b237591d19f06128d92 100644 --- a/tensorflow/stream_executor/cuda/cuda_event.cc +++ b/tensorflow/stream_executor/cuda/cuda_event.cc @@ -19,8 +19,7 @@ limitations under the License. #include "tensorflow/stream_executor/cuda/cuda_stream.h" #include "tensorflow/stream_executor/lib/statusor.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { CUDAEvent::CUDAEvent(CUDAExecutor* parent) @@ -68,5 +67,4 @@ const CUevent& CUDAEvent::cuda_event() { } } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/cuda/cuda_event.h b/tensorflow/stream_executor/cuda/cuda_event.h index 56667e65d38199fd4c340147c4e40a17c5bb2b2d..f62344672ed624f1ed60b5452d33b6f8273f2b47 100644 --- a/tensorflow/stream_executor/cuda/cuda_event.h +++ b/tensorflow/stream_executor/cuda/cuda_event.h @@ -21,8 +21,7 @@ limitations under the License. #include "tensorflow/stream_executor/event.h" #include "tensorflow/stream_executor/lib/status.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { // CUDAEvent wraps a CUevent in the platform-independent EventInterface @@ -58,7 +57,6 @@ class CUDAEvent : public internal::EventInterface { }; } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_EVENT_H_ diff --git a/tensorflow/stream_executor/cuda/cuda_fft.cc b/tensorflow/stream_executor/cuda/cuda_fft.cc index a922f14fb4af695877b449d2f960fae1a356a82f..5b34740f9f1f9067f949ad41bcae0b97d3d3c7f4 100644 --- a/tensorflow/stream_executor/cuda/cuda_fft.cc +++ b/tensorflow/stream_executor/cuda/cuda_fft.cc @@ -31,8 +31,7 @@ limitations under the License. #include "tensorflow/stream_executor/plugin_registry.h" #include "tensorflow/stream_executor/stream_executor_internal.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { PLUGIN_REGISTRY_DEFINE_PLUGIN_ID(kCuFftPlugin); @@ -44,7 +43,7 @@ namespace wrap { // manner on first use. This dynamic loading technique is used to avoid DSO // dependencies on vendor libraries which may or may not be available in the // deployed binary environment. -#define PERFTOOLS_GPUTOOLS_CUFFT_WRAP(__name) \ +#define STREAM_EXECUTOR_CUFFT_WRAP(__name) \ struct WrapperShim__##__name { \ template \ cufftResult operator()(CUDAExecutor *parent, Args... args) { \ @@ -68,7 +67,7 @@ namespace wrap { __macro(cufftGetSizeMany) \ __macro(cufftMakePlanMany) -CUFFT_ROUTINE_EACH(PERFTOOLS_GPUTOOLS_CUFFT_WRAP) +CUFFT_ROUTINE_EACH(STREAM_EXECUTOR_CUFFT_WRAP) } // namespace wrap @@ -514,62 +513,59 @@ bool CUDAFft::DoFftWithDirectionInternal(Stream *stream, fft::Plan *plan, return true; } -#define PERFTOOLS_GPUTOOLS_CUDA_DEFINE_FFT(__type, __fft_type1, __fft_type2, \ - __fft_type3) \ - bool CUDAFft::DoFft(Stream *stream, fft::Plan *plan, \ - const DeviceMemory> &input, \ - DeviceMemory> *output) { \ - return DoFftWithDirectionInternal( \ - stream, plan, wrap::cufftExec##__fft_type1, input, output); \ - } \ - bool CUDAFft::DoFft(Stream *stream, fft::Plan *plan, \ - const DeviceMemory<__type> &input, \ - DeviceMemory> *output) { \ - return DoFftInternal(stream, plan, wrap::cufftExec##__fft_type2, input, \ - output); \ - } \ - bool CUDAFft::DoFft(Stream *stream, fft::Plan *plan, \ - const DeviceMemory> &input, \ - DeviceMemory<__type> *output) { \ - return DoFftInternal(stream, plan, wrap::cufftExec##__fft_type3, input, \ - output); \ +#define STREAM_EXECUTOR_CUDA_DEFINE_FFT(__type, __fft_type1, __fft_type2, \ + __fft_type3) \ + bool CUDAFft::DoFft(Stream *stream, fft::Plan *plan, \ + const DeviceMemory> &input, \ + DeviceMemory> *output) { \ + return DoFftWithDirectionInternal( \ + stream, plan, wrap::cufftExec##__fft_type1, input, output); \ + } \ + bool CUDAFft::DoFft(Stream *stream, fft::Plan *plan, \ + const DeviceMemory<__type> &input, \ + DeviceMemory> *output) { \ + return DoFftInternal(stream, plan, wrap::cufftExec##__fft_type2, input, \ + output); \ + } \ + bool CUDAFft::DoFft(Stream *stream, fft::Plan *plan, \ + const DeviceMemory> &input, \ + DeviceMemory<__type> *output) { \ + return DoFftInternal(stream, plan, wrap::cufftExec##__fft_type3, input, \ + output); \ } -PERFTOOLS_GPUTOOLS_CUDA_DEFINE_FFT(float, C2C, R2C, C2R) -PERFTOOLS_GPUTOOLS_CUDA_DEFINE_FFT(double, Z2Z, D2Z, Z2D) +STREAM_EXECUTOR_CUDA_DEFINE_FFT(float, C2C, R2C, C2R) +STREAM_EXECUTOR_CUDA_DEFINE_FFT(double, Z2Z, D2Z, Z2D) -#undef PERFTOOLS_GPUTOOLS_CUDA_DEFINE_FFT +#undef STREAM_EXECUTOR_CUDA_DEFINE_FFT } // namespace cuda -} // namespace gputools -} // namespace perftools - -namespace gpu = ::perftools::gputools; - -REGISTER_MODULE_INITIALIZER(register_cufft, { - gpu::port::Status status = - gpu::PluginRegistry::Instance() - ->RegisterFactory( - gpu::cuda::kCudaPlatformId, gpu::cuda::kCuFftPlugin, "cuFFT", - [](gpu::internal::StreamExecutorInterface - *parent) -> gpu::fft::FftSupport * { - gpu::cuda::CUDAExecutor *cuda_executor = - dynamic_cast(parent); - if (cuda_executor == nullptr) { - LOG(ERROR) - << "Attempting to initialize an instance of the cuFFT " - << "support library with a non-CUDA StreamExecutor"; - return nullptr; - } - - return new gpu::cuda::CUDAFft(cuda_executor); - }); + +void initialize_cufft() { + port::Status status = + PluginRegistry::Instance()->RegisterFactory( + cuda::kCudaPlatformId, cuda::kCuFftPlugin, "cuFFT", + [](internal::StreamExecutorInterface *parent) -> fft::FftSupport * { + cuda::CUDAExecutor *cuda_executor = + dynamic_cast(parent); + if (cuda_executor == nullptr) { + LOG(ERROR) << "Attempting to initialize an instance of the cuFFT " + << "support library with a non-CUDA StreamExecutor"; + return nullptr; + } + + return new cuda::CUDAFft(cuda_executor); + }); if (!status.ok()) { LOG(ERROR) << "Unable to register cuFFT factory: " << status.error_message(); } - gpu::PluginRegistry::Instance()->SetDefaultFactory(gpu::cuda::kCudaPlatformId, - gpu::PluginKind::kFft, - gpu::cuda::kCuFftPlugin); -}); + PluginRegistry::Instance()->SetDefaultFactory( + cuda::kCudaPlatformId, PluginKind::kFft, cuda::kCuFftPlugin); +} + +} // namespace stream_executor + +REGISTER_MODULE_INITIALIZER(register_cufft, + { stream_executor::initialize_cufft(); }); diff --git a/tensorflow/stream_executor/cuda/cuda_fft.h b/tensorflow/stream_executor/cuda/cuda_fft.h index 04c7dfe501c451e4848bef68bed9685c079dd523..8171e61418a3185455e50ee76315eb2493c36c01 100644 --- a/tensorflow/stream_executor/cuda/cuda_fft.h +++ b/tensorflow/stream_executor/cuda/cuda_fft.h @@ -26,8 +26,7 @@ limitations under the License. #include "tensorflow/stream_executor/plugin_registry.h" #include "tensorflow/stream_executor/scratch_allocator.h" -namespace perftools { -namespace gputools { +namespace stream_executor { class Stream; @@ -133,7 +132,6 @@ class CUDAFft : public fft::FftSupport { }; } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_FFT_H_ diff --git a/tensorflow/stream_executor/cuda/cuda_gpu_executor.cc b/tensorflow/stream_executor/cuda/cuda_gpu_executor.cc index 5ecaf46b8cae3c1e1f312816e7e5aec8ff8ce306..7c87d33d21b58af477a541953b9030e581cfa46d 100644 --- a/tensorflow/stream_executor/cuda/cuda_gpu_executor.cc +++ b/tensorflow/stream_executor/cuda/cuda_gpu_executor.cc @@ -66,8 +66,7 @@ limitations under the License. extern bool FLAGS_check_gpu_leaks; bool FLAGS_prefer_cubin_to_ptx = true; -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { // Hook that can be used to CUBIN-ate PTX before it is loaded into the driver. @@ -1127,7 +1126,7 @@ DeviceDescription *CUDAExecutor::PopulateDeviceDescription() const { builder.set_name(device_name); } - for (size_t i = 0; i < ARRAYSIZE(kAllUnqueryableDeviceParams); i++) { + for (size_t i = 0; i < TF_ARRAYSIZE(kAllUnqueryableDeviceParams); i++) { const auto ¶ms = kAllUnqueryableDeviceParams[i]; if (params.cc_major == cc_major_ && params.cc_minor == cc_minor_) { builder.set_blocks_per_core_limit(params.blocks_per_core_limit); @@ -1168,17 +1167,14 @@ DeviceDescription *CUDAExecutor::PopulateDeviceDescription() const { } // namespace cuda -namespace gpu = ::perftools::gputools; - void initialize_cuda_gpu_executor() { - *gpu::internal::MakeCUDAExecutorImplementation() = []( - const gpu::PluginConfig &config) { - return new gpu::cuda::CUDAExecutor{config}; + *internal::MakeCUDAExecutorImplementation() = [](const PluginConfig &config) { + return new cuda::CUDAExecutor{config}; }; } -} // namespace gputools -} // namespace perftools +} // namespace stream_executor -REGISTER_MODULE_INITIALIZER( - cuda_gpu_executor, {perftools::gputools::initialize_cuda_gpu_executor();}); +REGISTER_MODULE_INITIALIZER(cuda_gpu_executor, { + stream_executor::initialize_cuda_gpu_executor(); +}); diff --git a/tensorflow/stream_executor/cuda/cuda_gpu_executor.h b/tensorflow/stream_executor/cuda/cuda_gpu_executor.h index dbbbcd476f096ff912d391604ba349f6cb979478..f686685474b35acfb54c327401500c42109006d0 100644 --- a/tensorflow/stream_executor/cuda/cuda_gpu_executor.h +++ b/tensorflow/stream_executor/cuda/cuda_gpu_executor.h @@ -35,8 +35,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/thread_annotations.h" #include "tensorflow/stream_executor/stream_executor_internal.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { // CUDA-platform implementation of the platform-agnostic @@ -273,7 +272,6 @@ class CUDAExecutor : public internal::StreamExecutorInterface { }; } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_GPU_EXECUTOR_H_ diff --git a/tensorflow/stream_executor/cuda/cuda_helpers.h b/tensorflow/stream_executor/cuda/cuda_helpers.h index 6a6134bf881646991065cba536e955ef7c29e88c..d55706c66a9b47abfe125eaaa09e4b0cc543622a 100644 --- a/tensorflow/stream_executor/cuda/cuda_helpers.h +++ b/tensorflow/stream_executor/cuda/cuda_helpers.h @@ -27,8 +27,7 @@ limitations under the License. #include "cuda/include/cuComplex.h" #include "cuda/include/cuda.h" -namespace perftools { -namespace gputools { +namespace stream_executor { template class DeviceMemory; @@ -101,7 +100,6 @@ inline cuDoubleComplex CUDAComplexValue(std::complex val) { } } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_HELPERS_H_ diff --git a/tensorflow/stream_executor/cuda/cuda_kernel.h b/tensorflow/stream_executor/cuda/cuda_kernel.h index 88d29fddd06ad7c07bf8e90f490db2f4458e3684..beaebe8f1233533053c97bbac7eb283deaf96a2c 100644 --- a/tensorflow/stream_executor/cuda/cuda_kernel.h +++ b/tensorflow/stream_executor/cuda/cuda_kernel.h @@ -40,8 +40,7 @@ limitations under the License. "CUDA runtime being included into CUDA GPU executor; should be driver only." #endif -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { // Wraps a CUfunction to implement the platform-independent KernelInterface. @@ -124,7 +123,6 @@ inline CUDAKernel *AsCUDAKernel(KernelBase *kernel) { } } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_KERNEL_H_ diff --git a/tensorflow/stream_executor/cuda/cuda_platform.cc b/tensorflow/stream_executor/cuda/cuda_platform.cc index 3a738461489212a026197bc58777883349ba4b54..7a6ef5a248f7f96535fb64a3b2ca8ca15bf997cc 100644 --- a/tensorflow/stream_executor/cuda/cuda_platform.cc +++ b/tensorflow/stream_executor/cuda/cuda_platform.cc @@ -24,8 +24,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/status.h" #include "tensorflow/stream_executor/lib/stringprintf.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { namespace { @@ -41,16 +40,16 @@ const DeviceOptions GetDeviceOptionsFromEnv() { std::getenv("TF_CUDA_PLATFORM_GPU_DEVICE_SCHEDULE"); if (gpu_schedule_string == nullptr) { - return perftools::gputools::DeviceOptions::Default(); + return DeviceOptions::Default(); } unsigned device_flags = 0; if (strcmp(kScheduleSpinString, gpu_schedule_string) == 0) { - device_flags = perftools::gputools::DeviceOptions::kScheduleSpin; + device_flags = DeviceOptions::kScheduleSpin; } else if (strcmp(kScheduleYieldString, gpu_schedule_string) == 0) { - device_flags = perftools::gputools::DeviceOptions::kScheduleYield; + device_flags = DeviceOptions::kScheduleYield; } else if (strcmp(kScheduleBlockingSyncString, gpu_schedule_string) == 0) { - device_flags = perftools::gputools::DeviceOptions::kScheduleBlockingSync; + device_flags = DeviceOptions::kScheduleBlockingSync; } else { LOG(QFATAL) << "Unknown option for environment variable " "TF_CUDA_PLATFORM_GPU_DEVICE_SCHEDULE " @@ -59,7 +58,7 @@ const DeviceOptions GetDeviceOptionsFromEnv() { << ", " << kScheduleYieldString << "}"; } - return perftools::gputools::DeviceOptions(device_flags); + return DeviceOptions(device_flags); } } // namespace @@ -202,11 +201,10 @@ static void InitializeCudaPlatform() { SE_CHECK_OK(MultiPlatformManager::RegisterPlatform(std::move(platform))); } -} // namespace gputools -} // namespace perftools +} // namespace stream_executor REGISTER_MODULE_INITIALIZER(cuda_platform, - perftools::gputools::InitializeCudaPlatform()); + stream_executor::InitializeCudaPlatform()); DECLARE_MODULE_INITIALIZER(multi_platform_manager); // Note that module initialization sequencing is not supported in the diff --git a/tensorflow/stream_executor/cuda/cuda_platform.h b/tensorflow/stream_executor/cuda/cuda_platform.h index dab25602d08f394a2f8504ef9affe92552e3311f..fc0e15d5a6a9142f064085d34fcfaedfb25f433a 100644 --- a/tensorflow/stream_executor/cuda/cuda_platform.h +++ b/tensorflow/stream_executor/cuda/cuda_platform.h @@ -31,8 +31,7 @@ limitations under the License. #include "tensorflow/stream_executor/stream_executor_pimpl.h" #include "tensorflow/stream_executor/trace_listener.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { // Opaque and unique identifier for the CUDA platform plugin. @@ -104,7 +103,6 @@ class CudaPlatform : public Platform { }; } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_PLATFORM_H_ diff --git a/tensorflow/stream_executor/cuda/cuda_platform_id.cc b/tensorflow/stream_executor/cuda/cuda_platform_id.cc index dfd11a9abe87852924d5a0f52454ebd432f566e4..a7bb304cc8cde52dc8aa4512e911de3f40188100 100644 --- a/tensorflow/stream_executor/cuda/cuda_platform_id.cc +++ b/tensorflow/stream_executor/cuda/cuda_platform_id.cc @@ -15,12 +15,10 @@ limitations under the License. #include "tensorflow/stream_executor/cuda/cuda_platform_id.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { PLATFORM_DEFINE_ID(kCudaPlatformId); } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/cuda/cuda_platform_id.h b/tensorflow/stream_executor/cuda/cuda_platform_id.h index c677724517c1394284fb8b723933fa4b71d48ceb..92bcfd8372262dc127b71d399b6249efa9eff5dc 100644 --- a/tensorflow/stream_executor/cuda/cuda_platform_id.h +++ b/tensorflow/stream_executor/cuda/cuda_platform_id.h @@ -18,8 +18,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { // Opaque and unique identifier for the cuda platform. @@ -30,7 +29,6 @@ namespace cuda { extern const Platform::Id kCudaPlatformId; } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_PLATFORM_ID_H_ diff --git a/tensorflow/stream_executor/cuda/cuda_rng.cc b/tensorflow/stream_executor/cuda/cuda_rng.cc index 8641b6022776048baf93723d914d68b4a73f881c..e289e7ced57b164412b0fd78b538c3b08c7db5aa 100644 --- a/tensorflow/stream_executor/cuda/cuda_rng.cc +++ b/tensorflow/stream_executor/cuda/cuda_rng.cc @@ -54,15 +54,14 @@ std::ostream &operator<<(std::ostream &in, const curandStatus_t &status) { } } -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { PLUGIN_REGISTRY_DEFINE_PLUGIN_ID(kCuRandPlugin); namespace wrap { -#define PERFTOOLS_GPUTOOLS_CURAND_WRAP(__name) \ +#define STREAM_EXECUTOR_CURAND_WRAP(__name) \ struct WrapperShim__##__name { \ template \ curandStatus_t operator()(CUDAExecutor *parent, Args... args) { \ @@ -71,15 +70,15 @@ namespace wrap { } \ } __name; -PERFTOOLS_GPUTOOLS_CURAND_WRAP(curandCreateGenerator); -PERFTOOLS_GPUTOOLS_CURAND_WRAP(curandDestroyGenerator); -PERFTOOLS_GPUTOOLS_CURAND_WRAP(curandSetStream); -PERFTOOLS_GPUTOOLS_CURAND_WRAP(curandGenerateUniform); -PERFTOOLS_GPUTOOLS_CURAND_WRAP(curandGenerateUniformDouble); -PERFTOOLS_GPUTOOLS_CURAND_WRAP(curandSetPseudoRandomGeneratorSeed); -PERFTOOLS_GPUTOOLS_CURAND_WRAP(curandSetGeneratorOffset); -PERFTOOLS_GPUTOOLS_CURAND_WRAP(curandGenerateNormal); -PERFTOOLS_GPUTOOLS_CURAND_WRAP(curandGenerateNormalDouble); +STREAM_EXECUTOR_CURAND_WRAP(curandCreateGenerator); +STREAM_EXECUTOR_CURAND_WRAP(curandDestroyGenerator); +STREAM_EXECUTOR_CURAND_WRAP(curandSetStream); +STREAM_EXECUTOR_CURAND_WRAP(curandGenerateUniform); +STREAM_EXECUTOR_CURAND_WRAP(curandGenerateUniformDouble); +STREAM_EXECUTOR_CURAND_WRAP(curandSetPseudoRandomGeneratorSeed); +STREAM_EXECUTOR_CURAND_WRAP(curandSetGeneratorOffset); +STREAM_EXECUTOR_CURAND_WRAP(curandGenerateNormal); +STREAM_EXECUTOR_CURAND_WRAP(curandGenerateNormalDouble); } // namespace wrap @@ -271,42 +270,40 @@ bool CUDARng::SetSeed(Stream *stream, const uint8 *seed, uint64 seed_bytes) { } } // namespace cuda -} // namespace gputools -} // namespace perftools - -namespace gpu = ::perftools::gputools; - -REGISTER_MODULE_INITIALIZER(register_curand, { - gpu::port::Status status = - gpu::PluginRegistry::Instance() - ->RegisterFactory( - gpu::cuda::kCudaPlatformId, gpu::cuda::kCuRandPlugin, "cuRAND", - [](gpu::internal::StreamExecutorInterface - *parent) -> gpu::rng::RngSupport * { - gpu::cuda::CUDAExecutor *cuda_executor = - dynamic_cast(parent); - if (cuda_executor == nullptr) { - LOG(ERROR) - << "Attempting to initialize an instance of the cuRAND " - << "support library with a non-CUDA StreamExecutor"; - return nullptr; - } - - gpu::cuda::CUDARng *rng = new gpu::cuda::CUDARng(cuda_executor); - if (!rng->Init()) { - // Note: Init() will log a more specific error. - delete rng; - return nullptr; - } - return rng; - }); + +void initialize_curand() { + port::Status status = + PluginRegistry::Instance()->RegisterFactory( + cuda::kCudaPlatformId, cuda::kCuRandPlugin, "cuRAND", + [](internal::StreamExecutorInterface *parent) -> rng::RngSupport * { + cuda::CUDAExecutor *cuda_executor = + dynamic_cast(parent); + if (cuda_executor == nullptr) { + LOG(ERROR) + << "Attempting to initialize an instance of the cuRAND " + << "support library with a non-CUDA StreamExecutor"; + return nullptr; + } + + cuda::CUDARng *rng = new cuda::CUDARng(cuda_executor); + if (!rng->Init()) { + // Note: Init() will log a more specific error. + delete rng; + return nullptr; + } + return rng; + }); if (!status.ok()) { LOG(ERROR) << "Unable to register cuRAND factory: " << status.error_message(); } - gpu::PluginRegistry::Instance()->SetDefaultFactory(gpu::cuda::kCudaPlatformId, - gpu::PluginKind::kRng, - gpu::cuda::kCuRandPlugin); -}); + PluginRegistry::Instance()->SetDefaultFactory( + cuda::kCudaPlatformId, PluginKind::kRng, cuda::kCuRandPlugin); +} + +} // namespace stream_executor + +REGISTER_MODULE_INITIALIZER(register_curand, + { stream_executor::initialize_curand(); }); diff --git a/tensorflow/stream_executor/cuda/cuda_rng.h b/tensorflow/stream_executor/cuda/cuda_rng.h index 5bbfd0b37a0dafb525e2ed57ccec2daa0a0c52d8..57ef398aaa88da7de769c49820325c6c9feb4d70 100644 --- a/tensorflow/stream_executor/cuda/cuda_rng.h +++ b/tensorflow/stream_executor/cuda/cuda_rng.h @@ -24,8 +24,7 @@ limitations under the License. typedef struct curandGenerator_st *curandGenerator_t; -namespace perftools { -namespace gputools { +namespace stream_executor { class Stream; template @@ -98,7 +97,6 @@ class CUDARng : public rng::RngSupport { }; } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_RNG_H_ diff --git a/tensorflow/stream_executor/cuda/cuda_stream.cc b/tensorflow/stream_executor/cuda/cuda_stream.cc index 3eb37a7d84189cacb54d1e45fb0113030a9402f2..b5aa7694f7e1d8d47f3252d3ba679292155119b5 100644 --- a/tensorflow/stream_executor/cuda/cuda_stream.cc +++ b/tensorflow/stream_executor/cuda/cuda_stream.cc @@ -19,8 +19,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/status.h" #include "tensorflow/stream_executor/stream.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { bool CUDAStream::Init() { @@ -59,5 +58,4 @@ CUstream AsCUDAStreamValue(Stream *stream) { } } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/cuda/cuda_stream.h b/tensorflow/stream_executor/cuda/cuda_stream.h index 7358243dc4087194aa67da9958546c0487e95b8a..02edff643117fc2e3c6e6f74d2932f3f4c00c66d 100644 --- a/tensorflow/stream_executor/cuda/cuda_stream.h +++ b/tensorflow/stream_executor/cuda/cuda_stream.h @@ -23,8 +23,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/thread_annotations.h" #include "tensorflow/stream_executor/stream_executor_internal.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { class CUDAExecutor; @@ -89,7 +88,6 @@ CUDAStream *AsCUDAStream(Stream *stream); CUstream AsCUDAStreamValue(Stream *stream); } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_STREAM_H_ diff --git a/tensorflow/stream_executor/cuda/cuda_timer.cc b/tensorflow/stream_executor/cuda/cuda_timer.cc index 4bd5503348f4dc92a0ce3c18aaf0128174a94121..991a12a23d632bd9fb4c97a340e244f6ffb4c7d3 100644 --- a/tensorflow/stream_executor/cuda/cuda_timer.cc +++ b/tensorflow/stream_executor/cuda/cuda_timer.cc @@ -20,23 +20,24 @@ limitations under the License. #include "tensorflow/stream_executor/cuda/cuda_stream.h" #include "tensorflow/stream_executor/lib/status.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { bool CUDATimer::Init() { CHECK(start_event_ == nullptr && stop_event_ == nullptr); CudaContext* context = parent_->cuda_context(); - if (!CUDADriver::CreateEvent(context, &start_event_, - CUDADriver::EventFlags::kDefault) - .ok()) { + port::Status status = CUDADriver::CreateEvent( + context, &start_event_, CUDADriver::EventFlags::kDefault); + if (!status.ok()) { + LOG(ERROR) << status; return false; } - if (!CUDADriver::CreateEvent(context, &stop_event_, - CUDADriver::EventFlags::kDefault) - .ok()) { - port::Status status = CUDADriver::DestroyEvent(context, &start_event_); + status = CUDADriver::CreateEvent(context, &stop_event_, + CUDADriver::EventFlags::kDefault); + if (!status.ok()) { + LOG(ERROR) << status; + status = CUDADriver::DestroyEvent(context, &start_event_); if (!status.ok()) { LOG(ERROR) << status; } @@ -71,18 +72,23 @@ float CUDATimer::GetElapsedMilliseconds() const { return elapsed_milliseconds; } -bool CUDATimer::Start(CUDAStream *stream) { - return CUDADriver::RecordEvent(parent_->cuda_context(), start_event_, - stream->cuda_stream()) - .ok(); +bool CUDATimer::Start(CUDAStream* stream) { + port::Status status = CUDADriver::RecordEvent( + parent_->cuda_context(), start_event_, stream->cuda_stream()); + if (!status.ok()) { + LOG(ERROR) << status; + } + return status.ok(); } -bool CUDATimer::Stop(CUDAStream *stream) { - return CUDADriver::RecordEvent(parent_->cuda_context(), stop_event_, - stream->cuda_stream()) - .ok(); +bool CUDATimer::Stop(CUDAStream* stream) { + port::Status status = CUDADriver::RecordEvent( + parent_->cuda_context(), stop_event_, stream->cuda_stream()); + if (!status.ok()) { + LOG(ERROR) << status; + } + return status.ok(); } } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/cuda/cuda_timer.h b/tensorflow/stream_executor/cuda/cuda_timer.h index 2abc55ec9411d1cc90ddcd46fe27d19472e7ef49..70554ec93120fcb0251ba0995a1ce9d6e5997016 100644 --- a/tensorflow/stream_executor/cuda/cuda_timer.h +++ b/tensorflow/stream_executor/cuda/cuda_timer.h @@ -23,8 +23,7 @@ limitations under the License. #include "tensorflow/stream_executor/cuda/cuda_driver.h" #include "tensorflow/stream_executor/cuda/cuda_gpu_executor.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { class CUDAExecutor; @@ -60,13 +59,13 @@ class CUDATimer : public internal::TimerInterface { // events. float GetElapsedMilliseconds() const; - // See perftools::gputools::Timer::Microseconds(). + // See Timer::Microseconds(). // TODO(leary) make this into an error code interface... uint64 Microseconds() const override { return GetElapsedMilliseconds() * 1e3; } - // See perftools::GPUTools::Timer::Nanoseconds(). + // See Timer::Nanoseconds(). uint64 Nanoseconds() const override { return GetElapsedMilliseconds() * 1e6; } private: @@ -85,7 +84,6 @@ struct TimerDeleter { }; } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_TIMER_H_ diff --git a/tensorflow/stream_executor/cuda/cudnn_version.cc b/tensorflow/stream_executor/cuda/cudnn_version.cc index 5591801aae2526d528289f9b2267d864cf766045..e8fcc0361850a561928d09f29f78fb57071c24b2 100644 --- a/tensorflow/stream_executor/cuda/cudnn_version.cc +++ b/tensorflow/stream_executor/cuda/cudnn_version.cc @@ -15,8 +15,7 @@ limitations under the License. #include "tensorflow/stream_executor/cuda/cudnn_version.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { bool IsSourceCompatibleWithCudnnLibrary(CudnnVersion source_version, @@ -38,5 +37,4 @@ bool IsSourceCompatibleWithCudnnLibrary(CudnnVersion source_version, } } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/cuda/cudnn_version.h b/tensorflow/stream_executor/cuda/cudnn_version.h index 2ed02e1700ced5087bfebacb6314cbc8771e3612..6464e7f8e8755b5b46b90a4b35d50509eb0cfde7 100644 --- a/tensorflow/stream_executor/cuda/cudnn_version.h +++ b/tensorflow/stream_executor/cuda/cudnn_version.h @@ -20,8 +20,7 @@ limitations under the License. #include "tensorflow/core/lib/strings/strcat.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { struct CudnnVersion { @@ -46,7 +45,6 @@ bool IsSourceCompatibleWithCudnnLibrary(CudnnVersion source_version, CudnnVersion loaded_version); } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDNN_VERSION_H_ diff --git a/tensorflow/stream_executor/cuda/cudnn_version_test.cc b/tensorflow/stream_executor/cuda/cudnn_version_test.cc index 42b3dc8cc67b085d7c0d483f2922ef4f19df0092..7d4c6399d040e9bcddff5d98d202ab00fdeffa58 100644 --- a/tensorflow/stream_executor/cuda/cudnn_version_test.cc +++ b/tensorflow/stream_executor/cuda/cudnn_version_test.cc @@ -17,8 +17,7 @@ limitations under the License. #include "tensorflow/core/platform/test.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace cuda { namespace { @@ -70,5 +69,4 @@ TEST(IsSourceCompatibleWithCudnnLibraryTest, Basic) { } // namespace } // namespace cuda -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/device_description.cc b/tensorflow/stream_executor/device_description.cc index 52f5319a3b16c771ce89843a963841b25df5467e..8ca0677f8a3b3173fd023bd120526a490c2e7878 100644 --- a/tensorflow/stream_executor/device_description.cc +++ b/tensorflow/stream_executor/device_description.cc @@ -21,8 +21,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/mathutil.h" #include "tensorflow/stream_executor/lib/strcat.h" -namespace perftools { -namespace gputools { +namespace stream_executor { static const uint64 kUninitializedUint64 = -1ULL; /* static */ const char *DeviceDescription::kUndefinedString = ""; @@ -234,6 +233,4 @@ uint64 CalculateRegisterLimitForTargetOccupancy( return 0; } - -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/device_description.h b/tensorflow/stream_executor/device_description.h index fcf0928096ed1f1bdf0499efb92af2bc9cb0eaa2..7f99d81ef3bc5e21d5ea225ecfbc5a97bbd01ef5 100644 --- a/tensorflow/stream_executor/device_description.h +++ b/tensorflow/stream_executor/device_description.h @@ -27,8 +27,7 @@ limitations under the License. #include "tensorflow/stream_executor/launch_dim.h" #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace internal { class DeviceDescriptionBuilder; } // namespace internal @@ -388,7 +387,6 @@ uint64 CalculateRegisterLimitForTargetOccupancy( const DeviceDescription &device_description, uint64 shared_memory_per_block, const ThreadDim &thread_dims, uint64 target_blocks_per_core); -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_DEVICE_DESCRIPTION_H_ diff --git a/tensorflow/stream_executor/device_memory.h b/tensorflow/stream_executor/device_memory.h index 4c92b7dc78523e894aaf954fc44b72e89891d153..5a5334e0f5f6e8744b92188de14d7fea0f2ff9a0 100644 --- a/tensorflow/stream_executor/device_memory.h +++ b/tensorflow/stream_executor/device_memory.h @@ -32,6 +32,16 @@ limitations under the License. namespace perftools { namespace gputools { +// Temporarily pull stream_executor into perftools::gputools while we migrate +// code to the new namespace. TODO(b/77980417): Remove this once we've +// completed the migration. +using namespace stream_executor; // NOLINT[build/namespaces] + +} // namespace gputools +} // namespace perftools + +namespace stream_executor { + class StreamExecutor; // void*-analogous device memory allocation. For the typed variation, see @@ -280,7 +290,6 @@ static_assert(sizeof(Float2) == 2 * sizeof(float), "Float2 must be packed"); static_assert(sizeof(Float4) == 4 * sizeof(float), "Float4 must be packed"); static_assert(sizeof(Double2) == 2 * sizeof(double), "Double2 must be packed"); -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_DEVICE_MEMORY_H_ diff --git a/tensorflow/stream_executor/device_options.h b/tensorflow/stream_executor/device_options.h index 169325e7d127edb83c896e1b18e5bf88641b46b1..2646950f42e986d0c732f212af0c24c403b142e5 100644 --- a/tensorflow/stream_executor/device_options.h +++ b/tensorflow/stream_executor/device_options.h @@ -25,8 +25,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/port.h" #include "tensorflow/stream_executor/platform/logging.h" -namespace perftools { -namespace gputools { +namespace stream_executor { // Indicates a set of options for a device's usage, which generally must be // provided at StreamExecutor device-initialization time. @@ -84,7 +83,6 @@ struct DeviceOptions { unsigned flags_; }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_DEVICE_OPTIONS_H_ diff --git a/tensorflow/stream_executor/dnn.cc b/tensorflow/stream_executor/dnn.cc index 0a3c4bcf503b36f3ca7811970e94fdd7930b7c1a..6edb572820104592c3e9e506bbc66ec008672ec0 100644 --- a/tensorflow/stream_executor/dnn.cc +++ b/tensorflow/stream_executor/dnn.cc @@ -18,8 +18,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/strcat.h" #include "tensorflow/stream_executor/lib/stringprintf.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace dnn { bool DnnSupport::GetConvolveAlgorithms( @@ -554,5 +553,4 @@ string NormalizeDescriptor::ToShortString() const { } } // namespace dnn -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/dnn.h b/tensorflow/stream_executor/dnn.h index 3c47d2c2e852055988eb6d99e2081ba935920704..39f21d8b105058caa28af1b32f600383821b46c1 100644 --- a/tensorflow/stream_executor/dnn.h +++ b/tensorflow/stream_executor/dnn.h @@ -38,8 +38,7 @@ namespace Eigen { struct half; } // namespace Eigen -namespace perftools { -namespace gputools { +namespace stream_executor { class HostBuffer; class Stream; @@ -876,6 +875,22 @@ enum class ElementwiseOperation { kAdd, kMultiply }; string ElementwiseOperationString(ElementwiseOperation op); +// A simple class representing the version of the backing library, to +// workaround the "too perfect forwarding" issue in gcc6+ compilers. +// See PR#16309 and issue #18402 for links discussing the issue. +class VersionInfo { + public: + VersionInfo(int major = 0, int minor = 0, int patch = 0) + : major_(major), minor_(minor), patch_(patch) {} + int major_version() { return major_; } + int minor_version() { return minor_; } + int patch() { return patch_; } + private: + int major_; + int minor_; + int patch_; +}; + // Suite of operations typically used for implementing Deep/Convolutional Neural // Nets. Note: A false return value of an operation indicates the // implementation is not available. @@ -886,8 +901,8 @@ class DnnSupport { virtual port::Status Init() = 0; - // Gets the version of the backing library, as a {major, minor, patch} tuple. - virtual port::StatusOr> GetVersion() { + // Gets the version of the backing library, as a VersionInfo object. + virtual port::StatusOr GetVersion() { return port::UnimplementedError( "DnnSupport::GetVersion not implemented on this platform."); } @@ -2285,7 +2300,6 @@ class DnnSupport { }; } // namespace dnn -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_DNN_H_ diff --git a/tensorflow/stream_executor/dso_loader.cc b/tensorflow/stream_executor/dso_loader.cc index 95168836278add5d6592ff0c3d0f7245e6f6bc5b..114143b3abef00e757da3263449454fb1908fd53 100644 --- a/tensorflow/stream_executor/dso_loader.cc +++ b/tensorflow/stream_executor/dso_loader.cc @@ -37,8 +37,7 @@ limitations under the License. #include "cuda/cuda_config.h" #endif -namespace perftools { -namespace gputools { +namespace stream_executor { namespace internal { string GetCudaVersion() { return TF_CUDA_VERSION; } @@ -291,5 +290,4 @@ static std::vector* CreatePrimordialRpaths() { } } // namespace internal -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/dso_loader.h b/tensorflow/stream_executor/dso_loader.h index 354c7b50b8209755991827b3c36afac790cb952b..9ee081cb3d64e8878fa9d7f0c33da7f6827da620 100644 --- a/tensorflow/stream_executor/dso_loader.h +++ b/tensorflow/stream_executor/dso_loader.h @@ -28,8 +28,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform.h" #include "tensorflow/stream_executor/platform/mutex.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace internal { // Permits StreamExecutor code to dynamically load a pre-determined set of @@ -114,7 +113,6 @@ class CachedDsoLoader { }; } // namespace internal -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_DSO_LOADER_H_ diff --git a/tensorflow/stream_executor/event.cc b/tensorflow/stream_executor/event.cc index c423a453e9f1cb4f1d484427c11eacf959da7a8d..50a6edd80bd39004e32f09bcde36fbc8a8b59ad9 100644 --- a/tensorflow/stream_executor/event.cc +++ b/tensorflow/stream_executor/event.cc @@ -19,8 +19,7 @@ limitations under the License. #include "tensorflow/stream_executor/stream_executor_pimpl.h" #include "tensorflow/stream_executor/stream.h" -namespace perftools { -namespace gputools { +namespace stream_executor { Event::Event(StreamExecutor* stream_exec) : stream_exec_(stream_exec), @@ -48,5 +47,4 @@ Event::Status Event::PollForStatus() { return stream_exec_->PollForEventStatus(this); } -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/event.h b/tensorflow/stream_executor/event.h index a06c26ea5191d53c7d8ffc1e2b460cb594dd6f4a..1f37262c78d82f72f8818f35db273e87a47bdc1c 100644 --- a/tensorflow/stream_executor/event.h +++ b/tensorflow/stream_executor/event.h @@ -20,8 +20,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace internal { class EventInterface; @@ -76,7 +75,6 @@ class Event { SE_DISALLOW_COPY_AND_ASSIGN(Event); }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_EVENT_H_ diff --git a/tensorflow/stream_executor/executor_cache.cc b/tensorflow/stream_executor/executor_cache.cc index d1a8aae167455a7dc728999fbbaf1a119cf6a101..0b3ad7ebbc905ab1bd7e3298bb560af0e3868656 100644 --- a/tensorflow/stream_executor/executor_cache.cc +++ b/tensorflow/stream_executor/executor_cache.cc @@ -17,8 +17,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/stringprintf.h" -namespace perftools { -namespace gputools { +namespace stream_executor { port::StatusOr ExecutorCache::GetOrCreate( const StreamExecutorConfig& config, @@ -104,5 +103,4 @@ ExecutorCache::Entry::~Entry() { configurations.clear(); } -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/executor_cache.h b/tensorflow/stream_executor/executor_cache.h index 12f2275f6d8ecde4819d56ab9411c1087b76ddc5..bbeeaed787c27a62e2b0b22a75779b5350187d83 100644 --- a/tensorflow/stream_executor/executor_cache.h +++ b/tensorflow/stream_executor/executor_cache.h @@ -24,8 +24,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/mutex.h" #include "tensorflow/stream_executor/stream_executor_pimpl.h" -namespace perftools { -namespace gputools { +namespace stream_executor { // Utility class to allow Platform objects to manage cached StreamExecutors. // Thread-safe. @@ -76,7 +75,6 @@ class ExecutorCache { SE_DISALLOW_COPY_AND_ASSIGN(ExecutorCache); }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_EXECUTOR_CACHE_H_ diff --git a/tensorflow/stream_executor/fft.h b/tensorflow/stream_executor/fft.h index 6b1728829abdeb5c4e20534675801a437341d732..814efb2e923cb25833515f86e1b3f6f0bec437da 100644 --- a/tensorflow/stream_executor/fft.h +++ b/tensorflow/stream_executor/fft.h @@ -48,8 +48,7 @@ limitations under the License. #include #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { class Stream; template @@ -210,7 +209,7 @@ class FftSupport { // Macro used to quickly declare overrides for abstract virtuals in the // fft::FftSupport base class. Assumes that it's emitted somewhere inside the -// ::perftools::gputools namespace. +// ::stream_executor namespace. #define TENSORFLOW_STREAM_EXECUTOR_GPU_FFT_SUPPORT_OVERRIDES \ std::unique_ptr Create1dPlan(Stream *stream, uint64 num_x, \ fft::Type type, bool in_place_fft) \ @@ -265,7 +264,6 @@ class FftSupport { DeviceMemory *output) override; } // namespace fft -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_FFT_H_ diff --git a/tensorflow/stream_executor/host/host_gpu_executor.cc b/tensorflow/stream_executor/host/host_gpu_executor.cc index 542f521ef778c3a69ec9adba74405131e07bcf1a..2c4819651acaa2c6ee99c720b2c3d80e5c2ea1a9 100644 --- a/tensorflow/stream_executor/host/host_gpu_executor.cc +++ b/tensorflow/stream_executor/host/host_gpu_executor.cc @@ -28,8 +28,7 @@ limitations under the License. bool FLAGS_stream_executor_cpu_real_clock_rate = false; -namespace perftools { -namespace gputools { +namespace stream_executor { namespace host { HostStream *AsHostStream(Stream *stream) { @@ -266,5 +265,4 @@ rng::RngSupport *HostExecutor::CreateRng() { } } // namespace host -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/host/host_gpu_executor.h b/tensorflow/stream_executor/host/host_gpu_executor.h index e2c0e6d6b77130bd190b026f1eaff68d21dbf632..0c3991c151d5bbb84333240aa1ca4f9e1587330d 100644 --- a/tensorflow/stream_executor/host/host_gpu_executor.h +++ b/tensorflow/stream_executor/host/host_gpu_executor.h @@ -28,8 +28,7 @@ limitations under the License. #include "tensorflow/stream_executor/stream_executor.h" #include "tensorflow/stream_executor/stream_executor_internal.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace host { // An implementation of StreamExecutor that does no communication or interaction @@ -210,7 +209,6 @@ class HostExecutor : public internal::StreamExecutorInterface { }; } // namespace host -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_HOST_HOST_GPU_EXECUTOR_H_ diff --git a/tensorflow/stream_executor/host/host_platform.cc b/tensorflow/stream_executor/host/host_platform.cc index 2cb7d3696758365804786800e6a46e8bca7d80e7..00a17a05edea24d6485a948d72856537021c99ab 100644 --- a/tensorflow/stream_executor/host/host_platform.cc +++ b/tensorflow/stream_executor/host/host_platform.cc @@ -26,10 +26,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/status_macros.h" #include "tensorflow/stream_executor/lib/stringprintf.h" -namespace gpu = ::perftools::gputools; - -namespace perftools { -namespace gputools { +namespace stream_executor { namespace host { HostPlatform::HostPlatform() : name_("Host") {} @@ -93,16 +90,15 @@ void HostPlatform::UnregisterTraceListener(TraceListener* listener) { } static void InitializeHostPlatform() { - std::unique_ptr platform(new gpu::host::HostPlatform); - SE_CHECK_OK(gpu::MultiPlatformManager::RegisterPlatform(std::move(platform))); + std::unique_ptr platform(new host::HostPlatform); + SE_CHECK_OK(MultiPlatformManager::RegisterPlatform(std::move(platform))); } } // namespace host -} // namespace gputools -} // namespace perftools +} // namespace stream_executor -REGISTER_MODULE_INITIALIZER( - host_platform, perftools::gputools::host::InitializeHostPlatform()); +REGISTER_MODULE_INITIALIZER(host_platform, + stream_executor::host::InitializeHostPlatform()); DECLARE_MODULE_INITIALIZER(multi_platform_manager); // Note that module initialization sequencing is not supported in the diff --git a/tensorflow/stream_executor/host/host_platform.h b/tensorflow/stream_executor/host/host_platform.h index 0faec6c8b789766804c262fc5cf29c950fd31d3c..c6f46a2cc4028644237f206f5cd076a49e964702 100644 --- a/tensorflow/stream_executor/host/host_platform.h +++ b/tensorflow/stream_executor/host/host_platform.h @@ -33,8 +33,7 @@ limitations under the License. #include "tensorflow/stream_executor/stream_executor_pimpl.h" #include "tensorflow/stream_executor/trace_listener.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace host { // Host (CPU) platform plugin, registered as a singleton value via module @@ -79,7 +78,6 @@ class HostPlatform : public Platform { }; } // namespace host -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_HOST_HOST_PLATFORM_H_ diff --git a/tensorflow/stream_executor/host/host_platform_id.cc b/tensorflow/stream_executor/host/host_platform_id.cc index 69a203f2985b67ecd67346e6dc394d961377e035..2256bccec3f4ef73bc667128aeb1fbf90b1dcd90 100644 --- a/tensorflow/stream_executor/host/host_platform_id.cc +++ b/tensorflow/stream_executor/host/host_platform_id.cc @@ -15,12 +15,10 @@ limitations under the License. #include "tensorflow/stream_executor/host/host_platform_id.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace host { PLATFORM_DEFINE_ID(kHostPlatformId); } // namespace host -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/host/host_platform_id.h b/tensorflow/stream_executor/host/host_platform_id.h index 61d84ea2e2facae3a18798db15d529109b6dfc93..18d1f282f1f921bc9bb2508225398c98ffeb89bc 100644 --- a/tensorflow/stream_executor/host/host_platform_id.h +++ b/tensorflow/stream_executor/host/host_platform_id.h @@ -18,8 +18,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace host { // Opaque and unique identifier for the host platform. @@ -30,7 +29,6 @@ namespace host { extern const Platform::Id kHostPlatformId; } // namespace host -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_HOST_HOST_PLATFORM_ID_H_ diff --git a/tensorflow/stream_executor/host/host_stream.cc b/tensorflow/stream_executor/host/host_stream.cc index 5961c3151600b778b4cbca7443b5913329368431..5a7d3b3dd49275edd5242c30b38bb4f505042816 100644 --- a/tensorflow/stream_executor/host/host_stream.cc +++ b/tensorflow/stream_executor/host/host_stream.cc @@ -17,8 +17,7 @@ limitations under the License. // the HostExecutor implementation. #include "tensorflow/stream_executor/host/host_stream.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace host { HostStream::HostStream() @@ -53,5 +52,4 @@ void HostStream::BlockUntilDone() { } // namespace host -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/host/host_stream.h b/tensorflow/stream_executor/host/host_stream.h index 9894d17febcae24089bfe2948eb96bad701978d7..5d7b8a378268c3226a61fa43e738f209e84b30e9 100644 --- a/tensorflow/stream_executor/host/host_stream.h +++ b/tensorflow/stream_executor/host/host_stream.h @@ -24,8 +24,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/threadpool.h" #include "tensorflow/stream_executor/stream_executor_internal.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace host { class HostStream : public internal::StreamInterface { @@ -52,7 +51,6 @@ class HostStream : public internal::StreamInterface { }; } // namespace host -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_HOST_HOST_STREAM_H_ diff --git a/tensorflow/stream_executor/host/host_timer.cc b/tensorflow/stream_executor/host/host_timer.cc index d84d825c92a6e8e06b004305c9d4cb032883f5a3..e138daf0e11f179c8690306b5cfef3ed95ae4cb2 100644 --- a/tensorflow/stream_executor/host/host_timer.cc +++ b/tensorflow/stream_executor/host/host_timer.cc @@ -19,8 +19,7 @@ limitations under the License. #include "tensorflow/stream_executor/stream.h" #include "tensorflow/stream_executor/stream_executor_internal.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace host { using std::chrono::duration_cast; @@ -46,5 +45,4 @@ void HostTimer::StartNow() { start_time_ = clock::now(); } void HostTimer::StopNow() { duration_ = clock::now() - start_time_; } } // namespace host -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/host/host_timer.h b/tensorflow/stream_executor/host/host_timer.h index 17af7c0521d2ee956d6f1ee92dd2229b647785bd..5954b8023beb4fe78d37cd41c0379804a8dfa8aa 100644 --- a/tensorflow/stream_executor/host/host_timer.h +++ b/tensorflow/stream_executor/host/host_timer.h @@ -20,8 +20,7 @@ limitations under the License. #include "tensorflow/stream_executor/stream_executor_internal.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace host { class HostTimer : public internal::TimerInterface { @@ -57,7 +56,6 @@ class HostTimer : public internal::TimerInterface { }; } // namespace host -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_HOST_HOST_TIMER_H_ diff --git a/tensorflow/stream_executor/host_buffer.h b/tensorflow/stream_executor/host_buffer.h index 8fa542e9ff8851551d424da6be073d0d9959af67..20299da5172f20b9b73c31b6491806dc57b1d2f0 100644 --- a/tensorflow/stream_executor/host_buffer.h +++ b/tensorflow/stream_executor/host_buffer.h @@ -18,8 +18,7 @@ limitations under the License. #include "tensorflow/stream_executor/dnn.h" -namespace perftools { -namespace gputools { +namespace stream_executor { // A HostBuffer is a block of memory in host memory containing the data for a // dnn::BatchDescriptor using a device-dependent memory layout. @@ -42,7 +41,6 @@ class HostBuffer { const dnn::BatchDescriptor descriptor_; }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_HOST_BUFFER_H_ diff --git a/tensorflow/stream_executor/kernel.cc b/tensorflow/stream_executor/kernel.cc index 636199cfa2762b7c42dd350dfd294762e3159299..d1aa596b73da3d7f8cdf4ddd6911d956239ed46d 100644 --- a/tensorflow/stream_executor/kernel.cc +++ b/tensorflow/stream_executor/kernel.cc @@ -27,8 +27,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/logging.h" #include "tensorflow/stream_executor/stream_executor.h" -namespace perftools { -namespace gputools { +namespace stream_executor { bool KernelMetadata::registers_per_thread(int *registers_per_thread) const { if (has_registers_per_thread_) { @@ -103,5 +102,4 @@ void KernelBase::set_name(port::StringPiece name) { demangled_name_ = port::Demangle(stubless_name.data()); } -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/kernel.h b/tensorflow/stream_executor/kernel.h index 5358eac1ae070efb2bead75c73208e9d283b498c..2216884b873cda98f09782866f23c06088b73e09 100644 --- a/tensorflow/stream_executor/kernel.h +++ b/tensorflow/stream_executor/kernel.h @@ -64,7 +64,7 @@ limitations under the License. // // Users typically won't need to type out the TypedKernel signature in full, it // will be typedef'd by automatically generated code; for example, see -// perftools::gputools::executor_sample::VecReduceAddKernel. +// stream_executor::executor_sample::VecReduceAddKernel. #ifndef TENSORFLOW_STREAM_EXECUTOR_KERNEL_H_ #define TENSORFLOW_STREAM_EXECUTOR_KERNEL_H_ @@ -82,8 +82,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/stringpiece.h" #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { class DeviceMemoryBase; template @@ -639,8 +638,8 @@ struct KernelInvocationChecker { // NOTE: if you encounter an error here, you can see the mismatch by looking // at the end of the last error message, which will be of the form: // - // ...::Compatible &, - // perftools::gputools::DeviceMemory, true, + // ...::Compatible &, + // stream_executor::DeviceMemory, true, // 0>' // requested here // @@ -711,7 +710,6 @@ struct KernelParamsOk, Args...> { std::tuple, std::tuple>::CheckAllNoStaticAssert(); }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_KERNEL_H_ diff --git a/tensorflow/stream_executor/kernel_cache_config.h b/tensorflow/stream_executor/kernel_cache_config.h index 9d7ab1b79f6c94ce2fe64f14f7ebe3a9db44daf4..e63d6c6a0c05e8cd681a4e12f73b2468e170dea5 100644 --- a/tensorflow/stream_executor/kernel_cache_config.h +++ b/tensorflow/stream_executor/kernel_cache_config.h @@ -18,8 +18,7 @@ limitations under the License. #ifndef TENSORFLOW_STREAM_EXECUTOR_KERNEL_CACHE_CONFIG_H_ #define TENSORFLOW_STREAM_EXECUTOR_KERNEL_CACHE_CONFIG_H_ -namespace perftools { -namespace gputools { +namespace stream_executor { // This enum represents potential configurations of L1/shared memory when // running a particular kernel. These values represent user preference, and @@ -38,7 +37,6 @@ enum class KernelCacheConfig { kPreferEqual, }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_KERNEL_CACHE_CONFIG_H_ diff --git a/tensorflow/stream_executor/kernel_spec.cc b/tensorflow/stream_executor/kernel_spec.cc index 0404c573f0164782bb056b124ef60092c5985a5a..6a1f0a591ff0879821c6636e346dbc20105bcf7e 100644 --- a/tensorflow/stream_executor/kernel_spec.cc +++ b/tensorflow/stream_executor/kernel_spec.cc @@ -15,9 +15,7 @@ limitations under the License. #include "tensorflow/stream_executor/kernel_spec.h" - -namespace perftools { -namespace gputools { +namespace stream_executor { KernelLoaderSpec::KernelLoaderSpec(port::StringPiece kernelname) : kernelname_(kernelname.ToString()) {} @@ -247,5 +245,4 @@ MultiKernelLoaderSpec *MultiKernelLoaderSpec::AddCudaCompressedPtxInMemory( MultiKernelLoaderSpec::MultiKernelLoaderSpec(size_t arity) : arity_(arity) {} -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/kernel_spec.h b/tensorflow/stream_executor/kernel_spec.h index 3811bd833e70b567002327cbe348574b37cb1126..7cc23bb4e64b45268f6bb00d9ea9ee4a686a0e25 100644 --- a/tensorflow/stream_executor/kernel_spec.h +++ b/tensorflow/stream_executor/kernel_spec.h @@ -56,8 +56,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/mutex.h" #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { // Describes how to load a kernel on a target platform. // @@ -374,7 +373,6 @@ class MultiKernelLoaderSpec { size_t arity_; }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_KERNEL_SPEC_H_ diff --git a/tensorflow/stream_executor/launch_dim.h b/tensorflow/stream_executor/launch_dim.h index b95462667eea2c86d92060f3973e660f31d2d2dc..68f2f748407a87ec9cf3bdd411bf96c1a64b5681 100644 --- a/tensorflow/stream_executor/launch_dim.h +++ b/tensorflow/stream_executor/launch_dim.h @@ -21,7 +21,7 @@ limitations under the License. // a single PC in a unit called a warp. There is a maximum number of threads // that can execute in a shared-context entity called a block. Presently, that // number is 1024 -- again, something that should not be relied on from this -// comment, but checked via perftools::gputools::DeviceDescription. +// comment, but checked via stream_executor::DeviceDescription. // // For additional information, see // http://docs.nvidia.com/cuda/kepler-tuning-guide/#device-utilization-and-occupancy @@ -40,8 +40,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/strcat.h" #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { // Basic type that represents a 3-dimensional index space. struct Dim3D { @@ -74,7 +73,6 @@ struct BlockDim : public Dim3D { } }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LAUNCH_DIM_H_ diff --git a/tensorflow/stream_executor/lib/array_slice.h b/tensorflow/stream_executor/lib/array_slice.h index bef61bb2fc5611e84bc36320fa98f4df919372b7..8e3c4ca047bd467628551d8b8d351fa3baa956f7 100644 --- a/tensorflow/stream_executor/lib/array_slice.h +++ b/tensorflow/stream_executor/lib/array_slice.h @@ -18,14 +18,23 @@ limitations under the License. #include "tensorflow/core/lib/gtl/array_slice.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { using tensorflow::gtl::ArraySlice; using tensorflow::gtl::MutableArraySlice; } // namespace port +} // namespace stream_executor + +namespace perftools { +namespace gputools { + +// Temporarily pull stream_executor into perftools::gputools while we migrate +// code to the new namespace. TODO(b/77980417): Remove this once we've +// completed the migration. +using namespace stream_executor; // NOLINT[build/namespaces] + } // namespace gputools } // namespace perftools diff --git a/tensorflow/stream_executor/lib/casts.h b/tensorflow/stream_executor/lib/casts.h index 2261944e252f6a757a154e64be0b10484a693125..ec562e804fae51ac09e336b2e03b8ab0d7f1ca0e 100644 --- a/tensorflow/stream_executor/lib/casts.h +++ b/tensorflow/stream_executor/lib/casts.h @@ -13,15 +13,14 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -// IWYU pragma: private, include "perftools/gputools/executor/stream_executor.h" +// IWYU pragma: private, include "third_party/tensorflow/stream_executor/stream_executor.h" #ifndef TENSORFLOW_STREAM_EXECUTOR_LIB_CASTS_H_ #define TENSORFLOW_STREAM_EXECUTOR_LIB_CASTS_H_ #include -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { // port::bit_cast is a template function that implements the @@ -96,7 +95,6 @@ inline Dest bit_cast(const Source& source) { } } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_CASTS_H_ diff --git a/tensorflow/stream_executor/lib/demangle.cc b/tensorflow/stream_executor/lib/demangle.cc index fa2b4fa005ca1a715b0593c7c270c62bcf936d21..adb6b4f2d11e537433ae86ad694d6e1de9171942 100644 --- a/tensorflow/stream_executor/lib/demangle.cc +++ b/tensorflow/stream_executor/lib/demangle.cc @@ -27,8 +27,7 @@ limitations under the License. #include #endif -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { // The API reference of abi::__cxa_demangle() can be found in @@ -49,5 +48,4 @@ string Demangle(const char *mangled) { } } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/lib/demangle.h b/tensorflow/stream_executor/lib/demangle.h index 30be522557808e4048c7a1f712d0b4884c1043ba..af16fa7d8cbcce5c17a7b25021546176570558bf 100644 --- a/tensorflow/stream_executor/lib/demangle.h +++ b/tensorflow/stream_executor/lib/demangle.h @@ -18,14 +18,12 @@ limitations under the License. #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { string Demangle(const char* mangled); } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_DEMANGLE_H_ diff --git a/tensorflow/stream_executor/lib/env.h b/tensorflow/stream_executor/lib/env.h index c9a22ebd55894f0ddf6654887e55289030871501..776eba04080e9e993b61997bc3e7532202ac4c5b 100644 --- a/tensorflow/stream_executor/lib/env.h +++ b/tensorflow/stream_executor/lib/env.h @@ -21,8 +21,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/stringpiece.h" #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { using tensorflow::Env; @@ -37,7 +36,6 @@ inline Status FileExists(const port::StringPiece& filename) { } } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_ENV_H_ diff --git a/tensorflow/stream_executor/lib/error.h b/tensorflow/stream_executor/lib/error.h index 89df70cb5e582508448b6886926b0c8905a14c2e..c659f5fc140d29e17fc0bf11ce707b01b2cb2e74 100644 --- a/tensorflow/stream_executor/lib/error.h +++ b/tensorflow/stream_executor/lib/error.h @@ -13,21 +13,19 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -// IWYU pragma: private, include "perftools/gputools/executor/stream_executor.h" +// IWYU pragma: private, include "third_party/tensorflow/stream_executor/stream_executor.h" #ifndef TENSORFLOW_STREAM_EXECUTOR_LIB_ERROR_H_ #define TENSORFLOW_STREAM_EXECUTOR_LIB_ERROR_H_ #include "tensorflow/core/lib/core/error_codes.pb.h" // IWYU pragma: export -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { namespace error = tensorflow::error; } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_ERROR_H_ diff --git a/tensorflow/stream_executor/lib/human_readable.h b/tensorflow/stream_executor/lib/human_readable.h index f918c180d98c0fc74d554e3c7e95995e6720d356..893865f6dadd1076a4daef6d820f667b57ab9023 100644 --- a/tensorflow/stream_executor/lib/human_readable.h +++ b/tensorflow/stream_executor/lib/human_readable.h @@ -22,8 +22,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/stringprintf.h" #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { class HumanReadableNumBytes { @@ -67,7 +66,6 @@ class HumanReadableNumBytes { }; } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_HUMAN_READABLE_H_ diff --git a/tensorflow/stream_executor/lib/initialize.h b/tensorflow/stream_executor/lib/initialize.h index 9a09318a6cbd215e396968bc1968be07a3029129..688b0214694478e9be1b1d14e58fda94367f547b 100644 --- a/tensorflow/stream_executor/lib/initialize.h +++ b/tensorflow/stream_executor/lib/initialize.h @@ -26,8 +26,7 @@ limitations under the License. #undef DECLARE_MODULE_INITIALIZER #undef REGISTER_MODULE_INITIALIZER_SEQUENCE -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { class Initializer { @@ -49,20 +48,18 @@ class Initializer { }; } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor -#define REGISTER_INITIALIZER(type, name, body) \ - static void google_init_##type##_##name() { body; } \ - perftools::gputools::port::Initializer google_initializer_##type##_##name( \ +#define REGISTER_INITIALIZER(type, name, body) \ + static void google_init_##type##_##name() { body; } \ + ::stream_executor::port::Initializer google_initializer_##type##_##name( \ google_init_##type##_##name) #define REGISTER_MODULE_INITIALIZER(name, body) \ REGISTER_INITIALIZER(module, name, body) -#define DECLARE_INITIALIZER(type, name) \ - extern perftools::gputools::port::Initializer \ - google_initializer_##type##_##name +#define DECLARE_INITIALIZER(type, name) \ + extern ::stream_executor::port::Initializer google_initializer_##type##_##name #define DECLARE_MODULE_INITIALIZER(name) DECLARE_INITIALIZER(module, name) diff --git a/tensorflow/stream_executor/lib/inlined_vector.h b/tensorflow/stream_executor/lib/inlined_vector.h index 55a1e3ad102c2fe25c9b4338a07029330bbf4ef6..40bdddb180f1135ef2fb3bfc51b06d9e68a3aff9 100644 --- a/tensorflow/stream_executor/lib/inlined_vector.h +++ b/tensorflow/stream_executor/lib/inlined_vector.h @@ -18,14 +18,12 @@ limitations under the License. #include "tensorflow/core/lib/gtl/inlined_vector.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { using tensorflow::gtl::InlinedVector; } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_INLINED_VECTOR_H_ diff --git a/tensorflow/stream_executor/lib/mathutil.h b/tensorflow/stream_executor/lib/mathutil.h index e8310d55ddae20f1a92c4184d408215d7ec8664a..c225dc5f3cc9725be4b01039c10103ed6e62372e 100644 --- a/tensorflow/stream_executor/lib/mathutil.h +++ b/tensorflow/stream_executor/lib/mathutil.h @@ -25,8 +25,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/logging.h" #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { class MathUtil { @@ -97,7 +96,6 @@ IntegralType MathUtil::CeilOrFloorOfRatio(IntegralType numerator, } } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_MATHUTIL_H_ diff --git a/tensorflow/stream_executor/lib/notification.h b/tensorflow/stream_executor/lib/notification.h index 9bb3e170dc7c13ecaa636887315b4bc5cfc82f2b..472d8c9845c2ae3efbc5ab06234a7c084d044a96 100644 --- a/tensorflow/stream_executor/lib/notification.h +++ b/tensorflow/stream_executor/lib/notification.h @@ -18,14 +18,12 @@ limitations under the License. #include "tensorflow/core/platform/notification.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { using tensorflow::Notification; } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_NOTIFICATION_H_ diff --git a/tensorflow/stream_executor/lib/numbers.cc b/tensorflow/stream_executor/lib/numbers.cc index 11a65e198d60545e3bfb76a8cbf141c4dba1b58e..b670c42ec84a86d49a4938da3fc6387af529e04b 100644 --- a/tensorflow/stream_executor/lib/numbers.cc +++ b/tensorflow/stream_executor/lib/numbers.cc @@ -17,8 +17,7 @@ limitations under the License. #include -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { bool safe_strto32(const char* str, int32* value) { @@ -38,5 +37,4 @@ bool safe_strto32(const string& str, int32* value) { } } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/lib/numbers.h b/tensorflow/stream_executor/lib/numbers.h index 4a8692b7461d513bbac287281b134a4fbac54707..2f48281d2d67f0fa9b4fb523c4986c5dc8acfe33 100644 --- a/tensorflow/stream_executor/lib/numbers.h +++ b/tensorflow/stream_executor/lib/numbers.h @@ -18,8 +18,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { // Convert strings to floating point values. @@ -28,7 +27,6 @@ namespace port { bool safe_strto32(const string& str, int32* value); } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_NUMBERS_H_ diff --git a/tensorflow/stream_executor/lib/path.cc b/tensorflow/stream_executor/lib/path.cc index f2591f47f7bf75004ee00019d70e118cd9d32f0e..56e08c316f95758cd2a05d18d59c06f3d7727959 100644 --- a/tensorflow/stream_executor/lib/path.cc +++ b/tensorflow/stream_executor/lib/path.cc @@ -16,8 +16,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/path.h" #include "tensorflow/stream_executor/lib/strcat.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { namespace internal { @@ -58,5 +57,4 @@ string JoinPathImpl(std::initializer_list paths) { } // namespace internal } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/lib/path.h b/tensorflow/stream_executor/lib/path.h index 93053dbcb6862edbaf93a27a9414ad7e4bb10be2..325f04ff47552e052d81c96bec74e816378b8254 100644 --- a/tensorflow/stream_executor/lib/path.h +++ b/tensorflow/stream_executor/lib/path.h @@ -20,8 +20,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/stringpiece.h" #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { using tensorflow::io::Dirname; @@ -56,7 +55,6 @@ inline string JoinPath(const T&... args) { } } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_PATH_H_ diff --git a/tensorflow/stream_executor/lib/process_state.cc b/tensorflow/stream_executor/lib/process_state.cc index 3d856187f05b3601525b048d3997088e145bdb81..72d71e62116726d458f709b75e50fcc0947df00e 100644 --- a/tensorflow/stream_executor/lib/process_state.cc +++ b/tensorflow/stream_executor/lib/process_state.cc @@ -25,8 +25,7 @@ limitations under the License. #endif #include -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { string Hostname() { @@ -54,5 +53,4 @@ bool GetCurrentDirectory(string* dir) { } } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/lib/process_state.h b/tensorflow/stream_executor/lib/process_state.h index 205e726d95ce80216a8c260c270d32b990fd5b78..248218c759e65fc2d6ee09b2f1c9e309567b1048 100644 --- a/tensorflow/stream_executor/lib/process_state.h +++ b/tensorflow/stream_executor/lib/process_state.h @@ -18,15 +18,13 @@ limitations under the License. #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { string Hostname(); bool GetCurrentDirectory(string* dir); } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_PROCESS_STATE_H_ diff --git a/tensorflow/stream_executor/lib/ptr_util.h b/tensorflow/stream_executor/lib/ptr_util.h index 3d5e56faf745e484649079cee916bc26c8b2c449..3f89794688ca390a3d1e475fbe69828be1914167 100644 --- a/tensorflow/stream_executor/lib/ptr_util.h +++ b/tensorflow/stream_executor/lib/ptr_util.h @@ -18,8 +18,7 @@ limitations under the License. #include -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { // Trait to select overloads and return types for MakeUnique. @@ -59,8 +58,17 @@ typename MakeUniqueResult::invalid MakeUnique(Args&&... /* args */) = delete; // NOLINT } // namespace port +} // namespace stream_executor + +namespace perftools { +namespace gputools { + +// Temporarily pull stream_executor into perftools::gputools while we migrate +// code to the new namespace. TODO(jlebar): Remove this once we've completed +// the migration. +using namespace stream_executor; // NOLINT[build/namespaces] + } // namespace gputools } // namespace perftools - #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_PTR_UTIL_H_ diff --git a/tensorflow/stream_executor/lib/stacktrace.h b/tensorflow/stream_executor/lib/stacktrace.h index ba7e5317f0f4a852c353102d29411752f69b11a9..a15b0f30261f1a102c0f7cf8b28b5fc9b5426e25 100644 --- a/tensorflow/stream_executor/lib/stacktrace.h +++ b/tensorflow/stream_executor/lib/stacktrace.h @@ -19,14 +19,12 @@ limitations under the License. #include "tensorflow/core/platform/stacktrace.h" #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { using tensorflow::CurrentStackTrace; } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_STACKTRACE_H_ diff --git a/tensorflow/stream_executor/lib/status.h b/tensorflow/stream_executor/lib/status.h index 8c289e1927fcf9b49851389e44d3fa09fdfea3ae..407b71b405bc8a73e5aebcd18b043420b074b708 100644 --- a/tensorflow/stream_executor/lib/status.h +++ b/tensorflow/stream_executor/lib/status.h @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -// IWYU pragma: private, include "perftools/gputools/executor/stream_executor.h" +// IWYU pragma: private, include "third_party/tensorflow/stream_executor/stream_executor.h" #ifndef TENSORFLOW_STREAM_EXECUTOR_LIB_STATUS_H_ #define TENSORFLOW_STREAM_EXECUTOR_LIB_STATUS_H_ @@ -23,15 +23,14 @@ limitations under the License. #include "tensorflow/stream_executor/lib/stringpiece.h" #include "tensorflow/stream_executor/platform/logging.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { using Status = tensorflow::Status; #define SE_CHECK_OK(val) TF_CHECK_OK(val) #define SE_ASSERT_OK(val) \ - ASSERT_EQ(::perftools::gputools::port::Status::OK(), (val)) + ASSERT_EQ(::stream_executor::port::Status::OK(), (val)) // Define some canonical error helpers. inline Status UnimplementedError(StringPiece message) { @@ -45,6 +44,16 @@ inline Status FailedPreconditionError(StringPiece message) { } } // namespace port +} // namespace stream_executor + +namespace perftools { +namespace gputools { + +// Temporarily pull stream_executor into perftools::gputools while we migrate +// code to the new namespace. TODO(b/77980417): Remove this once we've +// completed the migration. +using namespace stream_executor; // NOLINT[build/namespaces] + } // namespace gputools } // namespace perftools diff --git a/tensorflow/stream_executor/lib/statusor.h b/tensorflow/stream_executor/lib/statusor.h index 138738ecab54986fd7d5cd76839d59da55623b1f..dab59096740102b94c0ff63c089b83ce052ea264 100644 --- a/tensorflow/stream_executor/lib/statusor.h +++ b/tensorflow/stream_executor/lib/statusor.h @@ -13,242 +13,21 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -// IWYU pragma: private, include "perftools/gputools/executor/stream_executor.h" -// -// StatusOr is the union of a Status object and a T -// object. StatusOr models the concept of an object that is either a -// usable value, or an error Status explaining why such a value is -// not present. To this end, StatusOr does not allow its Status -// value to be Status::OK. Further, StatusOr does not allow the -// contained pointer to be NULL. -// -// The primary use-case for StatusOr is as the return value of a -// function which may fail. -// -// Example client usage for a StatusOr, where T is not a pointer: -// -// StatusOr result = DoBigCalculationThatCouldFail(); -// if (result.ok()) { -// float answer = result.ValueOrDie(); -// printf("Big calculation yielded: %f", answer); -// } else { -// LOG(ERROR) << result.status(); -// } -// -// Example client usage for a StatusOr: -// -// StatusOr result = FooFactory::MakeNewFoo(arg); -// if (result.ok()) { -// std::unique_ptr foo(result.ValueOrDie()); -// foo->DoSomethingCool(); -// } else { -// LOG(ERROR) << result.status(); -// } -// -// Example client usage for a StatusOr>: -// -// StatusOr> result = FooFactory::MakeNewFoo(arg); -// if (result.ok()) { -// std::unique_ptr foo = std::move(result.ValueOrDie()); -// foo->DoSomethingCool(); -// } else { -// LOG(ERROR) << result.status(); -// } -// -// Example factory implementation returning StatusOr: -// -// StatusOr FooFactory::MakeNewFoo(int arg) { -// if (arg <= 0) { -// return Status(port::error::INVALID_ARGUMENT, -// "Arg must be positive"); -// } else { -// return new Foo(arg); -// } -// } -// +// IWYU pragma: private, include "third_party/tensorflow/stream_executor/stream_executor.h" #ifndef TENSORFLOW_STREAM_EXECUTOR_LIB_STATUSOR_H_ #define TENSORFLOW_STREAM_EXECUTOR_LIB_STATUSOR_H_ -#include -#include "tensorflow/stream_executor/platform/port.h" -#include -#include +#include "tensorflow/compiler/xla/statusor.h" -#include "tensorflow/stream_executor/lib/error.h" -#include "tensorflow/stream_executor/lib/status.h" -#include "tensorflow/stream_executor/platform/logging.h" -#include "tensorflow/stream_executor/platform/port.h" - -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { -template -class StatusOr { - template friend class StatusOr; - - public: - // Construct a new StatusOr with Status::UNKNOWN status - StatusOr() : status_(error::UNKNOWN, "") {} - - // Construct a new StatusOr with the given non-ok status. After calling - // this constructor, calls to ValueOrDie() is invalid. - // - // NOTE: Not explicit - we want to use StatusOr as a return - // value, so it is convenient and sensible to be able to do 'return - // Status()' when the return type is StatusOr. - // - // REQUIRES: status != Status::OK. - // In optimized builds, passing Status::OK here will have the effect - // of passing PosixErrorSpace::EINVAL as a fallback. - StatusOr(const Status& status); // NOLINT - - // Construct a new StatusOr with the given value. If T is a plain pointer, - // value must not be NULL. After calling this constructor, calls to - // ValueOrDie() will succeed, and calls to status() will return OK. - // - // NOTE: Not explicit - we want to use StatusOr as a return type - // so it is convenient and sensible to be able to do 'return T()' - // when the return type is StatusOr. - // - // REQUIRES: if T is a plain pointer, value != NULL. - // In optimized builds, passing a NULL pointer here will have - // the effect of passing PosixErrorSpace::EINVAL as a fallback. - StatusOr(const T& value); // NOLINT - - // Conversion copy constructor, T must be copy constructible from U - template - StatusOr(const StatusOr& other) // NOLINT - : status_(other.status_), - value_(other.value_) {} - - // Conversion assignment operator, T must be assignable from U - template - StatusOr& operator=(const StatusOr& other) { - status_ = other.status_; - value_ = other.value_; - return *this; - } - - // Rvalue-reference overloads of the other constructors and assignment - // operators, to support move-only types and avoid unnecessary copying. - StatusOr(T&& value); // NOLINT - - // Move conversion operator to avoid unnecessary copy. - // T must be assignable from U. - // Not marked with explicit so the implicit conversion can happen. - template - StatusOr(StatusOr&& other) // NOLINT - : status_(std::move(other.status_)), - value_(std::move(other.value_)) {} - - // Move assignment operator to avoid unnecessary copy. - // T must be assignable from U - template - StatusOr& operator=(StatusOr&& other) { - status_ = std::move(other.status_); - value_ = std::move(other.value_); - return *this; - } - - // Returns a reference to our status. If this contains a T, then - // returns Status::OK. - const Status& status() const { return status_; } - - // Returns this->status().ok() - bool ok() const { return status_.ok(); } - - // Returns a reference to our current value, requires that this->ok(). - // If you need to initialize a T object from the stored value, - // ConsumeValueOrDie() may be more efficient. - const T& ValueOrDie() const; - T& ValueOrDie(); - - // Returns our current value, requires this->ok(). Use this if - // you would otherwise want to say std::move(s.ValueOrDie()), for example - // if you need to initialize a T object from the stored value and you don't - // need subsequent access to the stored value. It uses T's move constructor, - // if it has one, so it will work with move-only types, and will often be - // more efficient than ValueOrDie, but may leave the stored value - // in an arbitrary valid state. - T ConsumeValueOrDie(); - - private: - Status status_; - T value_; - - void CheckValueNotNull(const T& value); - - template - struct IsNull { - // For non-pointer U, a reference can never be NULL. - static inline bool IsValueNull(const U& t) { return false; } - }; - - template - struct IsNull { - static inline bool IsValueNull(const U* t) { return t == NULL; } - }; -}; - -//////////////////////////////////////////////////////////////////////////////// -// Implementation details for StatusOr - -template -StatusOr::StatusOr(const T& value) - : status_(), value_(value) { - CheckValueNotNull(value); -} - -template -const T& StatusOr::ValueOrDie() const { - TF_CHECK_OK(status_); - return value_; -} - -template -T& StatusOr::ValueOrDie() { - TF_CHECK_OK(status_); - return value_; -} - -template -T StatusOr::ConsumeValueOrDie() { - TF_CHECK_OK(status_); - return std::move(value_); -} - -template -StatusOr::StatusOr(const Status& status) - : status_(status) { - assert(!status.ok()); - if (status.ok()) { - status_ = - Status(error::INTERNAL, - "Status::OK is not a valid constructor argument to StatusOr"); - } -} - -template -StatusOr::StatusOr(T&& value) - : status_() { - CheckValueNotNull(value); - value_ = std::move(value); -} - +// Use XLA's StatusOr so we don't duplicate code. template -void StatusOr::CheckValueNotNull(const T& value) { - assert(!IsNull::IsValueNull(value)); - if (IsNull::IsValueNull(value)) { - status_ = - Status(error::INTERNAL, - "NULL is not a valid constructor argument to StatusOr"); - } -} +using StatusOr = ::xla::StatusOr; } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_STATUSOR_H_ diff --git a/tensorflow/stream_executor/lib/str_util.h b/tensorflow/stream_executor/lib/str_util.h index 5dd3d06affa424f0919f107c76ba40feeb165122..a81c6668184c158b61eff4d4b77bb7ad07165c8e 100644 --- a/tensorflow/stream_executor/lib/str_util.h +++ b/tensorflow/stream_executor/lib/str_util.h @@ -19,8 +19,7 @@ limitations under the License. #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/stream_executor/lib/stringpiece.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { using tensorflow::str_util::Join; @@ -38,7 +37,6 @@ inline string StripSuffixString(port::StringPiece str, port::StringPiece suffix) using tensorflow::str_util::Lowercase; } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_STR_UTIL_H_ diff --git a/tensorflow/stream_executor/lib/strcat.h b/tensorflow/stream_executor/lib/strcat.h index 424cb75f0e805a840c5043a42e7cbf3f9c7d195c..c959e4df5b2d6e3954e2edf7bbafc9efb8259dc6 100644 --- a/tensorflow/stream_executor/lib/strcat.h +++ b/tensorflow/stream_executor/lib/strcat.h @@ -13,22 +13,20 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -// IWYU pragma: private, include "perftools/gputools/executor/stream_executor.h" +// IWYU pragma: private, include "third_party/tensorflow/stream_executor/stream_executor.h" #ifndef TENSORFLOW_STREAM_EXECUTOR_LIB_STRCAT_H_ #define TENSORFLOW_STREAM_EXECUTOR_LIB_STRCAT_H_ #include "tensorflow/core/lib/strings/strcat.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { using tensorflow::strings::StrCat; using tensorflow::strings::StrAppend; } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_STRCAT_H_ diff --git a/tensorflow/stream_executor/lib/stringpiece.h b/tensorflow/stream_executor/lib/stringpiece.h index 97ee0c9206480202db887018c21c1e41482a04e9..b80de5df306117e315b282e53ec871d60bd6d517 100644 --- a/tensorflow/stream_executor/lib/stringpiece.h +++ b/tensorflow/stream_executor/lib/stringpiece.h @@ -19,14 +19,12 @@ limitations under the License. #include "tensorflow/core/lib/core/stringpiece.h" #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { using tensorflow::StringPiece; } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_STRINGPIECE_H_ diff --git a/tensorflow/stream_executor/lib/stringprintf.h b/tensorflow/stream_executor/lib/stringprintf.h index 504de25a681c4ec880c04a46aad89ddeeb622c71..2f65ed9c6a86516cdb6f8aeaaf9c48b3622cb613 100644 --- a/tensorflow/stream_executor/lib/stringprintf.h +++ b/tensorflow/stream_executor/lib/stringprintf.h @@ -18,15 +18,13 @@ limitations under the License. #include "tensorflow/core/lib/strings/stringprintf.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { using tensorflow::strings::Printf; using tensorflow::strings::Appendf; } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_STRINGPRINTF_H_ diff --git a/tensorflow/stream_executor/lib/thread_options.h b/tensorflow/stream_executor/lib/thread_options.h index bd7f63714e25ec88b6d54bc683fe96a313589485..079cf757acd98d994b1351423205ef99a64478f2 100644 --- a/tensorflow/stream_executor/lib/thread_options.h +++ b/tensorflow/stream_executor/lib/thread_options.h @@ -18,14 +18,12 @@ limitations under the License. #include "tensorflow/core/platform/env.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { using tensorflow::ThreadOptions; } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_THREAD_OPTIONS_H_ diff --git a/tensorflow/stream_executor/lib/threadpool.h b/tensorflow/stream_executor/lib/threadpool.h index 35630c5106a1b732d1100c416b2abc7fe6404c55..220068ade11b98e2442cdc1c68b2aaeb422ae6b3 100644 --- a/tensorflow/stream_executor/lib/threadpool.h +++ b/tensorflow/stream_executor/lib/threadpool.h @@ -21,14 +21,12 @@ limitations under the License. #include "tensorflow/stream_executor/lib/notification.h" #include "tensorflow/stream_executor/lib/thread_options.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace port { using tensorflow::thread::ThreadPool; } // namespace port -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_LIB_THREADPOOL_H_ diff --git a/tensorflow/stream_executor/multi_platform_manager.cc b/tensorflow/stream_executor/multi_platform_manager.cc index f9f3737a06dad3f146ef9fc8e2ec50160b3a01b5..5b51398d8cab5df7c7514bc3bedf87f5c33c6e5a 100644 --- a/tensorflow/stream_executor/multi_platform_manager.cc +++ b/tensorflow/stream_executor/multi_platform_manager.cc @@ -20,8 +20,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/str_util.h" #include "tensorflow/stream_executor/lib/stringprintf.h" -namespace perftools { -namespace gputools { +namespace stream_executor { /* static */ mutex MultiPlatformManager::platforms_mutex_{LINKER_INITIALIZED}; @@ -132,8 +131,7 @@ MultiPlatformManager::InitializePlatformWithId( GetPlatformByIdMap()->clear(); } -} // namespace gputools -} // namespace perftools +} // namespace stream_executor REGISTER_MODULE_INITIALIZER( multi_platform_manager, diff --git a/tensorflow/stream_executor/multi_platform_manager.h b/tensorflow/stream_executor/multi_platform_manager.h index 438653ee20bdb1fd83cd9e75c4bcd35af277cc28..672855d5fb6152b8bfb00d8f0faa64873711646d 100644 --- a/tensorflow/stream_executor/multi_platform_manager.h +++ b/tensorflow/stream_executor/multi_platform_manager.h @@ -22,8 +22,8 @@ limitations under the License. // In your BUILD rule, add a dependency on a platform plugin that you'd like // to use, such as: // -// //perftools/gputools/executor/cuda:cuda_platform -// //perftools/gputools/executor/opencl:opencl_platform +// //third_party/tensorflow/stream_executor/cuda:cuda_platform +// //third_party/tensorflow/stream_executor/opencl:opencl_platform // // This will register platform plugins that can be discovered via this // interface. Sample API usage: @@ -56,10 +56,10 @@ limitations under the License. // And similarly, for standard interfaces (BLAS, RNG, etc.) you can add // dependencies on support libraries, e.g.: // -// //perftools/gputools/executor/cuda:pluton_blas_plugin -// //perftools/gputools/executor/cuda:cudnn_plugin -// //perftools/gputools/executor/cuda:cublas_plugin -// //perftools/gputools/executor/cuda:curand_plugin +// //third_party/tensorflow/stream_executor/cuda:pluton_blas_plugin +// //third_party/tensorflow/stream_executor/cuda:cudnn_plugin +// //third_party/tensorflow/stream_executor/cuda:cublas_plugin +// //third_party/tensorflow/stream_executor/cuda:curand_plugin #ifndef TENSORFLOW_STREAM_EXECUTOR_MULTI_PLATFORM_MANAGER_H_ #define TENSORFLOW_STREAM_EXECUTOR_MULTI_PLATFORM_MANAGER_H_ @@ -75,8 +75,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/port.h" #include "tensorflow/stream_executor/platform/thread_annotations.h" -namespace perftools { -namespace gputools { +namespace stream_executor { // Manages multiple platforms that may be present on the current machine. class MultiPlatformManager { @@ -181,7 +180,6 @@ class MultiPlatformManager { SE_DISALLOW_COPY_AND_ASSIGN(MultiPlatformManager); }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_MULTI_PLATFORM_MANAGER_H_ diff --git a/tensorflow/stream_executor/platform.cc b/tensorflow/stream_executor/platform.cc index 4cdc22bd16a3ea66037696f6a9d70bcb86ef5ebb..777abced8634410440fd5c2772a3b981a1489dd6 100644 --- a/tensorflow/stream_executor/platform.cc +++ b/tensorflow/stream_executor/platform.cc @@ -22,8 +22,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/logging.h" #include "tensorflow/stream_executor/stream_executor_pimpl.h" -namespace perftools { -namespace gputools { +namespace stream_executor { string PlatformKindString(PlatformKind kind) { switch (kind) { @@ -135,5 +134,4 @@ port::Status Platform::EnablePeerAccess() { return port::Status::OK(); } -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/platform.h b/tensorflow/stream_executor/platform.h index 54f8aa86c269ff0d32648e1d4629179cafd5be76..5cb7047b6f39483f237b5bb249906d9ce8a06b9e 100644 --- a/tensorflow/stream_executor/platform.h +++ b/tensorflow/stream_executor/platform.h @@ -29,8 +29,7 @@ limitations under the License. #include "tensorflow/stream_executor/plugin.h" #include "tensorflow/stream_executor/trace_listener.h" -namespace perftools { -namespace gputools { +namespace stream_executor { class StreamExecutor; @@ -106,7 +105,7 @@ class Platform { namespace { \ int plugin_id_value; \ } \ - const perftools::gputools::Platform::Id ID_VAR_NAME = &plugin_id_value; + const ::stream_executor::Platform::Id ID_VAR_NAME = &plugin_id_value; // Returns a key uniquely identifying this platform. virtual Id id() const = 0; @@ -205,7 +204,6 @@ class Platform { SE_DISALLOW_COPY_AND_ASSIGN(Platform); }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_PLATFORM_H_ diff --git a/tensorflow/stream_executor/platform/default/mutex.h b/tensorflow/stream_executor/platform/default/mutex.h index 62de0cbce0b534020166c47a64c10c595fb3d5ae..c9f5a7c609e5bbe59ea456e30d575b991aa37b65 100644 --- a/tensorflow/stream_executor/platform/default/mutex.h +++ b/tensorflow/stream_executor/platform/default/mutex.h @@ -18,8 +18,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/mutex.h" -namespace perftools { -namespace gputools { +namespace stream_executor { #undef mutex_lock #undef tf_shared_lock @@ -35,7 +34,6 @@ using tensorflow::tf_shared_lock; #define tf_shared_lock(x) \ static_assert(0, "tf_shared_lock_decl_missing_var_name"); -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_PLATFORM_DEFAULT_MUTEX_H_ diff --git a/tensorflow/stream_executor/platform/port.h b/tensorflow/stream_executor/platform/port.h index 6603df487878e62271a144b14d78518044c66c81..57ad965ef1101d7d86a4f9cf0a52016e2b3ae713 100644 --- a/tensorflow/stream_executor/platform/port.h +++ b/tensorflow/stream_executor/platform/port.h @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -// IWYU pragma: private, include "perftools/gputools/executor/stream_executor.h" +// IWYU pragma: private, include "third_party/tensorflow/stream_executor/stream_executor.h" #ifndef TENSORFLOW_STREAM_EXECUTOR_PLATFORM_PORT_H_ #define TENSORFLOW_STREAM_EXECUTOR_PLATFORM_PORT_H_ @@ -22,8 +22,7 @@ limitations under the License. #include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/types.h" -namespace perftools { -namespace gputools { +namespace stream_executor { using tensorflow::int8; using tensorflow::int16; @@ -39,19 +38,12 @@ using tensorflow::uint64; using std::string; #endif -#if !defined(COMPILER_MSVC) -#define ARRAYSIZE(a) \ - ((sizeof(a) / sizeof(*(a))) / \ - static_cast(!(sizeof(a) % sizeof(*(a))))) -#endif - using tensorflow::LinkerInitialized; using tensorflow::LINKER_INITIALIZED; #define SE_FALLTHROUGH_INTENDED TF_FALLTHROUGH_INTENDED -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #define SE_DISALLOW_COPY_AND_ASSIGN TF_DISALLOW_COPY_AND_ASSIGN #define SE_MUST_USE_RESULT TF_MUST_USE_RESULT diff --git a/tensorflow/stream_executor/plugin.cc b/tensorflow/stream_executor/plugin.cc index 6424658e22f930a7adfa4dd7fe330308e62e00cf..cfbc52ff17b3a699e97d073d0d0c77bba811db72 100644 --- a/tensorflow/stream_executor/plugin.cc +++ b/tensorflow/stream_executor/plugin.cc @@ -15,8 +15,7 @@ limitations under the License. #include "tensorflow/stream_executor/plugin.h" -namespace perftools { -namespace gputools { +namespace stream_executor { // Mostly-arbitrary ID only used as a sentinel "not otherwise initialized" // value. This value should never [need to] be specified aside by initialization @@ -51,5 +50,4 @@ PluginConfig& PluginConfig::SetRng(PluginId rng) { return *this; } -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/plugin.h b/tensorflow/stream_executor/plugin.h index 0b88b86e2b1cf8cbd3dddfa5ca3ae9cdc779a952..0505412e7ac412dda8cda1c5c4589148241c5b38 100644 --- a/tensorflow/stream_executor/plugin.h +++ b/tensorflow/stream_executor/plugin.h @@ -16,8 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_STREAM_EXECUTOR_PLUGIN_H_ #define TENSORFLOW_STREAM_EXECUTOR_PLUGIN_H_ -namespace perftools { -namespace gputools { +namespace stream_executor { // A plugin ID is a unique identifier for each registered plugin type. typedef void* PluginId; @@ -83,7 +82,6 @@ class PluginConfig { PluginId blas_, dnn_, fft_, rng_; }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_PLUGIN_H_ diff --git a/tensorflow/stream_executor/plugin_registry.cc b/tensorflow/stream_executor/plugin_registry.cc index 54761139eaf3ac5e2edb5ccc45dabe56fa1384d4..7812703efd8b4f143c195225c2c12502378bf5f2 100644 --- a/tensorflow/stream_executor/plugin_registry.cc +++ b/tensorflow/stream_executor/plugin_registry.cc @@ -19,8 +19,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/stringprintf.h" #include "tensorflow/stream_executor/multi_platform_manager.h" -namespace perftools { -namespace gputools { +namespace stream_executor { const PluginId kNullPlugin = nullptr; @@ -244,5 +243,4 @@ EMIT_PLUGIN_SPECIALIZATIONS(DnnFactory, dnn, "DNN"); EMIT_PLUGIN_SPECIALIZATIONS(FftFactory, fft, "FFT"); EMIT_PLUGIN_SPECIALIZATIONS(RngFactory, rng, "RNG"); -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/plugin_registry.h b/tensorflow/stream_executor/plugin_registry.h index 8636a49ce6861d24eb5567bc8d8b268b840a310f..49628ecd246e1a785ef3583f5a38567f553352f5 100644 --- a/tensorflow/stream_executor/plugin_registry.h +++ b/tensorflow/stream_executor/plugin_registry.h @@ -28,8 +28,7 @@ limitations under the License. #include "tensorflow/stream_executor/plugin.h" #include "tensorflow/stream_executor/rng.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace internal { class StreamExecutorInterface; @@ -160,7 +159,6 @@ class PluginRegistry { SE_DISALLOW_COPY_AND_ASSIGN(PluginRegistry); }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_PLUGIN_REGISTRY_H_ diff --git a/tensorflow/stream_executor/rng.cc b/tensorflow/stream_executor/rng.cc index 1c05005067ce7341994aa7b8f1cb9c4df734ce2c..b0efad91084b7fbfed70ecca9bf879810de2ec1d 100644 --- a/tensorflow/stream_executor/rng.cc +++ b/tensorflow/stream_executor/rng.cc @@ -17,8 +17,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/logging.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace rng { bool RngSupport::CheckSeed(const uint8 *seed, uint64 seed_bytes) { @@ -47,5 +46,4 @@ const int RngSupport::kMaxSeedBytes; #endif } // namespace rng -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/rng.h b/tensorflow/stream_executor/rng.h index 36d0fdd454f9a98b749fb2a8f9ceb14eec9d0fe1..acbf8fce4caa8cbbc9b2288725ea7bf07ee84af4 100644 --- a/tensorflow/stream_executor/rng.h +++ b/tensorflow/stream_executor/rng.h @@ -22,8 +22,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/logging.h" #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { class Stream; template @@ -89,7 +88,6 @@ class RngSupport { }; } // namespace rng -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_RNG_H_ diff --git a/tensorflow/stream_executor/scratch_allocator.cc b/tensorflow/stream_executor/scratch_allocator.cc index 0c1db414f2e9617070c8c12e2a334ed7977ae9f1..8fc4c4c509cfd35c8a733b01a8d85965fe524def 100644 --- a/tensorflow/stream_executor/scratch_allocator.cc +++ b/tensorflow/stream_executor/scratch_allocator.cc @@ -18,8 +18,7 @@ limitations under the License. #include "tensorflow/stream_executor/lib/status_macros.h" #include "tensorflow/stream_executor/stream.h" -namespace perftools { -namespace gputools { +namespace stream_executor { ScratchAllocator::~ScratchAllocator() {} @@ -38,5 +37,4 @@ port::StatusOr> OneTimeScratchAllocator::AllocateBytes( return temporary_->device_memory(); } -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/scratch_allocator.h b/tensorflow/stream_executor/scratch_allocator.h index 94d5ede161323ff39d7c75589958026b66f22cc5..2aed2c4437381bd29750aae1120a8783b61bfb9f 100644 --- a/tensorflow/stream_executor/scratch_allocator.h +++ b/tensorflow/stream_executor/scratch_allocator.h @@ -23,8 +23,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/port.h" #include "tensorflow/stream_executor/temporary_device_memory.h" -namespace perftools { -namespace gputools { +namespace stream_executor { class Stream; @@ -77,7 +76,6 @@ class OneTimeScratchAllocator : public ScratchAllocator { SE_DISALLOW_COPY_AND_ASSIGN(OneTimeScratchAllocator); }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_SCRATCH_ALLOCATOR_H_ diff --git a/tensorflow/stream_executor/shared_memory_config.h b/tensorflow/stream_executor/shared_memory_config.h index de556cb73404d5d919b779e5d712f465ff2bc821..7cbeb3bcd91816ece725afecc2e85e7f200c6045 100644 --- a/tensorflow/stream_executor/shared_memory_config.h +++ b/tensorflow/stream_executor/shared_memory_config.h @@ -19,8 +19,7 @@ limitations under the License. #ifndef TENSORFLOW_STREAM_EXECUTOR_SHARED_MEMORY_CONFIG_H_ #define TENSORFLOW_STREAM_EXECUTOR_SHARED_MEMORY_CONFIG_H_ -namespace perftools { -namespace gputools { +namespace stream_executor { // SharedMemoryConfig enum describes potential widths of shared memory banks for // a device or kernel. @@ -30,7 +29,6 @@ enum class SharedMemoryConfig { kEightByte, // Sets shared memory banks to be eight bytes wide. }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_SHARED_MEMORY_CONFIG_H_ diff --git a/tensorflow/stream_executor/stream.cc b/tensorflow/stream_executor/stream.cc index fe498507a8019c3b0994a98fb52301a1e9e52764..f59d9a13acffa405f8627416d87d812fe4e62c8e 100644 --- a/tensorflow/stream_executor/stream.cc +++ b/tensorflow/stream_executor/stream.cc @@ -28,8 +28,7 @@ limitations under the License. #include "tensorflow/stream_executor/stream_executor_internal.h" #include "tensorflow/stream_executor/stream_executor_pimpl.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace { // Code to turn parameters to functions on stream into strings that @@ -5192,5 +5191,4 @@ port::Status Stream::BlockHostUntilDone() { return first_error; } -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/stream.h b/tensorflow/stream_executor/stream.h index 4af426001f27a7983ce3a0832e53a1ce2c0ebd8e..d4a81440e9651e739545a3e9ccaceb1877949c75 100644 --- a/tensorflow/stream_executor/stream.h +++ b/tensorflow/stream_executor/stream.h @@ -38,8 +38,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/thread_annotations.h" #include "tensorflow/stream_executor/temporary_memory_manager.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace host { class HostBlas; @@ -2098,7 +2097,6 @@ struct Quantization { dnn::QuantizedActivationMode::k32Bit; }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_STREAM_H_ diff --git a/tensorflow/stream_executor/stream_executor.h b/tensorflow/stream_executor/stream_executor.h index 2995dccf469a2977b9593f2045062d568e569dc5..d63d485df5668a5cbb5c61525401947bebeab616 100644 --- a/tensorflow/stream_executor/stream_executor.h +++ b/tensorflow/stream_executor/stream_executor.h @@ -35,4 +35,15 @@ limitations under the License. #include "tensorflow/stream_executor/stream_executor_pimpl.h" // IWYU pragma: export #include "tensorflow/stream_executor/timer.h" // IWYU pragma: export +namespace perftools { +namespace gputools { + +// Temporarily pull stream_executor into perftools::gputools while we migrate +// code to the new namespace. TODO(b/77980417): Remove this once we've +// completed the migration. +using namespace stream_executor; // NOLINT[build/namespaces] + +} // namespace gputools +} // namespace perftools + #endif // TENSORFLOW_STREAM_EXECUTOR_STREAM_EXECUTOR_H_ diff --git a/tensorflow/stream_executor/stream_executor_internal.cc b/tensorflow/stream_executor/stream_executor_internal.cc index 273d970b6fa4a581381689191b183a30f4f2bcd3..8297228e6fecddffa8fc68a1a028456dc8e75a65 100644 --- a/tensorflow/stream_executor/stream_executor_internal.cc +++ b/tensorflow/stream_executor/stream_executor_internal.cc @@ -15,8 +15,7 @@ limitations under the License. #include "tensorflow/stream_executor/stream_executor_internal.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace internal { // -- CUDA @@ -38,5 +37,4 @@ StreamExecutorFactory* MakeOpenCLExecutorImplementation() { StreamExecutorFactory MakeHostExecutorImplementation; } // namespace internal -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/stream_executor_internal.h b/tensorflow/stream_executor/stream_executor_internal.h index 37ef182e1445a85dd0a97eac02ba064a26dc0f1d..2584c92f0c5a1129e2f10aa7148161a8d2d40c50 100644 --- a/tensorflow/stream_executor/stream_executor_internal.h +++ b/tensorflow/stream_executor/stream_executor_internal.h @@ -45,8 +45,7 @@ limitations under the License. #include "tensorflow/stream_executor/trace_listener.h" #include "tensorflow/stream_executor/lib/inlined_vector.h" -namespace perftools { -namespace gputools { +namespace stream_executor { class Stream; class Timer; @@ -343,7 +342,6 @@ extern StreamExecutorFactory MakeHostExecutorImplementation; } // namespace internal -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_STREAM_EXECUTOR_INTERNAL_H_ diff --git a/tensorflow/stream_executor/stream_executor_pimpl.cc b/tensorflow/stream_executor/stream_executor_pimpl.cc index f55fa68402953b0e084c7f688b9481cfa0ec1b7e..2e1adeb31e46d923eedfc1fab3c21271e5b76c2c 100644 --- a/tensorflow/stream_executor/stream_executor_pimpl.cc +++ b/tensorflow/stream_executor/stream_executor_pimpl.cc @@ -39,8 +39,7 @@ namespace { bool FLAGS_check_device_leaks = false; } // namespace -namespace perftools { -namespace gputools { +namespace stream_executor { namespace { string StackTraceIfVLOG10() { @@ -788,5 +787,4 @@ internal::StreamExecutorInterface *StreamExecutor::implementation() { return implementation_->GetUnderlyingExecutor(); } -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/stream_executor_pimpl.h b/tensorflow/stream_executor/stream_executor_pimpl.h index 69d0374d73d116ed8e19ace2329f6edb360bbea5..39af7115d8f697428fe7134ad0704a4ebcf033fe 100644 --- a/tensorflow/stream_executor/stream_executor_pimpl.h +++ b/tensorflow/stream_executor/stream_executor_pimpl.h @@ -37,8 +37,7 @@ limitations under the License. #include "tensorflow/stream_executor/stream_executor_internal.h" #include "tensorflow/stream_executor/trace_listener.h" -namespace perftools { -namespace gputools { +namespace stream_executor { // Structure used for device memory leak checking. struct AllocRecord { @@ -95,7 +94,7 @@ class StreamExecutor { // Parameters: // spec: The MultiKernelLoaderSpec is usually generated as a compile-time // constant into an appropriate namespace. For example, see - // perftools::gputools::executor_sample::kKernelLoaderSpecs, from which a + // stream_executor::executor_sample::kKernelLoaderSpecs, from which a // MultiKernelLoaderSpec is selected. // kernel: Outparam that the kernel is loaded into. A given Kernel // instantiation should not be loaded into more than once. @@ -803,7 +802,6 @@ inline Stream &Stream::ThenLaunch(ThreadDim thread_dims, BlockDim block_dims, return *this; } -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_STREAM_EXECUTOR_PIMPL_H_ diff --git a/tensorflow/stream_executor/temporary_device_memory.cc b/tensorflow/stream_executor/temporary_device_memory.cc index c33166b2246b51ac045349d9e2c0794a3f3868a1..f113ce9be57b9bd954df0273fa6f9af1dc7d9546 100644 --- a/tensorflow/stream_executor/temporary_device_memory.cc +++ b/tensorflow/stream_executor/temporary_device_memory.cc @@ -17,8 +17,7 @@ limitations under the License. #include "tensorflow/stream_executor/stream.h" -namespace perftools { -namespace gputools { +namespace stream_executor { TemporaryDeviceMemoryBase::~TemporaryDeviceMemoryBase() { parent_->temporary_memory_manager()->MarkFinalized(device_memory_, @@ -64,5 +63,4 @@ TemporaryDeviceMemoryBase::TemporaryDeviceMemoryBase( DCHECK(IsAllocated()); } -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/temporary_device_memory.h b/tensorflow/stream_executor/temporary_device_memory.h index 2255e7ffd716d8f6d6aad2c447e05c6dfca766df..77be8599a2de59f9a8287337210ad0c6bf76ece7 100644 --- a/tensorflow/stream_executor/temporary_device_memory.h +++ b/tensorflow/stream_executor/temporary_device_memory.h @@ -43,8 +43,7 @@ limitations under the License. #include "tensorflow/stream_executor/device_memory.h" -namespace perftools { -namespace gputools { +namespace stream_executor { class Stream; namespace internal { @@ -132,7 +131,6 @@ class TemporaryDeviceMemory : public TemporaryDeviceMemoryBase { } }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_TEMPORARY_DEVICE_MEMORY_H_ diff --git a/tensorflow/stream_executor/temporary_memory_manager.cc b/tensorflow/stream_executor/temporary_memory_manager.cc index 449ab7d3f0bff5dff0c55b716753ca0befd4062f..420dbb0933db3adf69a425f6223250e66f960261 100644 --- a/tensorflow/stream_executor/temporary_memory_manager.cc +++ b/tensorflow/stream_executor/temporary_memory_manager.cc @@ -21,8 +21,7 @@ limitations under the License. #include "tensorflow/stream_executor/stream.h" #include "tensorflow/stream_executor/stream_executor_pimpl.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace internal { void TemporaryMemoryManager::ForceDeallocateAll() { @@ -124,5 +123,4 @@ TemporaryMemoryManager::AllocateArrayBase(uint64 element_count, } } // namespace internal -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/temporary_memory_manager.h b/tensorflow/stream_executor/temporary_memory_manager.h index 2e6fbd9d62a041a5b4f2d2c8111d30ba0da69ce5..faf13380dc20afa9675a1b8d51aad3906d9c7bbe 100644 --- a/tensorflow/stream_executor/temporary_memory_manager.h +++ b/tensorflow/stream_executor/temporary_memory_manager.h @@ -31,8 +31,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/thread_annotations.h" #include "tensorflow/stream_executor/temporary_device_memory.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace internal { // Record used inside the TemporaryMemoryManager as metadata for a given device @@ -147,7 +146,6 @@ TemporaryMemoryManager::AllocateArray(uint64 element_count) { } } // namespace internal -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_TEMPORARY_MEMORY_MANAGER_H_ diff --git a/tensorflow/stream_executor/timer.cc b/tensorflow/stream_executor/timer.cc index 41d7e4359d4145b34d094a366824501d5a606940..a29791a1049d57cbd85604e4bfce154b57155009 100644 --- a/tensorflow/stream_executor/timer.cc +++ b/tensorflow/stream_executor/timer.cc @@ -21,8 +21,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/logging.h" #include "tensorflow/stream_executor/stream_executor.h" -namespace perftools { -namespace gputools { +namespace stream_executor { Timer::Timer(StreamExecutor *parent) : parent_(parent), @@ -34,5 +33,4 @@ uint64 Timer::Microseconds() const { return implementation_->Microseconds(); } uint64 Timer::Nanoseconds() const { return implementation_->Nanoseconds(); } -} // namespace gputools -} // namespace perftools +} // namespace stream_executor diff --git a/tensorflow/stream_executor/timer.h b/tensorflow/stream_executor/timer.h index 0a37caa0f2fa0b5f489d1e495368c807df70f51f..fba7dd8f589ad49685bdf5fd8faa1ee18fb55113 100644 --- a/tensorflow/stream_executor/timer.h +++ b/tensorflow/stream_executor/timer.h @@ -20,8 +20,7 @@ limitations under the License. #include "tensorflow/stream_executor/platform/port.h" -namespace perftools { -namespace gputools { +namespace stream_executor { namespace internal { class TimerInterface; @@ -69,7 +68,6 @@ class Timer { SE_DISALLOW_COPY_AND_ASSIGN(Timer); }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_TIMER_H_ diff --git a/tensorflow/stream_executor/trace_listener.h b/tensorflow/stream_executor/trace_listener.h index d1e87c348b1f867009fdb6b741d984b2f58cef21..0e874a1d47b4daa12b501f87ad1c84afe84b0b6b 100644 --- a/tensorflow/stream_executor/trace_listener.h +++ b/tensorflow/stream_executor/trace_listener.h @@ -23,8 +23,7 @@ limitations under the License. #include "tensorflow/stream_executor/launch_dim.h" #include "tensorflow/stream_executor/lib/status.h" -namespace perftools { -namespace gputools { +namespace stream_executor { class Stream; @@ -69,7 +68,6 @@ class TraceListener { const port::Status* result) {} }; -} // namespace gputools -} // namespace perftools +} // namespace stream_executor #endif // TENSORFLOW_STREAM_EXECUTOR_TRACE_LISTENER_H_ diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl index b286834ded3b32dbc7804c32befd3bd8d4bbe660..51e856bed0e995fc6e1e367401c93b5f3648c361 100644 --- a/tensorflow/tensorflow.bzl +++ b/tensorflow/tensorflow.bzl @@ -163,7 +163,6 @@ def if_override_eigen_strong_inline(a): def get_win_copts(is_external=False): WINDOWS_COPTS = [ - "/D__VERSION__=\\\"MSVC\\\"", "/DPLATFORM_WINDOWS", "/DEIGEN_HAS_C99_MATH", "/DTENSORFLOW_USE_EIGEN_THREADPOOL", diff --git a/tensorflow/tools/api/generator/create_python_api.py b/tensorflow/tools/api/generator/create_python_api.py index c7748f5b7a7c1013f11e053e7a36ddfd9594c6ea..c06a39bfbdf06cb478007a884f38728755845f19 100644 --- a/tensorflow/tools/api/generator/create_python_api.py +++ b/tensorflow/tools/api/generator/create_python_api.py @@ -160,7 +160,7 @@ def get_api_init_text(): # we want to traverse over TensorFlow Python modules. for module in sys.modules.values(): # Only look at tensorflow modules. - if (not module or not hasattr(module, "__name__") or + if (not module or not hasattr(module, '__name__') or 'tensorflow.' not in module.__name__): continue # Do not generate __init__.py files for contrib modules for now. diff --git a/tensorflow/tools/api/golden/tensorflow.data.-dataset.pbtxt b/tensorflow/tools/api/golden/tensorflow.data.-dataset.pbtxt index 0900adaf762df1415c8db63c3879ca2fabc28d9f..cbbd077c97ba7f304e13f7883254f5a12668255e 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\', \'shuffle\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'file_pattern\', \'shuffle\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\', \'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 7b16ac90c925beb25e065d26e73ee2a54b06d9dc..9a56ae8675c4fdae7424859fb5105574d83da494 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\', \'shuffle\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'file_pattern\', \'shuffle\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\', \'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 9cf5f2ae2057ab4a16131527cf2ef2fa6ada28e5..e5ec824bb8959f77f99ce3527f168bacf8de200a 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 @@ -65,7 +65,7 @@ tf_class { } member_method { name: "list_files" - argspec: "args=[\'file_pattern\', \'shuffle\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'file_pattern\', \'shuffle\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\', \'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 8c3d6691439e619c906996a3ddaea4317c4a9597..008239789c788692a36445dc2c51e26113b6a0c5 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\', \'shuffle\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'file_pattern\', \'shuffle\', \'seed\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " } member_method { name: "make_initializable_iterator" diff --git a/tensorflow/tools/api/golden/tensorflow.distributions.bijectors.-bijector.pbtxt b/tensorflow/tools/api/golden/tensorflow.distributions.bijectors.-bijector.pbtxt deleted file mode 100644 index 11565bd3e4178202fa82e2e079d1035190dbd6ec..0000000000000000000000000000000000000000 --- a/tensorflow/tools/api/golden/tensorflow.distributions.bijectors.-bijector.pbtxt +++ /dev/null @@ -1,65 +0,0 @@ -path: "tensorflow.distributions.bijectors.Bijector" -tf_class { - is_instance: "" - is_instance: "" - member { - name: "dtype" - mtype: "" - } - member { - name: "event_ndims" - mtype: "" - } - member { - name: "graph_parents" - mtype: "" - } - member { - name: "is_constant_jacobian" - mtype: "" - } - member { - name: "name" - mtype: "" - } - member { - name: "validate_args" - mtype: "" - } - member_method { - name: "__init__" - argspec: "args=[\'self\', \'event_ndims\', \'graph_parents\', \'is_constant_jacobian\', \'validate_args\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'False\', \'False\', \'None\', \'None\'], " - } - member_method { - name: "forward" - argspec: "args=[\'self\', \'x\', \'name\'], varargs=None, keywords=None, defaults=[\'forward\'], " - } - member_method { - name: "forward_event_shape" - argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "forward_event_shape_tensor" - argspec: "args=[\'self\', \'input_shape\', \'name\'], varargs=None, keywords=None, defaults=[\'forward_event_shape_tensor\'], " - } - member_method { - name: "forward_log_det_jacobian" - argspec: "args=[\'self\', \'x\', \'name\'], varargs=None, keywords=None, defaults=[\'forward_log_det_jacobian\'], " - } - member_method { - name: "inverse" - argspec: "args=[\'self\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'inverse\'], " - } - member_method { - name: "inverse_event_shape" - argspec: "args=[\'self\', \'output_shape\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "inverse_event_shape_tensor" - argspec: "args=[\'self\', \'output_shape\', \'name\'], varargs=None, keywords=None, defaults=[\'inverse_event_shape_tensor\'], " - } - member_method { - name: "inverse_log_det_jacobian" - argspec: "args=[\'self\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'inverse_log_det_jacobian\'], " - } -} diff --git a/tensorflow/tools/api/golden/tensorflow.distributions.bijectors.-identity.pbtxt b/tensorflow/tools/api/golden/tensorflow.distributions.bijectors.-identity.pbtxt deleted file mode 100644 index 1e5fe624eb838e188594d03b656c12890db344a1..0000000000000000000000000000000000000000 --- a/tensorflow/tools/api/golden/tensorflow.distributions.bijectors.-identity.pbtxt +++ /dev/null @@ -1,66 +0,0 @@ -path: "tensorflow.distributions.bijectors.Identity" -tf_class { - is_instance: "" - is_instance: "" - is_instance: "" - member { - name: "dtype" - mtype: "" - } - member { - name: "event_ndims" - mtype: "" - } - member { - name: "graph_parents" - mtype: "" - } - member { - name: "is_constant_jacobian" - mtype: "" - } - member { - name: "name" - mtype: "" - } - member { - name: "validate_args" - mtype: "" - } - member_method { - name: "__init__" - argspec: "args=[\'self\', \'validate_args\', \'event_ndims\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'0\', \'identity\'], " - } - member_method { - name: "forward" - argspec: "args=[\'self\', \'x\', \'name\'], varargs=None, keywords=None, defaults=[\'forward\'], " - } - member_method { - name: "forward_event_shape" - argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "forward_event_shape_tensor" - argspec: "args=[\'self\', \'input_shape\', \'name\'], varargs=None, keywords=None, defaults=[\'forward_event_shape_tensor\'], " - } - member_method { - name: "forward_log_det_jacobian" - argspec: "args=[\'self\', \'x\', \'name\'], varargs=None, keywords=None, defaults=[\'forward_log_det_jacobian\'], " - } - member_method { - name: "inverse" - argspec: "args=[\'self\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'inverse\'], " - } - member_method { - name: "inverse_event_shape" - argspec: "args=[\'self\', \'output_shape\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "inverse_event_shape_tensor" - argspec: "args=[\'self\', \'output_shape\', \'name\'], varargs=None, keywords=None, defaults=[\'inverse_event_shape_tensor\'], " - } - member_method { - name: "inverse_log_det_jacobian" - argspec: "args=[\'self\', \'y\', \'name\'], varargs=None, keywords=None, defaults=[\'inverse_log_det_jacobian\'], " - } -} diff --git a/tensorflow/tools/api/golden/tensorflow.distributions.bijectors.pbtxt b/tensorflow/tools/api/golden/tensorflow.distributions.bijectors.pbtxt deleted file mode 100644 index 1d0144f36ec332740889dc8caa5add8f41960d92..0000000000000000000000000000000000000000 --- a/tensorflow/tools/api/golden/tensorflow.distributions.bijectors.pbtxt +++ /dev/null @@ -1,11 +0,0 @@ -path: "tensorflow.distributions.bijectors" -tf_module { - member { - name: "Bijector" - mtype: "" - } - member { - name: "Identity" - mtype: "" - } -} diff --git a/tensorflow/tools/api/golden/tensorflow.distributions.pbtxt b/tensorflow/tools/api/golden/tensorflow.distributions.pbtxt index 2fba7c506ed9d2490e7c19c1746d3f4e9645424f..90b60ef074dd2eaf911291e6c725b98e2891e728 100644 --- a/tensorflow/tools/api/golden/tensorflow.distributions.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.distributions.pbtxt @@ -68,10 +68,6 @@ tf_module { name: "Uniform" mtype: "" } - member { - name: "bijectors" - mtype: "" - } member_method { name: "kl_divergence" argspec: "args=[\'distribution_a\', \'distribution_b\', \'allow_nan_stats\', \'name\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], " diff --git a/tensorflow/tools/api/golden/tensorflow.estimator.-run-config.pbtxt b/tensorflow/tools/api/golden/tensorflow.estimator.-run-config.pbtxt index 05e603efb7cbad8c4c42a7a15074d2634af8d21c..c8da55d8021b7659446d0771a089b7b605d86c4f 100644 --- a/tensorflow/tools/api/golden/tensorflow.estimator.-run-config.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.estimator.-run-config.pbtxt @@ -6,6 +6,10 @@ tf_class { name: "cluster_spec" mtype: "" } + member { + name: "device_fn" + mtype: "" + } member { name: "evaluation_master" mtype: "" @@ -84,7 +88,7 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'model_dir\', \'tf_random_seed\', \'save_summary_steps\', \'save_checkpoints_steps\', \'save_checkpoints_secs\', \'session_config\', \'keep_checkpoint_max\', \'keep_checkpoint_every_n_hours\', \'log_step_count_steps\', \'train_distribute\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'100\', \'\', \'\', \'None\', \'5\', \'10000\', \'100\', \'None\'], " + argspec: "args=[\'self\', \'model_dir\', \'tf_random_seed\', \'save_summary_steps\', \'save_checkpoints_steps\', \'save_checkpoints_secs\', \'session_config\', \'keep_checkpoint_max\', \'keep_checkpoint_every_n_hours\', \'log_step_count_steps\', \'train_distribute\', \'device_fn\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'100\', \'\', \'\', \'None\', \'5\', \'10000\', \'100\', \'None\', \'None\'], " } member_method { name: "replace" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.-model.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.-model.pbtxt index 7be2f4f61f6b9637f372591e49efc0c93c7a8c0a..cdf2da712f38170c4b82a682d4630036bbf93925 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.-model.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.-model.pbtxt @@ -3,7 +3,6 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -74,10 +69,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "state_updates" mtype: "" @@ -128,7 +119,7 @@ tf_class { } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -260,7 +251,7 @@ tf_class { } member_method { name: "test_on_batch" - argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " } member_method { name: "to_json" @@ -272,6 +263,6 @@ tf_class { } member_method { name: "train_on_batch" - argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " + argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " } } diff --git a/tensorflow/tools/api/golden/tensorflow.keras.-sequential.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.-sequential.pbtxt index bf361cf8054571c0b056e1373acb838aaea87173..5c2c29e60fe11aa29cf2a27cac055bb7aea49dde 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.-sequential.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.-sequential.pbtxt @@ -4,7 +4,6 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +14,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -75,10 +70,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "state_updates" mtype: "" @@ -133,7 +124,7 @@ tf_class { } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -277,7 +268,7 @@ tf_class { } member_method { name: "test_on_batch" - argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " } member_method { name: "to_json" @@ -289,6 +280,6 @@ tf_class { } member_method { name: "train_on_batch" - argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " + argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " } } diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-activation.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-activation.pbtxt index db8f626b98b70fd99f38e696aa16c72e74e86e25..96272d1b7d61430188bbbf2680bd2beb9f1e9675 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-activation.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-activation.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.Activation" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 809b3a5430449176a0d7423ec7f4499ceb620890..8fd55c8686de77ec764e9d564c78d0df4f545915 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-activity-regularization.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-activity-regularization.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.ActivityRegularization" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-add.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-add.pbtxt index 68d41bb6cc258ca87d4664ac0fb9d5649f89ebaf..47d1532c3c8cf248f6e9a9a35e10b9559286263f 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-add.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-add.pbtxt @@ -3,7 +3,6 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 970b777e514194db4ac49fe58bea737b35436217..797d422a90a5ad21d0014a0003b11d281c25e579 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-alpha-dropout.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-alpha-dropout.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.AlphaDropout" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 529c64ab293d596012aefd42e0695bd1eb7e44d1..269be1455b6bf3bbe325f3928584960578e3793d 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 @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.AveragePooling1D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -93,7 +83,7 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'pool_size\', \'strides\', \'padding\'], varargs=None, keywords=kwargs, defaults=[\'2\', \'None\', \'valid\'], " + argspec: "args=[\'self\', \'pool_size\', \'strides\', \'padding\', \'data_format\'], varargs=None, keywords=kwargs, defaults=[\'2\', \'None\', \'valid\', \'None\'], " } member_method { name: "add_loss" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 7e7c330d74fe3b71ecd0eb87e34719e47ae70784..344813621534dbb5de3719088c06313e55519dd7 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 @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.AveragePooling2D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 ada8466d7473072b1878861ab36ec40b07fa1914..979008d0edb7d0f9d9c1246e1fcc7d2e2871d28c 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 @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.AveragePooling3D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-average.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-average.pbtxt index 2a5c1cd530a7a532f6cdd3c184f4ee7eb88d23d3..0ffdffd4cdee14fbfeb68f2575300632ab21d7a6 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-average.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-average.pbtxt @@ -3,7 +3,6 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 9a2cb29815d59f3761ea25e9ea36ff6489c85b88..6b00f110eea2aaa551bffe8ec225042c5469210e 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 @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.AvgPool1D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -93,7 +83,7 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'pool_size\', \'strides\', \'padding\'], varargs=None, keywords=kwargs, defaults=[\'2\', \'None\', \'valid\'], " + argspec: "args=[\'self\', \'pool_size\', \'strides\', \'padding\', \'data_format\'], varargs=None, keywords=kwargs, defaults=[\'2\', \'None\', \'valid\', \'None\'], " } member_method { name: "add_loss" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 f5e991ea42e5ee2723b64574d4598dc8463f1c8c..caff5a2f1db61c6958980446d3bd54009776e1a4 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 @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.AvgPool2D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 31732214a62524017e39776cdfb9ab629746e8ae..4a7239492177aae2ee098fc3033904d3d1a31ba4 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 @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.AvgPool3D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 422eddf10db6763e10405dba5537ca161d1b8994..9804394fa53d6b3c1ff136a73212863143bbdb39 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-batch-normalization.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-batch-normalization.pbtxt @@ -1,9 +1,7 @@ path: "tensorflow.keras.layers.BatchNormalization" tf_class { is_instance: "" - is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -92,7 +82,7 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'axis\', \'momentum\', \'epsilon\', \'center\', \'scale\', \'beta_initializer\', \'gamma_initializer\', \'moving_mean_initializer\', \'moving_variance_initializer\', \'beta_regularizer\', \'gamma_regularizer\', \'beta_constraint\', \'gamma_constraint\'], varargs=None, keywords=kwargs, defaults=[\'-1\', \'0.99\', \'0.001\', \'True\', \'True\', \'zeros\', \'ones\', \'zeros\', \'ones\', \'None\', \'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'axis\', \'momentum\', \'epsilon\', \'center\', \'scale\', \'beta_initializer\', \'gamma_initializer\', \'moving_mean_initializer\', \'moving_variance_initializer\', \'beta_regularizer\', \'gamma_regularizer\', \'beta_constraint\', \'gamma_constraint\', \'renorm\', \'renorm_clipping\', \'renorm_momentum\', \'fused\', \'trainable\', \'virtual_batch_size\', \'adjustment\', \'name\'], varargs=None, keywords=kwargs, defaults=[\'-1\', \'0.99\', \'0.001\', \'True\', \'True\', \'zeros\', \'ones\', \'zeros\', \'ones\', \'None\', \'None\', \'None\', \'None\', \'False\', \'None\', \'0.99\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "add_loss" @@ -104,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-bidirectional.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-bidirectional.pbtxt index 9053a37916314198842bc21b0608a9b69a64c264..5e5b04c7c695c6d88e7f42b77290a582be087763 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-bidirectional.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-bidirectional.pbtxt @@ -3,7 +3,6 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -18,10 +17,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -70,10 +65,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable" mtype: "" @@ -112,11 +103,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-concatenate.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-concatenate.pbtxt index 3d536d2182fc4480a2ee5fba177543ca21fbd5ac..b8eb4079b9eea3b054f9c2ad4298f6a1669de79e 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-concatenate.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-concatenate.pbtxt @@ -3,7 +3,6 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 a535f1817021f07c29e19fca4b3a7e7e60f59fb1..3fdb101425d0f010f77cd70bc3721af269a36b0a 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 @@ -4,7 +4,6 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -47,10 +46,6 @@ tf_class { name: "filters" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -139,10 +134,6 @@ tf_class { name: "recurrent_regularizer" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "states" mtype: "" @@ -193,11 +184,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 801a0339720919f8b3f6beee0f045d58b2c0a371..0be42471e35eeba224376b24aa846db69e011274 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv1-d.pbtxt @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.Conv1D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -93,7 +83,7 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'filters\', \'kernel_size\', \'strides\', \'padding\', \'dilation_rate\', \'activation\', \'use_bias\', \'kernel_initializer\', \'bias_initializer\', \'kernel_regularizer\', \'bias_regularizer\', \'activity_regularizer\', \'kernel_constraint\', \'bias_constraint\'], varargs=None, keywords=kwargs, defaults=[\'1\', \'valid\', \'1\', \'None\', \'True\', \'glorot_uniform\', \'zeros\', \'None\', \'None\', \'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'filters\', \'kernel_size\', \'strides\', \'padding\', \'data_format\', \'dilation_rate\', \'activation\', \'use_bias\', \'kernel_initializer\', \'bias_initializer\', \'kernel_regularizer\', \'bias_regularizer\', \'activity_regularizer\', \'kernel_constraint\', \'bias_constraint\'], varargs=None, keywords=kwargs, defaults=[\'1\', \'valid\', \'channels_last\', \'1\', \'None\', \'True\', \'glorot_uniform\', \'zeros\', \'None\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "add_loss" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 13352e264a5305190717bb973a3f2bce4d7f4fff..39ba31a70942811cbd36caf33c9bda90a5449703 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 @@ -1,11 +1,9 @@ path: "tensorflow.keras.layers.Conv2DTranspose" tf_class { is_instance: "" - is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -16,10 +14,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -68,10 +62,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -106,11 +96,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 f400e4a15c362037e85ac375cee98bb5f6358669..26d9d8c476f4e429ffe112cd490ad277b478c65b 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv2-d.pbtxt @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.Conv2D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 b3a9f573b8ba652d2544b21f36f65fe81a6ebb50..43611017fa37077c7ff05e690dfd50a0e6e5ae1e 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 @@ -1,11 +1,9 @@ path: "tensorflow.keras.layers.Conv3DTranspose" tf_class { is_instance: "" - is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -16,10 +14,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -68,10 +62,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -106,11 +96,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 a9be09c0abd19aeb4df30116ef2befc3948bfbf4..fa4925ab99d719104c5ca1a0003c25d85f78f3df 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-conv3-d.pbtxt @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.Conv3D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 be1ef5eb928d16cc6bf78c289aa20d815c728b23..c5c5d5e7c083dca63972cb7033e675aa483d8039 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution1-d.pbtxt @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.Convolution1D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -93,7 +83,7 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'filters\', \'kernel_size\', \'strides\', \'padding\', \'dilation_rate\', \'activation\', \'use_bias\', \'kernel_initializer\', \'bias_initializer\', \'kernel_regularizer\', \'bias_regularizer\', \'activity_regularizer\', \'kernel_constraint\', \'bias_constraint\'], varargs=None, keywords=kwargs, defaults=[\'1\', \'valid\', \'1\', \'None\', \'True\', \'glorot_uniform\', \'zeros\', \'None\', \'None\', \'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'filters\', \'kernel_size\', \'strides\', \'padding\', \'data_format\', \'dilation_rate\', \'activation\', \'use_bias\', \'kernel_initializer\', \'bias_initializer\', \'kernel_regularizer\', \'bias_regularizer\', \'activity_regularizer\', \'kernel_constraint\', \'bias_constraint\'], varargs=None, keywords=kwargs, defaults=[\'1\', \'valid\', \'channels_last\', \'1\', \'None\', \'True\', \'glorot_uniform\', \'zeros\', \'None\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "add_loss" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 30034f7eaf6d9073695353e5c8d9ead0cc8de7cc..36dc2d2e9a70fe7a1d32352a14dabbade2a2efc2 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 @@ -1,11 +1,9 @@ path: "tensorflow.keras.layers.Convolution2DTranspose" tf_class { is_instance: "" - is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -16,10 +14,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -68,10 +62,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -106,11 +96,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 189b38054c004facfeeff8ad2ae87848b89040f2..23ec74370bab010d2b5ba257502a32c8ca7e4a57 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution2-d.pbtxt @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.Convolution2D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 a76d85c629c1fe620dafd62a0f0e05e9009109e2..0e4089c5785ccdcafbb8b3dc1ca75ffbe49d9434 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 @@ -1,11 +1,9 @@ path: "tensorflow.keras.layers.Convolution3DTranspose" tf_class { is_instance: "" - is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -16,10 +14,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -68,10 +62,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -106,11 +96,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 782195d4ad5883d8c0ea6a657cc10258f2080a55..23ddbe1a925e33432727d13dc875972136083056 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-convolution3-d.pbtxt @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.Convolution3D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 2cb7a39ea595e1ff699b96554cb135377d20a488..e04ab6bea85baa1252c9c43f891f9ed5a9dedb87 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping1-d.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.Cropping1D" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 80803306992bba3b601824a93cb3086ef3947369..655314afffd1e4fb0987c913eb69edc4254c77ef 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping2-d.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.Cropping2D" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 678f40bbc23db15ff7c1138169478fb4412a449d..d5215f1330a2fd0db4696a7b743d931e164a5a9f 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-cropping3-d.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.Cropping3D" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-dense.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-dense.pbtxt index fac826109b6a32305ece86c4990f08afe2236ce8..310a3c3b918684803c86d8e2ea1731604a041cdb 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-dense.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-dense.pbtxt @@ -1,9 +1,7 @@ path: "tensorflow.keras.layers.Dense" tf_class { is_instance: "" - is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-depthwise-conv2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-depthwise-conv2-d.pbtxt index b38716aa2cd9cedd94dbfe1468c3988e419db714..2d67b5f720209e3648e69659d44c3d8c4e639231 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-depthwise-conv2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-depthwise-conv2-d.pbtxt @@ -2,10 +2,8 @@ path: "tensorflow.keras.layers.DepthwiseConv2D" tf_class { is_instance: "" is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -16,10 +14,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -68,10 +62,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -106,11 +96,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-dot.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-dot.pbtxt index 285d544af2d69d564afdec748598b39b6b95670f..0e493a7f2bffc772e9cc9cd5afdc4c092fb92118 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-dot.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-dot.pbtxt @@ -3,7 +3,6 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-dropout.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-dropout.pbtxt index b77976974cccb96fc2373c093d2bdf279560c46f..14726b4b6cecc39da98af4220211a1c0351b2ff1 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-dropout.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-dropout.pbtxt @@ -1,9 +1,7 @@ path: "tensorflow.keras.layers.Dropout" tf_class { is_instance: "" - is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 b07714d3f2d158496e0482f8611e55ea0fb0fd51..32a50455ed8a2a146a81fde0d884ebc867b8d0ed 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.ELU" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-embedding.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-embedding.pbtxt index e67d4ddfc47077d62319ab097e5333a373cbfc80..2f615d81124a7ef5e1bd7181a10abfb1b7a8df24 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-embedding.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-embedding.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.Embedding" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-flatten.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-flatten.pbtxt index b2a668e5a88d312656f48ddd0e9f7aa9f6306991..82dc878a8c7f7f011df4dd3fa0445217fa250a98 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-flatten.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-flatten.pbtxt @@ -1,9 +1,7 @@ path: "tensorflow.keras.layers.Flatten" tf_class { is_instance: "" - is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 4274b8d42552c4421acaee5e18393ea6093acb7d..d79d02b95433a7399f27ff1354cba315f8a2c3ae 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.GRUCell" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 8d9f06083cfc271ebacbe714bca0037a94753fd7..1d38ae64bb86d11ecd352371608a11f6736bb0ac 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 @@ -3,7 +3,6 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -34,10 +33,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "implementation" mtype: "" @@ -126,10 +121,6 @@ tf_class { name: "reset_after" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "states" mtype: "" @@ -176,11 +167,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 f4f1a5d51c5d5689918af4facf907f79d9ca71ec..135de9cd95141a93d1a45b80ada66e339c484c89 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-gaussian-dropout.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-gaussian-dropout.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.GaussianDropout" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 e502df5e177d422403d0643c18a9588afb9d9713..5db6e433ee02b8050822f76bb762329055c11aa4 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-gaussian-noise.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-gaussian-noise.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.GaussianNoise" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 9c8d5bfcd8966384230e7d5cdcc1cac53a0eab9a..bf0dba0a925b4fb3a88173f4a6dfcc565c6edd91 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 @@ -1,9 +1,8 @@ path: "tensorflow.keras.layers.GlobalAveragePooling1D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 8dd65f1f248daaf120780f19050c45d297b7902e..6da98036094228603f12f65f66d958e2e4b9daeb 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 @@ -1,9 +1,8 @@ path: "tensorflow.keras.layers.GlobalAveragePooling2D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 5e30571cc730ee23767a044036b590460deec00b..345593dec812a240251c0c07da759e131fefaad9 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 @@ -1,9 +1,8 @@ path: "tensorflow.keras.layers.GlobalAveragePooling3D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 ba90fa454696d1cb4e77d80a2dc77ff65def4714..5d3be9085e51520c48a8d33b1dde6480a0039bad 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 @@ -1,9 +1,8 @@ path: "tensorflow.keras.layers.GlobalAvgPool1D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 8823857758307c208527b144c0cc73b566f2f115..0b79a87e0507ecc795f14a63684dd5b5d7dafc1e 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 @@ -1,9 +1,8 @@ path: "tensorflow.keras.layers.GlobalAvgPool2D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 500ced852ba6b19502769ba9052f2e364af7e283..68cdbac652f74b72d1cb769fbefbee750025767d 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 @@ -1,9 +1,8 @@ path: "tensorflow.keras.layers.GlobalAvgPool3D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 cf2717ed46b56e639fb774c1e922648e1653ec0d..d5872b444fa9ae617e6ad55bc39f43ec4be7d92f 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 @@ -1,9 +1,8 @@ path: "tensorflow.keras.layers.GlobalMaxPool1D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 a86ff1a46997f19b11e6ef03be432b45687a2df2..4b0cf9a5d38f868d9fdd16a042daea6d0f56fdf8 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 @@ -1,9 +1,8 @@ path: "tensorflow.keras.layers.GlobalMaxPool2D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 e01cc7c1b09ad6a40380613d54b771c6a1c89c1c..4c1adb2131f204121cf74c9a77d346902632fef2 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 @@ -1,9 +1,8 @@ path: "tensorflow.keras.layers.GlobalMaxPool3D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 259c1fb37c787f5318570b7aca6935d2f0ed997f..815f1cf580562e62db99862e51ba02e2b2051b57 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 @@ -1,9 +1,8 @@ path: "tensorflow.keras.layers.GlobalMaxPooling1D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 0c41bf97f763f1e40e8fac714709ccac1483a00b..e027dd6cc282b75a6a6c2a78878fa57d6706b547 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 @@ -1,9 +1,8 @@ path: "tensorflow.keras.layers.GlobalMaxPooling2D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 bec8817aa393ba2d8a6410408938402366cbb01d..c647b24a23258b96a3c1780fd2f3e499658cfe7d 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 @@ -1,9 +1,8 @@ path: "tensorflow.keras.layers.GlobalMaxPooling3D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 17be86222901c0f5a9a18c0e5f1c5bcac6c06a17..75d70734b4144627f02a8f619991356e38889389 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-input-layer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-input-layer.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.InputLayer" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-input-spec.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-input-spec.pbtxt index 3aeef347ae1f96a3ef40493cc6b722a887e81786..29edabe0483a21d7db35eec04d6ae7a855a82da1 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-input-spec.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-input-spec.pbtxt @@ -1,6 +1,6 @@ path: "tensorflow.keras.layers.InputSpec" tf_class { - is_instance: "" + is_instance: "" is_instance: "" member_method { name: "__init__" 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 6d2a8c56196d9b3c80f570c7f1d3ac803253fff6..0ed383a3554f81c3db490cf5d242546a14b64d15 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.LSTMCell" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 490b5b618c65e28f1ae2e01e8d35e7f3973cc180..6d14c9c8f69286552271ca6dab5271a5af48593f 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 @@ -3,7 +3,6 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -34,10 +33,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "implementation" mtype: "" @@ -122,10 +117,6 @@ tf_class { name: "recurrent_regularizer" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "states" mtype: "" @@ -176,11 +167,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-lambda.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-lambda.pbtxt index 21a65b838af35e2f540eacab823513e7bf54b434..ddf96aba34bf574f4b9046ed932d8f136864f157 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-lambda.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-lambda.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.Lambda" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-layer.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-layer.pbtxt index 127b04738e70c11b2dc1071cf174cf5de23c5133..aca282d62427f5c8186fa3ac86daebd6fba09ce5 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-layer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-layer.pbtxt @@ -1,7 +1,6 @@ path: "tensorflow.keras.layers.Layer" tf_class { is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -12,10 +11,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -64,10 +59,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -90,7 +81,7 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None" + argspec: "args=[\'self\', \'trainable\', \'name\', \'dtype\'], varargs=None, keywords=kwargs, defaults=[\'True\', \'None\', \'None\'], " } member_method { name: "add_loss" @@ -102,11 +93,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 87e49f2ed5b5d73aee5e9aa2511485b1f3f4bcd9..b9c53b43c87bd8c1047663001159d7286360a008 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.LeakyReLU" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 1aa3aad3246b83931a47e69a4aa76fdf2b5aee22..2ee566d03b4631ebcff6bcb0b93ec274849c5b67 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.LocallyConnected1D" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 5e9dc7d4774c651a186a4e320d0cfd088e87b6b3..db0d0e816a6d863f13c7eb085edf269d71e2f252 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.LocallyConnected2D" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-masking.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-masking.pbtxt index 0d101e5b68cdb2cdf24ed472c724cfc885e3d95d..82008b89d038e26155b8ae952c2943557ed8c35e 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-masking.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-masking.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.Masking" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 c85cd49ac8ce2c1fc0759671865b7174cd1c1480..31a34a17d04129c8dfe8ec6f98b9fdbc110e13f4 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 @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.MaxPool1D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -93,7 +83,7 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'pool_size\', \'strides\', \'padding\'], varargs=None, keywords=kwargs, defaults=[\'2\', \'None\', \'valid\'], " + argspec: "args=[\'self\', \'pool_size\', \'strides\', \'padding\', \'data_format\'], varargs=None, keywords=kwargs, defaults=[\'2\', \'None\', \'valid\', \'None\'], " } member_method { name: "add_loss" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 4f59e330c92f96101c65a9a24f66196e84587ccb..70d24ac75c4f850374fa8cffb881652f97e97d22 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 @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.MaxPool2D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 c0ea0eb0505d20e70d641f2a646a060d7dbfabda..55b16564b30919fb48d97862e2e8cfc0fdda8de5 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 @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.MaxPool3D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 ca37ae51314516ae67c7725eb2ccd3d25154e2ac..a230b74c383eff93443d07e48f6818352302d9e7 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 @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.MaxPooling1D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -93,7 +83,7 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'pool_size\', \'strides\', \'padding\'], varargs=None, keywords=kwargs, defaults=[\'2\', \'None\', \'valid\'], " + argspec: "args=[\'self\', \'pool_size\', \'strides\', \'padding\', \'data_format\'], varargs=None, keywords=kwargs, defaults=[\'2\', \'None\', \'valid\', \'None\'], " } member_method { name: "add_loss" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 3ede2378347f5eddb0e8fae775a0200ea484d3f8..d98f7c39f546048b1483617b24f004d5c9759d14 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 @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.MaxPooling2D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 d87e25a7ba8e7cce615431723b53a0106c2b5279..b2e96a4203758270afb8c225a05a481dbc329a84 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 @@ -1,10 +1,8 @@ path: "tensorflow.keras.layers.MaxPooling3D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-maximum.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-maximum.pbtxt index e4df7b48ae6b41400375920a48ef8577bb69376e..0c45bbdf171ad9831e51d2b1ba952fab9eb3d0e5 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-maximum.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-maximum.pbtxt @@ -3,7 +3,6 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-multiply.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-multiply.pbtxt index 6bf7c77743c31b6d74df35d827e9d5bc9a25d303..6423d83418aa40b57afb3d5ff22f4ec605183587 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-multiply.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-multiply.pbtxt @@ -3,7 +3,6 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -104,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 c14be132b7e406c99841576be8d8fa9ab99aa816..6e17081375b7fbdcb000dfcbc0cd48ab072fe6e5 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.PReLU" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-permute.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-permute.pbtxt index 72ffbceae01da900778dba1ec14e646aa17b39e5..d01d371da596256873c3799b99c45db01c674ff9 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-permute.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-permute.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.Permute" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 d3e780c8b22ed580f61ffc3d9b2bad7278391402..d3f5508640e268df484af5adaf66621fa3d92d5b 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.RNN" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "states" mtype: "" @@ -107,11 +98,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 a27980a9d17397e558a4b732e3dc332a0c1e8432..44e1007f5420bbb8feda891901c138cb776c071e 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-repeat-vector.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-repeat-vector.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.RepeatVector" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-reshape.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-reshape.pbtxt index 67f991276c6908ff54fd516e84533542a5f60528..8fc3ec33310f531a6ca6948dc20db67543cde69a 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-reshape.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-reshape.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.Reshape" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 fccea5e8af5ab81e712669ff1b2567d8bde8607e..457d27749504167f2865bc272251bd89b5d3297e 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 @@ -1,11 +1,9 @@ path: "tensorflow.keras.layers.SeparableConv1D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -16,10 +14,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -68,10 +62,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -106,11 +96,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 d20663bdb0bc2eea323d35b1e3d4d27122f50472..54eda8ee2121d4fadd73a33010202850d743cc65 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 @@ -1,11 +1,9 @@ path: "tensorflow.keras.layers.SeparableConv2D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -16,10 +14,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -68,10 +62,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -106,11 +96,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 889fa0a1b58bbd3babd293b7b1b45915a9ee3ca4..711196554698b79c64e1c67bc875ce33c70563c6 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 @@ -1,11 +1,9 @@ path: "tensorflow.keras.layers.SeparableConvolution1D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -16,10 +14,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -68,10 +62,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -106,11 +96,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 c850f3fedc814b20f0f95cc3cf4fd5c973446b5b..815e34a48de542feb078d0002dbebbbf4d199e63 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 @@ -1,11 +1,9 @@ path: "tensorflow.keras.layers.SeparableConvolution2D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -16,10 +14,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -68,10 +62,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -106,11 +96,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 526d88ccba60eb25c68432e5baa03fd3a878f718..6614760e5e72556ac61c3788065b3faf2d286800 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.SimpleRNNCell" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 7fddae34472411f49d42b4d65d12034d056ec818..bfcfd71ecddfb618ceb53969e8782f535786009b 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 @@ -3,7 +3,6 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -34,10 +33,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -114,10 +109,6 @@ tf_class { name: "recurrent_regularizer" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "states" mtype: "" @@ -164,11 +155,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-softmax.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-softmax.pbtxt index 5b9b62fc970238e49e6d4849285606d0a7908b23..9c4618c4e91ae288284b35c4a0c6bbbfe604d91d 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-softmax.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-softmax.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.Softmax" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 769da30999993fad05ae0f7c04e256e6cf01a774..9a0a19d2d52f34e61c273bd4d4a27c46940dd5bf 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 @@ -2,9 +2,7 @@ path: "tensorflow.keras.layers.SpatialDropout1D" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 fca2e42a1519fcf3a9f0ec996c50b148b2df05fd..446f7122a6a2dc1d2f7377cef00d7b5b9a89cd3a 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 @@ -2,9 +2,7 @@ path: "tensorflow.keras.layers.SpatialDropout2D" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 36e8de09a967c5940bf8078234f5980a78ec8009..52a0485b5cea2be6f2f4d9b0ee31cb2388adcef7 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 @@ -2,9 +2,7 @@ path: "tensorflow.keras.layers.SpatialDropout3D" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -67,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -105,11 +95,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 a96f16fae99af9c30959d228202055e9aebfaf58..c82e7a192dfca8ca38832a53f9135125b4c34286 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.StackedRNNCells" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "state_size" mtype: "" @@ -107,11 +98,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 e1cbd0e150ed890ae57c1725249d1340fc2cb663..9ccf251a18034371607cfdc6091f2282136feec9 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.ThresholdedReLU" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 f0d35728fb1c42d563ff0598dd84da51a766a764..e080a07799fe1b7ec5f73b4e7bf78053a2c9dd3c 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-time-distributed.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-time-distributed.pbtxt @@ -3,7 +3,6 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -66,10 +61,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable" mtype: "" @@ -108,11 +99,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 74efaea6ddb22ec2fe9d41558978c183b0e06671..5fadca0b8386951976e8a6330bacf0dcf169e2d1 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.UpSampling1D" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 dc5bd5fd5319f9bbd601a3c4083ae566b47e1aaa..2d395bf7e87b8835bcf63d792d68dc3ac4083051 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.UpSampling2D" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 e01ccfb74aead591f1018cdcbb1c888767ecdb20..18d58ec3b23f7d4fd13ff45f4c1d4d95e7722ed5 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.UpSampling3D" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.layers.-wrapper.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.layers.-wrapper.pbtxt index 7e6f90f7623677244865ac285c134dc79f7b9b69..6223cb2f3c1230a60f3cf3dd57a0e803cf4f15d9 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.layers.-wrapper.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.layers.-wrapper.pbtxt @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.Wrapper" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable" mtype: "" @@ -107,11 +98,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 4d0d402dad442ccf52267f5ce40b05400afbfbc7..e71bba6a7f1df713cd13b4a0249d001d56bb31b0 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.ZeroPadding1D" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 b353a529bcf8e543d334fee57fca26ebc83036a4..aba6d8cb1f43bb070f9b17b5290afc5ce30246b5 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.ZeroPadding2D" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" 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 9fe1256e616dbca4f35101df160dc55bc68bfa8a..ce545ecc954e84e702fbd24047370a3417dc0fb1 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 @@ -2,7 +2,6 @@ path: "tensorflow.keras.layers.ZeroPadding3D" tf_class { is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -13,10 +12,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -65,10 +60,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "trainable_variables" mtype: "" @@ -103,11 +94,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" diff --git a/tensorflow/tools/api/golden/tensorflow.keras.models.-model.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.models.-model.pbtxt index 8ccf15f9ab0fcfa59907ff05a962a84d3d86ccb4..b3f3f169227c4548131d86165427655524f50d9d 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.models.-model.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.models.-model.pbtxt @@ -3,7 +3,6 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -14,10 +13,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -74,10 +69,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "state_updates" mtype: "" @@ -128,7 +119,7 @@ tf_class { } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -260,7 +251,7 @@ tf_class { } member_method { name: "test_on_batch" - argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " } member_method { name: "to_json" @@ -272,6 +263,6 @@ tf_class { } member_method { name: "train_on_batch" - argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " + argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " } } diff --git a/tensorflow/tools/api/golden/tensorflow.keras.models.-sequential.pbtxt b/tensorflow/tools/api/golden/tensorflow.keras.models.-sequential.pbtxt index be12b0bd2ec509ff394eaa3f43db0b54badd7fba..4ac6811bace97c25d62a2596523efdb1f789f7b2 100644 --- a/tensorflow/tools/api/golden/tensorflow.keras.models.-sequential.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.keras.models.-sequential.pbtxt @@ -4,7 +4,6 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" - is_instance: "" is_instance: "" is_instance: "" member { @@ -15,10 +14,6 @@ tf_class { name: "dtype" mtype: "" } - member { - name: "graph" - mtype: "" - } member { name: "inbound_nodes" mtype: "" @@ -75,10 +70,6 @@ tf_class { name: "output_shape" mtype: "" } - member { - name: "scope_name" - mtype: "" - } member { name: "state_updates" mtype: "" @@ -133,7 +124,7 @@ tf_class { } member_method { name: "add_weight" - argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], " + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'partitioner\', \'use_resource\', \'getter\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -277,7 +268,7 @@ tf_class { } member_method { name: "test_on_batch" - argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " } member_method { name: "to_json" @@ -289,6 +280,6 @@ tf_class { } member_method { name: "train_on_batch" - argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " + argspec: "args=[\'self\', \'x\', \'y\', \'sample_weight\', \'class_weight\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'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 1c4f550d7f05b8be33326cb39d7a5f3bf663f5e6..38fd78a5a828c7d0da98c97fdc01f504397c6fe6 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling1-d.pbtxt @@ -1,8 +1,10 @@ path: "tensorflow.layers.AveragePooling1D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -25,6 +27,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -53,6 +59,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -95,7 +105,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -109,6 +123,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -117,10 +135,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -133,6 +163,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -141,4 +175,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } 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 d2db0952693f2989e6a9e8748a254eb4db483206..86a524cc91e10616cd049cec93843e419ec670c7 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling2-d.pbtxt @@ -1,8 +1,10 @@ path: "tensorflow.layers.AveragePooling2D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -25,6 +27,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -53,6 +59,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -95,7 +105,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -109,6 +123,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -117,10 +135,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -133,6 +163,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -141,4 +175,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } 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 34d9a9df281c09a2e2030daf74a2ceb8066085bb..8a811fe4561ac3790e43a5553ca04fca002e420d 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-average-pooling3-d.pbtxt @@ -1,8 +1,10 @@ path: "tensorflow.layers.AveragePooling3D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -25,6 +27,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -53,6 +59,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -95,7 +105,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -109,6 +123,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -117,10 +135,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -133,6 +163,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -141,4 +175,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-batch-normalization.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-batch-normalization.pbtxt index 21ad0efecf88c42a3a679910ddfe095585a7933a..3923e706be7a352d770bf309aecd1fadb2a05e81 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-batch-normalization.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-batch-normalization.pbtxt @@ -1,7 +1,9 @@ path: "tensorflow.layers.BatchNormalization" tf_class { is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -24,6 +26,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -52,6 +58,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -94,7 +104,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -108,6 +122,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\', \'training\'], varargs=None, keywords=None, defaults=[\'False\'], " } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -116,10 +134,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -132,6 +162,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -140,4 +174,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-conv1-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-conv1-d.pbtxt index ed38747c7671a267bb640ecb96a4c5fcc46c5edf..7a0a8a2a51295d9009f44e2ea126e8a4694147af 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-conv1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-conv1-d.pbtxt @@ -1,8 +1,10 @@ path: "tensorflow.layers.Conv1D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -25,6 +27,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -53,6 +59,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -95,7 +105,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -109,6 +123,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -117,10 +135,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -133,6 +163,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -141,4 +175,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } 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 ff453c6059477c20528fc768d93c65d208cdfc4a..7ed3a652519a6429ff429925c29f3c296a6d2958 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-conv2-d-transpose.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-conv2-d-transpose.pbtxt @@ -1,9 +1,11 @@ path: "tensorflow.layers.Conv2DTranspose" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -26,6 +28,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -54,6 +60,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -96,7 +106,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -110,6 +124,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -118,10 +136,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -134,6 +164,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -142,4 +176,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-conv2-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-conv2-d.pbtxt index 5583bd22dce18b0a0593b73bde509818b63b3f29..23831aa74f1c3dea99e6da407e5a63693f94e37a 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-conv2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-conv2-d.pbtxt @@ -1,8 +1,10 @@ path: "tensorflow.layers.Conv2D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -25,6 +27,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -53,6 +59,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -95,7 +105,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -109,6 +123,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -117,10 +135,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -133,6 +163,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -141,4 +175,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } 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 63f0c32a7c8f7e530c76c64fa619102bc12f9ad9..9d41a6b09900d984706accd70a353cc26585d9b5 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-conv3-d-transpose.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-conv3-d-transpose.pbtxt @@ -1,9 +1,11 @@ path: "tensorflow.layers.Conv3DTranspose" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -26,6 +28,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -54,6 +60,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -96,7 +106,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -110,6 +124,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -118,10 +136,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -134,6 +164,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -142,4 +176,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-conv3-d.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-conv3-d.pbtxt index b77726252ccca30a7c6555fb569eb65b69e34998..865fe08e63c81222395d125938c8830c02030733 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-conv3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-conv3-d.pbtxt @@ -1,8 +1,10 @@ path: "tensorflow.layers.Conv3D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -25,6 +27,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -53,6 +59,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -95,7 +105,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -109,6 +123,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -117,10 +135,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -133,6 +163,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -141,4 +175,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-dense.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-dense.pbtxt index 92db9f6dcd2f77c4253eb77df4a26fb632b2a766..ee164aae204d3f6c09af79f7fbac825ce470098d 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-dense.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-dense.pbtxt @@ -1,7 +1,9 @@ path: "tensorflow.layers.Dense" tf_class { is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -24,6 +26,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -52,6 +58,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -94,7 +104,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -108,6 +122,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -116,10 +134,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -132,6 +162,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -140,4 +174,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-dropout.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-dropout.pbtxt index 80fa846a24c9162d8521bdb4f098b9cd8e34aedb..8167dc79cdf9f83b9b97557638bf0702a1be86da 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-dropout.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-dropout.pbtxt @@ -1,7 +1,9 @@ path: "tensorflow.layers.Dropout" tf_class { is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -24,6 +26,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -52,6 +58,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -94,7 +104,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -108,6 +122,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\', \'training\'], varargs=None, keywords=None, defaults=[\'False\'], " } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -116,10 +134,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -132,6 +162,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -140,4 +174,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-flatten.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-flatten.pbtxt index f63213b3dde40aa54b165c1c269c26fd2cd9e3b4..efa4419692993aa9975c1af2e647288ae9f38eba 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-flatten.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-flatten.pbtxt @@ -1,7 +1,9 @@ path: "tensorflow.layers.Flatten" tf_class { is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -24,6 +26,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -52,6 +58,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -94,7 +104,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -108,6 +122,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -116,10 +134,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -132,6 +162,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -140,4 +174,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-input-spec.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-input-spec.pbtxt index 7c1d05cd2bfc8be978e82566e7a3086040978b4a..2ff89f0a6faef905bcafdcb36121f506e9a9977a 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-input-spec.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-input-spec.pbtxt @@ -1,6 +1,6 @@ path: "tensorflow.layers.InputSpec" tf_class { - is_instance: "" + is_instance: "" is_instance: "" member_method { name: "__init__" diff --git a/tensorflow/tools/api/golden/tensorflow.layers.-layer.pbtxt b/tensorflow/tools/api/golden/tensorflow.layers.-layer.pbtxt index 4e45b2d513bb72bb47433d72c310d6a34fbc0c01..b3a6dfdffa28d3628e09e6aab823534ba84edf16 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: "" is_instance: "" member { @@ -23,6 +24,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -51,6 +56,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -81,7 +90,7 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'trainable\', \'name\', \'dtype\', \'activity_regularizer\'], varargs=None, keywords=kwargs, defaults=[\'True\', \'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'trainable\', \'name\', \'dtype\'], varargs=None, keywords=kwargs, defaults=[\'True\', \'None\', \'None\'], " } member_method { name: "add_loss" @@ -93,7 +102,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -107,6 +120,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -115,10 +132,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -131,6 +160,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -139,4 +172,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } 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 19ec33fce775caa634e71e2295ac945a6f70ade9..cef396489dde24698cd9a63b6247292958cfec4e 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling1-d.pbtxt @@ -1,8 +1,10 @@ path: "tensorflow.layers.MaxPooling1D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -25,6 +27,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -53,6 +59,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -95,7 +105,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -109,6 +123,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -117,10 +135,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -133,6 +163,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -141,4 +175,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } 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 76180c333a21c592a3b53bb445df9b12d3596552..565f0c7a79661f77d5987d671266ff69268b03b2 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling2-d.pbtxt @@ -1,8 +1,10 @@ path: "tensorflow.layers.MaxPooling2D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -25,6 +27,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -53,6 +59,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -95,7 +105,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -109,6 +123,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -117,10 +135,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -133,6 +163,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -141,4 +175,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } 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 ded75c8ff09efc6746ddd2284f53d2c021cc473c..595ce2eeadfc95fe44895ffd976024aee80ee948 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling3-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-max-pooling3-d.pbtxt @@ -1,8 +1,10 @@ path: "tensorflow.layers.MaxPooling3D" tf_class { is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -25,6 +27,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -53,6 +59,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -95,7 +105,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -109,6 +123,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -117,10 +135,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -133,6 +163,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -141,4 +175,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } 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 3dbfa5453f8e0ebb02429df9c4cbdf98de6b8ced..ccca96f72248e390ca65db061d836ee58c8e3205 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-separable-conv1-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-separable-conv1-d.pbtxt @@ -1,9 +1,11 @@ path: "tensorflow.layers.SeparableConv1D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -26,6 +28,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -54,6 +60,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -96,7 +106,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -110,6 +124,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -118,10 +136,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -134,6 +164,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -142,4 +176,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } } 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 ab171df1d1650e19836018f3316e6919f6d36def..1c99c9618254cadb1b4e95c7223ca9361e4fa861 100644 --- a/tensorflow/tools/api/golden/tensorflow.layers.-separable-conv2-d.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.layers.-separable-conv2-d.pbtxt @@ -1,9 +1,11 @@ path: "tensorflow.layers.SeparableConv2D" tf_class { is_instance: "" - is_instance: "" - is_instance: "" + is_instance: "" + is_instance: "" + is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -26,6 +28,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -54,6 +60,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -96,7 +106,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -110,6 +124,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -118,10 +136,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -134,6 +164,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -142,4 +176,12 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=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 9c71a24d0500e2091e0ae94cc4dd7ed6b788a54f..e1abd43ab54f9d0c131fd94cfeefcdd7ccdb00f1 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 @@ -4,6 +4,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -26,6 +27,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -54,6 +59,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -92,7 +101,7 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'num_units\', \'forget_bias\', \'state_is_tuple\', \'activation\', \'reuse\', \'name\'], varargs=None, keywords=None, defaults=[\'1.0\', \'True\', \'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'num_units\', \'forget_bias\', \'state_is_tuple\', \'activation\', \'reuse\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'1.0\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "add_loss" @@ -104,7 +113,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -118,6 +131,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\', \'state\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -126,10 +143,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -142,6 +171,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -150,6 +183,14 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "zero_state" argspec: "args=[\'self\', \'batch_size\', \'dtype\'], varargs=None, keywords=None, defaults=None" 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 9e19f96b7452616956fb7fd3ca62d8f4b25a2122..93e7e40199884fa77ae7da6d15113b8b4ca125ab 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 @@ -4,6 +4,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -26,6 +27,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -54,6 +59,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -92,7 +101,7 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'num_units\', \'activation\', \'reuse\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'num_units\', \'activation\', \'reuse\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], " } member_method { name: "add_loss" @@ -104,7 +113,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -118,6 +131,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\', \'state\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -126,10 +143,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -142,6 +171,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -150,6 +183,14 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "zero_state" argspec: "args=[\'self\', \'batch_size\', \'dtype\'], varargs=None, keywords=None, defaults=None" 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 7540aa62861895a7c41840476d4edb79785a77a9..3c3e38229738fec3b25f437a73f3a9d216d970af 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: "" is_instance: "" member { @@ -25,6 +26,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -53,6 +58,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -103,7 +112,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -117,6 +130,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -125,10 +142,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -141,6 +170,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -149,6 +182,14 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "zero_state" argspec: "args=[\'self\', \'batch_size\', \'dtype\'], varargs=None, keywords=None, defaults=None" 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 fc1ff386690f9c7acb11d4cc0770e394f78350ad..db16660f1145b55c824c698653094977dd6c718b 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: "" is_instance: "" member { @@ -25,6 +26,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -53,6 +58,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -107,7 +116,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -121,6 +134,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -129,10 +146,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -145,6 +174,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -153,6 +186,14 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "zero_state" argspec: "args=[\'self\', \'batch_size\', \'dtype\'], varargs=None, keywords=None, defaults=None" 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 751122cfff3bf9c55dd9fa264fdf2e1960940724..465fc1cd9c886acaeb8ac49ff54417b9fd101b04 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 @@ -4,6 +4,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -26,6 +27,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -54,6 +59,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -92,7 +101,7 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'num_units\', \'activation\', \'reuse\', \'kernel_initializer\', \'bias_initializer\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'num_units\', \'activation\', \'reuse\', \'kernel_initializer\', \'bias_initializer\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "add_loss" @@ -104,7 +113,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -118,6 +131,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\', \'state\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -126,10 +143,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -142,6 +171,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -150,6 +183,14 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "zero_state" argspec: "args=[\'self\', \'batch_size\', \'dtype\'], varargs=None, keywords=None, defaults=None" 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 4b6313f395fd8fd4ec2af78365117620263e7a55..38a387d55a4d34524cdc8af1f7ce1617eb120373 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 @@ -4,6 +4,7 @@ tf_class { is_instance: "" is_instance: "" is_instance: "" + is_instance: "" is_instance: "" is_instance: "" member { @@ -26,6 +27,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -54,6 +59,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -92,7 +101,7 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'num_units\', \'use_peepholes\', \'cell_clip\', \'initializer\', \'num_proj\', \'proj_clip\', \'num_unit_shards\', \'num_proj_shards\', \'forget_bias\', \'state_is_tuple\', \'activation\', \'reuse\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'1.0\', \'True\', \'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'num_units\', \'use_peepholes\', \'cell_clip\', \'initializer\', \'num_proj\', \'proj_clip\', \'num_unit_shards\', \'num_proj_shards\', \'forget_bias\', \'state_is_tuple\', \'activation\', \'reuse\', \'name\', \'dtype\'], varargs=None, keywords=None, defaults=[\'False\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'1.0\', \'True\', \'None\', \'None\', \'None\', \'None\'], " } member_method { name: "add_loss" @@ -104,7 +113,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -118,6 +131,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\', \'state\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -126,10 +143,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -142,6 +171,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -150,6 +183,14 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "zero_state" argspec: "args=[\'self\', \'batch_size\', \'dtype\'], varargs=None, keywords=None, defaults=None" 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 00e8c71140596ecea237ce05a09feff1fbb49001..b9e3d934759accd885036fa4c5a7013ef64736f3 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: "" is_instance: "" member { @@ -25,6 +26,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -53,6 +58,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -103,7 +112,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -117,6 +130,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\', \'state\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -125,10 +142,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -141,6 +170,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -149,6 +182,14 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "zero_state" argspec: "args=[\'self\', \'batch_size\', \'dtype\'], varargs=None, keywords=None, defaults=None" 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 3852f90dd6c4a254e20e789bdeb7796d61cef6bc..75b5898c591cbe2b761c0f709159c5489cb8f76a 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: "" is_instance: "" member { @@ -24,6 +25,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -52,6 +57,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -90,7 +99,7 @@ tf_class { } member_method { name: "__init__" - argspec: "args=[\'self\', \'trainable\', \'name\', \'dtype\', \'activity_regularizer\'], varargs=None, keywords=kwargs, defaults=[\'True\', \'None\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'trainable\', \'name\', \'dtype\'], varargs=None, keywords=kwargs, defaults=[\'True\', \'None\', \'None\'], " } member_method { name: "add_loss" @@ -102,7 +111,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -116,6 +129,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -124,10 +141,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -140,6 +169,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -148,6 +181,14 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "zero_state" argspec: "args=[\'self\', \'batch_size\', \'dtype\'], varargs=None, keywords=None, defaults=None" 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 8f3f0f7506ef49014b31cd4bc04f1cb1e0d696fc..fee0dc63b997f328a4e3d44040c4056de4128eb7 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: "" is_instance: "" member { @@ -25,6 +26,10 @@ tf_class { name: "input" mtype: "" } + member { + name: "input_mask" + mtype: "" + } member { name: "input_shape" mtype: "" @@ -53,6 +58,10 @@ tf_class { name: "output" mtype: "" } + member { + name: "output_mask" + mtype: "" + } member { name: "output_shape" mtype: "" @@ -103,7 +112,11 @@ 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\'], varargs=args, keywords=kwargs, defaults=None" + } + member_method { + name: "add_weight" + argspec: "args=[\'self\', \'name\', \'shape\', \'dtype\', \'initializer\', \'regularizer\', \'trainable\', \'constraint\', \'use_resource\', \'partitioner\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\'], " } member_method { name: "apply" @@ -117,6 +130,10 @@ tf_class { name: "call" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=kwargs, defaults=None" } + member_method { + name: "compute_mask" + argspec: "args=[\'self\', \'inputs\', \'mask\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "compute_output_shape" argspec: "args=[\'self\', \'input_shape\'], varargs=None, keywords=None, defaults=None" @@ -125,10 +142,22 @@ tf_class { name: "count_params" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "from_config" + argspec: "args=[\'cls\', \'config\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "get_config" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_input_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_input_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -141,6 +170,10 @@ tf_class { name: "get_output_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_output_mask_at" + argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "get_output_shape_at" argspec: "args=[\'self\', \'node_index\'], varargs=None, keywords=None, defaults=None" @@ -149,6 +182,14 @@ tf_class { name: "get_updates_for" argspec: "args=[\'self\', \'inputs\'], varargs=None, keywords=None, defaults=None" } + member_method { + name: "get_weights" + argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "set_weights" + argspec: "args=[\'self\', \'weights\'], varargs=None, keywords=None, defaults=None" + } member_method { name: "zero_state" argspec: "args=[\'self\', \'batch_size\', \'dtype\'], varargs=None, keywords=None, defaults=None" diff --git a/tensorflow/tools/api/golden/tensorflow.pbtxt b/tensorflow/tools/api/golden/tensorflow.pbtxt index afa3b78eb7fb3618edce06bbff288c37fdf71015..0b12bc060efc2ecfa4819e77cbb0d26c5a6dc419 100644 --- a/tensorflow/tools/api/golden/tensorflow.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.pbtxt @@ -912,9 +912,13 @@ tf_module { name: "decode_base64" argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " } + member_method { + name: "decode_compressed" + argspec: "args=[\'bytes\', \'compression_type\', \'name\'], varargs=None, keywords=None, defaults=[\'\', \'None\'], " + } member_method { name: "decode_csv" - argspec: "args=[\'records\', \'record_defaults\', \'field_delim\', \'use_quote_delim\', \'name\', \'na_value\'], varargs=None, keywords=None, defaults=[\',\', \'True\', \'None\', \'\'], " + argspec: "args=[\'records\', \'record_defaults\', \'field_delim\', \'use_quote_delim\', \'name\', \'na_value\', \'select_cols\'], varargs=None, keywords=None, defaults=[\',\', \'True\', \'None\', \'\', \'None\'], " } member_method { name: "decode_json_example" @@ -1976,6 +1980,10 @@ tf_module { name: "string_split" argspec: "args=[\'source\', \'delimiter\', \'skip_empty\'], varargs=None, keywords=None, defaults=[\' \', \'True\'], " } + member_method { + name: "string_strip" + argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " + } member_method { name: "string_to_hash_bucket" argspec: "args=[\'string_tensor\', \'num_buckets\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], " diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-file-writer.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-file-writer.pbtxt index dcf747971b7b8bf243502b2388da635705b8ee3e..6b65b0ace3cf7740ab03390841c941592000d127 100644 --- a/tensorflow/tools/api/golden/tensorflow.summary.-file-writer.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.summary.-file-writer.pbtxt @@ -5,7 +5,7 @@ tf_class { is_instance: "" member_method { name: "__init__" - argspec: "args=[\'self\', \'logdir\', \'graph\', \'max_queue\', \'flush_secs\', \'graph_def\', \'filename_suffix\'], varargs=None, keywords=None, defaults=[\'None\', \'10\', \'120\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'logdir\', \'graph\', \'max_queue\', \'flush_secs\', \'graph_def\', \'filename_suffix\', \'session\'], varargs=None, keywords=None, defaults=[\'None\', \'10\', \'120\', \'None\', \'None\', \'None\'], " } member_method { name: "add_event" diff --git a/tensorflow/tools/api/golden/tensorflow.train.-checkpoint-saver-hook.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-checkpoint-saver-hook.pbtxt index c3037baa8c951ecd9b60267ee7cc8674ead88dbe..327799729c9e7d9ce2931f8951db995195ff35e9 100644 --- a/tensorflow/tools/api/golden/tensorflow.train.-checkpoint-saver-hook.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.train.-checkpoint-saver-hook.pbtxt @@ -5,7 +5,7 @@ tf_class { is_instance: "" member_method { name: "__init__" - argspec: "args=[\'self\', \'checkpoint_dir\', \'save_secs\', \'save_steps\', \'saver\', \'checkpoint_basename\', \'scaffold\', \'listeners\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'model.ckpt\', \'None\', \'None\'], " + argspec: "args=[\'self\', \'checkpoint_dir\', \'save_secs\', \'save_steps\', \'saver\', \'checkpoint_basename\', \'scaffold\', \'listeners\', \'steps_per_run\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'model.ckpt\', \'None\', \'None\', \'1\'], " } member_method { name: "after_create_session" diff --git a/tensorflow/tools/api/golden/tensorflow.train.-checkpoint.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-checkpoint.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..17f393d27c4b917a78902a4d6f13630311cb44ef --- /dev/null +++ b/tensorflow/tools/api/golden/tensorflow.train.-checkpoint.pbtxt @@ -0,0 +1,23 @@ +path: "tensorflow.train.Checkpoint" +tf_class { + is_instance: "" + is_instance: "" + is_instance: "" + is_instance: "" + member { + name: "save_counter" + mtype: "" + } + member_method { + name: "__init__" + argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None" + } + member_method { + name: "restore" + argspec: "args=[\'self\', \'save_path\'], varargs=None, keywords=None, defaults=None" + } + member_method { + name: "save" + argspec: "args=[\'self\', \'file_prefix\', \'session\'], varargs=None, keywords=None, defaults=[\'None\'], " + } +} diff --git a/tensorflow/tools/api/golden/tensorflow.train.-scaffold.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-scaffold.pbtxt index 62b956c5ef7dc54e92431f25ec948e341c0e1f24..38cc98b48e78aa93f7614a9baff236f7b119f99d 100644 --- a/tensorflow/tools/api/golden/tensorflow.train.-scaffold.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.train.-scaffold.pbtxt @@ -38,6 +38,10 @@ tf_class { name: "__init__" argspec: "args=[\'self\', \'init_op\', \'init_feed_dict\', \'init_fn\', \'ready_op\', \'ready_for_local_init_op\', \'local_init_op\', \'summary_op\', \'saver\', \'copy_from_scaffold\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\'], " } + member_method { + name: "default_local_init_op" + argspec: "args=[], varargs=None, keywords=None, defaults=None" + } member_method { name: "finalize" argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" diff --git a/tensorflow/tools/api/golden/tensorflow.train.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.pbtxt index bec72e1e609c3e32ca8366396b9b1cb577feab9d..9fb18e77afd7c9c989ad5e967be291406e7239aa 100644 --- a/tensorflow/tools/api/golden/tensorflow.train.pbtxt +++ b/tensorflow/tools/api/golden/tensorflow.train.pbtxt @@ -20,6 +20,10 @@ tf_module { name: "BytesList" mtype: "" } + member { + name: "Checkpoint" + mtype: "" + } member { name: "CheckpointSaverHook" mtype: "" diff --git a/tensorflow/tools/ci_build/Dockerfile.rbe.cpu b/tensorflow/tools/ci_build/Dockerfile.rbe.cpu index 6f0798b1afc34bc08df6f3f8f467a329fcf0fe9b..3bc52b9ed611a0f0a4a269a2864d5b349ee9232c 100644 --- a/tensorflow/tools/ci_build/Dockerfile.rbe.cpu +++ b/tensorflow/tools/ci_build/Dockerfile.rbe.cpu @@ -1,4 +1,4 @@ -FROM launcher.gcr.io/google/rbe-debian8:r322167 +FROM launcher.gcr.io/google/rbe-debian8:r327695 LABEL maintainer="Yu Yi " # Copy install scripts diff --git a/tensorflow/tools/ci_build/builds/test_user_ops.sh b/tensorflow/tools/ci_build/builds/test_user_ops.sh index caa3a40817c80b27271f76de0a95a743cb2916f6..c342367bacea9d2ba8152d928b93bf61cf60d0e7 100755 --- a/tensorflow/tools/ci_build/builds/test_user_ops.sh +++ b/tensorflow/tools/ci_build/builds/test_user_ops.sh @@ -213,27 +213,34 @@ USER_OP=$(echo "${USER_OP_SO}" | sed -e 's/\.so//') echo "Invoking user op ${USER_OP} defined in file ${USER_OP_SO} "\ "via pip installation" -ORIG_OUTPUT=$("${PYTHON_BIN_PATH}" -c "import tensorflow as tf; print(tf.Session('').run(tf.load_op_library('./${USER_OP_SO}').${USER_OP}(${OP_INPUT})))") - -# Format OUTPUT for analysis -if [[ -z $(echo "${ORIG_OUTPUT}" | grep -o ',') ]]; then - if [[ ${IS_MAC} == "1" ]]; then - OUTPUT=$(echo "${ORIG_OUTPUT}" | sed -E -e 's/[ \t]+/,/g') +function run_op() { + local ORIG_OUTPUT=$1 + local ADDITIONAL_LOG=$2 + + # Format OUTPUT for analysis + if [[ -z $(echo "${ORIG_OUTPUT}" | grep -o ',') ]]; then + if [[ ${IS_MAC} == "1" ]]; then + local OUTPUT=$(echo "${ORIG_OUTPUT}" | sed -E -e 's/[ \t]+/,/g') + else + local OUTPUT=$(echo "${ORIG_OUTPUT}" | sed -r -e 's/[ \t]+/,/g') + fi else - OUTPUT=$(echo "${ORIG_OUTPUT}" | sed -r -e 's/[ \t]+/,/g') + local OUTPUT="${ORIG_OUTPUT}" fi -else - OUTPUT="${ORIG_OUTPUT}" -fi -EQUALS_EXPECTED=$("${PYTHON_BIN_PATH}" -c "print(${OUTPUT} == ${EXPECTED_OUTPUT})") + local EQUALS_EXPECTED=$("${PYTHON_BIN_PATH}" -c "print(${OUTPUT} == ${EXPECTED_OUTPUT})") -if [[ "${EQUALS_EXPECTED}" != "True" ]]; then - die "FAILED: Output from user op (${OUTPUT}) does not match expected "\ -"output ${EXPECTED_OUTPUT}" -else - echo "Output from user op (${OUTPUT}) matches expected output" -fi + if [[ "${EQUALS_EXPECTED}" != "True" ]]; then + local ERROR="FAILED: Output from user op (${OUTPUT}) does not match expected "\ + "output ${EXPECTED_OUTPUT}"${ADDITIONAL_LOG} + die ${ERROR} + else + echo "Output from user op (${OUTPUT}) matches expected output" + fi +} + +run_op $("${PYTHON_BIN_PATH}" -c "import tensorflow as tf; print(tf.Session('').run(tf.load_op_library('./${USER_OP_SO}').${USER_OP}(${OP_INPUT})))") +run_op $("${PYTHON_BIN_PATH}" -c "import tensorflow as tf; tf.enable_eager_execution(); print(tf.load_op_library('./${USER_OP_SO}').${USER_OP}(${OP_INPUT}))") " in eager mode" popd diff --git a/tensorflow/tools/ci_build/install/install_pip_packages.sh b/tensorflow/tools/ci_build/install/install_pip_packages.sh index d406b83a6246d18c335fb52cea1256d7809fa61a..5aaf544afdcb881fc00553cc78f916752f4527ac 100755 --- a/tensorflow/tools/ci_build/install/install_pip_packages.sh +++ b/tensorflow/tools/ci_build/install/install_pip_packages.sh @@ -16,10 +16,15 @@ set -e -# We don't apt-get install so that we can install a newer version of pip. Not -# needed after we upgrade to Ubuntu 16.04 -easy_install -U pip -easy_install3 -U pip +# We don't apt-get install so that we can install a newer version of pip. +# Only needed for Ubuntu 14.04 ,and not needed for Ubuntu 16.04 / Debian 8,9 +if $(cat /etc/*-release | grep -q 14.04); then + easy_install -U pip==9.0.3 + easy_install3 -U pip==9.0.3 +else + pip2 install --upgrade pip==9.0.3 + pip3 install --upgrade pip==9.0.3 +fi # Install pip packages from whl files to avoid the time-consuming process of # building from source. @@ -60,8 +65,13 @@ rm -rf /usr/lib/python3/dist-packages/six* # numpy needs to be installed from source to fix segfaults. See: # https://github.com/tensorflow/tensorflow/issues/6968 # This workaround isn't needed for Ubuntu 16.04 or later. -pip2 install --no-binary=:all: --upgrade numpy==1.12.0 -pip3 install --no-binary=:all: --upgrade numpy==1.12.0 +if $(cat /etc/*-release | grep -q 14.04); then + pip2 install --no-binary=:all: --upgrade numpy==1.12.0 + pip3 install --no-binary=:all: --upgrade numpy==1.12.0 +else + pip2 install --upgrade numpy==1.12.0 + pip3 install --upgrade numpy==1.12.0 +fi pip2 install scipy==0.18.1 pip3 install scipy==0.18.1 diff --git a/tensorflow/tools/ci_build/install/install_pip_packages_remote.sh b/tensorflow/tools/ci_build/install/install_pip_packages_remote.sh index 39a6d557d185d8564a79315fc738a054325aa0bc..0beabcf5ef83000bb0b0c5d3dc077a3d2bbe118a 100755 --- a/tensorflow/tools/ci_build/install/install_pip_packages_remote.sh +++ b/tensorflow/tools/ci_build/install/install_pip_packages_remote.sh @@ -20,10 +20,8 @@ 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 +pip2 install --upgrade setuptools +pip3 install --upgrade setuptools # The rest of the pip packages will be installed in # `install_pip_packages.sh` diff --git a/tensorflow/tools/ci_build/linux/cpu/run_mkl.sh b/tensorflow/tools/ci_build/linux/cpu/run_mkl.sh index dbf376be6f7467466b5383c57dca539f9a470a38..2a9f29518886d5dcb96f1afc8b12d6be307aa965 100755 --- a/tensorflow/tools/ci_build/linux/cpu/run_mkl.sh +++ b/tensorflow/tools/ci_build/linux/cpu/run_mkl.sh @@ -30,7 +30,10 @@ export PYTHON_BIN_PATH=`which python2` yes "" | $PYTHON_BIN_PATH configure.py # Run bazel test command. Double test timeouts to avoid flakes. +# Setting KMP_BLOCKTIME to 0 lets OpenMP threads to sleep right after parallel execution +# in an MKL primitive. This reduces the effects of an oversubscription of OpenMP threads +# caused by executing multiple tests concurrently. bazel test --test_tag_filters=-no_oss,-oss_serial,-gpu,-benchmark-test --test_lang_filters=py -k \ --jobs=${N_JOBS} --test_timeout 300,450,1200,3600 --build_tests_only \ - --config=mkl --config=opt --test_output=errors -- \ + --config=mkl --test_env=KMP_BLOCKTIME=0 --config=opt --test_output=errors -- \ //tensorflow/... -//tensorflow/compiler/... -//tensorflow/contrib/... diff --git a/tensorflow/tools/ci_build/xla/linux/gpu/run_py3.sh b/tensorflow/tools/ci_build/xla/linux/gpu/run_py3.sh index a94a627dfb632cf01518c2022fd01b168afb4a7e..a410c10b61b9f3f2cf8fc00074237a2bcfcbbf78 100755 --- a/tensorflow/tools/ci_build/xla/linux/gpu/run_py3.sh +++ b/tensorflow/tools/ci_build/xla/linux/gpu/run_py3.sh @@ -35,7 +35,7 @@ echo "build --distinct_host_configuration=false" >> .tf_configure.bazelrc bazel clean # Run bazel test command. Double test timeouts to avoid flakes. -bazel test --config=cuda --test_tag_filters=-no_gpu,-benchmark-test -k \ +bazel test --config=cuda --test_tag_filters=-no_gpu,-benchmark-test,-no_oss -k \ --jobs=${N_JOBS} --test_timeout 300,450,1200,3600 \ --build_tests_only --test_output=errors --local_test_jobs=8 \ --run_under=//tensorflow/tools/ci_build/gpu_build:parallel_gpu_execute \ diff --git a/tensorflow/tools/docs/BUILD b/tensorflow/tools/docs/BUILD index 0c1fd0cf9dc4e91601ac50648757613ec08a7f38..58b5ef8345c9de83e2d50cd01fe11e11f51fe298 100644 --- a/tensorflow/tools/docs/BUILD +++ b/tensorflow/tools/docs/BUILD @@ -103,10 +103,11 @@ py_test( data = ["//tensorflow:docs_src"], srcs_version = "PY2AND3", tags = [ - # No reason to run sanitizers for this test. + # No reason to run sanitizers or fastbuild for this test. "noasan", "nomsan", "notsan", + "optonly", ], deps = [ ":generate_lib", diff --git a/tensorflow/tools/docs/generate.py b/tensorflow/tools/docs/generate.py index c750539a76a85ec122fb87f44160422158966314..fc93085e3e0316cf274f4d9b325d6af0ea3a2f83 100644 --- a/tensorflow/tools/docs/generate.py +++ b/tensorflow/tools/docs/generate.py @@ -43,10 +43,6 @@ if __name__ == '__main__': flags = doc_generator.parse_known_args() - # Suppress documentation of some symbols that users should never use. - del tf.layers.Layer.inbound_nodes - del tf.layers.Layer.outbound_nodes - # tf_debug is not imported with tf, it's a separate module altogether doc_generator.set_py_modules([('tf', tf), ('tfdbg', tf_debug)]) diff --git a/tensorflow/tools/docs/generate_lib.py b/tensorflow/tools/docs/generate_lib.py index 6185c9ddd7b2c044728daa6f14befee3c8add036..111d54d8205f805cc24d21c610acc81610b8d47d 100644 --- a/tensorflow/tools/docs/generate_lib.py +++ b/tensorflow/tools/docs/generate_lib.py @@ -311,6 +311,10 @@ def build_doc_index(src_dir): continue title_parser = _GetMarkdownTitle() title_parser.process(os.path.join(dirpath, base_name)) + if title_parser.title is None: + msg = ('`{}` has no markdown title (# title)'.format( + os.path.join(dirpath, base_name))) + raise ValueError(msg) key_parts = os.path.join(suffix, base_name[:-3]).split('/') if key_parts[-1] == 'index': key_parts = key_parts[:-1] diff --git a/tensorflow/tools/docs/parser.py b/tensorflow/tools/docs/parser.py index cec23b1a36ca7e6a94d850c393271fa8616a0717..fb0bd2c2ff438aa9b3fa04719c447a2f3a91a95e 100644 --- a/tensorflow/tools/docs/parser.py +++ b/tensorflow/tools/docs/parser.py @@ -657,12 +657,14 @@ def _get_arg_spec(func): argspec_defaults.pop(i-first_default_arg) else: first_default_arg -= 1 - return tf_inspect.FullArgSpec(args=argspec_args, - varargs=argspec.varargs, - varkw=argspec.varkw, - defaults=tuple(argspec_defaults), - kwonlyargs=[], kwonlydefaults=None, - annotations={}) + return tf_inspect.FullArgSpec( + args=argspec_args, + varargs=argspec.varargs, + varkw=argspec.varkw, + defaults=tuple(argspec_defaults), + kwonlyargs=[], + kwonlydefaults=None, + annotations={}) else: # Regular function or method, getargspec will work fine. return tf_inspect.getfullargspec(func) @@ -672,7 +674,7 @@ def _remove_first_line_indent(string): return '\n'.join([line[indent:] for line in string.split('\n')]) -PAREN_NUMBER_RE = re.compile("^\(([0-9.e-]+)\)") +PAREN_NUMBER_RE = re.compile(r'^\(([0-9.e-]+)\)') def _generate_signature(func, reverse_index): @@ -1145,10 +1147,11 @@ class _ClassPageInfo(object): for short_name in parser_config.tree[self.full_name]: # Remove builtin members that we never want to document. - if short_name in ['__class__', '__base__', '__weakref__', '__doc__', - '__module__', '__dict__', '__abstractmethods__', - '__slots__', '__getnewargs__', '__str__', - '__repr__', '__hash__']: + if short_name in [ + '__class__', '__base__', '__weakref__', '__doc__', '__module__', + '__dict__', '__abstractmethods__', '__slots__', '__getnewargs__', + '__str__', '__repr__', '__hash__' + ]: continue child_name = '.'.join([self.full_name, short_name]) @@ -1193,7 +1196,8 @@ class _ClassPageInfo(object): # obvious what they do, don't include them in the docs if there's no # docstring. if not child_doc.brief.strip() and short_name in [ - '__del__', '__copy__']: + '__del__', '__copy__' + ]: print('Skipping %s, defined in %s, no docstring.' % (child_name, defining_class)) continue diff --git a/tensorflow/tools/docs/parser_test.py b/tensorflow/tools/docs/parser_test.py index d7757d78ed42a5307635776b3595153465ac8d71..274d48ef66071a4e6a5ebea65087f18382fea6a2 100644 --- a/tensorflow/tools/docs/parser_test.py +++ b/tensorflow/tools/docs/parser_test.py @@ -408,67 +408,98 @@ class ParserTest(googletest.TestCase): # pylint: disable=protected-access # Make sure everything works for regular functions. - expected = tf_inspect.FullArgSpec(args=['arg1', 'arg2', 'kwarg1', 'kwarg2'], - varargs=None, varkw=None, defaults=(1, 2), - kwonlyargs=[], kwonlydefaults=None, - annotations={}) + expected = tf_inspect.FullArgSpec( + args=['arg1', 'arg2', 'kwarg1', 'kwarg2'], + varargs=None, + varkw=None, + defaults=(1, 2), + kwonlyargs=[], + kwonlydefaults=None, + annotations={}) self.assertEqual(expected, parser._get_arg_spec(test_function_for_partial1)) # Make sure doing nothing works. - expected = tf_inspect.FullArgSpec(args=['arg1', 'arg2', 'kwarg1', 'kwarg2'], - varargs=None, varkw=None, defaults=(1, 2), - kwonlyargs=[], kwonlydefaults=None, - annotations={}) + expected = tf_inspect.FullArgSpec( + args=['arg1', 'arg2', 'kwarg1', 'kwarg2'], + varargs=None, + varkw=None, + defaults=(1, 2), + kwonlyargs=[], + kwonlydefaults=None, + annotations={}) partial = functools.partial(test_function_for_partial1) self.assertEqual(expected, parser._get_arg_spec(partial)) # Make sure setting args from the front works. - expected = tf_inspect.FullArgSpec(args=['arg2', 'kwarg1', 'kwarg2'], - varargs=None, varkw=None, defaults=(1, 2), - kwonlyargs=[], kwonlydefaults=None, - annotations={}) + expected = tf_inspect.FullArgSpec( + args=['arg2', 'kwarg1', 'kwarg2'], + varargs=None, + varkw=None, + defaults=(1, 2), + kwonlyargs=[], + kwonlydefaults=None, + annotations={}) partial = functools.partial(test_function_for_partial1, 1) self.assertEqual(expected, parser._get_arg_spec(partial)) - expected = tf_inspect.FullArgSpec(args=['kwarg2'], - varargs=None, varkw=None, defaults=(2,), - kwonlyargs=[], kwonlydefaults=None, - annotations={}) + expected = tf_inspect.FullArgSpec( + args=['kwarg2'], + varargs=None, + varkw=None, + defaults=(2,), + kwonlyargs=[], + kwonlydefaults=None, + annotations={}) partial = functools.partial(test_function_for_partial1, 1, 2, 3) self.assertEqual(expected, parser._get_arg_spec(partial)) # Make sure setting kwargs works. - expected = tf_inspect.FullArgSpec(args=['arg1', 'arg2', 'kwarg2'], - varargs=None, varkw=None, defaults=(2,), - kwonlyargs=[], kwonlydefaults=None, - annotations={}) + expected = tf_inspect.FullArgSpec( + args=['arg1', 'arg2', 'kwarg2'], + varargs=None, + varkw=None, + defaults=(2,), + kwonlyargs=[], + kwonlydefaults=None, + annotations={}) partial = functools.partial(test_function_for_partial1, kwarg1=0) self.assertEqual(expected, parser._get_arg_spec(partial)) - expected = tf_inspect.FullArgSpec(args=['arg1', 'arg2', 'kwarg1'], - varargs=None, varkw=None, defaults=(1,), - kwonlyargs=[], kwonlydefaults=None, - annotations={}) + expected = tf_inspect.FullArgSpec( + args=['arg1', 'arg2', 'kwarg1'], + varargs=None, + varkw=None, + defaults=(1,), + kwonlyargs=[], + kwonlydefaults=None, + annotations={}) partial = functools.partial(test_function_for_partial1, kwarg2=0) self.assertEqual(expected, parser._get_arg_spec(partial)) - expected = tf_inspect.FullArgSpec(args=['arg1'], - varargs=None, varkw=None, defaults=(), - kwonlyargs=[], kwonlydefaults=None, - annotations={}) + expected = tf_inspect.FullArgSpec( + args=['arg1'], + varargs=None, + varkw=None, + defaults=(), + kwonlyargs=[], + kwonlydefaults=None, + annotations={}) partial = functools.partial(test_function_for_partial1, arg2=0, kwarg1=0, kwarg2=0) self.assertEqual(expected, parser._get_arg_spec(partial)) # Make sure *args, *kwargs is accounted for. - expected = tf_inspect.FullArgSpec(args=[], - varargs='my_args', varkw='my_kwargs', - defaults=(), - kwonlyargs=[], kwonlydefaults=None, - annotations={}) + expected = tf_inspect.FullArgSpec( + args=[], + varargs='my_args', + varkw='my_kwargs', + defaults=(), + kwonlyargs=[], + kwonlydefaults=None, + annotations={}) partial = functools.partial(test_function_for_partial2, 0, 1) self.assertEqual(expected, parser._get_arg_spec(partial)) - + # pylint: enable=protected-access def testSaveReferenceResolver(self): diff --git a/tensorflow/tools/git/gen_git_source.py b/tensorflow/tools/git/gen_git_source.py index 6ec162e4a993187aa7113325c6bed1a88f9e247c..a8d878cf16cb9aae1c8639d97cb2802f685c2a8a 100755 --- a/tensorflow/tools/git/gen_git_source.py +++ b/tensorflow/tools/git/gen_git_source.py @@ -125,7 +125,7 @@ def configure(src_base_path, gen_path, debug=False): try: # In python 3.5, symlink function exists even on Windows. But requires # Windows Admin privileges, otherwise an OSError will be thrown. - if hasattr(os, 'symlink'): + if hasattr(os, "symlink"): os.symlink(src, os.path.join(gen_path, target)) else: shutil.copy2(src, os.path.join(gen_path, target)) @@ -176,7 +176,7 @@ def get_git_version(git_base_path, git_tag_override): val = bytes( "-".join([git_tag_override, "0", abbrev_commit])) return val if val else unknown_label - except subprocess.CalledProcessError: + except (subprocess.CalledProcessError, OSError): return unknown_label @@ -192,7 +192,15 @@ def write_version_info(filename, git_version): contents = """/* Generated by gen_git_source.py */ #include const char* tf_git_version() {return "%s";} -const char* tf_compiler_version() {return __VERSION__;} +const char* tf_compiler_version() { +#ifdef _MSC_VER +#define STRINGIFY(x) #x +#define TOSTRING(x) STRINGIFY(x) + return "MSVC " TOSTRING(_MSC_FULL_VER); +#else + return __VERSION__; +#endif +} const int tf_cxx11_abi_flag() { #ifdef _GLIBCXX_USE_CXX11_ABI return _GLIBCXX_USE_CXX11_ABI; diff --git a/tensorflow/tools/git/gen_git_source.sh b/tensorflow/tools/git/gen_git_source.sh index db20bb00e84b47bd15244e70b925f59e62731deb..cd128af6b36f2f99e5cf91961476d30384227e9b 100755 --- a/tensorflow/tools/git/gen_git_source.sh +++ b/tensorflow/tools/git/gen_git_source.sh @@ -28,7 +28,15 @@ fi cat < ${OUTPUT_FILENAME} #include const char* tf_git_version() {return "${GIT_VERSION}";} -const char* tf_compiler_version() {return __VERSION__;} +const char* tf_compiler_version() { +#ifdef _MSC_VER +#define STRINGIFY(x) #x +#define TOSTRING(x) STRINGIFY(x) + return "MSVC " TOSTRING(_MSC_FULL_VER); +#else + return __VERSION__; +#endif +} const int tf_cxx11_abi_flag() { #ifdef _GLIBCXX_USE_CXX11_ABI return _GLIBCXX_USE_CXX11_ABI; diff --git a/tensorflow/tools/graph_transforms/transform_graph.cc b/tensorflow/tools/graph_transforms/transform_graph.cc index 28387c2b48c06ecffd2afa0705a8dea5bc368460..8ce8f5e24b9f002e50d456c8ccab8a6414fca724 100644 --- a/tensorflow/tools/graph_transforms/transform_graph.cc +++ b/tensorflow/tools/graph_transforms/transform_graph.cc @@ -24,6 +24,9 @@ limitations under the License. #include "tensorflow/core/util/command_line_flags.h" #include "tensorflow/tools/graph_transforms/file_utils.h" #include "tensorflow/tools/graph_transforms/transform_utils.h" +#if !defined(PLATFORM_WINDOWS) +#include +#endif namespace tensorflow { namespace graph_transforms { @@ -130,16 +133,64 @@ Status ParseTransformParameters(const string& transforms_string, return Status::OK(); } +std::string ExpandPath(const std::string& path_string) { +#if defined(PLATFORM_WINDOWS) + return path_string; +#else + if (path_string.empty() || path_string[0] != '~') { + return path_string; + } + + const char* home = NULL; + std::string::size_type prefix = path_string.find_first_of('/'); + if (path_string.length() == 1 || prefix == 1) { + // The value of $HOME, e.g., ~/foo + home = getenv("HOME"); + if (!home) { + // If HOME is not available, get uid + struct passwd* pw = getpwuid(getuid()); + if (pw) { + home = pw->pw_dir; + } + } + } else { + // The value of ~user, e.g., ~user/foo + std::string user(path_string, 1, (prefix == std::string::npos) + ? std::string::npos + : prefix - 1); + struct passwd* pw = getpwnam(user.c_str()); + if (pw) { + home = pw->pw_dir; + } + } + + if (!home) { + return path_string; + } + + string path(home); + if (prefix == std::string::npos) { + return path; + } + + if (path.length() == 0 || path[path.length() - 1] != '/') { + path += '/'; + } + path += path_string.substr(prefix + 1); + return path; +#endif +} + int ParseFlagsAndTransformGraph(int argc, char* argv[], bool init_main) { - string in_graph = ""; - string out_graph = ""; + string in_graph_string = ""; + string out_graph_string = ""; string inputs_string = ""; string outputs_string = ""; string transforms_string = ""; bool output_as_text = false; std::vector flag_list = { - Flag("in_graph", &in_graph, "input graph file name"), - Flag("out_graph", &out_graph, "output graph file name"), + Flag("in_graph", &in_graph_string, "input graph file name"), + Flag("out_graph", &out_graph_string, "output graph file name"), Flag("inputs", &inputs_string, "inputs"), Flag("outputs", &outputs_string, "outputs"), Flag("transforms", &transforms_string, "list of transforms"), @@ -166,11 +217,11 @@ int ParseFlagsAndTransformGraph(int argc, char* argv[], bool init_main) { LOG(ERROR) << "Unknown argument " << argv[1] << ".\n" << usage; return -1; } - if (in_graph.empty()) { + if (in_graph_string.empty()) { LOG(ERROR) << "in_graph graph can't be empty.\n" << usage; return -1; } - if (out_graph.empty()) { + if (out_graph_string.empty()) { LOG(ERROR) << "out_graph graph can't be empty.\n" << usage; return -1; } @@ -179,6 +230,9 @@ int ParseFlagsAndTransformGraph(int argc, char* argv[], bool init_main) { return -1; } + string in_graph = ExpandPath(in_graph_string); + string out_graph = ExpandPath(out_graph_string); + std::vector inputs = str_util::Split(inputs_string, ','); std::vector outputs = str_util::Split(outputs_string, ','); TransformParameters transform_params; @@ -197,7 +251,7 @@ int ParseFlagsAndTransformGraph(int argc, char* argv[], bool init_main) { GraphDef graph_def; Status load_status = LoadTextOrBinaryGraphFile(in_graph, &graph_def); if (!load_status.ok()) { - LOG(ERROR) << "Loading graph '" << in_graph << "' failed with " + LOG(ERROR) << "Loading graph '" << in_graph_string << "' failed with " << load_status.error_message(); LOG(ERROR) << usage; return -1; @@ -219,7 +273,7 @@ int ParseFlagsAndTransformGraph(int argc, char* argv[], bool init_main) { save_status = WriteBinaryProto(Env::Default(), out_graph, graph_def); } if (!save_status.ok()) { - LOG(ERROR) << "Saving graph '" << out_graph << "' failed with " + LOG(ERROR) << "Saving graph '" << out_graph_string << "' failed with " << save_status.error_message(); return -1; } diff --git a/tensorflow/tools/lib_package/BUILD b/tensorflow/tools/lib_package/BUILD index 0ede8c63704ac4a474eb0d19e17cf5f365abca77..569b6678cabddd17eaf21921d668221f1a865625 100644 --- a/tensorflow/tools/lib_package/BUILD +++ b/tensorflow/tools/lib_package/BUILD @@ -118,6 +118,7 @@ genrule( "@com_googlesource_code_re2//:LICENSE", "@cub_archive//:LICENSE.TXT", "@curl//:COPYING", + "@double_conversion//:LICENSE", "@eigen_archive//:COPYING.MPL2", "@farmhash_archive//:COPYING", "@fft2d//:fft/readme.txt", @@ -155,6 +156,7 @@ genrule( "@com_googlesource_code_re2//:LICENSE", "@cub_archive//:LICENSE.TXT", "@curl//:COPYING", + "@double_conversion//:LICENSE", "@eigen_archive//:COPYING.MPL2", "@farmhash_archive//:COPYING", "@fft2d//:fft/readme.txt", diff --git a/tensorflow/tools/lib_package/README.md b/tensorflow/tools/lib_package/README.md index 700814826035a3274739f1d3784140ea0a59b36f..cb6aef262456a202b93a2787626688b7712e2352 100644 --- a/tensorflow/tools/lib_package/README.md +++ b/tensorflow/tools/lib_package/README.md @@ -35,8 +35,8 @@ The following commands: bazel test --config opt //tensorflow/tools/lib_package:libtensorflow_test bazel build --config opt \ //tensorflow/tools/lib_package:libtensorflow_jni.tar.gz \ - //tensorflow/tools/lib_package:libtensorflow.jar \ - //tensorflow/tools/lib_package:libtensorflow-src.jar + //tensorflow/java:libtensorflow.jar \ + //tensorflow/java:libtensorflow-src.jar ``` test and produce the following: @@ -44,9 +44,9 @@ test and produce the following: - The native library (`libtensorflow_jni.so`) packaged in an archive at: `bazel-bin/tensorflow/tools/lib_package/libtensorflow_jni.tar.gz` - The Java archive at: - `bazel-bin/tensorflow/tools/lib_package/libtensorflow.jar` + `bazel-bin/tensorflow/java/libtensorflow.jar` - The Java archive for Java sources at: - `bazel-bin/tensorflow/tools/lib_package/libtensorflow-src.jar` + `bazel-bin/tensorflow/java/libtensorflow-src.jar` ## Release diff --git a/tensorflow/tools/pip_package/BUILD b/tensorflow/tools/pip_package/BUILD index 2ef105755f2dde0c1181ab3eaaa02ae553beca88..7b508f87ab75e27cfdb3c41002873a8d38117f13 100644 --- a/tensorflow/tools/pip_package/BUILD +++ b/tensorflow/tools/pip_package/BUILD @@ -66,7 +66,6 @@ COMMON_PIP_DEPS = [ "//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", @@ -129,6 +128,7 @@ filegroup( "@com_googlesource_code_re2//:LICENSE", "@cub_archive//:LICENSE.TXT", "@curl//:COPYING", + "@double_conversion//:LICENSE", "@eigen_archive//:COPYING.MPL2", "@farmhash_archive//:COPYING", "@fft2d//:fft/readme.txt", diff --git a/tensorflow/tools/pip_package/setup.py b/tensorflow/tools/pip_package/setup.py index 6da3223d339057708cf82073e1666e6879c67e26..b88d023cbcaa9d6444ae14824227ef13145db4f1 100644 --- a/tensorflow/tools/pip_package/setup.py +++ b/tensorflow/tools/pip_package/setup.py @@ -22,7 +22,9 @@ import os import re import sys -from setuptools import find_packages, setup, Command +from setuptools import Command +from setuptools import find_packages +from setuptools import setup from setuptools.command.install import install as InstallCommandBase from setuptools.dist import Distribution @@ -97,7 +99,9 @@ TEST_PACKAGES = [ 'scipy >= 0.15.1', ] + class BinaryDistribution(Distribution): + def has_ext_modules(self): return True @@ -179,9 +183,9 @@ class InstallHeaders(Command): def find_files(pattern, root): """Return all the files matching pattern below root dir.""" - for path, _, files in os.walk(root): + for dirpath, _, files in os.walk(root): for filename in fnmatch.filter(files, pattern): - yield os.path.join(path, filename) + yield os.path.join(dirpath, filename) matches = ['../' + x for x in find_files('*', 'external') if '.py' not in x] diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl index 48728ac131a638f68223255632a3fc2f85c6c92b..884dcc5d0c0b36207b7a76ec4225b39b9533d262 100644 --- a/tensorflow/workspace.bzl +++ b/tensorflow/workspace.bzl @@ -50,31 +50,31 @@ def tf_workspace(path_prefix="", tf_repo_name=""): mkl_repository( name = "mkl_linux", urls = [ - "https://mirror.bazel.build/github.com/intel/mkl-dnn/releases/download/v0.12/mklml_lnx_2018.0.1.20171227.tgz", - "https://github.com/intel/mkl-dnn/releases/download/v0.12/mklml_lnx_2018.0.1.20171227.tgz", + "https://mirror.bazel.build/github.com/intel/mkl-dnn/releases/download/v0.13/mklml_lnx_2018.0.2.20180127.tgz", + "https://github.com/intel/mkl-dnn/releases/download/v0.13/mklml_lnx_2018.0.2.20180127.tgz", ], - sha256 = "feacc3d82565c1231470359b42c696236fae873704e0b013436afba5fd4fd30f", - strip_prefix = "mklml_lnx_2018.0.1.20171227", + sha256 = "74844bd77294742bf2396ff040369d1aa4cdd9e826fcd38cf8398ae83564d146", + strip_prefix = "mklml_lnx_2018.0.2.20180127", build_file = clean_dep("//third_party/mkl:mkl.BUILD") ) mkl_repository( name = "mkl_windows", urls = [ - "https://mirror.bazel.build/github.com/intel/mkl-dnn/releases/download/v0.12/mklml_win_2018.0.1.20171227.zip", - "https://github.com/intel/mkl-dnn/releases/download/v0.12/mklml_win_2018.0.1.20171227.zip" + "https://mirror.bazel.build/github.com/intel/mkl-dnn/releases/download/v0.13/mklml_win_2018.0.2.20180127.zip", + "https://github.com/intel/mkl-dnn/releases/download/v0.13/mklml_win_2018.0.2.20180127.zip" ], - sha256 = "24bae8d7b22b431a654acadea43f2243c46ae6b1e5a73a4a936825f31d284ee4", - strip_prefix = "mklml_win_2018.0.1.20171227", + sha256 = "d8fbf0faa0684bffa3548005d05fe5cfe56ff9dbc0e15e7612d7ac01055a6ded", + strip_prefix = "mklml_win_2018.0.2.20180127", build_file = clean_dep("//third_party/mkl:mkl.BUILD") ) mkl_repository( name = "mkl_darwin", urls = [ - "https://mirror.bazel.build/github.com/intel/mkl-dnn/releases/download/v0.12/mklml_mac_2018.0.1.20171227.tgz", - "https://github.com/intel/mkl-dnn/releases/download/v0.12/mklml_mac_2018.0.1.20171227.tgz" + "https://mirror.bazel.build/github.com/intel/mkl-dnn/releases/download/v0.13/mklml_mac_2018.0.2.20180127.tgz", + "https://github.com/intel/mkl-dnn/releases/download/v0.13/mklml_mac_2018.0.2.20180127.tgz" ], - sha256 = "0e954ec6fd3dc5e37f64c4043f6b5613dd687558da3df1028b3b7c29ff5cf77f", - strip_prefix = "mklml_mac_2018.0.1.20171227", + sha256 = "aa740d71e14562bfea56e6829e6dc186e7487cbcf6748a88dec73826b7ec1943", + strip_prefix = "mklml_mac_2018.0.2.20180127", build_file = clean_dep("//third_party/mkl:mkl.BUILD") ) @@ -85,11 +85,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""): tf_http_archive( name = "mkl_dnn", urls = [ - "https://mirror.bazel.build/github.com/intel/mkl-dnn/archive/v0.12.tar.gz", - "https://github.com/intel/mkl-dnn/archive/v0.12.tar.gz", + "https://mirror.bazel.build/github.com/intel/mkl-dnn/archive/v0.13.tar.gz", + "https://github.com/intel/mkl-dnn/archive/v0.13.tar.gz", ], - sha256 = "86fa2a8c12a56e3b725945acedeaa82492746be02545aba6d710f097e013e19e", - strip_prefix = "mkl-dnn-0.12", + sha256 = "d2cfd93a70cfe86ebe054477c530c9b5c1218b70f75856eb6d1956c68ee89e8f", + strip_prefix = "mkl-dnn-0.13", build_file = clean_dep("//third_party/mkl_dnn:mkldnn.BUILD"), ) @@ -167,11 +167,12 @@ def tf_workspace(path_prefix="", tf_repo_name=""): tf_http_archive( name = "gemmlowp", urls = [ - "https://mirror.bazel.build/github.com/google/gemmlowp/archive/7c7c744640ddc3d0af18fb245b4d23228813a71b.zip", - "https://github.com/google/gemmlowp/archive/7c7c744640ddc3d0af18fb245b4d23228813a71b.zip", + # TODO (yongtang): uncomment once mirror.bazel.build is propagated. + # "https://mirror.bazel.build/github.com/google/gemmlowp/archive/38ebac7b059e84692f53e5938f97a9943c120d98.zip", + "https://github.com/google/gemmlowp/archive/38ebac7b059e84692f53e5938f97a9943c120d98.zip", ], - sha256 = "b852cc90259a7357c8a323f108f2cec6e85979fc3b18b5590b99e0130044b2cf", - strip_prefix = "gemmlowp-7c7c744640ddc3d0af18fb245b4d23228813a71b", + sha256 = "b87faa7294dfcc5d678f22a59d2c01ca94ea1e2a3b488c38a95a67889ed0a658", + strip_prefix = "gemmlowp-38ebac7b059e84692f53e5938f97a9943c120d98", ) tf_http_archive( @@ -210,11 +211,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""): tf_http_archive( name = "jpeg", urls = [ - "https://mirror.bazel.build/github.com/libjpeg-turbo/libjpeg-turbo/archive/1.5.1.tar.gz", - "https://github.com/libjpeg-turbo/libjpeg-turbo/archive/1.5.1.tar.gz", + "https://mirror.bazel.build/github.com/libjpeg-turbo/libjpeg-turbo/archive/1.5.3.tar.gz", + "https://github.com/libjpeg-turbo/libjpeg-turbo/archive/1.5.3.tar.gz", ], - sha256 = "c15a9607892113946379ccea3ca8b85018301b200754f209453ab21674268e77", - strip_prefix = "libjpeg-turbo-1.5.1", + sha256 = "1a17020f859cb12711175a67eab5c71fc1904e04b587046218e36106e07eabde", + strip_prefix = "libjpeg-turbo-1.5.3", build_file = clean_dep("//third_party/jpeg:jpeg.BUILD"), ) @@ -232,11 +233,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""): tf_http_archive( name = "org_sqlite", urls = [ - "https://mirror.bazel.build/www.sqlite.org/2017/sqlite-amalgamation-3200000.zip", - "http://www.sqlite.org/2017/sqlite-amalgamation-3200000.zip", + "https://mirror.bazel.build/www.sqlite.org/2018/sqlite-amalgamation-3230100.zip", + "https://www.sqlite.org/2018/sqlite-amalgamation-3230100.zip", ], - sha256 = "208780b3616f9de0aeb50822b7a8f5482f6515193859e91ed61637be6ad74fd4", - strip_prefix = "sqlite-amalgamation-3200000", + sha256 = "4239a1f69e5721d07d9a374eb84d594225229e54be4ee628da2995f4315d8dfc", + strip_prefix = "sqlite-amalgamation-3230100", build_file = clean_dep("//third_party:sqlite.BUILD"), ) @@ -427,11 +428,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""): tf_http_archive( name = "grpc", urls = [ - "https://mirror.bazel.build/github.com/grpc/grpc/archive/bd6bdf93279a39a8cd92978fd7c9d14eccd98fc2.tar.gz", - "https://github.com/grpc/grpc/archive/bd6bdf93279a39a8cd92978fd7c9d14eccd98fc2.tar.gz", + "https://mirror.bazel.build/github.com/grpc/grpc/archive/d184fa229d75d336aedea0041bd59cb93e7e267f.tar.gz", + "https://github.com/grpc/grpc/archive/d184fa229d75d336aedea0041bd59cb93e7e267f.tar.gz", ], - sha256 = "0a05bd355e4571b01d813dddffa38e57e689ac41b264dc9b1bd6ec66463ef5d6", - strip_prefix = "grpc-bd6bdf93279a39a8cd92978fd7c9d14eccd98fc2", + sha256 = "895b31310e718a61f7335759a778c068a6edde1c089883598a0830cbb7075673", + strip_prefix = "grpc-d184fa229d75d336aedea0041bd59cb93e7e267f", ) @@ -451,11 +452,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/7e78daafdd22f3f17720a103d29d89590534004e.tar.gz", - "https://github.com/llvm-mirror/llvm/archive/7e78daafdd22f3f17720a103d29d89590534004e.tar.gz", + "https://mirror.bazel.build/github.com/llvm-mirror/llvm/archive/c1e9b6f826c86c87a7e7173f1baf7e7df9f43e32.tar.gz", + "https://github.com/llvm-mirror/llvm/archive/c1e9b6f826c86c87a7e7173f1baf7e7df9f43e32.tar.gz", ], - sha256 = "a6d94bd9de23515a1e3792a830421e3885977ea43d03427cdbe68f98cb7e0045", - strip_prefix = "llvm-7e78daafdd22f3f17720a103d29d89590534004e", + sha256 = "92b7c01074f694a77b4d664951d1ec071e30ef19c61e673158e95fbb6e447b54", + strip_prefix = "llvm-c1e9b6f826c86c87a7e7173f1baf7e7df9f43e32", build_file = clean_dep("//third_party/llvm:llvm.BUILD"), ) @@ -693,6 +694,16 @@ def tf_workspace(path_prefix="", tf_repo_name=""): build_file = clean_dep("//third_party/flatbuffers:flatbuffers.BUILD"), ) + native.new_http_archive( + name = "double_conversion", + urls = [ + "https://github.com/google/double-conversion/archive/3992066a95b823efc8ccc1baf82a1cfc73f6e9b8.zip", + ], + sha256 = "2f7fbffac0d98d201ad0586f686034371a6d152ca67508ab611adc2386ad30de", + strip_prefix = "double-conversion-3992066a95b823efc8ccc1baf82a1cfc73f6e9b8", + build_file = clean_dep("//third_party:double_conversion.BUILD") + ) + tf_http_archive( name = "tflite_mobilenet", sha256 = "23f814d1c076bdf03715dfb6cab3713aa4fbdf040fd5448c43196bd2e97a4c1b", diff --git a/third_party/double_conversion.BUILD b/third_party/double_conversion.BUILD new file mode 100644 index 0000000000000000000000000000000000000000..9f905216c036bf5e48e1a1b94cd3dd61f3e53c41 --- /dev/null +++ b/third_party/double_conversion.BUILD @@ -0,0 +1,38 @@ +# Bazel(http://bazel.io) BUILD file + +licenses(["notice"]) + +exports_files(["LICENSE"]) + +cc_library( + name = "double-conversion", + srcs = [ + "double-conversion/bignum.cc", + "double-conversion/bignum-dtoa.cc", + "double-conversion/cached-powers.cc", + "double-conversion/diy-fp.cc", + "double-conversion/double-conversion.cc", + "double-conversion/fast-dtoa.cc", + "double-conversion/fixed-dtoa.cc", + "double-conversion/strtod.cc", + "double-conversion/utils.h", + ], + hdrs = [ + "double-conversion/bignum.h", + "double-conversion/bignum-dtoa.h", + "double-conversion/cached-powers.h", + "double-conversion/diy-fp.h", + "double-conversion/double-conversion.h", + "double-conversion/fast-dtoa.h", + "double-conversion/fixed-dtoa.h", + "double-conversion/ieee.h", + "double-conversion/strtod.h", + ], + includes = [ + ".", + ], + linkopts = [ + "-lm", + ], + visibility = ["//visibility:public"], +) diff --git a/third_party/llvm/llvm.BUILD b/third_party/llvm/llvm.BUILD index 075b46896ed868d8e2e1bcddf6d867974a248313..cbb1b2fe429e90088510661fe333f208f26ff061 100644 --- a/third_party/llvm/llvm.BUILD +++ b/third_party/llvm/llvm.BUILD @@ -2006,7 +2006,6 @@ cc_library( ]) + [ "include/llvm/BinaryFormat/MachO.def", "include/llvm/Support/VCSRevision.h", - "include/llvm/ExecutionEngine/ObjectMemoryBuffer.h", ], deps = [ ":config", @@ -2053,6 +2052,7 @@ cc_library( "include/llvm/Target/*.def", "include/llvm/Target/*.inc", "include/llvm/CodeGen/*.def", + "include/llvm/CodeGen/*.inc", ]), deps = [ ":analysis", diff --git a/third_party/repo.bzl b/third_party/repo.bzl index aa178fa8cab92d9d299e5ed09927d8572816a0af..36f5aa5bdee43a511abf5634af85643ac7e11cfc 100644 --- a/third_party/repo.bzl +++ b/third_party/repo.bzl @@ -17,6 +17,7 @@ _SINGLE_URL_WHITELIST = depset([ "arm_compiler", "ortools_archive", + "gemmlowp", ]) def _is_windows(ctx): @@ -68,7 +69,7 @@ def _apply_delete(ctx, paths): _execute_and_check_ret_code(ctx, cmd) def _tf_http_archive(ctx): - if ("mirror.bazel.build" not in ctx.attr.urls[0] or + if ("mirror.bazel.build" not in ctx.attr.urls[0] and (len(ctx.attr.urls) < 2 and ctx.attr.name not in _SINGLE_URL_WHITELIST)): fail("tf_http_archive(urls) must have redundant URLs. The " +