Beispiel #1
0
def aggregate_indexed_slices_gradients(grads):
    """Aggregates gradients containing `IndexedSlices`s."""
    if len(grads) < 1:
        return None
    elif len(grads) == 1:
        return grads[0]
    else:
        grads = [g for g in grads if g is not None]
        # If any gradient is a `Tensor`, sum them up and return a dense tensor
        # object.
        if any(isinstance(g, ops.Tensor) for g in grads):
            return math_ops.add_n(grads)

        # The following `_as_indexed_slices_list` casts ids of IndexedSlices into
        # int64. It is to make sure the inputs of `concat` all have same the data
        # type.
        grads = math_ops._as_indexed_slices_list(grads)  # pylint: disable=protected-access

        grads = [flatten_nested_indexed_slices(x) for x in grads]
        # Form IndexedSlices out of the concatenated values and indices.
        concat_grad = ops.IndexedSlices(
            array_ops.concat([x.values for x in grads], axis=0),
            array_ops.concat([x.indices for x in grads], axis=0),
            grads[0].dense_shape)

        return concat_grad
Beispiel #2
0
def _AggregateIndexedSlicesGradients(grads):
  """Aggregates gradients of type `IndexedSlices` by concatenation."""
  if len(grads) < 1:
    return None
  elif len(grads) == 1:
    return grads[0]
  else:
    grads = math_ops._as_indexed_slices_list(  # pylint: disable=protected-access
        [g for g in grads if g is not None])
    grads = [_HandleNestedIndexedSlices(x) for x in grads]  # pylint: disable=protected-access
    # Form IndexedSlices out of the concatenated values and indices.
    concat_grad = ops.IndexedSlices(
        array_ops.concat([x.values for x in grads], axis=0),
        array_ops.concat([x.indices for x in grads], axis=0),
        grads[0].dense_shape)

    return concat_grad
def _AggregateIndexedSlicesGradients(grads):
  """Aggregates gradients of type `IndexedSlices` by concatenation."""
  if len(grads) < 1:
    return None
  elif len(grads) == 1:
    return grads[0]
  else:
    grads = math_ops._as_indexed_slices_list(  # pylint: disable=protected-access
        [g for g in grads if g is not None])
    grads = [_HandleNestedIndexedSlices(x) for x in grads]  # pylint: disable=protected-access
    # Form IndexedSlices out of the concatenated values and indices.
    concat_grad = ops.IndexedSlices(
        array_ops.concat([x.values for x in grads], axis=0),
        array_ops.concat([x.indices for x in grads], axis=0),
        grads[0].dense_shape)

    return concat_grad
def merge(inputs, name=None):
  """Returns the value of an available element of `inputs`.

  This op tests each of the tensors in `inputs` in turn to determine if any of
  them is available. If it finds an available tensor, it returns it and its
  index in `inputs`.

  It is an error if more than one tensor in `inputs` is available. If no tensor
  in `inputs` is available, the returned tensor and index are not set.

  This op handles both `Tensor`s and `IndexedSlices`. If inputs has a mix of
  `Tensor`s and `IndexedSlices`, all inputs are converted to IndexedSlices
  before merging.

  Args:
    inputs: The input tensors, at most one of which is available.
    name: A name for this operation (optional).

  Returns:
    A tuple containing the chosen input tensor and its index in `inputs`.

  Raises:
    ValueError: If inputs are IndexedSlices and some but not all have a
      dense_shape property.
  """
  with ops.op_scope(inputs, name, "Merge") as name:
    inputs = [ops.convert_to_tensor_or_indexed_slices(inp) for inp in inputs]
    if all([isinstance(inp, ops.Tensor) for inp in inputs]):
      return gen_control_flow_ops._merge(inputs, name=name)
    else:
      inputs = math_ops._as_indexed_slices_list(inputs)
      values, _ = gen_control_flow_ops._merge([inp.values for inp in inputs],
                                              name=name)
      indices, chosen_index = gen_control_flow_ops._merge(
          [inp.indices for inp in inputs], name="indices")
      if any(inp.dense_shape for inp in inputs):
        if not all(inp.dense_shape for inp in inputs):
          raise ValueError("Either all merged IndexedSlices must have a "
                           "dense_shape, or none must have a dense_shape.")
        dense_shape, _ = gen_control_flow_ops._merge(
            [inp.dense_shape for inp in inputs], name="dense_shape")
      else:
        dense_shape = None
      return ops.IndexedSlices(values, indices, dense_shape), chosen_index
Beispiel #5
0
def _AggregatedGrads(grads, op, loop_state, aggregation_method=None):
  """Get the aggregated gradients for op.

  Args:
    grads: The map of memoized gradients.
    op: The op to get gradients for.
    loop_state: An object for maintaining the state of the while loops in the
                graph. It is of type ControlFlowState. None if the graph
                contains no while loops.
    aggregation_method: Specifies the method used to combine gradient terms.
      Accepted values are constants defined in the class `AggregationMethod`.

  Returns:
    A list of gradients, one per each output of `op`. If the gradients
      for a particular output is a list, this function aggregates it
      before returning.

  Raises:
    TypeError: if the incoming grads are not Tensors or IndexedSlices.
    ValueError: if the arguments are invalid.

  """
  if aggregation_method is None:
    aggregation_method = AggregationMethod.DEFAULT
  if aggregation_method not in [
      AggregationMethod.ADD_N, AggregationMethod.EXPERIMENTAL_TREE,
      AggregationMethod.EXPERIMENTAL_ACCUMULATE_N
  ]:
    raise ValueError("Invalid aggregation_method specified %s." %
                     aggregation_method)
  out_grads = _GetGrads(grads, op)
  for i, out_grad in enumerate(out_grads):
    if loop_state:
      if isinstance(out_grad, (ops.Tensor, ops.IndexedSlices)):
        assert control_flow_ops.IsLoopSwitch(op)
        continue
    # Grads have to be Tensors or IndexedSlices
    if (isinstance(out_grad, collections.Sequence) and not all([
        isinstance(g, (ops.Tensor, ops.IndexedSlices)) for g in out_grad
        if g is not None
    ])):
      raise TypeError("gradients have to be either all Tensors "
                      "or all IndexedSlices")
    # Aggregate multiple gradients, and convert [] to None.
    if out_grad:
      if len(out_grad) < 2:
        used = "nop"
        out_grads[i] = out_grad[0]
      elif all([isinstance(g, ops.Tensor) for g in out_grad if g is not None]):
        tensor_shape = _AccumulatorShape(out_grad)
        if (aggregation_method == AggregationMethod.EXPERIMENTAL_ACCUMULATE_N
            and len(out_grad) > 2 and tensor_shape.is_fully_defined()):
          # The benefit of using AccumulateN is that its inputs can be combined
          # in any order and this can allow the expression to be evaluated with
          # a smaller memory footprint.  When used with gpu_allocator_retry,
          # it is possible to compute a sum of terms which are much larger than
          # total GPU memory.
          # AccumulateN can currently only be used if we know the shape for
          # an accumulator variable.  If this is not known, or if we only have
          # 2 grads then we fall through to the "tree" case below.
          used = "accumulate_n"
          out_grads[i] = math_ops.accumulate_n(out_grad)
        elif aggregation_method in [
            AggregationMethod.EXPERIMENTAL_TREE,
            AggregationMethod.EXPERIMENTAL_ACCUMULATE_N
        ]:
          # Aggregate all gradients by doing pairwise sums: this may
          # reduce performance, but it can improve memory because the
          # gradients can be released earlier.
          #
          # TODO(vrv): Consider replacing this with a version of
          # tf.AddN() that eagerly frees its inputs as soon as they are
          # ready, so the order of this tree does not become a problem.
          used = "tree"
          with ops.name_scope(op.name + "_gradient_sum"):
            running_sum = out_grad[0]
            for grad in out_grad[1:]:
              running_sum = math_ops.add_n([running_sum, grad])
            out_grads[i] = running_sum
        else:
          used = "add_n"
          out_grads[i] = _MultiDeviceAddN(out_grad)
        logging.vlog(2, "  _AggregatedGrads %d x %s using %s",
                     len(out_grad), tensor_shape, used)
      else:
        out_grad = math_ops._as_indexed_slices_list(
            [g for g in out_grad if g is not None])
        out_grad = [_HandleNestedIndexedSlices(x) for x in out_grad]
        # Form IndexedSlices out of the concatenated values and
        # indices.
        out_grads[i] = ops.IndexedSlices(
            array_ops.concat_v2([x.values for x in out_grad], 0),
            array_ops.concat_v2([x.indices for x in out_grad], 0),
            out_grad[0].dense_shape)
    else:
      out_grads[i] = []
  return out_grads