def _fn_with_diet_vars(fn, args, params):
  """Call function with args; use diet variables according to params."""

  vs_ctr = []

  def grad_fn(inputs, variables, outputs, output_grads):
    """Custom gradient function."""
    del outputs  # recomputing below
    with common_layers.fn_device_dependency("diet_grad",
                                            output_grads[0].device) as out_dep:
      with tf.variable_scope(vs_ctr[0], reuse=True):
        outputs = fn(*inputs)

      variables = [common_layers.underlying_variable_ref(v) for v in variables]
      dequantized_variables = [
          params.dequantized[v.name][-1] for v in variables
      ]

      grads = tf.gradients(outputs, inputs + dequantized_variables,
                           output_grads)
      grad_inputs = grads[:len(inputs)]
      grad_variables = grads[len(inputs):]

      opt = _create_diet_optimizer(params)

      # Apply grad_variables here
      var_updates = []
      for v, dv in zip(variables, grad_variables):
        with tf.variable_scope(vs_ctr[0].name):
          opt.create_slots(v)
        update_op = opt.update_variable(v, dv)
        var_updates.append(update_op)

      with tf.control_dependencies(var_updates):
        grad_inputs = [tf.identity(dx) for dx in grad_inputs]

      out_dep.append(grad_inputs)

      return grad_inputs, [None] * len(variables)

  @common_layers.fn_with_custom_grad(grad_fn, use_global_vars=True)
  def forward(*inputs):
    with tf.variable_scope(
        None, default_name="diet",
        custom_getter=make_diet_var_getter(params)) as vs:
      vs_ctr.append(vs)
      outputs = fn(*inputs)
      return outputs

  with common_layers.fn_device_dependency("diet_forward",
                                          args[0].device) as out_dep:
    outputs = forward(*args)
    out_dep.append(outputs)
  return outputs
def _fn_with_diet_vars ( fn , args , params ) : vs_ctr = [ ] def grad_fn ( inputs , variables , outputs , output_grads ) : del outputs with common_layers . fn_device_dependency ( "diet_grad" , output_grads [ 0 ] . device ) as out_dep : with tf . variable_scope ( vs_ctr [ 0 ] , reuse = True ) : outputs = fn ( * inputs ) variables = [ common_layers . underlying_variable_ref ( v ) for v in variables ] dequantized_variables = [ params . dequantized [ v . name ] [ - 1 ] for v in variables ] grads = tf . gradients ( outputs , inputs + dequantized_variables , output_grads ) grad_inputs = grads [ : len ( inputs ) ] grad_variables = grads [ len ( inputs ) : ] opt = _create_diet_optimizer ( params ) var_updates = [ ] for v , dv in zip ( variables , grad_variables ) : with tf . variable_scope ( vs_ctr [ 0 ] . name ) : opt . create_slots ( v ) update_op = opt . update_variable ( v , dv ) var_updates . append ( update_op ) with tf . control_dependencies ( var_updates ) : grad_inputs = [ tf . identity ( dx ) for dx in grad_inputs ] out_dep . append ( grad_inputs ) return grad_inputs , [ None ] * len ( variables ) @ common_layers . fn_with_custom_grad ( grad_fn , use_global_vars = True ) def forward ( * inputs ) : with tf . variable_scope ( None , default_name = "diet" , custom_getter = make_diet_var_getter ( params ) ) as vs : vs_ctr . append ( vs ) outputs = fn ( * inputs ) return outputs with common_layers . fn_device_dependency ( "diet_forward" , args [ 0 ] . device ) as out_dep : outputs = forward ( * args ) out_dep . append ( outputs ) return outputs
Call function with args; use diet variables according to params.