Python tensorflow.python.ops.math_ops.unsorted_segment_sum() Examples
The following are 24
code examples of tensorflow.python.ops.math_ops.unsorted_segment_sum().
You can vote up the ones you like or vote down the ones you don't like,
and go to the original project or source file by following the links above each example.
You may also want to check out all available functions/classes of the module
tensorflow.python.ops.math_ops
, or try the search function
.
Example #1
Source File: data_flow_grad.py From lambda-packs with MIT License | 6 votes |
def _DynamicStitchGrads(op, grad): """Gradients for DynamicStitch.""" num_values = len(op.inputs) // 2 indices_grad = [None] * num_values def AsInt32(x): return (x if op.inputs[0].dtype == dtypes.int32 else math_ops.cast(x, dtypes.int32)) inputs = [AsInt32(op.inputs[i]) for i in xrange(num_values)] if isinstance(grad, ops.IndexedSlices): output_shape = array_ops.shape(op.outputs[0]) output_rows = output_shape[0] grad = math_ops.unsorted_segment_sum(grad.values, grad.indices, output_rows) values_grad = [array_ops.gather(grad, inp) for inp in inputs] return indices_grad + values_grad
Example #2
Source File: math_grad.py From lambda-packs with MIT License | 6 votes |
def _SegmentMinOrMaxGrad(op, grad, is_sorted): """Gradient for SegmentMin and (unsorted) SegmentMax. They share similar code.""" zeros = array_ops.zeros(array_ops.shape(op.inputs[0]), dtype=op.inputs[0].dtype) # Get the number of selected (minimum or maximum) elements in each segment. gathered_outputs = array_ops.gather(op.outputs[0], op.inputs[1]) is_selected = math_ops.equal(op.inputs[0], gathered_outputs) if is_sorted: num_selected = math_ops.segment_sum(math_ops.cast(is_selected, grad.dtype), op.inputs[1]) else: num_selected = math_ops.unsorted_segment_sum(math_ops.cast(is_selected, grad.dtype), op.inputs[1], op.inputs[2]) # Compute the gradient for each segment. The gradient for the ith segment is # divided evenly among the selected elements in that segment. weighted_grads = math_ops.div(grad, num_selected) gathered_grads = array_ops.gather(weighted_grads, op.inputs[1]) if is_sorted: return array_ops.where(is_selected, gathered_grads, zeros), None else: return array_ops.where(is_selected, gathered_grads, zeros), None, None
Example #3
Source File: optimizer.py From lambda-packs with MIT License | 6 votes |
def _deduplicate_indexed_slices(values, indices): """Sums `values` associated with any non-unique `indices`. Args: values: A `Tensor` with rank >= 1. indices: A one-dimensional integer `Tensor`, indexing into the first dimension of `values` (as in an IndexedSlices object). Returns: A tuple of (`summed_values`, `unique_indices`) where `unique_indices` is a de-duplicated version of `indices` and `summed_values` contains the sum of `values` slices associated with each unique index. """ unique_indices, new_index_positions = array_ops.unique(indices) summed_values = math_ops.unsorted_segment_sum( values, new_index_positions, array_ops.shape(unique_indices)[0]) return (summed_values, unique_indices)
Example #4
Source File: data_flow_grad.py From auto-alt-text-lambda-api with MIT License | 6 votes |
def _DynamicStitchGrads(op, grad): """Gradients for DynamicStitch.""" num_values = len(op.inputs) // 2 indices_grad = [None] * num_values def AsInt32(x): return (x if op.inputs[0].dtype == dtypes.int32 else math_ops.cast(x, dtypes.int32)) inputs = [AsInt32(op.inputs[i]) for i in xrange(num_values)] if isinstance(grad, ops.IndexedSlices): output_shape = array_ops.shape(op.outputs[0]) output_rows = output_shape[0] grad = math_ops.unsorted_segment_sum(grad.values, grad.indices, output_rows) values_grad = [array_ops.gather(grad, inp) for inp in inputs] return indices_grad + values_grad
Example #5
Source File: data_flow_grad.py From keras-lambda with MIT License | 6 votes |
def _DynamicStitchGrads(op, grad): """Gradients for DynamicStitch.""" num_values = len(op.inputs) // 2 indices_grad = [None] * num_values def AsInt32(x): return (x if op.inputs[0].dtype == dtypes.int32 else math_ops.cast(x, dtypes.int32)) inputs = [AsInt32(op.inputs[i]) for i in xrange(num_values)] if isinstance(grad, ops.IndexedSlices): output_shape = array_ops.shape(op.outputs[0]) output_rows = output_shape[0] grad = math_ops.unsorted_segment_sum(grad.values, grad.indices, output_rows) values_grad = [array_ops.gather(grad, inp) for inp in inputs] return indices_grad + values_grad
Example #6
Source File: optimizer.py From Serverless-Deep-Learning-with-TensorFlow-and-AWS-Lambda with MIT License | 6 votes |
def _deduplicate_indexed_slices(values, indices): """Sums `values` associated with any non-unique `indices`. Args: values: A `Tensor` with rank >= 1. indices: A one-dimensional integer `Tensor`, indexing into the first dimension of `values` (as in an IndexedSlices object). Returns: A tuple of (`summed_values`, `unique_indices`) where `unique_indices` is a de-duplicated version of `indices` and `summed_values` contains the sum of `values` slices associated with each unique index. """ unique_indices, new_index_positions = array_ops.unique(indices) summed_values = math_ops.unsorted_segment_sum( values, new_index_positions, array_ops.shape(unique_indices)[0]) return (summed_values, unique_indices)
Example #7
Source File: data_flow_grad.py From deep_image_model with Apache License 2.0 | 6 votes |
def _DynamicStitchGrads(op, grad): """Gradients for DynamicStitch.""" num_values = len(op.inputs) // 2 indices_grad = [None] * num_values def AsInt32(x): return (x if op.inputs[0].dtype == dtypes.int32 else math_ops.cast(x, dtypes.int32)) inputs = [AsInt32(op.inputs[i]) for i in xrange(num_values)] if isinstance(grad, ops.IndexedSlices): output_shape = array_ops.shape(op.outputs[0]) output_rows = output_shape[0] grad = math_ops.unsorted_segment_sum(grad.values, grad.indices, output_rows) values_grad = [array_ops.gather(grad, inp) for inp in inputs] return indices_grad + values_grad
Example #8
Source File: math_grad.py From Serverless-Deep-Learning-with-TensorFlow-and-AWS-Lambda with MIT License | 6 votes |
def _SegmentMinOrMaxGrad(op, grad, is_sorted): """Gradient for SegmentMin and (unsorted) SegmentMax. They share similar code.""" zeros = array_ops.zeros(array_ops.shape(op.inputs[0]), dtype=op.inputs[0].dtype) # Get the number of selected (minimum or maximum) elements in each segment. gathered_outputs = array_ops.gather(op.outputs[0], op.inputs[1]) is_selected = math_ops.equal(op.inputs[0], gathered_outputs) if is_sorted: num_selected = math_ops.segment_sum(math_ops.cast(is_selected, grad.dtype), op.inputs[1]) else: num_selected = math_ops.unsorted_segment_sum( math_ops.cast(is_selected, grad.dtype), op.inputs[1], op.inputs[2]) # Compute the gradient for each segment. The gradient for the ith segment is # divided evenly among the selected elements in that segment. weighted_grads = math_ops.div(grad, num_selected) gathered_grads = array_ops.gather(weighted_grads, op.inputs[1]) if is_sorted: return array_ops.where(is_selected, gathered_grads, zeros), None else: return array_ops.where(is_selected, gathered_grads, zeros), None, None
Example #9
Source File: data_flow_grad.py From Serverless-Deep-Learning-with-TensorFlow-and-AWS-Lambda with MIT License | 6 votes |
def _DynamicStitchGrads(op, grad): """Gradients for DynamicStitch and ParallelDynamicStitch.""" num_values = len(op.inputs) // 2 indices_grad = [None] * num_values def AsInt32(x): return (x if op.inputs[0].dtype == dtypes.int32 else math_ops.cast(x, dtypes.int32)) inputs = [AsInt32(op.inputs[i]) for i in xrange(num_values)] if isinstance(grad, ops.IndexedSlices): output_shape = array_ops.shape(op.outputs[0]) output_rows = output_shape[0] grad = math_ops.unsorted_segment_sum(grad.values, grad.indices, output_rows) values_grad = [array_ops.gather(grad, inp) for inp in inputs] return indices_grad + values_grad
Example #10
Source File: clustering_ops.py From keras-lambda with MIT License | 5 votes |
def _full_batch_training_op(self, inputs, cluster_idx_list, cluster_centers): """Creates an op for training for full batch case. Args: inputs: list of input Tensors. cluster_idx_list: A vector (or list of vectors). Each element in the vector corresponds to an input row in 'inp' and specifies the cluster id corresponding to the input. cluster_centers: Tensor Ref of cluster centers. Returns: An op for doing an update of mini-batch k-means. """ cluster_sums = [] cluster_counts = [] epsilon = constant_op.constant(1e-6, dtype=inputs[0].dtype) for inp, cluster_idx in zip(inputs, cluster_idx_list): with ops.colocate_with(inp): cluster_sums.append( math_ops.unsorted_segment_sum(inp, cluster_idx, self._num_clusters)) cluster_counts.append( math_ops.unsorted_segment_sum( array_ops.reshape( array_ops.ones( array_ops.reshape(array_ops.shape(inp)[0], [-1])), [-1, 1]), cluster_idx, self._num_clusters)) with ops.colocate_with(cluster_centers): new_clusters_centers = math_ops.add_n(cluster_sums) / (math_ops.cast( math_ops.add_n(cluster_counts), cluster_sums[0].dtype) + epsilon) if self._clusters_l2_normalized(): new_clusters_centers = nn_impl.l2_normalize(new_clusters_centers, dim=1) return state_ops.assign(cluster_centers, new_clusters_centers)
Example #11
Source File: gradients_impl.py From lambda-packs with MIT License | 5 votes |
def _IndexedSlicesToTensor(value, dtype=None, name=None, as_ref=False): """Converts an IndexedSlices object `value` to a Tensor. NOTE(mrry): This function is potentially expensive. Args: value: An ops.IndexedSlices object. dtype: The dtype of the Tensor to be returned. name: Optional name to use for the returned Tensor. as_ref: True if a ref is requested. Returns: A dense Tensor representing the values in the given IndexedSlices. Raises: ValueError: If the IndexedSlices does not have the same dtype. """ _ = as_ref if dtype and not dtype.is_compatible_with(value.dtype): raise ValueError( "Tensor conversion requested dtype %s for IndexedSlices with dtype %s" % (dtype.name, value.dtype.name)) if value.dense_shape is None: raise ValueError( "Tensor conversion requested for IndexedSlices without dense_shape: %s" % str(value)) # TODO(mrry): Consider adding static shape information to # IndexedSlices, to avoid using numpy here. dense_shape_value = tensor_util.constant_value(value.dense_shape) if dense_shape_value is not None: num_elements = np.prod(dense_shape_value) if num_elements >= _LARGE_SPARSE_NUM_ELEMENTS: warnings.warn( "Converting sparse IndexedSlices to a dense Tensor with %d elements. " "This may consume a large amount of memory." % num_elements) else: warnings.warn( "Converting sparse IndexedSlices to a dense Tensor of unknown shape. " "This may consume a large amount of memory.") return math_ops.unsorted_segment_sum( value.values, value.indices, value.dense_shape[0], name=name)
Example #12
Source File: optimizer.py From keras-lambda with MIT License | 5 votes |
def _apply_sparse_duplicate_indices(self, grad, var): """Add ops to apply sparse gradients to `var`, with repeated sparse indices. Optimizers which override this method must deal with IndexedSlices objects such as the following: IndexedSlicesValue(values=[1, 1], indices=[0, 0], dense_shape=[1]) The correct interpretation is: IndexedSlicesValue(values=[2], indices=[0], dense_shape=[1]) Many optimizers deal incorrectly with repeated indices when updating based on sparse gradients (e.g. summing squares rather than squaring the sum, or applying momentum terms multiple times). Adding first is always the correct behavior, so this is enforced here by reconstructing the IndexedSlices to have only unique indices, then calling _apply_sparse. Optimizers which deal correctly with repeated indices may instead override this method to avoid the overhead of summing indices. Args: grad: `IndexedSlices`. var: A `Variable` object. Returns: An `Operation`. """ unique_indices, new_index_positions = array_ops.unique(grad.indices) summed_values = math_ops.unsorted_segment_sum( grad.values, new_index_positions, array_ops.shape(unique_indices)[0]) gradient_no_duplicate_indices = ops.IndexedSlices( indices=unique_indices, values=summed_values, dense_shape=grad.dense_shape) return self._apply_sparse(gradient_no_duplicate_indices, var)
Example #13
Source File: math_grad.py From keras-lambda with MIT License | 5 votes |
def _SparseSegmentSumGrad(op, grad): """Gradient for SparseSegmentSum.""" input_rows = array_ops.shape(op.inputs[0])[0] return (math_ops.unsorted_segment_sum( array_ops.gather(grad, op.inputs[2]), op.inputs[1], input_rows), None, None)
Example #14
Source File: gradients_impl.py From keras-lambda with MIT License | 5 votes |
def _IndexedSlicesToTensor(value, dtype=None, name=None, as_ref=False): """Converts an IndexedSlices object `value` to a Tensor. NOTE(mrry): This function is potentially expensive. Args: value: An ops.IndexedSlices object. dtype: The dtype of the Tensor to be returned. name: Optional name to use for the returned Tensor. as_ref: True if a ref is requested. Returns: A dense Tensor representing the values in the given IndexedSlices. Raises: ValueError: If the IndexedSlices does not have the same dtype. """ _ = as_ref if dtype and not dtype.is_compatible_with(value.dtype): raise ValueError( "Tensor conversion requested dtype %s for IndexedSlices with dtype %s" % (dtype.name, value.dtype.name)) if value.dense_shape is None: raise ValueError( "Tensor conversion requested for IndexedSlices without dense_shape: %s" % str(value)) # TODO(mrry): Consider adding static shape information to # IndexedSlices, to avoid using numpy here. dense_shape_value = tensor_util.constant_value(value.dense_shape) if dense_shape_value is not None: num_elements = np.prod(dense_shape_value) if num_elements >= _LARGE_SPARSE_NUM_ELEMENTS: warnings.warn( "Converting sparse IndexedSlices to a dense Tensor with %d elements. " "This may consume a large amount of memory." % num_elements) else: warnings.warn( "Converting sparse IndexedSlices to a dense Tensor of unknown shape. " "This may consume a large amount of memory.") return math_ops.unsorted_segment_sum( value.values, value.indices, value.dense_shape[0], name=name)
Example #15
Source File: math_grad.py From Serverless-Deep-Learning-with-TensorFlow-and-AWS-Lambda with MIT License | 5 votes |
def _SparseSegmentSumGrad(op, grad): """Gradient for SparseSegmentSum.""" input_rows = array_ops.shape(op.inputs[0])[0] return (math_ops.unsorted_segment_sum( array_ops.gather(grad, op.inputs[2]), op.inputs[1], input_rows), None, None)
Example #16
Source File: gradients_impl.py From Serverless-Deep-Learning-with-TensorFlow-and-AWS-Lambda with MIT License | 5 votes |
def _IndexedSlicesToTensor(value, dtype=None, name=None, as_ref=False): """Converts an IndexedSlices object `value` to a Tensor. NOTE(mrry): This function is potentially expensive. Args: value: An ops.IndexedSlices object. dtype: The dtype of the Tensor to be returned. name: Optional name to use for the returned Tensor. as_ref: True if a ref is requested. Returns: A dense Tensor representing the values in the given IndexedSlices. Raises: ValueError: If the IndexedSlices does not have the same dtype. """ _ = as_ref if dtype and not dtype.is_compatible_with(value.dtype): raise ValueError( "Tensor conversion requested dtype %s for IndexedSlices with dtype %s" % (dtype.name, value.dtype.name)) if value.dense_shape is None: raise ValueError( "Tensor conversion requested for IndexedSlices without dense_shape: %s" % str(value)) # TODO(mrry): Consider adding static shape information to # IndexedSlices, to avoid using numpy here. dense_shape_value = tensor_util.constant_value(value.dense_shape) if dense_shape_value is not None: num_elements = np.prod(dense_shape_value) if num_elements >= _LARGE_SPARSE_NUM_ELEMENTS: warnings.warn( "Converting sparse IndexedSlices to a dense Tensor with %d elements. " "This may consume a large amount of memory." % num_elements) else: warnings.warn( "Converting sparse IndexedSlices to a dense Tensor of unknown shape. " "This may consume a large amount of memory.") return math_ops.unsorted_segment_sum( value.values, value.indices, value.dense_shape[0], name=name)
Example #17
Source File: math_grad.py From deep_image_model with Apache License 2.0 | 5 votes |
def _SparseSegmentSumGrad(op, grad): """Gradient for SparseSegmentSum.""" input_rows = array_ops.shape(op.inputs[0])[0] return (math_ops.unsorted_segment_sum( array_ops.gather(grad, op.inputs[2]), op.inputs[1], input_rows), None, None)
Example #18
Source File: gradients_impl.py From deep_image_model with Apache License 2.0 | 5 votes |
def _IndexedSlicesToTensor(value, dtype=None, name=None, as_ref=False): """Converts an IndexedSlices object `value` to a Tensor. NOTE(mrry): This function is potentially expensive. Args: value: An ops.IndexedSlices object. dtype: The dtype of the Tensor to be returned. name: Optional name to use for the returned Tensor. as_ref: True if a ref is requested. Returns: A dense Tensor representing the values in the given IndexedSlices. Raises: ValueError: If the IndexedSlices does not have the same dtype. """ _ = as_ref if dtype and not dtype.is_compatible_with(value.dtype): raise ValueError( "Tensor conversion requested dtype %s for IndexedSlices with dtype %s" % (dtype.name, value.dtype.name)) if value.dense_shape is None: raise ValueError( "Tensor conversion requested for IndexedSlices without dense_shape: %s" % str(value)) # TODO(mrry): Consider adding static shape information to # IndexedSlices, to avoid using numpy here. dense_shape_value = tensor_util.constant_value(value.dense_shape) if dense_shape_value is not None: num_elements = np.prod(dense_shape_value) if num_elements >= _LARGE_SPARSE_NUM_ELEMENTS: warnings.warn( "Converting sparse IndexedSlices to a dense Tensor with %d elements. " "This may consume a large amount of memory." % num_elements) else: warnings.warn( "Converting sparse IndexedSlices to a dense Tensor of unknown shape. " "This may consume a large amount of memory.") return math_ops.unsorted_segment_sum( value.values, value.indices, value.dense_shape[0], name=name)
Example #19
Source File: clustering_ops.py From auto-alt-text-lambda-api with MIT License | 5 votes |
def _full_batch_training_op(self, inputs, cluster_idx_list, cluster_centers): """Creates an op for training for full batch case. Args: inputs: list of input Tensors. cluster_idx_list: A vector (or list of vectors). Each element in the vector corresponds to an input row in 'inp' and specifies the cluster id corresponding to the input. cluster_centers: Tensor Ref of cluster centers. Returns: An op for doing an update of mini-batch k-means. """ cluster_sums = [] cluster_counts = [] epsilon = constant_op.constant(1e-6, dtype=inputs[0].dtype) for inp, cluster_idx in zip(inputs, cluster_idx_list): with ops.colocate_with(inp): cluster_sums.append( math_ops.unsorted_segment_sum(inp, cluster_idx, self._num_clusters)) cluster_counts.append( math_ops.unsorted_segment_sum( array_ops.reshape( array_ops.ones( array_ops.reshape(array_ops.shape(inp)[0], [-1])), [-1, 1]), cluster_idx, self._num_clusters)) with ops.colocate_with(cluster_centers): new_clusters_centers = math_ops.add_n(cluster_sums) / (math_ops.cast( math_ops.add_n(cluster_counts), cluster_sums[0].dtype) + epsilon) if self._clusters_l2_normalized(): new_clusters_centers = nn_impl.l2_normalize(new_clusters_centers, dim=1) return state_ops.assign(cluster_centers, new_clusters_centers)
Example #20
Source File: optimizer.py From auto-alt-text-lambda-api with MIT License | 5 votes |
def _apply_sparse_duplicate_indices(self, grad, var): """Add ops to apply sparse gradients to `var`, with repeated sparse indices. Optimizers which override this method must deal with IndexedSlices objects such as the following: IndexedSlicesValue(values=[1, 1], indices=[0, 0], dense_shape=[1]) The correct interpretation is: IndexedSlicesValue(values=[2], indices=[0], dense_shape=[1]) Many optimizers deal incorrectly with repeated indices when updating based on sparse gradients (e.g. summing squares rather than squaring the sum, or applying momentum terms multiple times). Adding first is always the correct behavior, so this is enforced here by reconstructing the IndexedSlices to have only unique indices, then calling _apply_sparse. Optimizers which deal correctly with repeated indices may instead override this method to avoid the overhead of summing indices. Args: grad: `IndexedSlices`. var: A `Variable` object. Returns: An `Operation`. """ unique_indices, new_index_positions = array_ops.unique(grad.indices) summed_values = math_ops.unsorted_segment_sum( grad.values, new_index_positions, array_ops.shape(unique_indices)[0]) gradient_no_duplicate_indices = ops.IndexedSlices( indices=unique_indices, values=summed_values, dense_shape=grad.dense_shape) return self._apply_sparse(gradient_no_duplicate_indices, var)
Example #21
Source File: math_grad.py From auto-alt-text-lambda-api with MIT License | 5 votes |
def _SparseSegmentSumGrad(op, grad): """Gradient for SparseSegmentSum.""" input_rows = array_ops.shape(op.inputs[0])[0] return (math_ops.unsorted_segment_sum( array_ops.gather(grad, op.inputs[2]), op.inputs[1], input_rows), None, None)
Example #22
Source File: gradients_impl.py From auto-alt-text-lambda-api with MIT License | 5 votes |
def _IndexedSlicesToTensor(value, dtype=None, name=None, as_ref=False): """Converts an IndexedSlices object `value` to a Tensor. NOTE(mrry): This function is potentially expensive. Args: value: An ops.IndexedSlices object. dtype: The dtype of the Tensor to be returned. name: Optional name to use for the returned Tensor. as_ref: True if a ref is requested. Returns: A dense Tensor representing the values in the given IndexedSlices. Raises: ValueError: If the IndexedSlices does not have the same dtype. """ _ = as_ref if dtype and not dtype.is_compatible_with(value.dtype): raise ValueError( "Tensor conversion requested dtype %s for IndexedSlices with dtype %s" % (dtype.name, value.dtype.name)) if value.dense_shape is None: raise ValueError( "Tensor conversion requested for IndexedSlices without dense_shape: %s" % str(value)) # TODO(mrry): Consider adding static shape information to # IndexedSlices, to avoid using numpy here. dense_shape_value = tensor_util.constant_value(value.dense_shape) if dense_shape_value is not None: num_elements = np.prod(dense_shape_value) if num_elements >= _LARGE_SPARSE_NUM_ELEMENTS: warnings.warn( "Converting sparse IndexedSlices to a dense Tensor with %d elements. " "This may consume a large amount of memory." % num_elements) else: warnings.warn( "Converting sparse IndexedSlices to a dense Tensor of unknown shape. " "This may consume a large amount of memory.") return math_ops.unsorted_segment_sum( value.values, value.indices, value.dense_shape[0], name=name)
Example #23
Source File: math_grad.py From lambda-packs with MIT License | 5 votes |
def _SparseSegmentSumGrad(op, grad): """Gradient for SparseSegmentSum.""" input_rows = array_ops.shape(op.inputs[0])[0] return (math_ops.unsorted_segment_sum( array_ops.gather(grad, op.inputs[2]), op.inputs[1], input_rows), None, None)
Example #24
Source File: sync_replicas_optimizer.py From deep_image_model with Apache License 2.0 | 4 votes |
def _aggregate_sparse_grad(self, grad, var, train_ops): """Aggregate sparse gradients. Args: grad: The sparse gradient to aggregate. var: The variable to apply this gradient to. train_ops: The train_ops for the worker to run. Returns: aggregated_grad: Aggregated grad. """ # Sparse gradients have to be inserted as one pair of (value, # indice) as an element instead of the whole "indexedslice" because # their shapes are not deterministic. sparse_grad_queue = (data_flow_ops.FIFOQueue( -1, (grad.values.dtype, grad.indices.dtype), shapes=(var.get_shape().as_list()[1:], ()), shared_name="sparse_grad_q_%s" % var.name)) self._sparse_grad_queues_and_devs.append((sparse_grad_queue, var.device)) # Sparse token is inserted after the "enqueue_many" finishes. This # is needed to make sure enough sparse gradients have been enqueued # before applying them to the variables. sparse_token_queue = (data_flow_ops.FIFOQueue( self._replicas_to_aggregate * 2, types_pb2.DT_INT32, shapes=(), shared_name="sparse_token_q_%s" % var.name)) self._one_element_queue_list.append((sparse_token_queue, var.device)) enqueue_spares_op = sparse_grad_queue.enqueue_many([grad.values, grad.indices]) with ops.control_dependencies([enqueue_spares_op]): train_ops.append(sparse_token_queue.enqueue((1,))) with ops.control_dependencies([sparse_token_queue.dequeue_many( self._replicas_to_aggregate)]): values, indices = sparse_grad_queue.dequeue_many(sparse_grad_queue.size()) concat_grad = ops.IndexedSlices(values, indices, grad.dense_shape) # Sum the gradients of the same variables in the sparse layers so # that each variable is only updated once. Note that with 2 # gradients g1 and g2 from 2 replicas for the same variable, # apply(g1+g2) is different from apply(g1) and then apply(g2) when # the optimizer is complex like Momentum or Adagrad. values = concat_grad.values indices = concat_grad.indices new_indices, indx = array_ops.unique(indices) num_indices = array_ops.shape(new_indices)[0] sum_values = math_ops.unsorted_segment_sum(values, indx, num_indices) return ops.IndexedSlices(sum_values, new_indices, concat_grad.dense_shape)