Release MobileDet code and model, and require tf_slim installation for OD API. (#8562)

* Merged commit includes the following changes:
311933687  by Sergio Guadarrama:

    Removes spurios use of tf.compat.v2, which results in spurious tf.compat.v1.compat.v2. Adds basic test to nasnet_utils.
    Replaces all remaining import tensorflow as tf with import tensorflow.compat.v1 as tf

--
311766063  by Sergio Guadarrama:

    Removes explicit tf.compat.v1 in all call sites (we already import tf.compat.v1, so this code was  doing tf.compat.v1.compat.v1). The existing code worked in latest version of tensorflow, 2.2, (and 1.15) but not in 1.14 or in 2.0.0a, this CL fixes it.

--
311624958  by Sergio Guadarrama:

    Updates README that doesn't render properly in github documentation

--
310980959  by Sergio Guadarrama:

    Moves research_models/slim off tf.contrib.slim/layers/framework to tf_slim

--
310263156  by Sergio Guadarrama:

    Adds model breakdown for MobilenetV3

--
308640516  by Sergio Guadarrama:

    Internal change

308244396  by Sergio Guadarrama:

    GroupNormalization support for MobilenetV3.

--
307475800  by Sergio Guadarrama:

    Internal change

--
302077708  by Sergio Guadarrama:

    Remove `disable_tf2` behavior from slim py_library targets

--
301208453  by Sergio Guadarrama:

    Automated refactoring to make code Python 3 compatible.

--
300816672  by Sergio Guadarrama:

    Internal change

299433840  by Sergio Guadarrama:

    Internal change

299221609  by Sergio Guadarrama:

    Explicitly disable Tensorflow v2 behaviors for all TF1.x binaries and tests

--
299179617  by Sergio Guadarrama:

    Internal change

299040784  by Sergio Guadarrama:

    Internal change

299036699  by Sergio Guadarrama:

    Internal change

298736510  by Sergio Guadarrama:

    Internal change

298732599  by Sergio Guadarrama:

    Internal change

298729507  by Sergio Guadarrama:

    Internal change

298253328  by Sergio Guadarrama:

    Internal change

297788346  by Sergio Guadarrama:

    Internal change

297785278  by Sergio Guadarrama:

    Internal change

297783127  by Sergio Guadarrama:

    Internal change

297725870  by Sergio Guadarrama:

    Internal change

297721811  by Sergio Guadarrama:

    Internal change

297711347  by Sergio Guadarrama:

    Internal change

297708059  by Sergio Guadarrama:

    Internal change

297701831  by Sergio Guadarrama:

    Internal change

297700038  by Sergio Guadarrama:

    Internal change

297670468  by Sergio Guadarrama:

    Internal change.

--
297350326  by Sergio Guadarrama:

    Explicitly replace "import tensorflow" with "tensorflow.compat.v1" for TF2.x migration

--
297201668  by Sergio Guadarrama:

    Explicitly replace "import tensorflow" with "tensorflow.compat.v1" for TF2.x migration

--
294483372  by Sergio Guadarrama:

    Internal change

PiperOrigin-RevId: 311933687

* Merged commit includes the following changes:
312578615  by Menglong Zhu:

    Modify the LSTM feature extractors to be python 3 compatible.

--
311264357  by Menglong Zhu:

    Removes contrib.slim

--
308957207  by Menglong Zhu:

    Automated refactoring to make code Python 3 compatible.

--
306976470  by yongzhe:

    Internal change

306777559  by Menglong Zhu:

    Internal change

--
299232507  by lzyuan:

    Internal update.

--
299221735  by lzyuan:

    Add small epsilon on max_range for quantize_op to prevent range collapse.

--

PiperOrigin-RevId: 312578615

* Merged commit includes the following changes:
310447280  by lzc:

    Internal changes.

--

PiperOrigin-RevId: 310447280
Co-authored-by: Sergio Guadarrama <sguada@google.com>
Co-authored-by: Menglong Zhu <menglong@google.com>

research/object_detection/utils/object_detection_evaluation_test.py

@@ -22,7 +22,7 @@ from absl.testing import parameterized
 import numpy as np
 import six
 from six.moves import range
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from object_detection import eval_util
 from object_detection.core import standard_fields
 from object_detection.utils import object_detection_evaluation

research/object_detection/utils/ops.py

@@ -24,20 +24,13 @@ import six
 
 from six.moves import range
 from six.moves import zip
-import tensorflow as tf
-
+import tensorflow.compat.v1 as tf
+import tf_slim as slim
 from object_detection.core import standard_fields as fields
 from object_detection.utils import shape_utils
 from object_detection.utils import spatial_transform_ops as spatial_ops
 from object_detection.utils import static_shape
 
-# pylint: disable=g-import-not-at-top
-try:
-  from tensorflow.contrib import framework as contrib_framework
-except ImportError:
-  # TF 2.0 doesn't ship with contrib.
-  pass
-# pylint: enable=g-import-not-at-top
 
 matmul_crop_and_resize = spatial_ops.matmul_crop_and_resize
 multilevel_roi_align = spatial_ops.multilevel_roi_align
@@ -595,7 +588,7 @@ def normalize_to_target(inputs,
       initial_norm = depth * [target_norm_value]
     else:
       initial_norm = target_norm_value
-    target_norm = contrib_framework.model_variable(
+    target_norm = slim.model_variable(
         name='weights',
         dtype=tf.float32,
         initializer=tf.constant(initial_norm, dtype=tf.float32),
@@ -833,7 +826,10 @@ def reframe_box_masks_to_image_masks(box_masks, boxes, image_height,
       boxes = tf.reshape(boxes, [-1, 2, 2])
       min_corner = tf.expand_dims(reference_boxes[:, 0:2], 1)
       max_corner = tf.expand_dims(reference_boxes[:, 2:4], 1)
-      transformed_boxes = (boxes - min_corner) / (max_corner - min_corner)
+      denom = max_corner - min_corner
+      # Prevent a divide by zero.
+      denom = tf.math.maximum(denom, 1e-4)
+      transformed_boxes = (boxes - min_corner) / denom
       return tf.reshape(transformed_boxes, [-1, 4])
 
     box_masks_expanded = tf.expand_dims(box_masks, axis=3)
@@ -841,6 +837,9 @@ def reframe_box_masks_to_image_masks(box_masks, boxes, image_height,
     unit_boxes = tf.concat(
         [tf.zeros([num_boxes, 2]), tf.ones([num_boxes, 2])], axis=1)
     reverse_boxes = transform_boxes_relative_to_boxes(unit_boxes, boxes)
+
+    # TODO(vighneshb) Use matmul_crop_and_resize so that the output shape
+    # is static. This will help us run and test on TPUs.
     return tf.image.crop_and_resize(
         image=box_masks_expanded,
         boxes=reverse_boxes,
@@ -1042,28 +1041,30 @@ def fpn_feature_levels(num_levels, unit_scale_index, image_ratio, boxes):
   return levels
 
 
-def bfloat16_to_float32_nested(tensor_nested):
+def bfloat16_to_float32_nested(input_nested):
   """Convert float32 tensors in a nested structure to bfloat16.
 
   Args:
-    tensor_nested: A Python dict, values being Tensor or Python list/tuple of
-      Tensor.
+    input_nested: A Python dict, values being Tensor or Python list/tuple of
+      Tensor or Non-Tensor.
 
   Returns:
     A Python dict with the same structure as `tensor_dict`,
     with all bfloat16 tensors converted to float32.
   """
-  if isinstance(tensor_nested, tf.Tensor):
-    if tensor_nested.dtype == tf.bfloat16:
-      return tf.cast(tensor_nested, dtype=tf.float32)
+  if isinstance(input_nested, tf.Tensor):
+    if input_nested.dtype == tf.bfloat16:
+      return tf.cast(input_nested, dtype=tf.float32)
     else:
-      return tensor_nested
-  elif isinstance(tensor_nested, (list, tuple)):
-    out_tensor_dict = [bfloat16_to_float32_nested(t) for t in tensor_nested]
-  elif isinstance(tensor_nested, dict):
+      return input_nested
+  elif isinstance(input_nested, (list, tuple)):
+    out_tensor_dict = [bfloat16_to_float32_nested(t) for t in input_nested]
+  elif isinstance(input_nested, dict):
     out_tensor_dict = {
-        k: bfloat16_to_float32_nested(v) for k, v in tensor_nested.items()
+        k: bfloat16_to_float32_nested(v) for k, v in input_nested.items()
     }
+  else:
+    return input_nested
   return out_tensor_dict
 
 
@@ -1101,3 +1102,28 @@ EqualizationLossConfig = collections.namedtuple('EqualizationLossConfig',
                                                 ['weight', 'exclude_prefixes'])
 
 
+
+
+def tile_context_tensors(tensor_dict):
+  """Tiles context fields to have num_frames along 0-th dimension."""
+
+  num_frames = tf.shape(tensor_dict[fields.InputDataFields.image])[0]
+
+  for key in tensor_dict:
+    if key not in fields.SEQUENCE_FIELDS:
+      original_tensor = tensor_dict[key]
+      tensor_shape = shape_utils.combined_static_and_dynamic_shape(
+          original_tensor)
+      tensor_dict[key] = tf.tile(
+          tf.expand_dims(original_tensor, 0),
+          tf.stack([num_frames] + [1] * len(tensor_shape), axis=0))
+  return tensor_dict
+
+
+def decode_image(tensor_dict):
+  """Decodes images in a tensor dict."""
+
+  tensor_dict[fields.InputDataFields.image] = tf.io.decode_image(
+      tensor_dict[fields.InputDataFields.image], channels=3)
+  tensor_dict[fields.InputDataFields.image].set_shape([None, None, 3])
+  return tensor_dict

research/object_detection/utils/ops_test.py

@@ -21,80 +21,75 @@ from __future__ import print_function
 import numpy as np
 import six
 from six.moves import range
-import tensorflow as tf
-
+import tensorflow.compat.v1 as tf
+import tf_slim as slim
 from object_detection.core import standard_fields as fields
 from object_detection.utils import ops
 from object_detection.utils import test_case
 
-# pylint: disable=g-import-not-at-top
-try:
-  from tensorflow.contrib import framework as contrib_framework
-  from tensorflow.contrib import slim
-except ImportError:
-  # TF 2.0 doesn't ship with contrib.
-  pass
-# pylint: enable=g-import-not-at-top
-
 
-class NormalizedToImageCoordinatesTest(tf.test.TestCase):
+class NormalizedToImageCoordinatesTest(test_case.TestCase):
 
   def test_normalized_to_image_coordinates(self):
-    normalized_boxes = tf.placeholder(tf.float32, shape=(None, 1, 4))
     normalized_boxes_np = np.array([[[0.0, 0.0, 1.0, 1.0]],
                                     [[0.5, 0.5, 1.0, 1.0]]])
-    image_shape = tf.convert_to_tensor([1, 4, 4, 3], dtype=tf.int32)
-    absolute_boxes = ops.normalized_to_image_coordinates(normalized_boxes,
-                                                         image_shape,
-                                                         parallel_iterations=2)
+
+    def graph_fn(normalized_boxes):
+      image_shape = tf.convert_to_tensor([1, 4, 4, 3], dtype=tf.int32)
+      absolute_boxes = ops.normalized_to_image_coordinates(
+          normalized_boxes, image_shape, parallel_iterations=2)
+      return absolute_boxes
 
     expected_boxes = np.array([[[0, 0, 4, 4]],
                                [[2, 2, 4, 4]]])
-    with self.test_session() as sess:
-      absolute_boxes = sess.run(absolute_boxes,
-                                feed_dict={normalized_boxes:
-                                           normalized_boxes_np})
 
+    absolute_boxes = self.execute(graph_fn, [normalized_boxes_np])
     self.assertAllEqual(absolute_boxes, expected_boxes)
 
 
-class ReduceSumTrailingDimensions(tf.test.TestCase):
+class ReduceSumTrailingDimensions(test_case.TestCase):
 
   def test_reduce_sum_trailing_dimensions(self):
-    input_tensor = tf.placeholder(tf.float32, shape=[None, None, None])
-    reduced_tensor = ops.reduce_sum_trailing_dimensions(input_tensor, ndims=2)
-    with self.test_session() as sess:
-      reduced_np = sess.run(reduced_tensor,
-                            feed_dict={input_tensor: np.ones((2, 2, 2),
-                                                             np.float32)})
+
+    def graph_fn(input_tensor):
+      reduced_tensor = ops.reduce_sum_trailing_dimensions(input_tensor, ndims=2)
+      return reduced_tensor
+
+    reduced_np = self.execute(graph_fn, [np.ones((2, 2, 2), np.float32)])
     self.assertAllClose(reduced_np, 2 * np.ones((2, 2), np.float32))
 
 
-class MeshgridTest(tf.test.TestCase):
+class MeshgridTest(test_case.TestCase):
 
   def test_meshgrid_numpy_comparison(self):
     """Tests meshgrid op with vectors, for which it should match numpy."""
+
     x = np.arange(4)
     y = np.arange(6)
+
+    def graph_fn():
+      xgrid, ygrid = ops.meshgrid(x, y)
+      return xgrid, ygrid
+
     exp_xgrid, exp_ygrid = np.meshgrid(x, y)
-    xgrid, ygrid = ops.meshgrid(x, y)
-    with self.test_session() as sess:
-      xgrid_output, ygrid_output = sess.run([xgrid, ygrid])
-      self.assertAllEqual(xgrid_output, exp_xgrid)
-      self.assertAllEqual(ygrid_output, exp_ygrid)
+    xgrid_output, ygrid_output = self.execute(graph_fn, [])
+    self.assertAllEqual(xgrid_output, exp_xgrid)
+    self.assertAllEqual(ygrid_output, exp_ygrid)
 
   def test_meshgrid_multidimensional(self):
     np.random.seed(18)
     x = np.random.rand(4, 1, 2).astype(np.float32)
     y = np.random.rand(2, 3).astype(np.float32)
 
-    xgrid, ygrid = ops.meshgrid(x, y)
-
     grid_shape = list(y.shape) + list(x.shape)
-    self.assertEqual(xgrid.get_shape().as_list(), grid_shape)
-    self.assertEqual(ygrid.get_shape().as_list(), grid_shape)
-    with self.test_session() as sess:
-      xgrid_output, ygrid_output = sess.run([xgrid, ygrid])
+
+    def graph_fn():
+      xgrid, ygrid = ops.meshgrid(x, y)
+      self.assertEqual(xgrid.get_shape().as_list(), grid_shape)
+      self.assertEqual(ygrid.get_shape().as_list(), grid_shape)
+      return xgrid, ygrid
+
+    xgrid_output, ygrid_output = self.execute(graph_fn, [])
 
     # Check the shape of the output grids
     self.assertEqual(xgrid_output.shape, tuple(grid_shape))
@@ -111,110 +106,150 @@ class MeshgridTest(tf.test.TestCase):
       self.assertEqual(ygrid_output[yind + xind], y[yind])
 
 
-class OpsTestFixedPadding(tf.test.TestCase):
+class OpsTestFixedPadding(test_case.TestCase):
 
   def test_3x3_kernel(self):
-    tensor = tf.constant([[[[0.], [0.]], [[0.], [0.]]]])
-    padded_tensor = ops.fixed_padding(tensor, 3)
-    with self.test_session() as sess:
-      padded_tensor_out = sess.run(padded_tensor)
+
+    def graph_fn():
+      tensor = tf.constant([[[[0.], [0.]], [[0.], [0.]]]])
+      padded_tensor = ops.fixed_padding(tensor, 3)
+      return padded_tensor
+
+    padded_tensor_out = self.execute(graph_fn, [])
     self.assertEqual((1, 4, 4, 1), padded_tensor_out.shape)
 
   def test_5x5_kernel(self):
-    tensor = tf.constant([[[[0.], [0.]], [[0.], [0.]]]])
-    padded_tensor = ops.fixed_padding(tensor, 5)
-    with self.test_session() as sess:
-      padded_tensor_out = sess.run(padded_tensor)
+
+    def graph_fn():
+      tensor = tf.constant([[[[0.], [0.]], [[0.], [0.]]]])
+      padded_tensor = ops.fixed_padding(tensor, 5)
+      return padded_tensor
+
+    padded_tensor_out = self.execute(graph_fn, [])
     self.assertEqual((1, 6, 6, 1), padded_tensor_out.shape)
 
   def test_3x3_atrous_kernel(self):
-    tensor = tf.constant([[[[0.], [0.]], [[0.], [0.]]]])
-    padded_tensor = ops.fixed_padding(tensor, 3, 2)
-    with self.test_session() as sess:
-      padded_tensor_out = sess.run(padded_tensor)
+
+    def graph_fn():
+      tensor = tf.constant([[[[0.], [0.]], [[0.], [0.]]]])
+      padded_tensor = ops.fixed_padding(tensor, 3, 2)
+      return padded_tensor
+
+    padded_tensor_out = self.execute(graph_fn, [])
     self.assertEqual((1, 6, 6, 1), padded_tensor_out.shape)
 
 
-class OpsTestPadToMultiple(tf.test.TestCase):
+class OpsTestPadToMultiple(test_case.TestCase):
 
   def test_zero_padding(self):
-    tensor = tf.constant([[[[0.], [0.]], [[0.], [0.]]]])
-    padded_tensor = ops.pad_to_multiple(tensor, 1)
-    with self.test_session() as sess:
-      padded_tensor_out = sess.run(padded_tensor)
+
+    def graph_fn():
+      tensor = tf.constant([[[[0.], [0.]], [[0.], [0.]]]])
+      padded_tensor = ops.pad_to_multiple(tensor, 1)
+      return padded_tensor
+
+    padded_tensor_out = self.execute(graph_fn, [])
     self.assertEqual((1, 2, 2, 1), padded_tensor_out.shape)
 
   def test_no_padding(self):
-    tensor = tf.constant([[[[0.], [0.]], [[0.], [0.]]]])
-    padded_tensor = ops.pad_to_multiple(tensor, 2)
-    with self.test_session() as sess:
-      padded_tensor_out = sess.run(padded_tensor)
+
+    def graph_fn():
+      tensor = tf.constant([[[[0.], [0.]], [[0.], [0.]]]])
+      padded_tensor = ops.pad_to_multiple(tensor, 2)
+      return padded_tensor
+
+    padded_tensor_out = self.execute(graph_fn, [])
     self.assertEqual((1, 2, 2, 1), padded_tensor_out.shape)
 
   def test_non_square_padding(self):
-    tensor = tf.constant([[[[0.], [0.]]]])
-    padded_tensor = ops.pad_to_multiple(tensor, 2)
-    with self.test_session() as sess:
-      padded_tensor_out = sess.run(padded_tensor)
+
+    def graph_fn():
+      tensor = tf.constant([[[[0.], [0.]]]])
+      padded_tensor = ops.pad_to_multiple(tensor, 2)
+      return padded_tensor
+
+    padded_tensor_out = self.execute(graph_fn, [])
     self.assertEqual((1, 2, 2, 1), padded_tensor_out.shape)
 
   def test_padding(self):
-    tensor = tf.constant([[[[0.], [0.]], [[0.], [0.]]]])
-    padded_tensor = ops.pad_to_multiple(tensor, 4)
-    with self.test_session() as sess:
-      padded_tensor_out = sess.run(padded_tensor)
+
+    def graph_fn():
+      tensor = tf.constant([[[[0.], [0.]], [[0.], [0.]]]])
+      padded_tensor = ops.pad_to_multiple(tensor, 4)
+      return padded_tensor
+
+    padded_tensor_out = self.execute(graph_fn, [])
     self.assertEqual((1, 4, 4, 1), padded_tensor_out.shape)
 
 
-class OpsTestPaddedOneHotEncoding(tf.test.TestCase):
+class OpsTestPaddedOneHotEncoding(test_case.TestCase):
 
   def test_correct_one_hot_tensor_with_no_pad(self):
-    indices = tf.constant([1, 2, 3, 5])
-    one_hot_tensor = ops.padded_one_hot_encoding(indices, depth=6, left_pad=0)
+
+    def graph_fn():
+      indices = tf.constant([1, 2, 3, 5])
+      one_hot_tensor = ops.padded_one_hot_encoding(indices, depth=6, left_pad=0)
+      return one_hot_tensor
+
     expected_tensor = np.array([[0, 1, 0, 0, 0, 0],
                                 [0, 0, 1, 0, 0, 0],
                                 [0, 0, 0, 1, 0, 0],
                                 [0, 0, 0, 0, 0, 1]], np.float32)
-    with self.test_session() as sess:
-      out_one_hot_tensor = sess.run(one_hot_tensor)
-      self.assertAllClose(out_one_hot_tensor, expected_tensor, rtol=1e-10,
-                          atol=1e-10)
+
+    # Executing on CPU only because output shape is not constant.
+    out_one_hot_tensor = self.execute_cpu(graph_fn, [])
+    self.assertAllClose(out_one_hot_tensor, expected_tensor, rtol=1e-10,
+                        atol=1e-10)
 
   def test_correct_one_hot_tensor_with_pad_one(self):
-    indices = tf.constant([1, 2, 3, 5])
-    one_hot_tensor = ops.padded_one_hot_encoding(indices, depth=6, left_pad=1)
+
+    def graph_fn():
+      indices = tf.constant([1, 2, 3, 5])
+      one_hot_tensor = ops.padded_one_hot_encoding(indices, depth=6, left_pad=1)
+      return one_hot_tensor
+
     expected_tensor = np.array([[0, 0, 1, 0, 0, 0, 0],
                                 [0, 0, 0, 1, 0, 0, 0],
                                 [0, 0, 0, 0, 1, 0, 0],
                                 [0, 0, 0, 0, 0, 0, 1]], np.float32)
-    with self.test_session() as sess:
-      out_one_hot_tensor = sess.run(one_hot_tensor)
-      self.assertAllClose(out_one_hot_tensor, expected_tensor, rtol=1e-10,
-                          atol=1e-10)
+    # Executing on CPU only because output shape is not constant.
+    out_one_hot_tensor = self.execute_cpu(graph_fn, [])
+    self.assertAllClose(out_one_hot_tensor, expected_tensor, rtol=1e-10,
+                        atol=1e-10)
 
   def test_correct_one_hot_tensor_with_pad_three(self):
-    indices = tf.constant([1, 2, 3, 5])
-    one_hot_tensor = ops.padded_one_hot_encoding(indices, depth=6, left_pad=3)
+
+    def graph_fn():
+      indices = tf.constant([1, 2, 3, 5])
+      one_hot_tensor = ops.padded_one_hot_encoding(indices, depth=6, left_pad=3)
+      return one_hot_tensor
+
     expected_tensor = np.array([[0, 0, 0, 0, 1, 0, 0, 0, 0],
                                 [0, 0, 0, 0, 0, 1, 0, 0, 0],
                                 [0, 0, 0, 0, 0, 0, 1, 0, 0],
                                 [0, 0, 0, 0, 0, 0, 0, 0, 1]], np.float32)
-    with self.test_session() as sess:
-      out_one_hot_tensor = sess.run(one_hot_tensor)
-      self.assertAllClose(out_one_hot_tensor, expected_tensor, rtol=1e-10,
-                          atol=1e-10)
+
+    # executing on CPU only because output shape is not constant.
+    out_one_hot_tensor = self.execute_cpu(graph_fn, [])
+    self.assertAllClose(out_one_hot_tensor, expected_tensor, rtol=1e-10,
+                        atol=1e-10)
 
   def test_correct_padded_one_hot_tensor_with_empty_indices(self):
+
     depth = 6
     pad = 2
-    indices = tf.constant([])
-    one_hot_tensor = ops.padded_one_hot_encoding(
-        indices, depth=depth, left_pad=pad)
+
+    def graph_fn():
+      indices = tf.constant([])
+      one_hot_tensor = ops.padded_one_hot_encoding(
+          indices, depth=depth, left_pad=pad)
+      return one_hot_tensor
+
     expected_tensor = np.zeros((0, depth + pad))
-    with self.test_session() as sess:
-      out_one_hot_tensor = sess.run(one_hot_tensor)
-      self.assertAllClose(out_one_hot_tensor, expected_tensor, rtol=1e-10,
-                          atol=1e-10)
+    # executing on CPU only because output shape is not constant.
+    out_one_hot_tensor = self.execute_cpu(graph_fn, [])
+    self.assertAllClose(out_one_hot_tensor, expected_tensor, rtol=1e-10,
+                        atol=1e-10)
 
   def test_return_none_on_zero_depth(self):
     indices = tf.constant([1, 2, 3, 4, 5])
@@ -242,28 +277,27 @@ class OpsTestPaddedOneHotEncoding(tf.test.TestCase):
       ops.padded_one_hot_encoding(indices, depth=0.1, left_pad=2)
 
 
-class OpsDenseToSparseBoxesTest(tf.test.TestCase):
+class OpsDenseToSparseBoxesTest(test_case.TestCase):
 
   def test_return_all_boxes_when_all_input_boxes_are_valid(self):
     num_classes = 4
     num_valid_boxes = 3
     code_size = 4
-    dense_location_placeholder = tf.placeholder(tf.float32,
-                                                shape=(num_valid_boxes,
-                                                       code_size))
-    dense_num_boxes_placeholder = tf.placeholder(tf.int32, shape=(num_classes))
-    box_locations, box_classes = ops.dense_to_sparse_boxes(
-        dense_location_placeholder, dense_num_boxes_placeholder, num_classes)
-    feed_dict = {dense_location_placeholder: np.random.uniform(
-        size=[num_valid_boxes, code_size]),
-                 dense_num_boxes_placeholder: np.array([1, 0, 0, 2],
-                                                       dtype=np.int32)}
-
-    expected_box_locations = feed_dict[dense_location_placeholder]
+
+    def graph_fn(dense_location, dense_num_boxes):
+      box_locations, box_classes = ops.dense_to_sparse_boxes(
+          dense_location, dense_num_boxes, num_classes)
+      return box_locations, box_classes
+
+    dense_location_np = np.random.uniform(size=[num_valid_boxes, code_size])
+    dense_num_boxes_np = np.array([1, 0, 0, 2], dtype=np.int32)
+
+    expected_box_locations = dense_location_np
     expected_box_classses = np.array([0, 3, 3])
-    with self.test_session() as sess:
-      box_locations, box_classes = sess.run([box_locations, box_classes],
-                                            feed_dict=feed_dict)
+
+    # Executing on CPU only since output shape is not constant.
+    box_locations, box_classes = self.execute_cpu(
+        graph_fn, [dense_location_np, dense_num_boxes_np])
 
     self.assertAllClose(box_locations, expected_box_locations, rtol=1e-6,
                         atol=1e-6)
@@ -275,29 +309,27 @@ class OpsDenseToSparseBoxesTest(tf.test.TestCase):
     num_boxes = 10
     code_size = 4
 
-    dense_location_placeholder = tf.placeholder(tf.float32, shape=(num_boxes,
-                                                                   code_size))
-    dense_num_boxes_placeholder = tf.placeholder(tf.int32, shape=(num_classes))
-    box_locations, box_classes = ops.dense_to_sparse_boxes(
-        dense_location_placeholder, dense_num_boxes_placeholder, num_classes)
-    feed_dict = {dense_location_placeholder: np.random.uniform(
-        size=[num_boxes, code_size]),
-                 dense_num_boxes_placeholder: np.array([1, 0, 0, 2],
-                                                       dtype=np.int32)}
-
-    expected_box_locations = (feed_dict[dense_location_placeholder]
-                              [:num_valid_boxes])
+    def graph_fn(dense_location, dense_num_boxes):
+      box_locations, box_classes = ops.dense_to_sparse_boxes(
+          dense_location, dense_num_boxes, num_classes)
+      return box_locations, box_classes
+
+    dense_location_np = np.random.uniform(size=[num_boxes, code_size])
+    dense_num_boxes_np = np.array([1, 0, 0, 2], dtype=np.int32)
+
+    expected_box_locations = dense_location_np[:num_valid_boxes]
     expected_box_classses = np.array([0, 3, 3])
-    with self.test_session() as sess:
-      box_locations, box_classes = sess.run([box_locations, box_classes],
-                                            feed_dict=feed_dict)
+
+    # Executing on CPU only since output shape is not constant.
+    box_locations, box_classes = self.execute_cpu(
+        graph_fn, [dense_location_np, dense_num_boxes_np])
 
     self.assertAllClose(box_locations, expected_box_locations, rtol=1e-6,
                         atol=1e-6)
     self.assertAllEqual(box_classes, expected_box_classses)
 
 
-class OpsTestIndicesToDenseVector(tf.test.TestCase):
+class OpsTestIndicesToDenseVector(test_case.TestCase):
 
   def test_indices_to_dense_vector(self):
     size = 10000
@@ -307,13 +339,14 @@ class OpsTestIndicesToDenseVector(tf.test.TestCase):
     expected_output = np.zeros(size, dtype=np.float32)
     expected_output[rand_indices] = 1.
 
-    tf_rand_indices = tf.constant(rand_indices)
-    indicator = ops.indices_to_dense_vector(tf_rand_indices, size)
+    def graph_fn():
+      tf_rand_indices = tf.constant(rand_indices)
+      indicator = ops.indices_to_dense_vector(tf_rand_indices, size)
+      return indicator
 
-    with self.test_session() as sess:
-      output = sess.run(indicator)
-      self.assertAllEqual(output, expected_output)
-      self.assertEqual(output.dtype, expected_output.dtype)
+    output = self.execute(graph_fn, [])
+    self.assertAllEqual(output, expected_output)
+    self.assertEqual(output.dtype, expected_output.dtype)
 
   def test_indices_to_dense_vector_size_at_inference(self):
     size = 5000
@@ -324,16 +357,15 @@ class OpsTestIndicesToDenseVector(tf.test.TestCase):
     expected_output = np.zeros(size, dtype=np.float32)
     expected_output[rand_indices] = 1.
 
-    tf_all_indices = tf.placeholder(tf.int32)
-    tf_rand_indices = tf.constant(rand_indices)
-    indicator = ops.indices_to_dense_vector(tf_rand_indices,
-                                            tf.shape(tf_all_indices)[0])
-    feed_dict = {tf_all_indices: all_indices}
+    def graph_fn(tf_all_indices):
+      tf_rand_indices = tf.constant(rand_indices)
+      indicator = ops.indices_to_dense_vector(tf_rand_indices,
+                                              tf.shape(tf_all_indices)[0])
+      return indicator
 
-    with self.test_session() as sess:
-      output = sess.run(indicator, feed_dict=feed_dict)
-      self.assertAllEqual(output, expected_output)
-      self.assertEqual(output.dtype, expected_output.dtype)
+    output = self.execute(graph_fn, [all_indices])
+    self.assertAllEqual(output, expected_output)
+    self.assertEqual(output.dtype, expected_output.dtype)
 
   def test_indices_to_dense_vector_int(self):
     size = 500
@@ -343,14 +375,15 @@ class OpsTestIndicesToDenseVector(tf.test.TestCase):
     expected_output = np.zeros(size, dtype=np.int64)
     expected_output[rand_indices] = 1
 
-    tf_rand_indices = tf.constant(rand_indices)
-    indicator = ops.indices_to_dense_vector(
-        tf_rand_indices, size, 1, dtype=tf.int64)
+    def graph_fn():
+      tf_rand_indices = tf.constant(rand_indices)
+      indicator = ops.indices_to_dense_vector(
+          tf_rand_indices, size, 1, dtype=tf.int64)
+      return indicator
 
-    with self.test_session() as sess:
-      output = sess.run(indicator)
-      self.assertAllEqual(output, expected_output)
-      self.assertEqual(output.dtype, expected_output.dtype)
+    output = self.execute(graph_fn, [])
+    self.assertAllEqual(output, expected_output)
+    self.assertEqual(output.dtype, expected_output.dtype)
 
   def test_indices_to_dense_vector_custom_values(self):
     size = 100
@@ -362,17 +395,18 @@ class OpsTestIndicesToDenseVector(tf.test.TestCase):
     expected_output = np.float32(np.ones(size) * default_value)
     expected_output[rand_indices] = indices_value
 
-    tf_rand_indices = tf.constant(rand_indices)
-    indicator = ops.indices_to_dense_vector(
-        tf_rand_indices,
-        size,
-        indices_value=indices_value,
-        default_value=default_value)
+    def graph_fn():
+      tf_rand_indices = tf.constant(rand_indices)
+      indicator = ops.indices_to_dense_vector(
+          tf_rand_indices,
+          size,
+          indices_value=indices_value,
+          default_value=default_value)
+      return indicator
 
-    with self.test_session() as sess:
-      output = sess.run(indicator)
-      self.assertAllClose(output, expected_output)
-      self.assertEqual(output.dtype, expected_output.dtype)
+    output = self.execute(graph_fn, [])
+    self.assertAllClose(output, expected_output)
+    self.assertEqual(output.dtype, expected_output.dtype)
 
   def test_indices_to_dense_vector_all_indices_as_input(self):
     size = 500
@@ -381,13 +415,14 @@ class OpsTestIndicesToDenseVector(tf.test.TestCase):
 
     expected_output = np.ones(size, dtype=np.float32)
 
-    tf_rand_indices = tf.constant(rand_indices)
-    indicator = ops.indices_to_dense_vector(tf_rand_indices, size)
+    def graph_fn():
+      tf_rand_indices = tf.constant(rand_indices)
+      indicator = ops.indices_to_dense_vector(tf_rand_indices, size)
+      return indicator
 
-    with self.test_session() as sess:
-      output = sess.run(indicator)
-      self.assertAllEqual(output, expected_output)
-      self.assertEqual(output.dtype, expected_output.dtype)
+    output = self.execute(graph_fn, [])
+    self.assertAllEqual(output, expected_output)
+    self.assertEqual(output.dtype, expected_output.dtype)
 
   def test_indices_to_dense_vector_empty_indices_as_input(self):
     size = 500
@@ -395,55 +430,58 @@ class OpsTestIndicesToDenseVector(tf.test.TestCase):
 
     expected_output = np.zeros(size, dtype=np.float32)
 
-    tf_rand_indices = tf.constant(rand_indices)
-    indicator = ops.indices_to_dense_vector(tf_rand_indices, size)
+    def graph_fn():
+      tf_rand_indices = tf.constant(rand_indices)
+      indicator = ops.indices_to_dense_vector(tf_rand_indices, size)
+      return indicator
 
-    with self.test_session() as sess:
-      output = sess.run(indicator)
-      self.assertAllEqual(output, expected_output)
-      self.assertEqual(output.dtype, expected_output.dtype)
+    output = self.execute(graph_fn, [])
+    self.assertAllEqual(output, expected_output)
+    self.assertEqual(output.dtype, expected_output.dtype)
 
 
-class GroundtruthFilterTest(tf.test.TestCase):
+class GroundtruthFilterTest(test_case.TestCase):
 
   def test_filter_groundtruth(self):
-    input_image = tf.placeholder(tf.float32, shape=(None, None, 3))
-    input_boxes = tf.placeholder(tf.float32, shape=(None, 4))
-    input_classes = tf.placeholder(tf.int32, shape=(None,))
-    input_is_crowd = tf.placeholder(tf.bool, shape=(None,))
-    input_area = tf.placeholder(tf.float32, shape=(None,))
-    input_difficult = tf.placeholder(tf.float32, shape=(None,))
-    input_label_types = tf.placeholder(tf.string, shape=(None,))
-    input_confidences = tf.placeholder(tf.float32, shape=(None,))
-    valid_indices = tf.placeholder(tf.int32, shape=(None,))
-    input_tensors = {
-        fields.InputDataFields.image: input_image,
-        fields.InputDataFields.groundtruth_boxes: input_boxes,
-        fields.InputDataFields.groundtruth_classes: input_classes,
-        fields.InputDataFields.groundtruth_is_crowd: input_is_crowd,
-        fields.InputDataFields.groundtruth_area: input_area,
-        fields.InputDataFields.groundtruth_difficult: input_difficult,
-        fields.InputDataFields.groundtruth_label_types: input_label_types,
-        fields.InputDataFields.groundtruth_confidences: input_confidences,
-    }
-    output_tensors = ops.retain_groundtruth(input_tensors, valid_indices)
-
-    image_tensor = np.random.rand(224, 224, 3)
-    feed_dict = {
-        input_image: image_tensor,
-        input_boxes:
-        np.array([[0.2, 0.4, 0.1, 0.8], [0.2, 0.4, 1.0, 0.8]], dtype=np.float),
-        input_classes: np.array([1, 2], dtype=np.int32),
-        input_is_crowd: np.array([False, True], dtype=np.bool),
-        input_area: np.array([32, 48], dtype=np.float32),
-        input_difficult: np.array([True, False], dtype=np.bool),
-        input_label_types:
-        np.array(['APPROPRIATE', 'INCORRECT'], dtype=np.string_),
-        input_confidences: np.array([0.99, 0.5], dtype=np.float32),
-        valid_indices: np.array([0], dtype=np.int32),
-    }
+
+    def graph_fn(input_image, input_boxes, input_classes, input_is_crowd,
+                 input_area, input_difficult, input_label_types,
+                 input_confidences, valid_indices):
+      input_tensors = {
+          fields.InputDataFields.image: input_image,
+          fields.InputDataFields.groundtruth_boxes: input_boxes,
+          fields.InputDataFields.groundtruth_classes: input_classes,
+          fields.InputDataFields.groundtruth_is_crowd: input_is_crowd,
+          fields.InputDataFields.groundtruth_area: input_area,
+          fields.InputDataFields.groundtruth_difficult: input_difficult,
+          fields.InputDataFields.groundtruth_label_types: input_label_types,
+          fields.InputDataFields.groundtruth_confidences: input_confidences,
+      }
+
+      output_tensors = ops.retain_groundtruth(input_tensors, valid_indices)
+      return output_tensors
+
+    input_image = np.random.rand(224, 224, 3)
+    input_boxes = np.array([[0.2, 0.4, 0.1, 0.8], [0.2, 0.4, 1.0, 0.8]],
+                           dtype=np.float32)
+    input_classes = np.array([1, 2], dtype=np.int32)
+    input_is_crowd = np.array([False, True], dtype=np.bool)
+    input_area = np.array([32, 48], dtype=np.float32)
+    input_difficult = np.array([True, False], dtype=np.bool)
+    input_label_types = np.array(['APPROPRIATE', 'INCORRECT'],
+                                 dtype=np.string_)
+    input_confidences = np.array([0.99, 0.5], dtype=np.float32)
+    valid_indices = np.array([0], dtype=np.int32)
+
+    # Strings are not supported on TPU.
+    output_tensors = self.execute_cpu(
+        graph_fn,
+        [input_image, input_boxes, input_classes, input_is_crowd, input_area,
+         input_difficult, input_label_types, input_confidences, valid_indices]
+    )
+
     expected_tensors = {
-        fields.InputDataFields.image: image_tensor,
+        fields.InputDataFields.image: input_image,
         fields.InputDataFields.groundtruth_boxes: [[0.2, 0.4, 0.1, 0.8]],
         fields.InputDataFields.groundtruth_classes: [1],
         fields.InputDataFields.groundtruth_is_crowd: [False],
@@ -452,35 +490,24 @@ class GroundtruthFilterTest(tf.test.TestCase):
         fields.InputDataFields.groundtruth_label_types: [six.b('APPROPRIATE')],
         fields.InputDataFields.groundtruth_confidences: [0.99],
     }
-    with self.test_session() as sess:
-      output_tensors = sess.run(output_tensors, feed_dict=feed_dict)
-      for key in [fields.InputDataFields.image,
-                  fields.InputDataFields.groundtruth_boxes,
-                  fields.InputDataFields.groundtruth_area,
-                  fields.InputDataFields.groundtruth_confidences]:
-        self.assertAllClose(expected_tensors[key], output_tensors[key])
-      for key in [fields.InputDataFields.groundtruth_classes,
-                  fields.InputDataFields.groundtruth_is_crowd,
-                  fields.InputDataFields.groundtruth_label_types]:
-        self.assertAllEqual(expected_tensors[key], output_tensors[key])
+    for key in [fields.InputDataFields.image,
+                fields.InputDataFields.groundtruth_boxes,
+                fields.InputDataFields.groundtruth_area,
+                fields.InputDataFields.groundtruth_confidences]:
+      self.assertAllClose(expected_tensors[key], output_tensors[key])
+
+    for key in [fields.InputDataFields.groundtruth_classes,
+                fields.InputDataFields.groundtruth_is_crowd,
+                fields.InputDataFields.groundtruth_label_types]:
+      self.assertAllEqual(expected_tensors[key], output_tensors[key])
 
   def test_filter_with_missing_fields(self):
-    input_boxes = tf.placeholder(tf.float32, shape=(None, 4))
-    input_classes = tf.placeholder(tf.int32, shape=(None,))
-    input_tensors = {
-        fields.InputDataFields.groundtruth_boxes: input_boxes,
-        fields.InputDataFields.groundtruth_classes: input_classes
-    }
-    valid_indices = tf.placeholder(tf.int32, shape=(None,))
-
-    feed_dict = {
-        input_boxes:
-        np.array([[0.2, 0.4, 0.1, 0.8], [0.2, 0.4, 1.0, 0.8]], dtype=np.float),
-        input_classes:
-        np.array([1, 2], dtype=np.int32),
-        valid_indices:
-        np.array([0], dtype=np.int32)
-    }
+
+    input_boxes = np.array([[0.2, 0.4, 0.1, 0.8], [0.2, 0.4, 1.0, 0.8]],
+                           dtype=np.float)
+    input_classes = np.array([1, 2], dtype=np.int32)
+    valid_indices = np.array([0], dtype=np.int32)
+
     expected_tensors = {
         fields.InputDataFields.groundtruth_boxes:
         [[0.2, 0.4, 0.1, 0.8]],
@@ -488,42 +515,46 @@ class GroundtruthFilterTest(tf.test.TestCase):
         [1]
     }
 
-    output_tensors = ops.retain_groundtruth(input_tensors, valid_indices)
-    with self.test_session() as sess:
-      output_tensors = sess.run(output_tensors, feed_dict=feed_dict)
-      for key in [fields.InputDataFields.groundtruth_boxes]:
-        self.assertAllClose(expected_tensors[key], output_tensors[key])
-      for key in [fields.InputDataFields.groundtruth_classes]:
-        self.assertAllEqual(expected_tensors[key], output_tensors[key])
+    def graph_fn(input_boxes, input_classes, valid_indices):
+      input_tensors = {
+          fields.InputDataFields.groundtruth_boxes: input_boxes,
+          fields.InputDataFields.groundtruth_classes: input_classes
+      }
+      output_tensors = ops.retain_groundtruth(input_tensors, valid_indices)
+      return output_tensors
+
+    output_tensors = self.execute(graph_fn, [input_boxes, input_classes,
+                                             valid_indices])
+
+    for key in [fields.InputDataFields.groundtruth_boxes]:
+      self.assertAllClose(expected_tensors[key], output_tensors[key])
+    for key in [fields.InputDataFields.groundtruth_classes]:
+      self.assertAllEqual(expected_tensors[key], output_tensors[key])
 
   def test_filter_with_empty_fields(self):
-    input_boxes = tf.placeholder(tf.float32, shape=(None, 4))
-    input_classes = tf.placeholder(tf.int32, shape=(None,))
-    input_is_crowd = tf.placeholder(tf.bool, shape=(None,))
-    input_area = tf.placeholder(tf.float32, shape=(None,))
-    input_difficult = tf.placeholder(tf.float32, shape=(None,))
-    input_confidences = tf.placeholder(tf.float32, shape=(None,))
-    valid_indices = tf.placeholder(tf.int32, shape=(None,))
-    input_tensors = {
-        fields.InputDataFields.groundtruth_boxes: input_boxes,
-        fields.InputDataFields.groundtruth_classes: input_classes,
-        fields.InputDataFields.groundtruth_is_crowd: input_is_crowd,
-        fields.InputDataFields.groundtruth_area: input_area,
-        fields.InputDataFields.groundtruth_difficult: input_difficult,
-        fields.InputDataFields.groundtruth_confidences: input_confidences,
-    }
-    output_tensors = ops.retain_groundtruth(input_tensors, valid_indices)
-
-    feed_dict = {
-        input_boxes:
-        np.array([[0.2, 0.4, 0.1, 0.8], [0.2, 0.4, 1.0, 0.8]], dtype=np.float),
-        input_classes: np.array([1, 2], dtype=np.int32),
-        input_is_crowd: np.array([False, True], dtype=np.bool),
-        input_area: np.array([], dtype=np.float32),
-        input_difficult: np.array([], dtype=np.float32),
-        input_confidences: np.array([0.99, 0.5], dtype=np.float32),
-        valid_indices: np.array([0], dtype=np.int32)
-    }
+
+    def graph_fn(input_boxes, input_classes, input_is_crowd, input_area,
+                 input_difficult, input_confidences, valid_indices):
+      input_tensors = {
+          fields.InputDataFields.groundtruth_boxes: input_boxes,
+          fields.InputDataFields.groundtruth_classes: input_classes,
+          fields.InputDataFields.groundtruth_is_crowd: input_is_crowd,
+          fields.InputDataFields.groundtruth_area: input_area,
+          fields.InputDataFields.groundtruth_difficult: input_difficult,
+          fields.InputDataFields.groundtruth_confidences: input_confidences,
+      }
+      output_tensors = ops.retain_groundtruth(input_tensors, valid_indices)
+      return output_tensors
+
+    input_boxes = np.array([[0.2, 0.4, 0.1, 0.8], [0.2, 0.4, 1.0, 0.8]],
+                           dtype=np.float)
+    input_classes = np.array([1, 2], dtype=np.int32)
+    input_is_crowd = np.array([False, True], dtype=np.bool)
+    input_area = np.array([], dtype=np.float32)
+    input_difficult = np.array([], dtype=np.float32)
+    input_confidences = np.array([0.99, 0.5], dtype=np.float32)
+    valid_indices = np.array([0], dtype=np.int32)
+
     expected_tensors = {
         fields.InputDataFields.groundtruth_boxes: [[0.2, 0.4, 0.1, 0.8]],
         fields.InputDataFields.groundtruth_classes: [1],
@@ -532,92 +563,87 @@ class GroundtruthFilterTest(tf.test.TestCase):
         fields.InputDataFields.groundtruth_difficult: [],
         fields.InputDataFields.groundtruth_confidences: [0.99],
     }
-    with self.test_session() as sess:
-      output_tensors = sess.run(output_tensors, feed_dict=feed_dict)
-      for key in [fields.InputDataFields.groundtruth_boxes,
-                  fields.InputDataFields.groundtruth_area,
-                  fields.InputDataFields.groundtruth_confidences]:
-        self.assertAllClose(expected_tensors[key], output_tensors[key])
-      for key in [fields.InputDataFields.groundtruth_classes,
-                  fields.InputDataFields.groundtruth_is_crowd]:
-        self.assertAllEqual(expected_tensors[key], output_tensors[key])
+    output_tensors = self.execute(graph_fn, [
+        input_boxes, input_classes, input_is_crowd, input_area,
+        input_difficult, input_confidences, valid_indices])
+
+    for key in [fields.InputDataFields.groundtruth_boxes,
+                fields.InputDataFields.groundtruth_area,
+                fields.InputDataFields.groundtruth_confidences]:
+      self.assertAllClose(expected_tensors[key], output_tensors[key])
+    for key in [fields.InputDataFields.groundtruth_classes,
+                fields.InputDataFields.groundtruth_is_crowd]:
+      self.assertAllEqual(expected_tensors[key], output_tensors[key])
 
   def test_filter_with_empty_groundtruth_boxes(self):
-    input_boxes = tf.placeholder(tf.float32, shape=(None, 4))
-    input_classes = tf.placeholder(tf.int32, shape=(None,))
-    input_is_crowd = tf.placeholder(tf.bool, shape=(None,))
-    input_area = tf.placeholder(tf.float32, shape=(None,))
-    input_difficult = tf.placeholder(tf.float32, shape=(None,))
-    input_confidences = tf.placeholder(tf.float32, shape=(None,))
-    valid_indices = tf.placeholder(tf.int32, shape=(None,))
-    input_tensors = {
-        fields.InputDataFields.groundtruth_boxes: input_boxes,
-        fields.InputDataFields.groundtruth_classes: input_classes,
-        fields.InputDataFields.groundtruth_is_crowd: input_is_crowd,
-        fields.InputDataFields.groundtruth_area: input_area,
-        fields.InputDataFields.groundtruth_difficult: input_difficult,
-        fields.InputDataFields.groundtruth_confidences: input_confidences,
-    }
-    output_tensors = ops.retain_groundtruth(input_tensors, valid_indices)
-
-    feed_dict = {
-        input_boxes: np.array([], dtype=np.float).reshape(0, 4),
-        input_classes: np.array([], dtype=np.int32),
-        input_is_crowd: np.array([], dtype=np.bool),
-        input_area: np.array([], dtype=np.float32),
-        input_difficult: np.array([], dtype=np.float32),
-        input_confidences: np.array([], dtype=np.float32),
-        valid_indices: np.array([], dtype=np.int32),
-    }
-    with self.test_session() as sess:
-      output_tensors = sess.run(output_tensors, feed_dict=feed_dict)
-      for key in input_tensors:
-        if key == fields.InputDataFields.groundtruth_boxes:
-          self.assertAllEqual([0, 4], output_tensors[key].shape)
-        else:
-          self.assertAllEqual([0], output_tensors[key].shape)
-
 
-class RetainGroundTruthWithPositiveClasses(tf.test.TestCase):
+    def graph_fn(input_boxes, input_classes, input_is_crowd, input_area,
+                 input_difficult, input_confidences, valid_indices):
+      input_tensors = {
+          fields.InputDataFields.groundtruth_boxes: input_boxes,
+          fields.InputDataFields.groundtruth_classes: input_classes,
+          fields.InputDataFields.groundtruth_is_crowd: input_is_crowd,
+          fields.InputDataFields.groundtruth_area: input_area,
+          fields.InputDataFields.groundtruth_difficult: input_difficult,
+          fields.InputDataFields.groundtruth_confidences: input_confidences,
+      }
+      output_tensors = ops.retain_groundtruth(input_tensors, valid_indices)
+      return output_tensors
+
+    input_boxes = np.array([], dtype=np.float).reshape(0, 4)
+    input_classes = np.array([], dtype=np.int32)
+    input_is_crowd = np.array([], dtype=np.bool)
+    input_area = np.array([], dtype=np.float32)
+    input_difficult = np.array([], dtype=np.float32)
+    input_confidences = np.array([], dtype=np.float32)
+    valid_indices = np.array([], dtype=np.int32)
+
+    output_tensors = self.execute(graph_fn, [input_boxes, input_classes,
+                                             input_is_crowd, input_area,
+                                             input_difficult,
+                                             input_confidences,
+                                             valid_indices])
+    for key in output_tensors:
+      if key == fields.InputDataFields.groundtruth_boxes:
+        self.assertAllEqual([0, 4], output_tensors[key].shape)
+      else:
+        self.assertAllEqual([0], output_tensors[key].shape)
+
+
+class RetainGroundTruthWithPositiveClasses(test_case.TestCase):
 
   def test_filter_groundtruth_with_positive_classes(self):
-    input_image = tf.placeholder(tf.float32, shape=(None, None, 3))
-    input_boxes = tf.placeholder(tf.float32, shape=(None, 4))
-    input_classes = tf.placeholder(tf.int32, shape=(None,))
-    input_is_crowd = tf.placeholder(tf.bool, shape=(None,))
-    input_area = tf.placeholder(tf.float32, shape=(None,))
-    input_difficult = tf.placeholder(tf.float32, shape=(None,))
-    input_label_types = tf.placeholder(tf.string, shape=(None,))
-    input_confidences = tf.placeholder(tf.float32, shape=(None,))
-    valid_indices = tf.placeholder(tf.int32, shape=(None,))
-    input_tensors = {
-        fields.InputDataFields.image: input_image,
-        fields.InputDataFields.groundtruth_boxes: input_boxes,
-        fields.InputDataFields.groundtruth_classes: input_classes,
-        fields.InputDataFields.groundtruth_is_crowd: input_is_crowd,
-        fields.InputDataFields.groundtruth_area: input_area,
-        fields.InputDataFields.groundtruth_difficult: input_difficult,
-        fields.InputDataFields.groundtruth_label_types: input_label_types,
-        fields.InputDataFields.groundtruth_confidences: input_confidences,
-    }
-    output_tensors = ops.retain_groundtruth_with_positive_classes(input_tensors)
-
-    image_tensor = np.random.rand(224, 224, 3)
-    feed_dict = {
-        input_image: image_tensor,
-        input_boxes:
-        np.array([[0.2, 0.4, 0.1, 0.8], [0.2, 0.4, 1.0, 0.8]], dtype=np.float),
-        input_classes: np.array([1, 0], dtype=np.int32),
-        input_is_crowd: np.array([False, True], dtype=np.bool),
-        input_area: np.array([32, 48], dtype=np.float32),
-        input_difficult: np.array([True, False], dtype=np.bool),
-        input_label_types:
-        np.array(['APPROPRIATE', 'INCORRECT'], dtype=np.string_),
-        input_confidences: np.array([0.99, 0.5], dtype=np.float32),
-        valid_indices: np.array([0], dtype=np.int32),
-    }
+
+    def graph_fn(input_image, input_boxes, input_classes, input_is_crowd,
+                 input_area, input_difficult, input_label_types,
+                 input_confidences):
+      input_tensors = {
+          fields.InputDataFields.image: input_image,
+          fields.InputDataFields.groundtruth_boxes: input_boxes,
+          fields.InputDataFields.groundtruth_classes: input_classes,
+          fields.InputDataFields.groundtruth_is_crowd: input_is_crowd,
+          fields.InputDataFields.groundtruth_area: input_area,
+          fields.InputDataFields.groundtruth_difficult: input_difficult,
+          fields.InputDataFields.groundtruth_label_types: input_label_types,
+          fields.InputDataFields.groundtruth_confidences: input_confidences,
+      }
+      output_tensors = ops.retain_groundtruth_with_positive_classes(
+          input_tensors)
+      return output_tensors
+
+    input_image = np.random.rand(224, 224, 3)
+    input_boxes = np.array([[0.2, 0.4, 0.1, 0.8], [0.2, 0.4, 1.0, 0.8]],
+                           dtype=np.float)
+    input_classes = np.array([1, 0], dtype=np.int32)
+    input_is_crowd = np.array([False, True], dtype=np.bool)
+    input_area = np.array([32, 48], dtype=np.float32)
+    input_difficult = np.array([True, False], dtype=np.bool)
+    input_label_types = np.array(['APPROPRIATE', 'INCORRECT'],
+                                 dtype=np.string_)
+    input_confidences = np.array([0.99, 0.5], dtype=np.float32)
+
     expected_tensors = {
-        fields.InputDataFields.image: image_tensor,
+        fields.InputDataFields.image: input_image,
         fields.InputDataFields.groundtruth_boxes: [[0.2, 0.4, 0.1, 0.8]],
         fields.InputDataFields.groundtruth_classes: [1],
         fields.InputDataFields.groundtruth_is_crowd: [False],
@@ -626,51 +652,70 @@ class RetainGroundTruthWithPositiveClasses(tf.test.TestCase):
         fields.InputDataFields.groundtruth_label_types: [six.b('APPROPRIATE')],
         fields.InputDataFields.groundtruth_confidences: [0.99],
     }
-    with self.test_session() as sess:
-      output_tensors = sess.run(output_tensors, feed_dict=feed_dict)
-      for key in [fields.InputDataFields.image,
-                  fields.InputDataFields.groundtruth_boxes,
-                  fields.InputDataFields.groundtruth_area,
-                  fields.InputDataFields.groundtruth_confidences]:
-        self.assertAllClose(expected_tensors[key], output_tensors[key])
-      for key in [fields.InputDataFields.groundtruth_classes,
-                  fields.InputDataFields.groundtruth_is_crowd,
-                  fields.InputDataFields.groundtruth_label_types]:
-        self.assertAllEqual(expected_tensors[key], output_tensors[key])
 
+    # Executing on CPU because string types are not supported on TPU.
+    output_tensors = self.execute_cpu(graph_fn,
+                                      [input_image, input_boxes,
+                                       input_classes, input_is_crowd,
+                                       input_area,
+                                       input_difficult, input_label_types,
+                                       input_confidences])
 
-class ReplaceNaNGroundtruthLabelScoresWithOnes(tf.test.TestCase):
+    for key in [fields.InputDataFields.image,
+                fields.InputDataFields.groundtruth_boxes,
+                fields.InputDataFields.groundtruth_area,
+                fields.InputDataFields.groundtruth_confidences]:
+      self.assertAllClose(expected_tensors[key], output_tensors[key])
+    for key in [fields.InputDataFields.groundtruth_classes,
+                fields.InputDataFields.groundtruth_is_crowd,
+                fields.InputDataFields.groundtruth_label_types]:
+      self.assertAllEqual(expected_tensors[key], output_tensors[key])
+
+
+class ReplaceNaNGroundtruthLabelScoresWithOnes(test_case.TestCase):
 
   def test_replace_nan_groundtruth_label_scores_with_ones(self):
-    label_scores = tf.constant([np.nan, 1.0, np.nan])
-    output_tensor = ops.replace_nan_groundtruth_label_scores_with_ones(
-        label_scores)
+
+    def graph_fn():
+      label_scores = tf.constant([np.nan, 1.0, np.nan])
+      output_tensor = ops.replace_nan_groundtruth_label_scores_with_ones(
+          label_scores)
+      return output_tensor
+
     expected_tensor = [1.0, 1.0, 1.0]
-    with self.test_session():
-      output_tensor = output_tensor.eval()
-      self.assertAllClose(expected_tensor, output_tensor)
+    output_tensor = self.execute(graph_fn, [])
+    self.assertAllClose(expected_tensor, output_tensor)
 
   def test_input_equals_output_when_no_nans(self):
+
     input_label_scores = [0.5, 1.0, 1.0]
-    label_scores_tensor = tf.constant(input_label_scores)
-    output_label_scores = ops.replace_nan_groundtruth_label_scores_with_ones(
-        label_scores_tensor)
-    with self.test_session():
-      output_label_scores = output_label_scores.eval()
-      self.assertAllClose(input_label_scores, output_label_scores)
+    def graph_fn():
+      label_scores_tensor = tf.constant(input_label_scores)
+      output_label_scores = ops.replace_nan_groundtruth_label_scores_with_ones(
+          label_scores_tensor)
+      return output_label_scores
+
+    output_label_scores = self.execute(graph_fn, [])
+
+    self.assertAllClose(input_label_scores, output_label_scores)
 
 
-class GroundtruthFilterWithCrowdBoxesTest(tf.test.TestCase):
+class GroundtruthFilterWithCrowdBoxesTest(test_case.TestCase):
 
   def test_filter_groundtruth_with_crowd_boxes(self):
-    input_tensors = {
-        fields.InputDataFields.groundtruth_boxes:
-        [[0.1, 0.2, 0.6, 0.8], [0.2, 0.4, 0.1, 0.8]],
-        fields.InputDataFields.groundtruth_classes: [1, 2],
-        fields.InputDataFields.groundtruth_is_crowd: [True, False],
-        fields.InputDataFields.groundtruth_area: [100.0, 238.7],
-        fields.InputDataFields.groundtruth_confidences: [0.5, 0.99],
-    }
+
+    def graph_fn():
+      input_tensors = {
+          fields.InputDataFields.groundtruth_boxes:
+          [[0.1, 0.2, 0.6, 0.8], [0.2, 0.4, 0.1, 0.8]],
+          fields.InputDataFields.groundtruth_classes: [1, 2],
+          fields.InputDataFields.groundtruth_is_crowd: [True, False],
+          fields.InputDataFields.groundtruth_area: [100.0, 238.7],
+          fields.InputDataFields.groundtruth_confidences: [0.5, 0.99],
+      }
+      output_tensors = ops.filter_groundtruth_with_crowd_boxes(
+          input_tensors)
+      return output_tensors
 
     expected_tensors = {
         fields.InputDataFields.groundtruth_boxes: [[0.2, 0.4, 0.1, 0.8]],
@@ -680,30 +725,32 @@ class GroundtruthFilterWithCrowdBoxesTest(tf.test.TestCase):
         fields.InputDataFields.groundtruth_confidences: [0.99],
     }
 
-    output_tensors = ops.filter_groundtruth_with_crowd_boxes(
-        input_tensors)
-    with self.test_session() as sess:
-      output_tensors = sess.run(output_tensors)
-      for key in [fields.InputDataFields.groundtruth_boxes,
-                  fields.InputDataFields.groundtruth_area,
-                  fields.InputDataFields.groundtruth_confidences]:
-        self.assertAllClose(expected_tensors[key], output_tensors[key])
-      for key in [fields.InputDataFields.groundtruth_classes,
-                  fields.InputDataFields.groundtruth_is_crowd]:
-        self.assertAllEqual(expected_tensors[key], output_tensors[key])
+    output_tensors = self.execute(graph_fn, [])
+    for key in [fields.InputDataFields.groundtruth_boxes,
+                fields.InputDataFields.groundtruth_area,
+                fields.InputDataFields.groundtruth_confidences]:
+      self.assertAllClose(expected_tensors[key], output_tensors[key])
+    for key in [fields.InputDataFields.groundtruth_classes,
+                fields.InputDataFields.groundtruth_is_crowd]:
+      self.assertAllEqual(expected_tensors[key], output_tensors[key])
 
 
-class GroundtruthFilterWithNanBoxTest(tf.test.TestCase):
+class GroundtruthFilterWithNanBoxTest(test_case.TestCase):
 
   def test_filter_groundtruth_with_nan_box_coordinates(self):
-    input_tensors = {
-        fields.InputDataFields.groundtruth_boxes:
-        [[np.nan, np.nan, np.nan, np.nan], [0.2, 0.4, 0.1, 0.8]],
-        fields.InputDataFields.groundtruth_classes: [1, 2],
-        fields.InputDataFields.groundtruth_is_crowd: [False, True],
-        fields.InputDataFields.groundtruth_area: [100.0, 238.7],
-        fields.InputDataFields.groundtruth_confidences: [0.5, 0.99],
-    }
+
+    def graph_fn():
+      input_tensors = {
+          fields.InputDataFields.groundtruth_boxes:
+          [[np.nan, np.nan, np.nan, np.nan], [0.2, 0.4, 0.1, 0.8]],
+          fields.InputDataFields.groundtruth_classes: [1, 2],
+          fields.InputDataFields.groundtruth_is_crowd: [False, True],
+          fields.InputDataFields.groundtruth_area: [100.0, 238.7],
+          fields.InputDataFields.groundtruth_confidences: [0.5, 0.99],
+      }
+      output_tensors = ops.filter_groundtruth_with_nan_box_coordinates(
+          input_tensors)
+      return output_tensors
 
     expected_tensors = {
         fields.InputDataFields.groundtruth_boxes: [[0.2, 0.4, 0.1, 0.8]],
@@ -713,30 +760,30 @@ class GroundtruthFilterWithNanBoxTest(tf.test.TestCase):
         fields.InputDataFields.groundtruth_confidences: [0.99],
     }
 
-    output_tensors = ops.filter_groundtruth_with_nan_box_coordinates(
-        input_tensors)
-    with self.test_session() as sess:
-      output_tensors = sess.run(output_tensors)
-      for key in [fields.InputDataFields.groundtruth_boxes,
-                  fields.InputDataFields.groundtruth_area,
-                  fields.InputDataFields.groundtruth_confidences]:
-        self.assertAllClose(expected_tensors[key], output_tensors[key])
-      for key in [fields.InputDataFields.groundtruth_classes,
-                  fields.InputDataFields.groundtruth_is_crowd]:
-        self.assertAllEqual(expected_tensors[key], output_tensors[key])
+    output_tensors = self.execute(graph_fn, [])
+    for key in [fields.InputDataFields.groundtruth_boxes,
+                fields.InputDataFields.groundtruth_area,
+                fields.InputDataFields.groundtruth_confidences]:
+      self.assertAllClose(expected_tensors[key], output_tensors[key])
+    for key in [fields.InputDataFields.groundtruth_classes,
+                fields.InputDataFields.groundtruth_is_crowd]:
+      self.assertAllEqual(expected_tensors[key], output_tensors[key])
 
 
-class GroundtruthFilterWithUnrecognizedClassesTest(tf.test.TestCase):
+class GroundtruthFilterWithUnrecognizedClassesTest(test_case.TestCase):
 
   def test_filter_unrecognized_classes(self):
-    input_tensors = {
-        fields.InputDataFields.groundtruth_boxes:
-        [[.3, .3, .5, .7], [0.2, 0.4, 0.1, 0.8]],
-        fields.InputDataFields.groundtruth_classes: [-1, 2],
-        fields.InputDataFields.groundtruth_is_crowd: [False, True],
-        fields.InputDataFields.groundtruth_area: [100.0, 238.7],
-        fields.InputDataFields.groundtruth_confidences: [0.5, 0.99],
-    }
+    def graph_fn():
+      input_tensors = {
+          fields.InputDataFields.groundtruth_boxes:
+          [[.3, .3, .5, .7], [0.2, 0.4, 0.1, 0.8]],
+          fields.InputDataFields.groundtruth_classes: [-1, 2],
+          fields.InputDataFields.groundtruth_is_crowd: [False, True],
+          fields.InputDataFields.groundtruth_area: [100.0, 238.7],
+          fields.InputDataFields.groundtruth_confidences: [0.5, 0.99],
+      }
+      output_tensors = ops.filter_unrecognized_classes(input_tensors)
+      return output_tensors
 
     expected_tensors = {
         fields.InputDataFields.groundtruth_boxes: [[0.2, 0.4, 0.1, 0.8]],
@@ -746,28 +793,30 @@ class GroundtruthFilterWithUnrecognizedClassesTest(tf.test.TestCase):
         fields.InputDataFields.groundtruth_confidences: [0.99],
     }
 
-    output_tensors = ops.filter_unrecognized_classes(input_tensors)
-    with self.test_session() as sess:
-      output_tensors = sess.run(output_tensors)
-      for key in [fields.InputDataFields.groundtruth_boxes,
-                  fields.InputDataFields.groundtruth_area,
-                  fields.InputDataFields.groundtruth_confidences]:
-        self.assertAllClose(expected_tensors[key], output_tensors[key])
-      for key in [fields.InputDataFields.groundtruth_classes,
-                  fields.InputDataFields.groundtruth_is_crowd]:
-        self.assertAllEqual(expected_tensors[key], output_tensors[key])
+    output_tensors = self.execute(graph_fn, [])
+    for key in [fields.InputDataFields.groundtruth_boxes,
+                fields.InputDataFields.groundtruth_area,
+                fields.InputDataFields.groundtruth_confidences]:
+      self.assertAllClose(expected_tensors[key], output_tensors[key])
+    for key in [fields.InputDataFields.groundtruth_classes,
+                fields.InputDataFields.groundtruth_is_crowd]:
+      self.assertAllEqual(expected_tensors[key], output_tensors[key])
 
 
-class OpsTestNormalizeToTarget(tf.test.TestCase):
+class OpsTestNormalizeToTarget(test_case.TestCase):
 
   def test_create_normalize_to_target(self):
+
+    if self.is_tf2():
+      self.skipTest('Skipping as variable names not supported in eager mode.')
+
     inputs = tf.random_uniform([5, 10, 12, 3])
     target_norm_value = 4.0
     dim = 3
     with self.test_session():
       output = ops.normalize_to_target(inputs, target_norm_value, dim)
       self.assertEqual(output.op.name, 'NormalizeToTarget/mul')
-      var_name = contrib_framework.get_variables()[0].name
+      var_name = slim.get_variables()[0].name
       self.assertEqual(var_name, 'NormalizeToTarget/weights:0')
 
   def test_invalid_dim(self):
@@ -788,94 +837,125 @@ class OpsTestNormalizeToTarget(tf.test.TestCase):
       ops.normalize_to_target(inputs, target_norm_value, dim)
 
   def test_correct_output_shape(self):
-    inputs = tf.random_uniform([5, 10, 12, 3])
-    target_norm_value = 4.0
-    dim = 3
-    with self.test_session():
+
+    if self.is_tf2():
+      self.skipTest('normalize_to_target not supported in eager mode because,'
+                    ' it requires creating variables.')
+
+    inputs = np.random.uniform(size=(5, 10, 12, 3)).astype(np.float32)
+    def graph_fn(inputs):
+      target_norm_value = 4.0
+      dim = 3
       output = ops.normalize_to_target(inputs, target_norm_value, dim)
-      self.assertEqual(output.get_shape().as_list(),
-                       inputs.get_shape().as_list())
+      return output
+
+    # Executing on CPU since creating a variable inside a conditional is not
+    # supported.
+    outputs = self.execute_cpu(graph_fn, [inputs])
+    self.assertEqual(outputs.shape, inputs.shape)
 
   def test_correct_initial_output_values(self):
-    inputs = tf.constant([[[[3, 4], [7, 24]],
-                           [[5, -12], [-1, 0]]]], tf.float32)
-    target_norm_value = 10.0
-    dim = 3
+
+    if self.is_tf2():
+      self.skipTest('normalize_to_target not supported in eager mode because,'
+                    ' it requires creating variables.')
+    def graph_fn():
+      inputs = tf.constant([[[[3, 4], [7, 24]],
+                             [[5, -12], [-1, 0]]]], tf.float32)
+      target_norm_value = 10.0
+      dim = 3
+      normalized_inputs = ops.normalize_to_target(inputs, target_norm_value,
+                                                  dim)
+      return normalized_inputs
+
     expected_output = [[[[30/5.0, 40/5.0], [70/25.0, 240/25.0]],
                         [[50/13.0, -120/13.0], [-10, 0]]]]
-    with self.test_session() as sess:
+    # Executing on CPU since creating a variable inside a conditional is not
+    # supported.
+    output = self.execute_cpu(graph_fn, [])
+    self.assertAllClose(output, expected_output)
+
+  def test_multiple_target_norm_values(self):
+
+    if self.is_tf2():
+      self.skipTest('normalize_to_target not supported in eager mode because,'
+                    ' it requires creating variables.')
+
+    def graph_fn():
+      inputs = tf.constant([[[[3, 4], [7, 24]],
+                             [[5, -12], [-1, 0]]]], tf.float32)
+      target_norm_value = [10.0, 20.0]
+      dim = 3
       normalized_inputs = ops.normalize_to_target(inputs, target_norm_value,
                                                   dim)
-      sess.run(tf.global_variables_initializer())
-      output = normalized_inputs.eval()
-      self.assertAllClose(output, expected_output)
+      return normalized_inputs
 
-  def test_multiple_target_norm_values(self):
-    inputs = tf.constant([[[[3, 4], [7, 24]],
-                           [[5, -12], [-1, 0]]]], tf.float32)
-    target_norm_value = [10.0, 20.0]
-    dim = 3
     expected_output = [[[[30/5.0, 80/5.0], [70/25.0, 480/25.0]],
                         [[50/13.0, -240/13.0], [-10, 0]]]]
-    with self.test_session() as sess:
-      normalized_inputs = ops.normalize_to_target(inputs, target_norm_value,
-                                                  dim)
-      sess.run(tf.global_variables_initializer())
-      output = normalized_inputs.eval()
-      self.assertAllClose(output, expected_output)
+
+    # Executing on CPU since creating a variable inside a conditional is not
+    # supported.
+    output = self.execute_cpu(graph_fn, [])
+    self.assertAllClose(output, expected_output)
 
 
-class OpsTestPositionSensitiveCropRegions(tf.test.TestCase):
+class OpsTestPositionSensitiveCropRegions(test_case.TestCase):
 
   def test_position_sensitive(self):
     num_spatial_bins = [3, 2]
     image_shape = [3, 2, 6]
 
-    # First channel is 1's, second channel is 2's, etc.
-    image = tf.constant(
-        list(range(1, 3 * 2 + 1)) * 6, dtype=tf.float32, shape=image_shape)
-    boxes = tf.random_uniform((2, 4))
-
     # The result for both boxes should be [[1, 2], [3, 4], [5, 6]]
     # before averaging.
     expected_output = np.array([3.5, 3.5]).reshape([2, 1, 1, 1])
 
     for crop_size_mult in range(1, 3):
       crop_size = [3 * crop_size_mult, 2 * crop_size_mult]
-      ps_crop_and_pool = ops.position_sensitive_crop_regions(
-          image, boxes, crop_size, num_spatial_bins, global_pool=True)
 
-      with self.test_session() as sess:
-        output = sess.run(ps_crop_and_pool)
-        self.assertAllClose(output, expected_output)
+      def graph_fn():
+        # First channel is 1's, second channel is 2's, etc.
+        image = tf.constant(
+            list(range(1, 3 * 2 + 1)) * 6, dtype=tf.float32, shape=image_shape)
+        boxes = tf.random_uniform((2, 4))
+
+        # pylint:disable=cell-var-from-loop
+        ps_crop_and_pool = ops.position_sensitive_crop_regions(
+            image, boxes, crop_size, num_spatial_bins, global_pool=True)
+        return ps_crop_and_pool
+
+      output = self.execute(graph_fn, [])
+      self.assertAllClose(output, expected_output)
 
   def test_position_sensitive_with_equal_channels(self):
     num_spatial_bins = [2, 2]
     image_shape = [3, 3, 4]
     crop_size = [2, 2]
 
-    image = tf.constant(
-        list(range(1, 3 * 3 + 1)), dtype=tf.float32, shape=[3, 3, 1])
-    tiled_image = tf.tile(image, [1, 1, image_shape[2]])
-    boxes = tf.random_uniform((3, 4))
-    box_ind = tf.constant([0, 0, 0], dtype=tf.int32)
-
-    # All channels are equal so position-sensitive crop and resize should
-    # work as the usual crop and resize for just one channel.
-    crop = tf.image.crop_and_resize(tf.expand_dims(image, axis=0), boxes,
-                                    box_ind, crop_size)
-    crop_and_pool = tf.reduce_mean(crop, [1, 2], keepdims=True)
-
-    ps_crop_and_pool = ops.position_sensitive_crop_regions(
-        tiled_image,
-        boxes,
-        crop_size,
-        num_spatial_bins,
-        global_pool=True)
-
-    with self.test_session() as sess:
-      expected_output, output = sess.run((crop_and_pool, ps_crop_and_pool))
-      self.assertAllClose(output, expected_output)
+    def graph_fn():
+      image = tf.constant(
+          list(range(1, 3 * 3 + 1)), dtype=tf.float32, shape=[3, 3, 1])
+      tiled_image = tf.tile(image, [1, 1, image_shape[2]])
+      boxes = tf.random_uniform((3, 4))
+      box_ind = tf.constant([0, 0, 0], dtype=tf.int32)
+
+      # All channels are equal so position-sensitive crop and resize should
+      # work as the usual crop and resize for just one channel.
+      crop = tf.image.crop_and_resize(tf.expand_dims(image, axis=0), boxes,
+                                      box_ind, crop_size)
+      crop_and_pool = tf.reduce_mean(crop, [1, 2], keepdims=True)
+
+      ps_crop_and_pool = ops.position_sensitive_crop_regions(
+          tiled_image,
+          boxes,
+          crop_size,
+          num_spatial_bins,
+          global_pool=True)
+
+      return crop_and_pool, ps_crop_and_pool
+
+    # Crop and resize op is not supported in TPUs.
+    expected_output, output = self.execute_cpu(graph_fn, [])
+    self.assertAllClose(output, expected_output)
 
   def test_raise_value_error_on_num_bins_less_than_one(self):
     num_spatial_bins = [1, -1]
@@ -907,24 +987,22 @@ class OpsTestPositionSensitiveCropRegions(tf.test.TestCase):
     image_shape = [1, 1, 5]
     crop_size = [2, 2]
 
-    image = tf.constant(1, dtype=tf.float32, shape=image_shape)
-    boxes = tf.constant([[0, 0, 1, 1]], dtype=tf.float32)
+    def graph_fn():
+      image = tf.constant(1, dtype=tf.float32, shape=image_shape)
+      boxes = tf.constant([[0, 0, 1, 1]], dtype=tf.float32)
+
+      return ops.position_sensitive_crop_regions(
+          image, boxes, crop_size, num_spatial_bins, global_pool=True)
 
     with self.assertRaisesRegexp(
         ValueError, 'Dimension size must be evenly divisible by 4 but is 5'):
-      ops.position_sensitive_crop_regions(
-          image, boxes, crop_size, num_spatial_bins, global_pool=True)
+      self.execute(graph_fn, [])
 
   def test_position_sensitive_with_global_pool_false(self):
     num_spatial_bins = [3, 2]
     image_shape = [3, 2, 6]
     num_boxes = 2
 
-    # First channel is 1's, second channel is 2's, etc.
-    image = tf.constant(
-        list(range(1, 3 * 2 + 1)) * 6, dtype=tf.float32, shape=image_shape)
-    boxes = tf.random_uniform((num_boxes, 4))
-
     expected_output = []
 
     # Expected output, when crop_size = [3, 2].
@@ -946,10 +1024,19 @@ class OpsTestPositionSensitiveCropRegions(tf.test.TestCase):
 
     for crop_size_mult in range(1, 3):
       crop_size = [3 * crop_size_mult, 2 * crop_size_mult]
-      ps_crop = ops.position_sensitive_crop_regions(
-          image, boxes, crop_size, num_spatial_bins, global_pool=False)
-      with self.test_session() as sess:
-        output = sess.run(ps_crop)
+      # First channel is 1's, second channel is 2's, etc.
+
+      def graph_fn():
+        # pylint:disable=cell-var-from-loop
+        image = tf.constant(
+            list(range(1, 3 * 2 + 1)) * 6, dtype=tf.float32, shape=image_shape)
+        boxes = tf.random_uniform((num_boxes, 4))
+
+        ps_crop = ops.position_sensitive_crop_regions(
+            image, boxes, crop_size, num_spatial_bins, global_pool=False)
+        return ps_crop
+
+      output = self.execute(graph_fn, [])
       self.assertAllClose(output, expected_output[crop_size_mult - 1])
 
   def test_position_sensitive_with_global_pool_false_and_do_global_pool(self):
@@ -957,11 +1044,6 @@ class OpsTestPositionSensitiveCropRegions(tf.test.TestCase):
     image_shape = [3, 2, 6]
     num_boxes = 2
 
-    # First channel is 1's, second channel is 2's, etc.
-    image = tf.constant(
-        list(range(1, 3 * 2 + 1)) * 6, dtype=tf.float32, shape=image_shape)
-    boxes = tf.random_uniform((num_boxes, 4))
-
     expected_output = []
 
     # Expected output, when crop_size = [3, 2].
@@ -988,16 +1070,22 @@ class OpsTestPositionSensitiveCropRegions(tf.test.TestCase):
     for crop_size_mult in range(1, 3):
       crop_size = [3 * crop_size_mult, 2 * crop_size_mult]
 
-      # Perform global_pooling after running the function with
-      # global_pool=False.
-      ps_crop = ops.position_sensitive_crop_regions(
-          image, boxes, crop_size, num_spatial_bins, global_pool=False)
-      ps_crop_and_pool = tf.reduce_mean(
-          ps_crop, reduction_indices=(1, 2), keepdims=True)
-
-      with self.test_session() as sess:
-        output = sess.run(ps_crop_and_pool)
-
+      def graph_fn():
+        # pylint:disable=cell-var-from-loop
+        # First channel is 1's, second channel is 2's, etc.
+        image = tf.constant(
+            list(range(1, 3 * 2 + 1)) * 6, dtype=tf.float32, shape=image_shape)
+        boxes = tf.random_uniform((num_boxes, 4))
+
+        # Perform global_pooling after running the function with
+        # global_pool=False.
+        ps_crop = ops.position_sensitive_crop_regions(
+            image, boxes, crop_size, num_spatial_bins, global_pool=False)
+        ps_crop_and_pool = tf.reduce_mean(
+            ps_crop, reduction_indices=(1, 2), keepdims=True)
+        return ps_crop_and_pool
+
+      output = self.execute(graph_fn, [])
       self.assertAllEqual(output, expected_output[crop_size_mult - 1])
 
   def test_raise_value_error_on_non_square_block_size(self):
@@ -1014,42 +1102,39 @@ class OpsTestPositionSensitiveCropRegions(tf.test.TestCase):
           image, boxes, crop_size, num_spatial_bins, global_pool=False)
 
 
-class OpsTestBatchPositionSensitiveCropRegions(tf.test.TestCase):
+class OpsTestBatchPositionSensitiveCropRegions(test_case.TestCase):
 
   def test_position_sensitive_with_single_bin(self):
     num_spatial_bins = [1, 1]
     image_shape = [2, 3, 3, 4]
     crop_size = [2, 2]
 
-    image = tf.random_uniform(image_shape)
-    boxes = tf.random_uniform((2, 3, 4))
-    box_ind = tf.constant([0, 0, 0, 1, 1, 1], dtype=tf.int32)
+    def graph_fn():
+      image = tf.random_uniform(image_shape)
+      boxes = tf.random_uniform((2, 3, 4))
+      box_ind = tf.constant([0, 0, 0, 1, 1, 1], dtype=tf.int32)
 
-    # When a single bin is used, position-sensitive crop and pool should be
-    # the same as non-position sensitive crop and pool.
-    crop = tf.image.crop_and_resize(image, tf.reshape(boxes, [-1, 4]), box_ind,
-                                    crop_size)
-    crop_and_pool = tf.reduce_mean(crop, [1, 2], keepdims=True)
-    crop_and_pool = tf.reshape(crop_and_pool, [2, 3, 1, 1, 4])
+      # When a single bin is used, position-sensitive crop and pool should be
+      # the same as non-position sensitive crop and pool.
+      crop = tf.image.crop_and_resize(image,
+                                      tf.reshape(boxes, [-1, 4]), box_ind,
+                                      crop_size)
+      crop_and_pool = tf.reduce_mean(crop, [1, 2], keepdims=True)
+      crop_and_pool = tf.reshape(crop_and_pool, [2, 3, 1, 1, 4])
 
-    ps_crop_and_pool = ops.batch_position_sensitive_crop_regions(
-        image, boxes, crop_size, num_spatial_bins, global_pool=True)
+      ps_crop_and_pool = ops.batch_position_sensitive_crop_regions(
+          image, boxes, crop_size, num_spatial_bins, global_pool=True)
+      return crop_and_pool, ps_crop_and_pool
 
-    with self.test_session() as sess:
-      expected_output, output = sess.run((crop_and_pool, ps_crop_and_pool))
-      self.assertAllClose(output, expected_output)
+    # Crop and resize is not supported on TPUs.
+    expected_output, output = self.execute_cpu(graph_fn, [])
+    self.assertAllClose(output, expected_output)
 
   def test_position_sensitive_with_global_pool_false_and_known_boxes(self):
     num_spatial_bins = [2, 2]
     image_shape = [2, 2, 2, 4]
     crop_size = [2, 2]
 
-    images = tf.constant(
-        list(range(1, 2 * 2 * 4 + 1)) * 2, dtype=tf.float32, shape=image_shape)
-
-    # First box contains whole image, and second box contains only first row.
-    boxes = tf.constant(np.array([[[0., 0., 1., 1.]],
-                                  [[0., 0., 0.5, 1.]]]), dtype=tf.float32)
     # box_ind = tf.constant([0, 1], dtype=tf.int32)
 
     expected_output = []
@@ -1069,105 +1154,147 @@ class OpsTestBatchPositionSensitiveCropRegions(tf.test.TestCase):
     )
     expected_output = np.stack(expected_output, axis=0)
 
-    ps_crop = ops.batch_position_sensitive_crop_regions(
-        images, boxes, crop_size, num_spatial_bins, global_pool=False)
+    def graph_fn():
+      images = tf.constant(
+          list(range(1, 2 * 2 * 4 + 1)) * 2, dtype=tf.float32,
+          shape=image_shape)
 
-    with self.test_session() as sess:
-      output = sess.run(ps_crop)
-      self.assertAllEqual(output, expected_output)
+      # First box contains whole image, and second box contains only first row.
+      boxes = tf.constant(np.array([[[0., 0., 1., 1.]],
+                                    [[0., 0., 0.5, 1.]]]), dtype=tf.float32)
+
+      ps_crop = ops.batch_position_sensitive_crop_regions(
+          images, boxes, crop_size, num_spatial_bins, global_pool=False)
+      return ps_crop
+
+    output = self.execute(graph_fn, [])
+    self.assertAllEqual(output, expected_output)
 
   def test_position_sensitive_with_global_pool_false_and_single_bin(self):
     num_spatial_bins = [1, 1]
     image_shape = [2, 3, 3, 4]
     crop_size = [1, 1]
 
-    images = tf.random_uniform(image_shape)
-    boxes = tf.random_uniform((2, 3, 4))
-    # box_ind = tf.constant([0, 0, 0, 1, 1, 1], dtype=tf.int32)
+    def graph_fn():
+      images = tf.random_uniform(image_shape)
+      boxes = tf.random_uniform((2, 3, 4))
+      # box_ind = tf.constant([0, 0, 0, 1, 1, 1], dtype=tf.int32)
 
-    # Since single_bin is used and crop_size = [1, 1] (i.e., no crop resize),
-    # the outputs are the same whatever the global_pool value is.
-    ps_crop_and_pool = ops.batch_position_sensitive_crop_regions(
-        images, boxes, crop_size, num_spatial_bins, global_pool=True)
-    ps_crop = ops.batch_position_sensitive_crop_regions(
-        images, boxes, crop_size, num_spatial_bins, global_pool=False)
+      # Since single_bin is used and crop_size = [1, 1] (i.e., no crop resize),
+      # the outputs are the same whatever the global_pool value is.
+      ps_crop_and_pool = ops.batch_position_sensitive_crop_regions(
+          images, boxes, crop_size, num_spatial_bins, global_pool=True)
+      ps_crop = ops.batch_position_sensitive_crop_regions(
+          images, boxes, crop_size, num_spatial_bins, global_pool=False)
+      return ps_crop_and_pool, ps_crop
 
-    with self.test_session() as sess:
-      pooled_output, unpooled_output = sess.run((ps_crop_and_pool, ps_crop))
-      self.assertAllClose(pooled_output, unpooled_output)
+    pooled_output, unpooled_output = self.execute(graph_fn, [])
+    self.assertAllClose(pooled_output, unpooled_output)
 
 
-class ReframeBoxMasksToImageMasksTest(tf.test.TestCase):
+# The following tests are only executed on CPU because the output
+# shape is not constant.
+class ReframeBoxMasksToImageMasksTest(test_case.TestCase):
 
   def testZeroImageOnEmptyMask(self):
-    box_masks = tf.constant([[[0, 0],
-                              [0, 0]]], dtype=tf.float32)
-    boxes = tf.constant([[0.0, 0.0, 1.0, 1.0]], dtype=tf.float32)
-    image_masks = ops.reframe_box_masks_to_image_masks(box_masks, boxes,
-                                                       image_height=4,
-                                                       image_width=4)
     np_expected_image_masks = np.array([[[0, 0, 0, 0],
                                          [0, 0, 0, 0],
                                          [0, 0, 0, 0],
                                          [0, 0, 0, 0]]], dtype=np.float32)
-    with self.test_session() as sess:
-      np_image_masks = sess.run(image_masks)
-      self.assertAllClose(np_image_masks, np_expected_image_masks)
+    def graph_fn():
+      box_masks = tf.constant([[[0, 0],
+                                [0, 0]]], dtype=tf.float32)
+      boxes = tf.constant([[0.0, 0.0, 1.0, 1.0]], dtype=tf.float32)
+      image_masks = ops.reframe_box_masks_to_image_masks(box_masks, boxes,
+                                                         image_height=4,
+                                                         image_width=4)
+      return image_masks
+
+    np_image_masks = self.execute_cpu(graph_fn, [])
+    self.assertAllClose(np_image_masks, np_expected_image_masks)
 
   def testZeroBoxMasks(self):
-    box_masks = tf.zeros([0, 3, 3], dtype=tf.float32)
-    boxes = tf.zeros([0, 4], dtype=tf.float32)
-    image_masks = ops.reframe_box_masks_to_image_masks(box_masks, boxes,
-                                                       image_height=4,
-                                                       image_width=4)
-    with self.test_session() as sess:
-      np_image_masks = sess.run(image_masks)
-      self.assertAllEqual(np_image_masks.shape, np.array([0, 4, 4]))
+
+    def graph_fn():
+      box_masks = tf.zeros([0, 3, 3], dtype=tf.float32)
+      boxes = tf.zeros([0, 4], dtype=tf.float32)
+      image_masks = ops.reframe_box_masks_to_image_masks(box_masks, boxes,
+                                                         image_height=4,
+                                                         image_width=4)
+      return image_masks
+
+    np_image_masks = self.execute_cpu(graph_fn, [])
+    self.assertAllEqual(np_image_masks.shape, np.array([0, 4, 4]))
+
+  def testBoxWithZeroArea(self):
+
+    def graph_fn():
+      box_masks = tf.zeros([1, 3, 3], dtype=tf.float32)
+      boxes = tf.constant([[0.1, 0.2, 0.1, 0.7]], dtype=tf.float32)
+      image_masks = ops.reframe_box_masks_to_image_masks(box_masks, boxes,
+                                                         image_height=4,
+                                                         image_width=4)
+      return image_masks
+
+    np_image_masks = self.execute_cpu(graph_fn, [])
+    self.assertAllEqual(np_image_masks.shape, np.array([1, 4, 4]))
 
   def testMaskIsCenteredInImageWhenBoxIsCentered(self):
-    box_masks = tf.constant([[[1, 1],
-                              [1, 1]]], dtype=tf.float32)
-    boxes = tf.constant([[0.25, 0.25, 0.75, 0.75]], dtype=tf.float32)
-    image_masks = ops.reframe_box_masks_to_image_masks(box_masks, boxes,
-                                                       image_height=4,
-                                                       image_width=4)
+
+    def graph_fn():
+      box_masks = tf.constant([[[1, 1],
+                                [1, 1]]], dtype=tf.float32)
+      boxes = tf.constant([[0.25, 0.25, 0.75, 0.75]], dtype=tf.float32)
+      image_masks = ops.reframe_box_masks_to_image_masks(box_masks, boxes,
+                                                         image_height=4,
+                                                         image_width=4)
+      return image_masks
+
     np_expected_image_masks = np.array([[[0, 0, 0, 0],
                                          [0, 1, 1, 0],
                                          [0, 1, 1, 0],
                                          [0, 0, 0, 0]]], dtype=np.float32)
-    with self.test_session() as sess:
-      np_image_masks = sess.run(image_masks)
-      self.assertAllClose(np_image_masks, np_expected_image_masks)
+    np_image_masks = self.execute_cpu(graph_fn, [])
+    self.assertAllClose(np_image_masks, np_expected_image_masks)
 
   def testMaskOffCenterRemainsOffCenterInImage(self):
-    box_masks = tf.constant([[[1, 0],
-                              [0, 1]]], dtype=tf.float32)
-    boxes = tf.constant([[0.25, 0.5, 0.75, 1.0]], dtype=tf.float32)
-    image_masks = ops.reframe_box_masks_to_image_masks(box_masks, boxes,
-                                                       image_height=4,
-                                                       image_width=4)
+
+    def graph_fn():
+      box_masks = tf.constant([[[1, 0],
+                                [0, 1]]], dtype=tf.float32)
+      boxes = tf.constant([[0.25, 0.5, 0.75, 1.0]], dtype=tf.float32)
+      image_masks = ops.reframe_box_masks_to_image_masks(box_masks, boxes,
+                                                         image_height=4,
+                                                         image_width=4)
+      return image_masks
+
     np_expected_image_masks = np.array([[[0, 0, 0, 0],
                                          [0, 0, 0.6111111, 0.16666669],
                                          [0, 0, 0.3888889, 0.83333337],
                                          [0, 0, 0, 0]]], dtype=np.float32)
-    with self.test_session() as sess:
-      np_image_masks = sess.run(image_masks)
-      self.assertAllClose(np_image_masks, np_expected_image_masks)
+    np_image_masks = self.execute_cpu(graph_fn, [])
+    self.assertAllClose(np_image_masks, np_expected_image_masks)
 
 
-class MergeBoxesWithMultipleLabelsTest(tf.test.TestCase):
+class MergeBoxesWithMultipleLabelsTest(test_case.TestCase):
 
   def testMergeBoxesWithMultipleLabels(self):
-    boxes = tf.constant(
-        [[0.25, 0.25, 0.75, 0.75], [0.0, 0.0, 0.5, 0.75],
-         [0.25, 0.25, 0.75, 0.75]],
-        dtype=tf.float32)
-    class_indices = tf.constant([0, 4, 2], dtype=tf.int32)
-    class_confidences = tf.constant([0.8, 0.2, 0.1], dtype=tf.float32)
-    num_classes = 5
-    merged_boxes, merged_classes, merged_confidences, merged_box_indices = (
-        ops.merge_boxes_with_multiple_labels(
-            boxes, class_indices, class_confidences, num_classes))
+
+    def graph_fn():
+      boxes = tf.constant(
+          [[0.25, 0.25, 0.75, 0.75], [0.0, 0.0, 0.5, 0.75],
+           [0.25, 0.25, 0.75, 0.75]],
+          dtype=tf.float32)
+      class_indices = tf.constant([0, 4, 2], dtype=tf.int32)
+      class_confidences = tf.constant([0.8, 0.2, 0.1], dtype=tf.float32)
+      num_classes = 5
+      merged_boxes, merged_classes, merged_confidences, merged_box_indices = (
+          ops.merge_boxes_with_multiple_labels(
+              boxes, class_indices, class_confidences, num_classes))
+
+      return (merged_boxes, merged_classes, merged_confidences,
+              merged_box_indices)
+
     expected_merged_boxes = np.array(
         [[0.25, 0.25, 0.75, 0.75], [0.0, 0.0, 0.5, 0.75]], dtype=np.float32)
     expected_merged_classes = np.array(
@@ -1175,28 +1302,32 @@ class MergeBoxesWithMultipleLabelsTest(tf.test.TestCase):
     expected_merged_confidences = np.array(
         [[0.8, 0, 0.1, 0, 0], [0, 0, 0, 0, 0.2]], dtype=np.float32)
     expected_merged_box_indices = np.array([0, 1], dtype=np.int32)
-    with self.test_session() as sess:
-      (np_merged_boxes, np_merged_classes, np_merged_confidences,
-       np_merged_box_indices) = sess.run(
-           [merged_boxes, merged_classes, merged_confidences,
-            merged_box_indices])
-      self.assertAllClose(np_merged_boxes, expected_merged_boxes)
-      self.assertAllClose(np_merged_classes, expected_merged_classes)
-      self.assertAllClose(np_merged_confidences, expected_merged_confidences)
-      self.assertAllClose(np_merged_box_indices, expected_merged_box_indices)
+
+    # Running on CPU only as tf.unique is not supported on TPU.
+    (np_merged_boxes, np_merged_classes, np_merged_confidences,
+     np_merged_box_indices) = self.execute_cpu(graph_fn, [])
+    self.assertAllClose(np_merged_boxes, expected_merged_boxes)
+    self.assertAllClose(np_merged_classes, expected_merged_classes)
+    self.assertAllClose(np_merged_confidences, expected_merged_confidences)
+    self.assertAllClose(np_merged_box_indices, expected_merged_box_indices)
 
   def testMergeBoxesWithMultipleLabelsCornerCase(self):
-    boxes = tf.constant(
-        [[0, 0, 1, 1], [0, 1, 1, 1], [1, 0, 1, 1], [1, 1, 1, 1],
-         [1, 1, 1, 1], [1, 0, 1, 1], [0, 1, 1, 1], [0, 0, 1, 1]],
-        dtype=tf.float32)
-    class_indices = tf.constant([0, 1, 2, 3, 2, 1, 0, 3], dtype=tf.int32)
-    class_confidences = tf.constant([0.1, 0.9, 0.2, 0.8, 0.3, 0.7, 0.4, 0.6],
-                                    dtype=tf.float32)
-    num_classes = 4
-    merged_boxes, merged_classes, merged_confidences, merged_box_indices = (
-        ops.merge_boxes_with_multiple_labels(
-            boxes, class_indices, class_confidences, num_classes))
+
+    def graph_fn():
+      boxes = tf.constant(
+          [[0, 0, 1, 1], [0, 1, 1, 1], [1, 0, 1, 1], [1, 1, 1, 1],
+           [1, 1, 1, 1], [1, 0, 1, 1], [0, 1, 1, 1], [0, 0, 1, 1]],
+          dtype=tf.float32)
+      class_indices = tf.constant([0, 1, 2, 3, 2, 1, 0, 3], dtype=tf.int32)
+      class_confidences = tf.constant([0.1, 0.9, 0.2, 0.8, 0.3, 0.7, 0.4, 0.6],
+                                      dtype=tf.float32)
+      num_classes = 4
+      merged_boxes, merged_classes, merged_confidences, merged_box_indices = (
+          ops.merge_boxes_with_multiple_labels(
+              boxes, class_indices, class_confidences, num_classes))
+      return (merged_boxes, merged_classes, merged_confidences,
+              merged_box_indices)
+
     expected_merged_boxes = np.array(
         [[0, 0, 1, 1], [0, 1, 1, 1], [1, 0, 1, 1], [1, 1, 1, 1]],
         dtype=np.float32)
@@ -1207,35 +1338,42 @@ class MergeBoxesWithMultipleLabelsTest(tf.test.TestCase):
         [[0.1, 0, 0, 0.6], [0.4, 0.9, 0, 0],
          [0, 0.7, 0.2, 0], [0, 0, 0.3, 0.8]], dtype=np.float32)
     expected_merged_box_indices = np.array([0, 1, 2, 3], dtype=np.int32)
-    with self.test_session() as sess:
-      (np_merged_boxes, np_merged_classes, np_merged_confidences,
-       np_merged_box_indices) = sess.run(
-           [merged_boxes, merged_classes, merged_confidences,
-            merged_box_indices])
-      self.assertAllClose(np_merged_boxes, expected_merged_boxes)
-      self.assertAllClose(np_merged_classes, expected_merged_classes)
-      self.assertAllClose(np_merged_confidences, expected_merged_confidences)
-      self.assertAllClose(np_merged_box_indices, expected_merged_box_indices)
+
+    # Running on CPU only as tf.unique is not supported on TPU.
+    (np_merged_boxes, np_merged_classes, np_merged_confidences,
+     np_merged_box_indices) = self.execute_cpu(graph_fn, [])
+
+    self.assertAllClose(np_merged_boxes, expected_merged_boxes)
+    self.assertAllClose(np_merged_classes, expected_merged_classes)
+    self.assertAllClose(np_merged_confidences, expected_merged_confidences)
+    self.assertAllClose(np_merged_box_indices, expected_merged_box_indices)
 
   def testMergeBoxesWithEmptyInputs(self):
-    boxes = tf.zeros([0, 4], dtype=tf.float32)
-    class_indices = tf.constant([], dtype=tf.int32)
-    class_confidences = tf.constant([], dtype=tf.float32)
-    num_classes = 5
-    merged_boxes, merged_classes, merged_confidences, merged_box_indices = (
-        ops.merge_boxes_with_multiple_labels(
-            boxes, class_indices, class_confidences, num_classes))
-    with self.test_session() as sess:
-      (np_merged_boxes, np_merged_classes, np_merged_confidences,
-       np_merged_box_indices) = sess.run(
-           [merged_boxes, merged_classes, merged_confidences,
-            merged_box_indices])
-      self.assertAllEqual(np_merged_boxes.shape, [0, 4])
-      self.assertAllEqual(np_merged_classes.shape, [0, 5])
-      self.assertAllEqual(np_merged_confidences.shape, [0, 5])
-      self.assertAllEqual(np_merged_box_indices.shape, [0])
+
+    def graph_fn():
+      boxes = tf.zeros([0, 4], dtype=tf.float32)
+      class_indices = tf.constant([], dtype=tf.int32)
+      class_confidences = tf.constant([], dtype=tf.float32)
+      num_classes = 5
+      merged_boxes, merged_classes, merged_confidences, merged_box_indices = (
+          ops.merge_boxes_with_multiple_labels(
+              boxes, class_indices, class_confidences, num_classes))
+      return (merged_boxes, merged_classes, merged_confidences,
+              merged_box_indices)
+
+    # Running on CPU only as tf.unique is not supported on TPU.
+    (np_merged_boxes, np_merged_classes, np_merged_confidences,
+     np_merged_box_indices) = self.execute_cpu(graph_fn, [])
+    self.assertAllEqual(np_merged_boxes.shape, [0, 4])
+    self.assertAllEqual(np_merged_classes.shape, [0, 5])
+    self.assertAllEqual(np_merged_confidences.shape, [0, 5])
+    self.assertAllEqual(np_merged_box_indices.shape, [0])
 
   def testMergeBoxesWithMultipleLabelsUsesInt64(self):
+
+    if self.is_tf2():
+      self.skipTest('Getting op names is not supported in eager mode.')
+
     boxes = tf.constant(
         [[0.25, 0.25, 0.75, 0.75], [0.0, 0.0, 0.5, 0.75],
          [0.25, 0.25, 0.75, 0.75]],
@@ -1345,20 +1483,18 @@ class MatmulGatherOnZerothAxis(test_case.TestCase):
     self.assertAllClose(gather_output, expected_output)
 
   def test_gather_with_dynamic_shape_input(self):
-    params_placeholder = tf.placeholder(tf.float32, shape=[None, 4])
-    indices_placeholder = tf.placeholder(tf.int32, shape=[None])
-    gather_result = ops.matmul_gather_on_zeroth_axis(
-        params_placeholder, indices_placeholder)
+
+    def graph_fn(params, indices):
+      return ops.matmul_gather_on_zeroth_axis(params, indices)
+
     params = np.array([[1, 2, 3, 4],
                        [5, 6, 7, 8],
                        [9, 10, 11, 12],
                        [0, 1, 0, 0]], dtype=np.float32)
     indices = np.array([0, 0, 0, 0, 0, 0])
     expected_output = np.array(6*[[1, 2, 3, 4]])
-    with self.test_session() as sess:
-      gather_output = sess.run(gather_result, feed_dict={
-          params_placeholder: params, indices_placeholder: indices})
-      self.assertAllClose(gather_output, expected_output)
+    gather_output = self.execute(graph_fn, [params, indices])
+    self.assertAllClose(gather_output, expected_output)
 
 
 class FpnFeatureLevelsTest(test_case.TestCase):
@@ -1419,22 +1555,27 @@ class TestBfloat16ToFloat32(test_case.TestCase):
 class TestGatherWithPaddingValues(test_case.TestCase):
 
   def test_gather_with_padding_values(self):
-    indices = tf.constant([1, -1, 0, -1])
-    input_tensor = tf.constant([[0, 0, 0.1, 0.1], [0, 0, 0.2, 0.2]],
-                               dtype=tf.float32)
     expected_gathered_tensor = [
         [0, 0, 0.2, 0.2],
         [0, 0, 0, 0],
         [0, 0, 0.1, 0.1],
         [0, 0, 0, 0],
     ]
-    gathered_tensor = ops.gather_with_padding_values(
-        input_tensor,
-        indices=indices,
-        padding_value=tf.zeros_like(input_tensor[0]))
-    self.assertEqual(gathered_tensor.dtype, tf.float32)
-    with self.test_session():
-      gathered_tensor_np = gathered_tensor.eval()
+
+    def graph_fn():
+      indices = tf.constant([1, -1, 0, -1])
+      input_tensor = tf.constant([[0, 0, 0.1, 0.1], [0, 0, 0.2, 0.2]],
+                                 dtype=tf.float32)
+
+      gathered_tensor = ops.gather_with_padding_values(
+          input_tensor,
+          indices=indices,
+          padding_value=tf.zeros_like(input_tensor[0]))
+      self.assertEqual(gathered_tensor.dtype, tf.float32)
+
+      return gathered_tensor
+
+    gathered_tensor_np = self.execute(graph_fn, [])
     self.assertAllClose(expected_gathered_tensor, gathered_tensor_np)
 
 

research/object_detection/utils/patch_ops.py

@@ -19,7 +19,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def get_patch_mask(y, x, patch_size, image_shape):

research/object_detection/utils/patch_ops_test.py

@@ -21,7 +21,7 @@ from __future__ import print_function
 
 from absl.testing import parameterized
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from object_detection.utils import patch_ops
 from object_detection.utils import test_case

research/object_detection/utils/per_image_evaluation_test.py

@@ -21,7 +21,7 @@ from __future__ import print_function
 
 import numpy as np
 from six.moves import range
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from object_detection.utils import per_image_evaluation
 

research/object_detection/utils/per_image_vrd_evaluation_test.py

@@ -19,7 +19,7 @@ from __future__ import division
 from __future__ import print_function
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from object_detection.utils import per_image_vrd_evaluation
 

research/object_detection/utils/shape_utils.py

@@ -20,7 +20,7 @@ from __future__ import division
 from __future__ import print_function
 
 from six.moves import zip
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from object_detection.utils import static_shape
 
@@ -69,17 +69,18 @@ def pad_tensor(t, length):
       is an integer, the first dimension of padded_t is set to length
       statically.
   """
-  t_rank = tf.rank(t)
-  t_shape = tf.shape(t)
-  t_d0 = t_shape[0]
-  pad_d0 = tf.expand_dims(length - t_d0, 0)
-  pad_shape = tf.cond(
-      tf.greater(t_rank, 1), lambda: tf.concat([pad_d0, t_shape[1:]], 0),
-      lambda: tf.expand_dims(length - t_d0, 0))
-  padded_t = tf.concat([t, tf.zeros(pad_shape, dtype=t.dtype)], 0)
-  if not _is_tensor(length):
-    padded_t = _set_dim_0(padded_t, length)
-  return padded_t
+
+  # Computing the padding statically makes the operation work with XLA.
+  rank = len(t.get_shape())
+  paddings = [[0 for _ in range(2)] for _ in range(rank)]
+  t_d0 = tf.shape(t)[0]
+
+  if isinstance(length, int) or len(length.get_shape()) == 0:  # pylint:disable=g-explicit-length-test
+    paddings[0][1] = length - t_d0
+  else:
+    paddings[0][1] = length[0] - t_d0
+
+  return tf.pad(t, paddings)
 
 
 def clip_tensor(t, length):

research/object_detection/utils/shape_utils_test.py

@@ -20,216 +20,215 @@ from __future__ import division
 from __future__ import print_function
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from object_detection.utils import shape_utils
+from object_detection.utils import test_case
 
-# pylint: disable=g-import-not-at-top
-try:
-  from tensorflow.contrib import framework as contrib_framework
-except ImportError:
-  # TF 2.0 doesn't ship with contrib.
-  pass
-# pylint: enable=g-import-not-at-top
 
-
-class UtilTest(tf.test.TestCase):
+class UtilTest(test_case.TestCase):
 
   def test_pad_tensor_using_integer_input(self):
-    t1 = tf.constant([1], dtype=tf.int32)
-    pad_t1 = shape_utils.pad_tensor(t1, 2)
-    t2 = tf.constant([[0.1, 0.2]], dtype=tf.float32)
-    pad_t2 = shape_utils.pad_tensor(t2, 2)
 
-    self.assertEqual(2, pad_t1.get_shape()[0])
-    self.assertEqual(2, pad_t2.get_shape()[0])
+    print('........pad tensor using interger input.')
+    def graph_fn():
+      t1 = tf.constant([1], dtype=tf.int32)
+      pad_t1 = shape_utils.pad_tensor(t1, 2)
+      t2 = tf.constant([[0.1, 0.2]], dtype=tf.float32)
+      pad_t2 = shape_utils.pad_tensor(t2, 2)
+
+      return pad_t1, pad_t2
+
+    pad_t1_result, pad_t2_result = self.execute(graph_fn, [])
 
-    with self.test_session() as sess:
-      pad_t1_result, pad_t2_result = sess.run([pad_t1, pad_t2])
-      self.assertAllEqual([1, 0], pad_t1_result)
-      self.assertAllClose([[0.1, 0.2], [0, 0]], pad_t2_result)
+    self.assertAllEqual([1, 0], pad_t1_result)
+    self.assertAllClose([[0.1, 0.2], [0, 0]], pad_t2_result)
 
   def test_pad_tensor_using_tensor_input(self):
-    t1 = tf.constant([1], dtype=tf.int32)
-    pad_t1 = shape_utils.pad_tensor(t1, tf.constant(2))
-    t2 = tf.constant([[0.1, 0.2]], dtype=tf.float32)
-    pad_t2 = shape_utils.pad_tensor(t2, tf.constant(2))
 
-    with self.test_session() as sess:
-      pad_t1_result, pad_t2_result = sess.run([pad_t1, pad_t2])
-      self.assertAllEqual([1, 0], pad_t1_result)
-      self.assertAllClose([[0.1, 0.2], [0, 0]], pad_t2_result)
+    def graph_fn():
+      t1 = tf.constant([1], dtype=tf.int32)
+      pad_t1 = shape_utils.pad_tensor(t1, tf.constant(2))
+      t2 = tf.constant([[0.1, 0.2]], dtype=tf.float32)
+      pad_t2 = shape_utils.pad_tensor(t2, tf.constant(2))
+
+      return pad_t1, pad_t2
+
+    pad_t1_result, pad_t2_result = self.execute(graph_fn, [])
+    self.assertAllEqual([1, 0], pad_t1_result)
+    self.assertAllClose([[0.1, 0.2], [0, 0]], pad_t2_result)
 
   def test_clip_tensor_using_integer_input(self):
-    t1 = tf.constant([1, 2, 3], dtype=tf.int32)
-    clip_t1 = shape_utils.clip_tensor(t1, 2)
-    t2 = tf.constant([[0.1, 0.2], [0.2, 0.4], [0.5, 0.8]], dtype=tf.float32)
-    clip_t2 = shape_utils.clip_tensor(t2, 2)
 
-    self.assertEqual(2, clip_t1.get_shape()[0])
-    self.assertEqual(2, clip_t2.get_shape()[0])
+    def graph_fn():
+      t1 = tf.constant([1, 2, 3], dtype=tf.int32)
+      clip_t1 = shape_utils.clip_tensor(t1, 2)
+      t2 = tf.constant([[0.1, 0.2], [0.2, 0.4], [0.5, 0.8]], dtype=tf.float32)
+      clip_t2 = shape_utils.clip_tensor(t2, 2)
 
-    with self.test_session() as sess:
-      clip_t1_result, clip_t2_result = sess.run([clip_t1, clip_t2])
-      self.assertAllEqual([1, 2], clip_t1_result)
-      self.assertAllClose([[0.1, 0.2], [0.2, 0.4]], clip_t2_result)
+      self.assertEqual(2, clip_t1.get_shape()[0])
+      self.assertEqual(2, clip_t2.get_shape()[0])
+
+      return clip_t1, clip_t2
+
+    clip_t1_result, clip_t2_result = self.execute(graph_fn, [])
+    self.assertAllEqual([1, 2], clip_t1_result)
+    self.assertAllClose([[0.1, 0.2], [0.2, 0.4]], clip_t2_result)
 
   def test_clip_tensor_using_tensor_input(self):
-    t1 = tf.constant([1, 2, 3], dtype=tf.int32)
-    clip_t1 = shape_utils.clip_tensor(t1, tf.constant(2))
-    t2 = tf.constant([[0.1, 0.2], [0.2, 0.4], [0.5, 0.8]], dtype=tf.float32)
-    clip_t2 = shape_utils.clip_tensor(t2, tf.constant(2))
 
-    with self.test_session() as sess:
-      clip_t1_result, clip_t2_result = sess.run([clip_t1, clip_t2])
-      self.assertAllEqual([1, 2], clip_t1_result)
-      self.assertAllClose([[0.1, 0.2], [0.2, 0.4]], clip_t2_result)
+    def graph_fn():
+      t1 = tf.constant([1, 2, 3], dtype=tf.int32)
+      clip_t1 = shape_utils.clip_tensor(t1, tf.constant(2))
+      t2 = tf.constant([[0.1, 0.2], [0.2, 0.4], [0.5, 0.8]], dtype=tf.float32)
+      clip_t2 = shape_utils.clip_tensor(t2, tf.constant(2))
+
+      return clip_t1, clip_t2
+
+    clip_t1_result, clip_t2_result = self.execute(graph_fn, [])
+    self.assertAllEqual([1, 2], clip_t1_result)
+    self.assertAllClose([[0.1, 0.2], [0.2, 0.4]], clip_t2_result)
 
   def test_pad_or_clip_tensor_using_integer_input(self):
-    t1 = tf.constant([1], dtype=tf.int32)
-    tt1 = shape_utils.pad_or_clip_tensor(t1, 2)
-    t2 = tf.constant([[0.1, 0.2]], dtype=tf.float32)
-    tt2 = shape_utils.pad_or_clip_tensor(t2, 2)
-
-    t3 = tf.constant([1, 2, 3], dtype=tf.int32)
-    tt3 = shape_utils.clip_tensor(t3, 2)
-    t4 = tf.constant([[0.1, 0.2], [0.2, 0.4], [0.5, 0.8]], dtype=tf.float32)
-    tt4 = shape_utils.clip_tensor(t4, 2)
-
-    self.assertEqual(2, tt1.get_shape()[0])
-    self.assertEqual(2, tt2.get_shape()[0])
-    self.assertEqual(2, tt3.get_shape()[0])
-    self.assertEqual(2, tt4.get_shape()[0])
-
-    with self.test_session() as sess:
-      tt1_result, tt2_result, tt3_result, tt4_result = sess.run(
-          [tt1, tt2, tt3, tt4])
-      self.assertAllEqual([1, 0], tt1_result)
-      self.assertAllClose([[0.1, 0.2], [0, 0]], tt2_result)
-      self.assertAllEqual([1, 2], tt3_result)
-      self.assertAllClose([[0.1, 0.2], [0.2, 0.4]], tt4_result)
+
+    def graph_fn():
+      t1 = tf.constant([1], dtype=tf.int32)
+      tt1 = shape_utils.pad_or_clip_tensor(t1, 2)
+      t2 = tf.constant([[0.1, 0.2]], dtype=tf.float32)
+      tt2 = shape_utils.pad_or_clip_tensor(t2, 2)
+
+      t3 = tf.constant([1, 2, 3], dtype=tf.int32)
+      tt3 = shape_utils.clip_tensor(t3, 2)
+      t4 = tf.constant([[0.1, 0.2], [0.2, 0.4], [0.5, 0.8]], dtype=tf.float32)
+      tt4 = shape_utils.clip_tensor(t4, 2)
+
+      self.assertEqual(2, tt1.get_shape()[0])
+      self.assertEqual(2, tt2.get_shape()[0])
+      self.assertEqual(2, tt3.get_shape()[0])
+      self.assertEqual(2, tt4.get_shape()[0])
+
+      return tt1, tt2, tt3, tt4
+
+    tt1_result, tt2_result, tt3_result, tt4_result = self.execute(graph_fn, [])
+    self.assertAllEqual([1, 0], tt1_result)
+    self.assertAllClose([[0.1, 0.2], [0, 0]], tt2_result)
+    self.assertAllEqual([1, 2], tt3_result)
+    self.assertAllClose([[0.1, 0.2], [0.2, 0.4]], tt4_result)
 
   def test_pad_or_clip_tensor_using_tensor_input(self):
-    t1 = tf.constant([1], dtype=tf.int32)
-    tt1 = shape_utils.pad_or_clip_tensor(t1, tf.constant(2))
-    t2 = tf.constant([[0.1, 0.2]], dtype=tf.float32)
-    tt2 = shape_utils.pad_or_clip_tensor(t2, tf.constant(2))
-
-    t3 = tf.constant([1, 2, 3], dtype=tf.int32)
-    tt3 = shape_utils.clip_tensor(t3, tf.constant(2))
-    t4 = tf.constant([[0.1, 0.2], [0.2, 0.4], [0.5, 0.8]], dtype=tf.float32)
-    tt4 = shape_utils.clip_tensor(t4, tf.constant(2))
-
-    with self.test_session() as sess:
-      tt1_result, tt2_result, tt3_result, tt4_result = sess.run(
-          [tt1, tt2, tt3, tt4])
-      self.assertAllEqual([1, 0], tt1_result)
-      self.assertAllClose([[0.1, 0.2], [0, 0]], tt2_result)
-      self.assertAllEqual([1, 2], tt3_result)
-      self.assertAllClose([[0.1, 0.2], [0.2, 0.4]], tt4_result)
-
-  def test_combines_static_dynamic_shape(self):
-    tensor = tf.placeholder(tf.float32, shape=(None, 2, 3))
-    combined_shape = shape_utils.combined_static_and_dynamic_shape(
-        tensor)
-    self.assertTrue(contrib_framework.is_tensor(combined_shape[0]))
-    self.assertListEqual(combined_shape[1:], [2, 3])
+
+    def graph_fn():
+      t1 = tf.constant([1], dtype=tf.int32)
+      tt1 = shape_utils.pad_or_clip_tensor(t1, tf.constant(2))
+      t2 = tf.constant([[0.1, 0.2]], dtype=tf.float32)
+      tt2 = shape_utils.pad_or_clip_tensor(t2, tf.constant(2))
+
+      t3 = tf.constant([1, 2, 3], dtype=tf.int32)
+      tt3 = shape_utils.clip_tensor(t3, tf.constant(2))
+      t4 = tf.constant([[0.1, 0.2], [0.2, 0.4], [0.5, 0.8]], dtype=tf.float32)
+      tt4 = shape_utils.clip_tensor(t4, tf.constant(2))
+
+      return tt1, tt2, tt3, tt4
+
+    tt1_result, tt2_result, tt3_result, tt4_result = self.execute(graph_fn, [])
+    self.assertAllEqual([1, 0], tt1_result)
+    self.assertAllClose([[0.1, 0.2], [0, 0]], tt2_result)
+    self.assertAllEqual([1, 2], tt3_result)
+    self.assertAllClose([[0.1, 0.2], [0.2, 0.4]], tt4_result)
+
+  def test_combined_static_dynamic_shape(self):
+
+    for n in [2, 3, 4]:
+      tensor = tf.zeros((n, 2, 3))
+      combined_shape = shape_utils.combined_static_and_dynamic_shape(
+          tensor)
+      self.assertListEqual(combined_shape[1:], [2, 3])
 
   def test_pad_or_clip_nd_tensor(self):
-    tensor_placeholder = tf.placeholder(tf.float32, [None, 5, 4, 7])
-    output_tensor = shape_utils.pad_or_clip_nd(
-        tensor_placeholder, [None, 3, 5, tf.constant(6)])
 
-    self.assertAllEqual(output_tensor.shape.as_list(), [None, 3, 5, None])
+    def graph_fn(input_tensor):
+      output_tensor = shape_utils.pad_or_clip_nd(
+          input_tensor, [None, 3, 5, tf.constant(6)])
 
-    with self.test_session() as sess:
-      output_tensor_np = sess.run(
-          output_tensor,
-          feed_dict={
-              tensor_placeholder: np.random.rand(2, 5, 4, 7),
-          })
+      return output_tensor
 
-    self.assertAllEqual(output_tensor_np.shape, [2, 3, 5, 6])
+    for n in [2, 3, 4, 5]:
+      input_np = np.zeros((n, 5, 4, 7))
+      output_tensor_np = self.execute(graph_fn, [input_np])
+      self.assertAllEqual(output_tensor_np.shape[1:], [3, 5, 6])
 
 
-class StaticOrDynamicMapFnTest(tf.test.TestCase):
+class StaticOrDynamicMapFnTest(test_case.TestCase):
 
   def test_with_dynamic_shape(self):
+
     def fn(input_tensor):
       return tf.reduce_sum(input_tensor)
-    input_tensor = tf.placeholder(tf.float32, shape=(None, 2))
-    map_fn_output = shape_utils.static_or_dynamic_map_fn(fn, input_tensor)
-
-    op_names = [op.name for op in tf.get_default_graph().get_operations()]
-    self.assertTrue(any(['map' == op_name[:3] for op_name in op_names]))
-
-    with self.test_session() as sess:
-      result1 = sess.run(
-          map_fn_output, feed_dict={
-              input_tensor: [[1, 2], [3, 1], [0, 4]]})
-      result2 = sess.run(
-          map_fn_output, feed_dict={
-              input_tensor: [[-1, 1], [0, 9]]})
-      self.assertAllEqual(result1, [3, 4, 4])
-      self.assertAllEqual(result2, [0, 9])
+
+    def graph_fn(input_tensor):
+      return shape_utils.static_or_dynamic_map_fn(fn, input_tensor)
+
+    # The input has different shapes, but due to how self.execute()
+    # works, the shape is known at graph compile time.
+    result1 = self.execute(
+        graph_fn, [np.array([[1, 2], [3, 1], [0, 4]]),])
+    result2 = self.execute(
+        graph_fn, [np.array([[-1, 1], [0, 9]]),])
+    self.assertAllEqual(result1, [3, 4, 4])
+    self.assertAllEqual(result2, [0, 9])
 
   def test_with_static_shape(self):
     def fn(input_tensor):
       return tf.reduce_sum(input_tensor)
-    input_tensor = tf.constant([[1, 2], [3, 1], [0, 4]], dtype=tf.float32)
-    map_fn_output = shape_utils.static_or_dynamic_map_fn(fn, input_tensor)
 
-    op_names = [op.name for op in tf.get_default_graph().get_operations()]
-    self.assertTrue(all(['map' != op_name[:3] for op_name in op_names]))
+    def graph_fn():
+      input_tensor = tf.constant([[1, 2], [3, 1], [0, 4]], dtype=tf.float32)
+      return shape_utils.static_or_dynamic_map_fn(fn, input_tensor)
 
-    with self.test_session() as sess:
-      result = sess.run(map_fn_output)
-      self.assertAllEqual(result, [3, 4, 4])
+    result = self.execute(graph_fn, [])
+    self.assertAllEqual(result, [3, 4, 4])
 
   def test_with_multiple_dynamic_shapes(self):
     def fn(elems):
       input_tensor, scalar_index_tensor = elems
       return tf.reshape(tf.slice(input_tensor, scalar_index_tensor, [1]), [])
 
-    input_tensor = tf.placeholder(tf.float32, shape=(None, 3))
-    scalar_index_tensor = tf.placeholder(tf.int32, shape=(None, 1))
-    map_fn_output = shape_utils.static_or_dynamic_map_fn(
-        fn, [input_tensor, scalar_index_tensor], dtype=tf.float32)
-
-    op_names = [op.name for op in tf.get_default_graph().get_operations()]
-    self.assertTrue(any(['map' == op_name[:3] for op_name in op_names]))
-
-    with self.test_session() as sess:
-      result1 = sess.run(
-          map_fn_output, feed_dict={
-              input_tensor: [[1, 2, 3], [4, 5, -1], [0, 6, 9]],
-              scalar_index_tensor: [[0], [2], [1]],
-          })
-      result2 = sess.run(
-          map_fn_output, feed_dict={
-              input_tensor: [[-1, 1, 0], [3, 9, 30]],
-              scalar_index_tensor: [[1], [0]]
-          })
-      self.assertAllEqual(result1, [1, -1, 6])
-      self.assertAllEqual(result2, [1, 3])
+    def graph_fn(input_tensor, scalar_index_tensor):
+      map_fn_output = shape_utils.static_or_dynamic_map_fn(
+          fn, [input_tensor, scalar_index_tensor], dtype=tf.float32)
+      return map_fn_output
+
+    # The input has different shapes, but due to how self.execute()
+    # works, the shape is known at graph compile time.
+
+    result1 = self.execute(
+        graph_fn, [
+            np.array([[1, 2, 3], [4, 5, -1], [0, 6, 9]]),
+            np.array([[0], [2], [1]]),
+        ])
+    result2 = self.execute(
+        graph_fn, [
+            np.array([[-1, 1, 0], [3, 9, 30]]),
+            np.array([[1], [0]])
+        ])
+    self.assertAllEqual(result1, [1, -1, 6])
+    self.assertAllEqual(result2, [1, 3])
 
   def test_with_multiple_static_shapes(self):
     def fn(elems):
       input_tensor, scalar_index_tensor = elems
       return tf.reshape(tf.slice(input_tensor, scalar_index_tensor, [1]), [])
 
-    input_tensor = tf.constant([[1, 2, 3], [4, 5, -1], [0, 6, 9]],
-                               dtype=tf.float32)
-    scalar_index_tensor = tf.constant([[0], [2], [1]], dtype=tf.int32)
-    map_fn_output = shape_utils.static_or_dynamic_map_fn(
-        fn, [input_tensor, scalar_index_tensor], dtype=tf.float32)
+    def graph_fn():
+      input_tensor = tf.constant([[1, 2, 3], [4, 5, -1], [0, 6, 9]],
+                                 dtype=tf.float32)
+      scalar_index_tensor = tf.constant([[0], [2], [1]], dtype=tf.int32)
+      map_fn_output = shape_utils.static_or_dynamic_map_fn(
+          fn, [input_tensor, scalar_index_tensor], dtype=tf.float32)
+      return map_fn_output
 
-    op_names = [op.name for op in tf.get_default_graph().get_operations()]
-    self.assertTrue(all(['map' != op_name[:3] for op_name in op_names]))
-
-    with self.test_session() as sess:
-      result = sess.run(map_fn_output)
-      self.assertAllEqual(result, [1, -1, 6])
+    result = self.execute(graph_fn, [])
+    self.assertAllEqual(result, [1, -1, 6])
 
   def test_fails_with_nested_input(self):
     def fn(input_tensor):
@@ -242,7 +241,7 @@ class StaticOrDynamicMapFnTest(tf.test.TestCase):
           fn, [input_tensor1, [input_tensor2]], dtype=tf.float32)
 
 
-class CheckMinImageShapeTest(tf.test.TestCase):
+class CheckMinImageShapeTest(test_case.TestCase):
 
   def test_check_min_image_dim_static_shape(self):
     input_tensor = tf.constant(np.zeros([1, 42, 42, 3]))
@@ -253,125 +252,151 @@ class CheckMinImageShapeTest(tf.test.TestCase):
       _ = shape_utils.check_min_image_dim(64, input_tensor)
 
   def test_check_min_image_dim_dynamic_shape(self):
-    input_placeholder = tf.placeholder(tf.float32, shape=[1, None, None, 3])
-    image_tensor = shape_utils.check_min_image_dim(33, input_placeholder)
 
-    with self.test_session() as sess:
-      sess.run(image_tensor,
-               feed_dict={input_placeholder: np.zeros([1, 42, 42, 3])})
-      with self.assertRaises(tf.errors.InvalidArgumentError):
-        sess.run(image_tensor,
-                 feed_dict={input_placeholder: np.zeros([1, 32, 32, 3])})
+    def graph_fn(input_tensor):
+      return shape_utils.check_min_image_dim(33, input_tensor)
+
+    self.execute(graph_fn,
+                 [np.zeros([1, 42, 42, 3])])
+    self.assertRaises(
+        ValueError, self.execute,
+        graph_fn, np.zeros([1, 32, 32, 3])
+    )
 
 
-class AssertShapeEqualTest(tf.test.TestCase):
+class AssertShapeEqualTest(test_case.TestCase):
 
   def test_unequal_static_shape_raises_exception(self):
     shape_a = tf.constant(np.zeros([4, 2, 2, 1]))
     shape_b = tf.constant(np.zeros([4, 2, 3, 1]))
-    with self.assertRaisesRegexp(
-        ValueError, 'Unequal shapes'):
-      shape_utils.assert_shape_equal(
-          shape_utils.combined_static_and_dynamic_shape(shape_a),
-          shape_utils.combined_static_and_dynamic_shape(shape_b))
+    self.assertRaisesRegex(
+        ValueError, 'Unequal shapes',
+        shape_utils.assert_shape_equal,
+        shape_utils.combined_static_and_dynamic_shape(shape_a),
+        shape_utils.combined_static_and_dynamic_shape(shape_b)
+    )
 
   def test_equal_static_shape_succeeds(self):
-    shape_a = tf.constant(np.zeros([4, 2, 2, 1]))
-    shape_b = tf.constant(np.zeros([4, 2, 2, 1]))
-    with self.test_session() as sess:
-      op = shape_utils.assert_shape_equal(
+
+    def graph_fn():
+      shape_a = tf.constant(np.zeros([4, 2, 2, 1]))
+      shape_b = tf.constant(np.zeros([4, 2, 2, 1]))
+
+      shape_utils.assert_shape_equal(
           shape_utils.combined_static_and_dynamic_shape(shape_a),
           shape_utils.combined_static_and_dynamic_shape(shape_b))
-      sess.run(op)
+
+      return tf.constant(0)
+
+    self.execute(graph_fn, [])
 
   def test_unequal_dynamic_shape_raises_tf_assert(self):
-    tensor_a = tf.placeholder(tf.float32, shape=[1, None, None, 3])
-    tensor_b = tf.placeholder(tf.float32, shape=[1, None, None, 3])
-    op = shape_utils.assert_shape_equal(
-        shape_utils.combined_static_and_dynamic_shape(tensor_a),
-        shape_utils.combined_static_and_dynamic_shape(tensor_b))
-    with self.test_session() as sess:
-      with self.assertRaises(tf.errors.InvalidArgumentError):
-        sess.run(op, feed_dict={tensor_a: np.zeros([1, 2, 2, 3]),
-                                tensor_b: np.zeros([1, 4, 4, 3])})
+
+    def graph_fn(tensor_a, tensor_b):
+      shape_utils.assert_shape_equal(
+          shape_utils.combined_static_and_dynamic_shape(tensor_a),
+          shape_utils.combined_static_and_dynamic_shape(tensor_b))
+      return tf.constant(0)
+
+    self.assertRaises(ValueError,
+                      self.execute, graph_fn,
+                      [np.zeros([1, 2, 2, 3]), np.zeros([1, 4, 4, 3])])
 
   def test_equal_dynamic_shape_succeeds(self):
-    tensor_a = tf.placeholder(tf.float32, shape=[1, None, None, 3])
-    tensor_b = tf.placeholder(tf.float32, shape=[1, None, None, 3])
-    op = shape_utils.assert_shape_equal(
-        shape_utils.combined_static_and_dynamic_shape(tensor_a),
-        shape_utils.combined_static_and_dynamic_shape(tensor_b))
-    with self.test_session() as sess:
-      sess.run(op, feed_dict={tensor_a: np.zeros([1, 2, 2, 3]),
-                              tensor_b: np.zeros([1, 2, 2, 3])})
+
+    def graph_fn(tensor_a, tensor_b):
+      shape_utils.assert_shape_equal(
+          shape_utils.combined_static_and_dynamic_shape(tensor_a),
+          shape_utils.combined_static_and_dynamic_shape(tensor_b)
+      )
+
+      return tf.constant(0)
+
+    self.execute(graph_fn, [np.zeros([1, 2, 2, 3]),
+                            np.zeros([1, 2, 2, 3])])
 
   def test_unequal_static_shape_along_first_dim_raises_exception(self):
     shape_a = tf.constant(np.zeros([4, 2, 2, 1]))
     shape_b = tf.constant(np.zeros([6, 2, 3, 1]))
-    with self.assertRaisesRegexp(
-        ValueError, 'Unequal first dimension'):
-      shape_utils.assert_shape_equal_along_first_dimension(
-          shape_utils.combined_static_and_dynamic_shape(shape_a),
-          shape_utils.combined_static_and_dynamic_shape(shape_b))
+
+    self.assertRaisesRegexp(
+        ValueError, 'Unequal first dimension',
+        shape_utils.assert_shape_equal_along_first_dimension,
+        shape_utils.combined_static_and_dynamic_shape(shape_a),
+        shape_utils.combined_static_and_dynamic_shape(shape_b)
+    )
 
   def test_equal_static_shape_along_first_dim_succeeds(self):
-    shape_a = tf.constant(np.zeros([4, 2, 2, 1]))
-    shape_b = tf.constant(np.zeros([4, 7, 2]))
-    with self.test_session() as sess:
-      op = shape_utils.assert_shape_equal_along_first_dimension(
+
+    def graph_fn():
+      shape_a = tf.constant(np.zeros([4, 2, 2, 1]))
+      shape_b = tf.constant(np.zeros([4, 7, 2]))
+      shape_utils.assert_shape_equal_along_first_dimension(
           shape_utils.combined_static_and_dynamic_shape(shape_a),
           shape_utils.combined_static_and_dynamic_shape(shape_b))
-      sess.run(op)
+      return tf.constant(0)
+
+    self.execute(graph_fn, [])
 
   def test_unequal_dynamic_shape_along_first_dim_raises_tf_assert(self):
-    tensor_a = tf.placeholder(tf.float32, shape=[None, None, None, 3])
-    tensor_b = tf.placeholder(tf.float32, shape=[None, None, 3])
-    op = shape_utils.assert_shape_equal_along_first_dimension(
-        shape_utils.combined_static_and_dynamic_shape(tensor_a),
-        shape_utils.combined_static_and_dynamic_shape(tensor_b))
-    with self.test_session() as sess:
-      with self.assertRaises(tf.errors.InvalidArgumentError):
-        sess.run(op, feed_dict={tensor_a: np.zeros([1, 2, 2, 3]),
-                                tensor_b: np.zeros([2, 4, 3])})
+
+    def graph_fn(tensor_a, tensor_b):
+      shape_utils.assert_shape_equal_along_first_dimension(
+          shape_utils.combined_static_and_dynamic_shape(tensor_a),
+          shape_utils.combined_static_and_dynamic_shape(tensor_b))
+
+      return tf.constant(0)
+
+    self.assertRaises(ValueError,
+                      self.execute, graph_fn,
+                      [np.zeros([1, 2, 2, 3]), np.zeros([2, 4, 3])])
 
   def test_equal_dynamic_shape_along_first_dim_succeeds(self):
-    tensor_a = tf.placeholder(tf.float32, shape=[None, None, None, 3])
-    tensor_b = tf.placeholder(tf.float32, shape=[None])
-    op = shape_utils.assert_shape_equal_along_first_dimension(
-        shape_utils.combined_static_and_dynamic_shape(tensor_a),
-        shape_utils.combined_static_and_dynamic_shape(tensor_b))
-    with self.test_session() as sess:
-      sess.run(op, feed_dict={tensor_a: np.zeros([5, 2, 2, 3]),
-                              tensor_b: np.zeros([5])})
+
+    def graph_fn(tensor_a, tensor_b):
+      shape_utils.assert_shape_equal_along_first_dimension(
+          shape_utils.combined_static_and_dynamic_shape(tensor_a),
+          shape_utils.combined_static_and_dynamic_shape(tensor_b))
+      return tf.constant(0)
+
+    self.execute(graph_fn, [np.zeros([5, 2, 2, 3]), np.zeros([5])])
 
 
-class FlattenExpandDimensionTest(tf.test.TestCase):
+class FlattenExpandDimensionTest(test_case.TestCase):
 
   def test_flatten_given_dims(self):
-    inputs = tf.random_uniform([5, 2, 10, 10, 3])
-    actual_flattened = shape_utils.flatten_dimensions(inputs, first=1, last=3)
-    expected_flattened = tf.reshape(inputs, [5, 20, 10, 3])
-    with self.test_session() as sess:
-      (actual_flattened_np,
-       expected_flattened_np) = sess.run([actual_flattened, expected_flattened])
+
+    def graph_fn():
+      inputs = tf.random_uniform([5, 2, 10, 10, 3])
+      actual_flattened = shape_utils.flatten_dimensions(inputs, first=1, last=3)
+      expected_flattened = tf.reshape(inputs, [5, 20, 10, 3])
+
+      return actual_flattened, expected_flattened
+
+    (actual_flattened_np,
+     expected_flattened_np) = self.execute(graph_fn, [])
     self.assertAllClose(expected_flattened_np, actual_flattened_np)
 
   def test_raises_value_error_incorrect_dimensions(self):
     inputs = tf.random_uniform([5, 2, 10, 10, 3])
-    with self.assertRaises(ValueError):
-      shape_utils.flatten_dimensions(inputs, first=0, last=6)
+    self.assertRaises(ValueError,
+                      shape_utils.flatten_dimensions, inputs,
+                      first=0, last=6)
 
   def test_flatten_first_two_dimensions(self):
-    inputs = tf.constant(
-        [
-            [[1, 2], [3, 4]],
-            [[5, 6], [7, 8]],
-            [[9, 10], [11, 12]]
-        ], dtype=tf.int32)
-    flattened_tensor = shape_utils.flatten_first_n_dimensions(
-        inputs, 2)
-    with self.test_session() as sess:
-      flattened_tensor_out = sess.run(flattened_tensor)
+
+    def graph_fn():
+      inputs = tf.constant(
+          [
+              [[1, 2], [3, 4]],
+              [[5, 6], [7, 8]],
+              [[9, 10], [11, 12]]
+          ], dtype=tf.int32)
+      flattened_tensor = shape_utils.flatten_first_n_dimensions(
+          inputs, 2)
+      return flattened_tensor
+
+    flattened_tensor_out = self.execute(graph_fn, [])
 
     expected_output = [[1, 2],
                        [3, 4],
@@ -382,20 +407,23 @@ class FlattenExpandDimensionTest(tf.test.TestCase):
     self.assertAllEqual(expected_output, flattened_tensor_out)
 
   def test_expand_first_dimension(self):
-    inputs = tf.constant(
-        [
-            [1, 2],
-            [3, 4],
-            [5, 6],
-            [7, 8],
-            [9, 10],
-            [11, 12]
-        ], dtype=tf.int32)
-    dims = [3, 2]
-    expanded_tensor = shape_utils.expand_first_dimension(
-        inputs, dims)
-    with self.test_session() as sess:
-      expanded_tensor_out = sess.run(expanded_tensor)
+
+    def graph_fn():
+      inputs = tf.constant(
+          [
+              [1, 2],
+              [3, 4],
+              [5, 6],
+              [7, 8],
+              [9, 10],
+              [11, 12]
+          ], dtype=tf.int32)
+      dims = [3, 2]
+      expanded_tensor = shape_utils.expand_first_dimension(
+          inputs, dims)
+      return expanded_tensor
+
+    expanded_tensor_out = self.execute(graph_fn, [])
 
     expected_output = [
         [[1, 2], [3, 4]],
@@ -404,19 +432,20 @@ class FlattenExpandDimensionTest(tf.test.TestCase):
     self.assertAllEqual(expected_output, expanded_tensor_out)
 
   def test_expand_first_dimension_with_incompatible_dims(self):
-    inputs_default = tf.constant(
-        [
-            [[1, 2]],
-            [[3, 4]],
-            [[5, 6]],
-        ], dtype=tf.int32)
-    inputs = tf.placeholder_with_default(inputs_default, [None, 1, 2])
-    dims = [3, 2]
-    expanded_tensor = shape_utils.expand_first_dimension(
-        inputs, dims)
-    with self.test_session() as sess:
-      with self.assertRaises(tf.errors.InvalidArgumentError):
-        sess.run(expanded_tensor)
+
+    def graph_fn():
+      inputs = tf.constant(
+          [
+              [[1, 2]],
+              [[3, 4]],
+              [[5, 6]],
+          ], dtype=tf.int32)
+      dims = [3, 2]
+      expanded_tensor = shape_utils.expand_first_dimension(
+          inputs, dims)
+      return expanded_tensor
+
+    self.assertRaises(ValueError, self.execute, graph_fn, [])
 
 
 if __name__ == '__main__':

research/object_detection/utils/spatial_transform_ops.py

@@ -18,7 +18,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def _coordinate_vector_1d(start, end, size, align_endpoints):

research/object_detection/utils/spatial_transform_ops_test.py

@@ -20,7 +20,7 @@ from __future__ import print_function
 
 import numpy as np
 from six.moves import range
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from object_detection.utils import spatial_transform_ops as spatial_ops
 from object_detection.utils import test_case

research/object_detection/utils/static_shape_test.py

@@ -19,7 +19,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from object_detection.utils import static_shape
 

research/object_detection/utils/target_assigner_utils.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 """Utility functions used by target assigner."""
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from object_detection.utils import shape_utils
 

research/object_detection/utils/target_assigner_utils_test.py

@@ -15,7 +15,7 @@
 """Tests for utils.target_assigner_utils."""
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from object_detection.utils import target_assigner_utils as ta_utils
 from object_detection.utils import test_case

research/object_detection/utils/test_case.py

@@ -18,8 +18,9 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 from six.moves import zip
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from tensorflow.python import tf2  # pylint: disable=import-outside-toplevel
+from object_detection.utils import tf_version
 if not tf2.enabled():
   from tensorflow.contrib import tpu as contrib_tpu  # pylint: disable=g-import-not-at-top, line-too-long
 
@@ -58,9 +59,9 @@ class TestCase(tf.test.TestCase):
 
   def is_tf2(self):
     """Returns whether TF2 is enabled."""
-    return tf2.enabled()
+    return tf_version.is_tf2()
 
-  def execute_tpu_tf1(self, compute_fn, inputs):
+  def execute_tpu_tf1(self, compute_fn, inputs, graph=None):
     """Executes compute_fn on TPU with Tensorflow 1.X.
 
     Args:
@@ -68,11 +69,13 @@ class TestCase(tf.test.TestCase):
         of input numpy tensors, performs computation and returns output numpy
         tensors.
       inputs: a list of numpy arrays to feed input to the `compute_fn`.
+      graph: (optional) If not None, provided `graph` is used for computation
+        instead of a brand new tf.Graph().
 
     Returns:
       A list of numpy arrays or a single numpy array.
     """
-    with self.test_session(graph=tf.Graph()) as sess:
+    with self.session(graph=(graph or tf.Graph())) as sess:
       placeholders = [tf.placeholder_with_default(v, v.shape) for v in inputs]
       def wrap_graph_fn(*args, **kwargs):
         results = compute_fn(*args, **kwargs)
@@ -117,7 +120,7 @@ class TestCase(tf.test.TestCase):
     tf.tpu.experimental.shutdown_tpu_system()
     return self.maybe_extract_single_output(outputs)
 
-  def execute_cpu_tf1(self, compute_fn, inputs):
+  def execute_cpu_tf1(self, compute_fn, inputs, graph=None):
     """Executes compute_fn on CPU with Tensorflow 1.X.
 
     Args:
@@ -125,13 +128,15 @@ class TestCase(tf.test.TestCase):
         of input numpy tensors, performs computation and returns output numpy
         tensors.
       inputs: a list of numpy arrays to feed input to the `compute_fn`.
+      graph: (optional) If not None, provided `graph` is used for computation
+        instead of a brand new tf.Graph().
 
     Returns:
       A list of numpy arrays or a single numpy array.
     """
     if self.is_tf2():
       raise ValueError('Required version Tenforflow 1.X is not available.')
-    with self.test_session(graph=tf.Graph()) as sess:
+    with self.session(graph=(graph or tf.Graph())) as sess:
       placeholders = [tf.placeholder_with_default(v, v.shape) for v in inputs]
       results = compute_fn(*placeholders)
       if (not (isinstance(results, dict) or isinstance(results, tf.Tensor)) and
@@ -163,7 +168,7 @@ class TestCase(tf.test.TestCase):
       return compute_fn(*tf_inputs)
     return self.maybe_extract_single_output(run())
 
-  def execute_cpu(self, compute_fn, inputs):
+  def execute_cpu(self, compute_fn, inputs, graph=None):
     """Executes compute_fn on CPU.
 
     Depending on the underlying TensorFlow installation (build deps) runs in
@@ -174,6 +179,8 @@ class TestCase(tf.test.TestCase):
         of input numpy tensors, performs computation and returns output numpy
         tensors.
       inputs: a list of numpy arrays to feed input to the `compute_fn`.
+      graph: (optional) If not None, provided `graph` is used for computation
+        instead of a brand new tf.Graph().
 
     Returns:
       A list of numpy arrays or a single tensor.
@@ -181,9 +188,9 @@ class TestCase(tf.test.TestCase):
     if self.is_tf2():
       return self.execute_cpu_tf2(compute_fn, inputs)
     else:
-      return self.execute_cpu_tf1(compute_fn, inputs)
+      return self.execute_cpu_tf1(compute_fn, inputs, graph)
 
-  def execute_tpu(self, compute_fn, inputs):
+  def execute_tpu(self, compute_fn, inputs, graph=None):
     """Executes compute_fn on TPU.
 
     Depending on the underlying TensorFlow installation (build deps) runs in
@@ -194,6 +201,8 @@ class TestCase(tf.test.TestCase):
         of input numpy tensors, performs computation and returns output numpy
         tensors.
       inputs: a list of numpy arrays to feed input to the `compute_fn`.
+      graph: (optional) If not None, provided `graph` is used for computation
+        instead of a brand new tf.Graph().
 
     Returns:
       A list of numpy arrays or a single tensor.
@@ -203,7 +212,7 @@ class TestCase(tf.test.TestCase):
     if self.is_tf2():
       return self.execute_tpu_tf2(compute_fn, inputs)
     else:
-      return self.execute_tpu_tf1(compute_fn, inputs)
+      return self.execute_tpu_tf1(compute_fn, inputs, graph)
 
   def execute_tf2(self, compute_fn, inputs):
     """Runs compute_fn with TensorFlow 2.0.
@@ -226,7 +235,7 @@ class TestCase(tf.test.TestCase):
     else:
       return self.execute_cpu_tf2(compute_fn, inputs)
 
-  def execute_tf1(self, compute_fn, inputs):
+  def execute_tf1(self, compute_fn, inputs, graph=None):
     """Runs compute_fn with TensorFlow 1.X.
 
     Executes on TPU if available, otherwise executes on CPU.
@@ -236,6 +245,8 @@ class TestCase(tf.test.TestCase):
         of input numpy tensors, performs computation and returns output numpy
         tensors.
       inputs: a list of numpy arrays to feed input to the `compute_fn`.
+      graph: (optional) If not None, provided `graph` is used for computation
+        instead of a brand new tf.Graph().
 
     Returns:
       A list of numpy arrays or a single tensor.
@@ -243,11 +254,11 @@ class TestCase(tf.test.TestCase):
     if self.is_tf2():
       raise ValueError('Required version Tenforflow 1.X is not available.')
     if self.has_tpu():
-      return self.execute_tpu_tf1(compute_fn, inputs)
+      return self.execute_tpu_tf1(compute_fn, inputs, graph)
     else:
-      return self.execute_cpu_tf1(compute_fn, inputs)
+      return self.execute_cpu_tf1(compute_fn, inputs, graph)
 
-  def execute(self, compute_fn, inputs):
+  def execute(self, compute_fn, inputs, graph=None):
     """Runs compute_fn with inputs and returns results.
 
     * Executes in either TF1.X or TF2.X style based on the TensorFlow version.
@@ -258,6 +269,8 @@ class TestCase(tf.test.TestCase):
         of input numpy tensors, performs computation and returns output numpy
         tensors.
       inputs: a list of numpy arrays to feed input to the `compute_fn`.
+      graph: (optional) If not None, provided `graph` is used for computation
+        instead of a brand new tf.Graph().
 
     Returns:
       A list of numpy arrays or a single tensor.
@@ -267,6 +280,6 @@ class TestCase(tf.test.TestCase):
     elif not self.has_tpu() and tf2.enabled():
       return self.execute_cpu_tf2(compute_fn, inputs)
     elif self.has_tpu() and not tf2.enabled():
-      return self.execute_tpu_tf1(compute_fn, inputs)
+      return self.execute_tpu_tf1(compute_fn, inputs, graph)
     else:
-      return self.execute_cpu_tf1(compute_fn, inputs)
+      return self.execute_cpu_tf1(compute_fn, inputs, graph)

research/object_detection/utils/test_case_test.py

@@ -15,7 +15,7 @@
 """Tests for google3.third_party.tensorflow_models.object_detection.utils.test_case."""
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from object_detection.utils import test_case
 
 

research/object_detection/utils/test_utils.py

@@ -21,7 +21,7 @@ from __future__ import print_function
 import numpy as np
 from six.moves import range
 from six.moves import zip
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from object_detection.core import anchor_generator
 from object_detection.core import box_coder
@@ -29,6 +29,7 @@ from object_detection.core import box_list
 from object_detection.core import box_predictor
 from object_detection.core import matcher
 from object_detection.utils import shape_utils
+from object_detection.utils import tf_version
 
 # Default size (both width and height) used for testing mask predictions.
 DEFAULT_MASK_SIZE = 5
@@ -233,3 +234,40 @@ def first_rows_close_as_set(a, b, k=None, rtol=1e-6, atol=1e-6):
       np.allclose(entry_a, entry_b, rtol, atol)
       for (entry_a, entry_b) in zip(a_sorted, b_sorted)
   ])
+
+
+class GraphContextOrNone(object):
+  """A new Graph context for TF1.X and None for TF2.X.
+
+  This is useful to write model tests that work with both TF1.X and TF2.X.
+
+  Example test using this pattern:
+
+  class ModelTest(test_case.TestCase):
+    def test_model(self):
+      with test_utils.GraphContextOrNone() as g:
+        model = Model()
+      def compute_fn():
+        out = model.predict()
+        return out['detection_boxes']
+      boxes = self.execute(compute_fn, [], graph=g)
+      self.assertAllClose(boxes, expected_boxes)
+  """
+
+  def __init__(self):
+    if tf_version.is_tf2():
+      self.graph = None
+    else:
+      self.graph = tf.Graph().as_default()
+
+  def __enter__(self):
+    if tf_version.is_tf2():
+      return None
+    else:
+      return self.graph.__enter__()
+
+  def __exit__(self, ttype, value, traceback):
+    if tf_version.is_tf2():
+      return False
+    else:
+      return self.graph.__exit__(ttype, value, traceback)

research/object_detection/utils/test_utils_test.py

@@ -20,12 +20,13 @@ from __future__ import division
 from __future__ import print_function
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
+from object_detection.utils import test_case
 from object_detection.utils import test_utils
 
 
-class TestUtilsTest(tf.test.TestCase):
+class TestUtilsTest(test_case.TestCase):
 
   def test_diagonal_gradient_image(self):
     """Tests if a good pyramid image is created."""
@@ -67,10 +68,10 @@ class TestUtilsTest(tf.test.TestCase):
     self.assertAllEqual(boxes[:, 0] < boxes[:, 2], true_column)
     self.assertAllEqual(boxes[:, 1] < boxes[:, 3], true_column)
 
-    self.assertTrue(boxes[:, 0].min() >= 0)
-    self.assertTrue(boxes[:, 1].min() >= 0)
-    self.assertTrue(boxes[:, 2].max() <= max_height)
-    self.assertTrue(boxes[:, 3].max() <= max_width)
+    self.assertGreaterEqual(boxes[:, 0].min(), 0)
+    self.assertGreaterEqual(boxes[:, 1].min(), 0)
+    self.assertLessEqual(boxes[:, 2].max(), max_height)
+    self.assertLessEqual(boxes[:, 3].max(), max_width)
 
   def test_first_rows_close_as_set(self):
     a = [1, 2, 3, 0, 0]

research/object_detection/utils/variables_helper.py

@@ -23,9 +23,9 @@ from __future__ import print_function
 import logging
 import re
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
+import tf_slim as slim
 
-from tensorflow.contrib import slim
 from tensorflow.python.ops import variables as tf_variables
 
 
@@ -47,6 +47,8 @@ def filter_variables(variables, filter_regex_list, invert=False):
   Returns:
     a list of filtered variables.
   """
+  if tf.executing_eagerly():
+    raise ValueError('Accessing variables is not supported in eager mode.')
   kept_vars = []
   variables_to_ignore_patterns = list([fre for fre in filter_regex_list if fre])
   for var in variables:
@@ -72,6 +74,8 @@ def multiply_gradients_matching_regex(grads_and_vars, regex_list, multiplier):
   Returns:
     grads_and_vars: A list of gradient to variable pairs (tuples).
   """
+  if tf.executing_eagerly():
+    raise ValueError('Accessing variables is not supported in eager mode.')
   variables = [pair[1] for pair in grads_and_vars]
   matching_vars = filter_variables(variables, regex_list, invert=True)
   for var in matching_vars:
@@ -93,6 +97,8 @@ def freeze_gradients_matching_regex(grads_and_vars, regex_list):
     grads_and_vars: A list of gradient to variable pairs (tuples) that do not
       contain the variables and gradients matching the regex.
   """
+  if tf.executing_eagerly():
+    raise ValueError('Accessing variables is not supported in eager mode.')
   variables = [pair[1] for pair in grads_and_vars]
   matching_vars = filter_variables(variables, regex_list, invert=True)
   kept_grads_and_vars = [pair for pair in grads_and_vars
@@ -123,6 +129,8 @@ def get_variables_available_in_checkpoint(variables,
   Raises:
     ValueError: if `variables` is not a list or dict.
   """
+  if tf.executing_eagerly():
+    raise ValueError('Accessing variables is not supported in eager mode.')
   if isinstance(variables, list):
     variable_names_map = {}
     for variable in variables:
@@ -170,6 +178,8 @@ def get_global_variables_safely():
   Returns:
     The result of tf.global_variables()
   """
+  if tf.executing_eagerly():
+    raise ValueError('Accessing variables is not supported in eager mode.')
   with tf.init_scope():
     if tf.executing_eagerly():
       raise ValueError("Global variables collection is not tracked when "

research/object_detection/utils/variables_helper_test.py

@@ -21,12 +21,13 @@ from __future__ import print_function
 
 import os
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
+from object_detection.utils import test_case
 from object_detection.utils import variables_helper
 
 
-class FilterVariablesTest(tf.test.TestCase):
+class FilterVariablesTest(test_case.TestCase):
 
   def _create_variables(self):
     return [tf.Variable(1.0, name='FeatureExtractor/InceptionV3/weights'),
@@ -37,26 +38,26 @@ class FilterVariablesTest(tf.test.TestCase):
   def test_return_all_variables_when_empty_regex(self):
     variables = self._create_variables()
     out_variables = variables_helper.filter_variables(variables, [''])
-    self.assertItemsEqual(out_variables, variables)
+    self.assertCountEqual(out_variables, variables)
 
   def test_return_variables_which_do_not_match_single_regex(self):
     variables = self._create_variables()
     out_variables = variables_helper.filter_variables(variables,
                                                       ['FeatureExtractor/.*'])
-    self.assertItemsEqual(out_variables, variables[2:])
+    self.assertCountEqual(out_variables, variables[2:])
 
   def test_return_variables_which_do_not_match_any_regex_in_list(self):
     variables = self._create_variables()
     out_variables = variables_helper.filter_variables(variables, [
         'FeatureExtractor.*biases', 'StackProposalGenerator.*biases'
     ])
-    self.assertItemsEqual(out_variables, [variables[0], variables[2]])
+    self.assertCountEqual(out_variables, [variables[0], variables[2]])
 
   def test_return_variables_matching_empty_regex_list(self):
     variables = self._create_variables()
     out_variables = variables_helper.filter_variables(
         variables, [''], invert=True)
-    self.assertItemsEqual(out_variables, [])
+    self.assertCountEqual(out_variables, [])
 
   def test_return_variables_matching_some_regex_in_list(self):
     variables = self._create_variables()
@@ -64,7 +65,7 @@ class FilterVariablesTest(tf.test.TestCase):
         variables,
         ['FeatureExtractor.*biases', 'StackProposalGenerator.*biases'],
         invert=True)
-    self.assertItemsEqual(out_variables, [variables[1], variables[3]])
+    self.assertCountEqual(out_variables, [variables[1], variables[3]])
 
 
 class MultiplyGradientsMatchingRegexTest(tf.test.TestCase):
@@ -90,7 +91,7 @@ class MultiplyGradientsMatchingRegexTest(tf.test.TestCase):
     with self.test_session() as sess:
       sess.run(init_op)
       output = sess.run(grads_and_vars)
-      self.assertItemsEqual(output, exp_output)
+      self.assertCountEqual(output, exp_output)
 
   def test_multiply_all_bias_variables(self):
     grads_and_vars = self._create_grads_and_vars()
@@ -103,10 +104,10 @@ class MultiplyGradientsMatchingRegexTest(tf.test.TestCase):
     with self.test_session() as sess:
       sess.run(init_op)
       output = sess.run(grads_and_vars)
-      self.assertItemsEqual(output, exp_output)
+      self.assertCountEqual(output, exp_output)
 
 
-class FreezeGradientsMatchingRegexTest(tf.test.TestCase):
+class FreezeGradientsMatchingRegexTest(test_case.TestCase):
 
   def _create_grads_and_vars(self):
     return [(tf.constant(1.0),
@@ -128,10 +129,10 @@ class FreezeGradientsMatchingRegexTest(tf.test.TestCase):
     with self.test_session() as sess:
       sess.run(init_op)
       output = sess.run(grads_and_vars)
-      self.assertItemsEqual(output, exp_output)
+      self.assertCountEqual(output, exp_output)
 
 
-class GetVariablesAvailableInCheckpointTest(tf.test.TestCase):
+class GetVariablesAvailableInCheckpointTest(test_case.TestCase):
 
   def test_return_all_variables_from_checkpoint(self):
     with tf.Graph().as_default():
@@ -147,7 +148,7 @@ class GetVariablesAvailableInCheckpointTest(tf.test.TestCase):
         saver.save(sess, checkpoint_path)
       out_variables = variables_helper.get_variables_available_in_checkpoint(
           variables, checkpoint_path)
-    self.assertItemsEqual(out_variables, variables)
+    self.assertCountEqual(out_variables, variables)
 
   def test_return_all_variables_from_checkpoint_with_partition(self):
     with tf.Graph().as_default():
@@ -165,7 +166,7 @@ class GetVariablesAvailableInCheckpointTest(tf.test.TestCase):
         saver.save(sess, checkpoint_path)
       out_variables = variables_helper.get_variables_available_in_checkpoint(
           variables, checkpoint_path)
-    self.assertItemsEqual(out_variables, variables)
+    self.assertCountEqual(out_variables, variables)
 
   def test_return_variables_available_in_checkpoint(self):
     checkpoint_path = os.path.join(self.get_temp_dir(), 'model.ckpt')
@@ -186,7 +187,7 @@ class GetVariablesAvailableInCheckpointTest(tf.test.TestCase):
       graph2_variables = graph1_variables + [tf.Variable(1.0, name='biases')]
       out_variables = variables_helper.get_variables_available_in_checkpoint(
           graph2_variables, checkpoint_path, include_global_step=False)
-    self.assertItemsEqual(out_variables, [weight_variable])
+    self.assertCountEqual(out_variables, [weight_variable])
 
   def test_return_variables_available_an_checkpoint_with_dict_inputs(self):
     checkpoint_path = os.path.join(self.get_temp_dir(), 'model.ckpt')
@@ -208,9 +209,9 @@ class GetVariablesAvailableInCheckpointTest(tf.test.TestCase):
       out_variables = variables_helper.get_variables_available_in_checkpoint(
           graph2_variables_dict, checkpoint_path)
 
-    self.assertTrue(isinstance(out_variables, dict))
-    self.assertItemsEqual(list(out_variables.keys()), ['ckpt_weights'])
-    self.assertTrue(out_variables['ckpt_weights'].op.name == 'weights')
+    self.assertIsInstance(out_variables, dict)
+    self.assertCountEqual(list(out_variables.keys()), ['ckpt_weights'])
+    self.assertEqual(out_variables['ckpt_weights'].op.name, 'weights')
 
   def test_return_variables_with_correct_sizes(self):
     checkpoint_path = os.path.join(self.get_temp_dir(), 'model.ckpt')
@@ -237,7 +238,7 @@ class GetVariablesAvailableInCheckpointTest(tf.test.TestCase):
 
     out_variables = variables_helper.get_variables_available_in_checkpoint(
         graph2_variables, checkpoint_path, include_global_step=True)
-    self.assertItemsEqual(out_variables, [bias_variable, global_step])
+    self.assertCountEqual(out_variables, [bias_variable, global_step])
 
 
 if __name__ == '__main__':