remove duplicate function

Author: fsx950223

efficientdet/utils.py

@@ -20,7 +20,6 @@
 import numpy as np
 import tensorflow.compat.v1 as tf
 import tensorflow.compat.v2 as tf2
-from tensorflow.python.eager import tape as tape_lib  # pylint:disable=g-direct-tensorflow-import
 from tensorflow.python.tpu import tpu_function  # pylint:disable=g-direct-tensorflow-import
 # pylint: disable=logging-format-interpolation
 
@@ -643,83 +642,12 @@ def build_model_with_precision(pp, mm, ii, *args, **kwargs):
   return outputs
 
 
-def _recompute_grad(f):
-  """An eager-compatible version of recompute_grad.
-
-  For f(*args, **kwargs), this supports gradients with respect to args or
-  kwargs, but kwargs are currently only supported in eager-mode.
-  Note that for keras layer and model objects, this is handled automatically.
-
-  Warning: If `f` was originally a tf.keras Model or Layer object, `g` will not
-  be able to access the member variables of that object, because `g` returns
-  through the wrapper function `inner`.  When recomputing gradients through
-  objects that inherit from tf2, we suggest keeping a reference to the
-  underlying object around for the purpose of accessing these variables.
-
-  Args:
-    f: function `f(*x)` that returns a `Tensor` or sequence of `Tensor` outputs.
-
-  Returns:
-   A function `g` that wraps `f`, but which recomputes `f` on the backwards
-   pass of a gradient call.
-  """
-
-  @tf.custom_gradient
-  def inner(*args, **kwargs):
-    """Inner function closure for calculating gradients."""
-    current_var_scope = tf.get_variable_scope()
-    with tape_lib.stop_recording():
-      result = f(*args, **kwargs)
-
-    def grad_wrapper(*wrapper_args, **grad_kwargs):
-      """Wrapper function to accomodate lack of kwargs in graph mode decorator."""
-
-      @tf.custom_gradient
-      def inner_recompute_grad(*dresult):
-        """Nested custom gradient function for computing grads in reverse and forward mode autodiff."""
-        # Gradient calculation for reverse mode autodiff.
-        variables = grad_kwargs.get('variables')
-        with tf.GradientTape() as t:
-          id_args = tf.nest.map_structure(tf.identity, args)
-          t.watch(id_args)
-          if variables is not None:
-            t.watch(variables)
-          with tf.control_dependencies(dresult):
-            with tf.variable_scope(current_var_scope):
-              result = f(*id_args, **kwargs)
-        kw_vars = []
-        if variables is not None:
-          kw_vars = list(variables)
-        grads = t.gradient(
-            result,
-            list(id_args) + kw_vars,
-            output_gradients=dresult,
-            unconnected_gradients=tf.UnconnectedGradients.ZERO)
-
-        def transpose(*t_args, **t_kwargs):
-          """Gradient function calculation for forward mode autodiff."""
-          # Just throw an error since gradients / activations are not stored on
-          # tape for recompute.
-          raise NotImplementedError(
-              'recompute_grad tried to transpose grad of {}. '
-              'Consider not using recompute_grad in forward mode'
-              'autodiff'.format(f.__name__))
-
-        return (grads[:len(id_args)], grads[len(id_args):]), transpose
-
-      return inner_recompute_grad(*wrapper_args)
-
-    return result, grad_wrapper
-
-  return inner
-
-
 def recompute_grad(recompute=False):
   """Decorator determine whether use gradient checkpoint."""
 
   def _wrapper(f):
     if recompute:
-      return _recompute_grad(f)
+      return tf.recompute_grad(f)
     return f
 
   return _wrapper