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
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
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