tensorflow · copybara-service · May 13, 2021 · Mar 5, 2021 · May 5, 2021
diff --git a/tensorflow/core/kernels/image/crop_and_resize_op.cc b/tensorflow/core/kernels/image/crop_and_resize_op.cc
@@ -33,6 +33,7 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/determinism.h"
 #include "tensorflow/core/util/work_sharder.h"
 
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
@@ -189,7 +190,7 @@ class CropAndResizeOp : public AsyncOpKernel {
 
       if (!status) {
         context->SetStatus(
-            errors::Internal("Failed launch CropAndResizeKernel."));
+            errors::Internal("Failed to launch CropAndResizeKernel."));
       }
     };
 
@@ -407,6 +408,14 @@ class CropAndResizeGradImageOp : public AsyncOpKernel {
         context, grads.dim_size(3) == depth,
         errors::InvalidArgument("image_size and grads are incompatible"), done);
 
+    if (std::is_same<Device, GPUDevice>::value) {
+      OP_REQUIRES_ASYNC(
+          context, !OpDeterminismRequired(), errors::Unimplemented(
+              "Deterministic GPU implementation of CropAndResizeBackpropImage"
+              " not available."),
+          done);
+    }
+
     // Allocate output tensor.
     Tensor* output = nullptr;
     OP_REQUIRES_OK_ASYNC(
@@ -426,7 +435,7 @@ class CropAndResizeGradImageOp : public AsyncOpKernel {
 
       if (!status) {
         context->SetStatus(errors::Internal(
-            "Failed launch CropAndResizeBackpropImage kernel."));
+            "Failed to launch CropAndResizeBackpropImage kernel."));
       }
     };
 
@@ -545,8 +554,13 @@ struct CropAndResizeBackpropImage<CPUDevice, T> {
 
     const DeviceBase::CpuWorkerThreads& worker_threads =
         *(context->device()->tensorflow_cpu_worker_threads());
-    Shard(worker_threads.num_threads, worker_threads.workers, num_boxes,
-          cost_per_box, CropAndResizeBackImgPerBox);
+
+    // Sharding introduces nondeterminism when the gradients associated with
+    // more than two crops backprop into the same element in the source image.
+    int max_threads = OpDeterminismRequired() ? 1 : worker_threads.num_threads;
+
+    Shard(max_threads, worker_threads.workers, num_boxes, cost_per_box,
+          CropAndResizeBackImgPerBox);
 
     return true;
   }
@@ -610,6 +624,14 @@ class CropAndResizeGradBoxesOp : public AsyncOpKernel {
         errors::InvalidArgument("boxes and grads have incompatible shape"),
         done);
 
+    if (std::is_same<Device, GPUDevice>::value) {
+      OP_REQUIRES_ASYNC(
+          context, !OpDeterminismRequired(), errors::Unimplemented(
+              "Deterministic GPU implementation of CropAndResizeBackpropBoxes"
+              " not available."),
+          done);
+    }
+
     // Allocate output tensor.
     Tensor* output = nullptr;
     OP_REQUIRES_OK_ASYNC(
@@ -628,7 +650,7 @@ class CropAndResizeGradBoxesOp : public AsyncOpKernel {
           box_index.tensor<int32, 1>(), output->tensor<float, 2>());
       if (!status) {
         context->SetStatus(errors::Internal(
-            "Failed launch CropAndResizeBackpropBoxes kernel."));
+            "Failed to launch CropAndResizeBackpropBoxes kernel."));
       }
     };
 

diff --git a/tensorflow/core/util/BUILD b/tensorflow/core/util/BUILD
@@ -147,6 +147,7 @@ filegroup(
         "bcast.h",
         "command_line_flags.h",
         "debug_events_writer.h",
+        "determinism.h",
         "device_name_utils.h",
         "dump_graph.h",
         "einsum_op_util.h",
@@ -214,6 +215,7 @@ filegroup(
         "bcast.cc",
         "command_line_flags.cc",
         "debug_events_writer.cc",
+        "determinism.cc",
         "device_name_utils.cc",
         "dump_graph.cc",
         "equal_graph_def.cc",

diff --git a/tensorflow/core/util/determinism.cc b/tensorflow/core/util/determinism.cc
@@ -0,0 +1,32 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES 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/util/determinism.h"
+#include "tensorflow/core/util/env_var.h"
+
+namespace tensorflow {
+
+bool OpDeterminismRequired() {
+  static bool op_determinism_required = [] {
+    bool deterministic_ops = false;
+    TF_CHECK_OK(tensorflow::ReadBoolFromEnvVar("TF_DETERMINISTIC_OPS",
+                                               /*default_val=*/false,
+                                               &deterministic_ops));
+    return deterministic_ops;
+  }();
+  return op_determinism_required;
+}
+
+}  // namespace tensorflow
diff --git a/tensorflow/core/util/determinism.h b/tensorflow/core/util/determinism.h
@@ -0,0 +1,27 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT 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_CORE_UTIL_DETERMINISM_H_
+#define TENSORFLOW_CORE_UTIL_DETERMINISM_H_
+
+namespace tensorflow {
+
+// TODO(duncanriach): use this to replace existing instances of
+// RequireDeterminism in various core/kernels.
+bool OpDeterminismRequired();
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_UTIL_DETERMINISM_H_
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
@@ -2925,6 +2925,7 @@ py_library(
     srcs_version = "PY3",
     deps = [
         ":client_testlib",
+        ":errors",
         ":gradients",
         ":image_ops",
         "//tensorflow/python/framework:for_generated_wrappers",

diff --git a/tensorflow/python/kernel_tests/sparse_xent_op_deterministic_test.py b/tensorflow/python/kernel_tests/sparse_xent_op_deterministic_test.py
@@ -65,7 +65,13 @@ def testExceptionThrowing(self):
 
 
 class SparseSoftmaxCrossEntropyWithLogitsDeterministicTest(test.TestCase):
-  """Test that SparseSoftmaxCrossEntropyWithLogits operates reproducibly."""
+  """Test that SparseSoftmaxCrossEntropyWithLogits operates reproducibly.
+
+  Note that the deterministic functionality currently tested by this class is
+  always activated (not enabled by TF_DETERMINISTIC_OPS), so this class does not
+  currently need to inherit from a base op test class (to ensure that the op
+  still functions correctly when determinism is enabled).
+  """
 
   def _randomInts(self, shape, high, dtype):
     return constant_op.constant(

diff --git a/tensorflow/python/kernel_tests/xent_op_deterministic_test.py b/tensorflow/python/kernel_tests/xent_op_deterministic_test.py
@@ -62,7 +62,13 @@ def testExceptionThrowing(self):
 
 
 class SoftmaxCrossEntropyWithLogitsDeterministicTest(test.TestCase):
-  """Test that SoftmaxCrossEntropyWithLogits operates reproducibly."""
+  """Test that SoftmaxCrossEntropyWithLogits operates reproducibly.
+
+  Note that the deterministic functionality currently tested by this class is
+  always activated (not enabled by TF_DETERMINISTIC_OPS), so this class does not
+  currently need to inherit from a base op test class (to ensure that the op
+  still functions correctly when determinism is enabled).
+  """
 
   def _randomFloats(self, shape, dtype, normalized_rows=False):
     a = (2 * np.random.random_sample(shape) - 1).astype(dtype)