Python tensorflow.unsorted_segment_sum() Examples

The following are 30 code examples of tensorflow.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 , or try the search function .
Example #1
Source File: segment_reduction_ops_test.py    From deep_image_model with Apache License 2.0 6 votes vote down vote up
def testGradient(self):
    num_cols = 2
    indices_flat = np.array([0, 4, 0, 8, 3, 8, 4, 7, 7, 3])
    num_segments = max(indices_flat) + 3
    for indices in indices_flat, indices_flat.reshape(5, 2):
      shape = indices.shape + (num_cols,)
      with self.test_session(use_gpu=self.use_gpu):
        tf_x, np_x = self._input(shape, dtype=tf.float64)
        s = tf.unsorted_segment_sum(data=tf_x,
                                    segment_ids=indices,
                                    num_segments=num_segments)
        jacob_t, jacob_n = tf.test.compute_gradient(
            tf_x,
            shape,
            s,
            [num_segments, num_cols],
            x_init_value=np_x.astype(np.double),
            delta=1)
      self.assertAllClose(jacob_t, jacob_n, rtol=1e-3, atol=1e-3) 
Example #2
Source File: expert_utils.py    From training_results_v0.5 with Apache License 2.0 6 votes vote down vote up
def combine(self, x):
    """Return the output from the experts.

    When one example goes to multiple experts, the outputs are summed.

    Args:
      x: a Tensor with shape [batch, num_experts, expert_capacity, depth]

    Returns:
      a `Tensor` with shape `[batch, length, depth]
    """
    depth = tf.shape(x)[-1]
    x *= tf.expand_dims(self._nonpadding, -1)
    ret = tf.unsorted_segment_sum(
        x, self._flat_indices, num_segments=self._batch * self._length)
    ret = tf.reshape(ret, [self._batch, self._length, depth])
    return ret 
Example #3
Source File: unsortedsegmentops.py    From dpu-utils with MIT License 6 votes vote down vote up
def unsorted_segment_log_softmax(logits, segment_ids, num_segments):
    """Perform an unsorted segment safe log_softmax."""
    # Note: if a segment is empty, the smallest value for the score will be returned,
    # which yields the correct behavior
    max_per_segment = tf.unsorted_segment_max(data=logits,
                                              segment_ids=segment_ids,
                                              num_segments=num_segments)
    scattered_maxes = tf.gather(params=max_per_segment,
                                indices=segment_ids)
    recentered_scores = logits - scattered_maxes
    exped_recentered_scores = tf.exp(recentered_scores)

    per_segment_sums = tf.unsorted_segment_sum(exped_recentered_scores, segment_ids, num_segments)
    per_segment_normalization_consts = tf.log(per_segment_sums)

    log_probs = recentered_scores - tf.gather(params=per_segment_normalization_consts, indices=segment_ids)
    return log_probs 
Example #4
Source File: training.py    From delft with Apache License 2.0 6 votes vote down vote up
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 = tf.unique(indices)
    summed_values = tf.unsorted_segment_sum(
      values, new_index_positions,
      tf.shape(unique_indices)[0])
    return (summed_values, unique_indices) 
Example #5
Source File: expert_utils.py    From training_results_v0.5 with Apache License 2.0 6 votes vote down vote up
def combine(self, x):
    """Return the output from the experts.

    When one example goes to multiple experts, the outputs are summed.

    Args:
      x: a Tensor with shape [batch, num_experts, expert_capacity, depth]

    Returns:
      a `Tensor` with shape `[batch, length, depth]
    """
    depth = tf.shape(x)[-1]
    x *= tf.expand_dims(self._nonpadding, -1)
    ret = tf.unsorted_segment_sum(
        x, self._flat_indices, num_segments=self._batch * self._length)
    ret = tf.reshape(ret, [self._batch, self._length, depth])
    return ret 
Example #6
Source File: attention.py    From jack with MIT License 6 votes vote down vote up
def apply_attention(attn_scores, states, length, is_self=False, with_sentinel=True, reuse=False, seq2_to_seq1=None):
    attn_scores += tf.expand_dims(misc.mask_for_lengths(length, tf.shape(attn_scores)[2]), 1)
    softmax = tf.nn.softmax if seq2_to_seq1 is None else lambda x: segment.segment_softmax(x, seq2_to_seq1)
    if is_self:
        # exclude attending to state itself
        attn_scores += tf.expand_dims(tf.diag(tf.fill([tf.shape(attn_scores)[1]], -1e6)), 0)
    if with_sentinel:
        with tf.variable_scope('sentinel', reuse=reuse):
            s = tf.get_variable('score', [1, 1, 1], tf.float32, tf.zeros_initializer())
        s = tf.tile(s, [tf.shape(attn_scores)[0], tf.shape(attn_scores)[1], 1])
        attn_probs = softmax(tf.concat([s, attn_scores], 2))
        attn_probs = attn_probs[:, :, 1:]
    else:
        attn_probs = softmax(attn_scores)
    attn_states = tf.einsum('abd,adc->abc', attn_probs, states)
    if seq2_to_seq1 is not None:
        attn_states = tf.unsorted_segment_sum(attn_states, seq2_to_seq1, tf.reduce_max(seq2_to_seq1) + 1)
    return attn_scores, attn_probs, attn_states 
Example #7
Source File: expert_utils.py    From training_results_v0.5 with Apache License 2.0 6 votes vote down vote up
def combine(self, expert_out, multiply_by_gates=True):
    """Sum together the expert output, weighted by the gates.

    The slice corresponding to a particular batch element `b` is computed
    as the sum over all experts `i` of the expert output, weighted by the
    corresponding gate values.  If `multiply_by_gates` is set to False, the
    gate values are ignored.

    Args:
      expert_out: a list of `num_experts` `Tensor`s, each with shape
        `[expert_batch_size_i, <extra_output_dims>]`.
      multiply_by_gates: a boolean

    Returns:
      a `Tensor` with shape `[batch_size, <extra_output_dims>]`.
    """
    # see comments on convert_gradient_to_tensor
    stitched = common_layers.convert_gradient_to_tensor(
        tf.concat(expert_out, 0))
    if multiply_by_gates:
      stitched *= tf.expand_dims(self._nonzero_gates, 1)
    combined = tf.unsorted_segment_sum(stitched, self._batch_index,
                                       tf.shape(self._gates)[0])
    return combined 
Example #8
Source File: segment.py    From jack with MIT License 6 votes vote down vote up
def segment_softmax(scores, segment_ids):
    """Given scores and a partition, converts scores to probs by performing
    softmax over all rows within a partition."""

    # Subtract max
    num_segments = tf.reduce_max(segment_ids) + 1
    if len(scores.get_shape()) == 2:
        max_per_partition = tf.unsorted_segment_max(tf.reduce_max(scores, axis=1), segment_ids, num_segments)
        scores -= tf.expand_dims(tf.gather(max_per_partition, segment_ids), axis=1)
    else:
        max_per_partition = tf.unsorted_segment_max(scores, segment_ids, num_segments)
        scores -= tf.gather(max_per_partition, segment_ids)

    # Compute probs
    scores_exp = tf.exp(scores)
    if len(scores.get_shape()) == 2:
        scores_exp_sum_per_partition = tf.unsorted_segment_sum(tf.reduce_sum(scores_exp, axis=1), segment_ids,
                                                               num_segments)
        probs = scores_exp / tf.expand_dims(tf.gather(scores_exp_sum_per_partition, segment_ids), axis=1)
    else:
        scores_exp_sum_per_partition = tf.unsorted_segment_sum(scores_exp, segment_ids, num_segments)
        probs = scores_exp / tf.gather(scores_exp_sum_per_partition, segment_ids)

    return probs 
Example #9
Source File: training.py    From ELMo_Chin with Apache License 2.0 6 votes vote down vote up
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 = tf.unique(indices)
    summed_values = tf.unsorted_segment_sum(
      values, new_index_positions,
      tf.shape(unique_indices)[0])
    return (summed_values, unique_indices) 
Example #10
Source File: training.py    From embedding with MIT License 6 votes vote down vote up
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 = tf.unique(indices)
    summed_values = tf.unsorted_segment_sum(
      values, new_index_positions,
      tf.shape(unique_indices)[0])
    return (summed_values, unique_indices) 
Example #11
Source File: expert_utils.py    From ASR with Apache License 2.0 6 votes vote down vote up
def Combine(self, expert_out, multiply_by_gates=True):
    """Sum together the expert output, weighted by the gates.

    The slice corresponding to a particular batch element `b` is computed
    as the sum over all experts `i` of the expert output, weighted by the
    corresponding gate values.  If `multiply_by_gates` is set to False, the
    gate values are ignored.

    Args:
      expert_out: a list of `num_experts` `Tensor`s, each with shape
        `[expert_batch_size_i, <extra_output_dims>]`.
      multiply_by_gates: a boolean

    Returns:
      a `Tensor` with shape `[batch_size, <extra_output_dims>]`.
    """
    # see comments on ConvertGradientToTensor
    stitched = ConvertGradientToTensor(tf.concat(expert_out, 0))
    if multiply_by_gates:
      stitched *= tf.expand_dims(self._nonzero_gates, 1)
    combined = tf.unsorted_segment_sum(stitched, self._batch_index,
                                       tf.shape(self._gates)[0])
    return combined 
Example #12
Source File: network_units.py    From DOTA_models with Apache License 2.0 6 votes vote down vote up
def embedding_lookup(embedding_matrix, indices, ids, weights, size):
  """Performs a weighted embedding lookup.

  Args:
    embedding_matrix: float Tensor from which to do the lookup.
    indices: int Tensor for the output rows of the looked up vectors.
    ids: int Tensor vectors to look up in the embedding_matrix.
    weights: float Tensor weights to apply to the looked up vectors.
    size: int number of output rows. Needed since some output rows may be
        empty.

  Returns:
    Weighted embedding vectors.
  """
  embeddings = tf.nn.embedding_lookup([embedding_matrix], ids)
  # TODO(googleuser): allow skipping weights.
  broadcast_weights_shape = tf.concat([tf.shape(weights), [1]], 0)
  embeddings *= tf.reshape(weights, broadcast_weights_shape)
  embeddings = tf.unsorted_segment_sum(embeddings, indices, size)
  return embeddings 
Example #13
Source File: network_units.py    From yolo_v2 with Apache License 2.0 6 votes vote down vote up
def embedding_lookup(embedding_matrix, indices, ids, weights, size):
  """Performs a weighted embedding lookup.

  Args:
    embedding_matrix: float Tensor from which to do the lookup.
    indices: int Tensor for the output rows of the looked up vectors.
    ids: int Tensor vectors to look up in the embedding_matrix.
    weights: float Tensor weights to apply to the looked up vectors.
    size: int number of output rows. Needed since some output rows may be
        empty.

  Returns:
    Weighted embedding vectors.
  """
  embeddings = tf.nn.embedding_lookup([embedding_matrix], ids)
  # TODO(googleuser): allow skipping weights.
  broadcast_weights_shape = tf.concat([tf.shape(weights), [1]], 0)
  embeddings *= tf.reshape(weights, broadcast_weights_shape)
  embeddings = tf.unsorted_segment_sum(embeddings, indices, size)
  return embeddings 
Example #14
Source File: segment_reduction_ops_test.py    From deep_image_model with Apache License 2.0 6 votes vote down vote up
def testValues(self):
    dtypes = [tf.float32,
              tf.float64,
              tf.int64,
              tf.int32,
              tf.complex64,
              tf.complex128]
    indices_flat = np.array([0, 4, 0, 8, 3, 8, 4, 7, 7, 3])
    num_segments = 12
    for indices in indices_flat, indices_flat.reshape(5, 2):
      shape = indices.shape + (2,)
      for dtype in dtypes:
        with self.test_session(use_gpu=self.use_gpu):
          tf_x, np_x = self._input(shape, dtype=dtype)
          np_ans = self._segmentReduce(indices,
                                       np_x,
                                       np.add,
                                       op2=None,
                                       num_out_rows=num_segments)
          s = tf.unsorted_segment_sum(data=tf_x,
                                      segment_ids=indices,
                                      num_segments=num_segments)
          tf_ans = s.eval()
        self._assertAllClose(indices, np_ans, tf_ans)
        self.assertShapeEqual(np_ans, s) 
Example #15
Source File: network_units.py    From hands-detection with MIT License 6 votes vote down vote up
def embedding_lookup(embedding_matrix, indices, ids, weights, size):
  """Performs a weighted embedding lookup.

  Args:
    embedding_matrix: float Tensor from which to do the lookup.
    indices: int Tensor for the output rows of the looked up vectors.
    ids: int Tensor vectors to look up in the embedding_matrix.
    weights: float Tensor weights to apply to the looked up vectors.
    size: int number of output rows. Needed since some output rows may be
        empty.

  Returns:
    Weighted embedding vectors.
  """
  embeddings = tf.nn.embedding_lookup([embedding_matrix], ids)
  # TODO(googleuser): allow skipping weights.
  broadcast_weights_shape = tf.concat([tf.shape(weights), [1]], 0)
  embeddings *= tf.reshape(weights, broadcast_weights_shape)
  embeddings = tf.unsorted_segment_sum(embeddings, indices, size)
  return embeddings 
Example #16
Source File: det_utils.py    From MobileNet with Apache License 2.0 6 votes vote down vote up
def find_dup(a):
  """ Find the duplicated elements in 1-D a tensor.
  Args:
    a: 1-D tensor.
    
  Return:
    more_than_one_vals: duplicated value in a.
    indexes_in_a: duplicated value's index in a.
    dups_in_a: duplicated value with duplicate in a.
  """
  unique_a_vals, unique_idx = tf.unique(a)
  count_a_unique = tf.unsorted_segment_sum(tf.ones_like(a),
                                           unique_idx,
                                           tf.shape(a)[0])

  more_than_one = tf.greater(count_a_unique, 1)
  more_than_one_idx = tf.squeeze(tf.where(more_than_one))
  more_than_one_vals = tf.squeeze(tf.gather(unique_a_vals, more_than_one_idx))

  not_duplicated, _ = tf.setdiff1d(a, more_than_one_vals)
  dups_in_a, indexes_in_a = tf.setdiff1d(a, not_duplicated)

  return more_than_one_vals, indexes_in_a, dups_in_a 
Example #17
Source File: network_units.py    From object_detection_kitti with Apache License 2.0 6 votes vote down vote up
def embedding_lookup(embedding_matrix, indices, ids, weights, size):
  """Performs a weighted embedding lookup.

  Args:
    embedding_matrix: float Tensor from which to do the lookup.
    indices: int Tensor for the output rows of the looked up vectors.
    ids: int Tensor vectors to look up in the embedding_matrix.
    weights: float Tensor weights to apply to the looked up vectors.
    size: int number of output rows. Needed since some output rows may be
        empty.

  Returns:
    Weighted embedding vectors.
  """
  embeddings = tf.nn.embedding_lookup([embedding_matrix], ids)
  # TODO(googleuser): allow skipping weights.
  broadcast_weights_shape = tf.concat([tf.shape(weights), [1]], 0)
  embeddings *= tf.reshape(weights, broadcast_weights_shape)
  embeddings = tf.unsorted_segment_sum(embeddings, indices, size)
  return embeddings 
Example #18
Source File: training.py    From bilm-tf with Apache License 2.0 6 votes vote down vote up
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 = tf.unique(indices)
    summed_values = tf.unsorted_segment_sum(
      values, new_index_positions,
      tf.shape(unique_indices)[0])
    return (summed_values, unique_indices) 
Example #19
Source File: expert_utils.py    From BERT with Apache License 2.0 6 votes vote down vote up
def combine(self, x):
    """Return the output from the experts.

    When one example goes to multiple experts, the outputs are summed.

    Args:
      x: a Tensor with shape [batch, num_experts, expert_capacity, depth]

    Returns:
      a `Tensor` with shape `[batch, length, depth]
    """
    depth = tf.shape(x)[-1]
    x *= tf.expand_dims(self._nonpadding, -1)
    ret = tf.unsorted_segment_sum(
        x, self._flat_indices, num_segments=self._batch * self._length)
    ret = tf.reshape(ret, [self._batch, self._length, depth])
    return ret 
Example #20
Source File: network_units.py    From Gun-Detector with Apache License 2.0 6 votes vote down vote up
def embedding_lookup(embedding_matrix, indices, ids, weights, size):
  """Performs a weighted embedding lookup.

  Args:
    embedding_matrix: float Tensor from which to do the lookup.
    indices: int Tensor for the output rows of the looked up vectors.
    ids: int Tensor vectors to look up in the embedding_matrix.
    weights: float Tensor weights to apply to the looked up vectors.
    size: int number of output rows. Needed since some output rows may be
        empty.

  Returns:
    Weighted embedding vectors.
  """
  embeddings = tf.nn.embedding_lookup([embedding_matrix], ids)
  # TODO(googleuser): allow skipping weights.
  broadcast_weights_shape = tf.concat([tf.shape(weights), [1]], 0)
  embeddings *= tf.reshape(weights, broadcast_weights_shape)
  embeddings = tf.unsorted_segment_sum(embeddings, indices, size)
  return embeddings 
Example #21
Source File: expert_utils.py    From BERT with Apache License 2.0 6 votes vote down vote up
def combine(self, expert_out, multiply_by_gates=True):
    """Sum together the expert output, weighted by the gates.

    The slice corresponding to a particular batch element `b` is computed
    as the sum over all experts `i` of the expert output, weighted by the
    corresponding gate values.  If `multiply_by_gates` is set to False, the
    gate values are ignored.

    Args:
      expert_out: a list of `num_experts` `Tensor`s, each with shape
        `[expert_batch_size_i, <extra_output_dims>]`.
      multiply_by_gates: a boolean

    Returns:
      a `Tensor` with shape `[batch_size, <extra_output_dims>]`.
    """
    # see comments on convert_gradient_to_tensor
    stitched = common_layers.convert_gradient_to_tensor(
        tf.concat(expert_out, 0))
    if multiply_by_gates:
      stitched *= tf.expand_dims(self._nonzero_gates, 1)
    combined = tf.unsorted_segment_sum(stitched, self._batch_index,
                                       tf.shape(self._gates)[0])
    return combined 
Example #22
Source File: training.py    From nlp_research with MIT License 6 votes vote down vote up
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 = tf.unique(indices)
    summed_values = tf.unsorted_segment_sum(
      values, new_index_positions,
      tf.shape(unique_indices)[0])
    return (summed_values, unique_indices) 
Example #23
Source File: expert_utils.py    From fine-lm with MIT License 6 votes vote down vote up
def combine(self, x):
    """Return the output from the experts.

    When one example goes to multiple experts, the outputs are summed.

    Args:
      x: a Tensor with shape [batch, num_experts, expert_capacity, depth]

    Returns:
      a `Tensor` with shape `[batch, length, depth]
    """
    depth = tf.shape(x)[-1]
    x *= tf.expand_dims(self._nonpadding, -1)
    ret = tf.unsorted_segment_sum(
        x, self._flat_indices, num_segments=self._batch * self._length)
    ret = tf.reshape(ret, [self._batch, self._length, depth])
    return ret 
Example #24
Source File: train_utils.py    From DeepPavlov with Apache License 2.0 6 votes vote down vote up
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 = tf.unique(indices)
    summed_values = tf.unsorted_segment_sum(values,
                                            new_index_positions,
                                            tf.shape(unique_indices)[0])
    return (summed_values, unique_indices) 
Example #25
Source File: expert_utils.py    From fine-lm with MIT License 6 votes vote down vote up
def combine(self, expert_out, multiply_by_gates=True):
    """Sum together the expert output, weighted by the gates.

    The slice corresponding to a particular batch element `b` is computed
    as the sum over all experts `i` of the expert output, weighted by the
    corresponding gate values.  If `multiply_by_gates` is set to False, the
    gate values are ignored.

    Args:
      expert_out: a list of `num_experts` `Tensor`s, each with shape
        `[expert_batch_size_i, <extra_output_dims>]`.
      multiply_by_gates: a boolean

    Returns:
      a `Tensor` with shape `[batch_size, <extra_output_dims>]`.
    """
    # see comments on convert_gradient_to_tensor
    stitched = common_layers.convert_gradient_to_tensor(
        tf.concat(expert_out, 0))
    if multiply_by_gates:
      stitched *= tf.expand_dims(self._nonzero_gates, 1)
    combined = tf.unsorted_segment_sum(stitched, self._batch_index,
                                       tf.shape(self._gates)[0])
    return combined 
Example #26
Source File: network_units.py    From object_detection_with_tensorflow with MIT License 6 votes vote down vote up
def embedding_lookup(embedding_matrix, indices, ids, weights, size):
  """Performs a weighted embedding lookup.

  Args:
    embedding_matrix: float Tensor from which to do the lookup.
    indices: int Tensor for the output rows of the looked up vectors.
    ids: int Tensor vectors to look up in the embedding_matrix.
    weights: float Tensor weights to apply to the looked up vectors.
    size: int number of output rows. Needed since some output rows may be
        empty.

  Returns:
    Weighted embedding vectors.
  """
  embeddings = tf.nn.embedding_lookup([embedding_matrix], ids)
  # TODO(googleuser): allow skipping weights.
  broadcast_weights_shape = tf.concat([tf.shape(weights), [1]], 0)
  embeddings *= tf.reshape(weights, broadcast_weights_shape)
  embeddings = tf.unsorted_segment_sum(embeddings, indices, size)
  return embeddings 
Example #27
Source File: expert_utils.py    From NMT_GAN with Apache License 2.0 6 votes vote down vote up
def Combine(self, expert_out, multiply_by_gates=True):
    """Sum together the expert output, weighted by the gates.

    The slice corresponding to a particular batch element `b` is computed
    as the sum over all experts `i` of the expert output, weighted by the
    corresponding gate values.  If `multiply_by_gates` is set to False, the
    gate values are ignored.

    Args:
      expert_out: a list of `num_experts` `Tensor`s, each with shape
        `[expert_batch_size_i, <extra_output_dims>]`.
      multiply_by_gates: a boolean

    Returns:
      a `Tensor` with shape `[batch_size, <extra_output_dims>]`.
    """
    # see comments on ConvertGradientToTensor
    stitched = ConvertGradientToTensor(tf.concat(expert_out, 0))
    if multiply_by_gates:
      stitched *= tf.expand_dims(self._nonzero_gates, 1)
    combined = tf.unsorted_segment_sum(stitched, self._batch_index,
                                       tf.shape(self._gates)[0])
    return combined 
Example #28
Source File: expert_utils.py    From NMT_GAN with Apache License 2.0 5 votes vote down vote up
def __init__(self, data_parallelism, model_parallelism, gates):
    """Constructs a Dispatcher.

    Args:
      data_parallelism: a Parallelism object.
      model_parallelism: a Parallelism object.
      gates: a list of 1d integer `Tensor`s, one per datashard.
        Says which expert to use for each batch element.

    Returns:
      a DistributedSingleDispatcher
    """
    gates = data_parallelism(tf.to_int32, gates)
    self._gates = gates
    self._data_parallelism = data_parallelism
    self._model_parallelism = model_parallelism

    # Compute the sizes number of examples going from each datashard to each
    # expert.
    def _PartSizes(gates):
      return tf.unsorted_segment_sum(
          tf.ones_like(gates), gates, model_parallelism.n)

    part_sizes_by_datashard = data_parallelism(_PartSizes, gates)
    self._part_sizes_by_expert = tf.unstack(
        tf.stack(part_sizes_by_datashard), num=model_parallelism.n, axis=1)

    # These indices will be used to combine the output on the datashards.
    def _StitchIndices(gates):
      return tf.dynamic_partition(
          tf.range(tf.size(gates)), gates, model_parallelism.n)

    self._stitch_indices = data_parallelism(_StitchIndices, gates) 
Example #29
Source File: chem_tensorflow_sparse.py    From gated-graph-neural-network-samples with MIT License 5 votes vote down vote up
def gated_regression(self, last_h, regression_gate, regression_transform):
        # last_h: [v x h]
        gate_input = tf.concat([last_h, self.placeholders['initial_node_representation']], axis=-1)  # [v x 2h]
        gated_outputs = tf.nn.sigmoid(regression_gate(gate_input)) * regression_transform(last_h)  # [v x 1]

        # Sum up all nodes per-graph
        graph_representations = tf.unsorted_segment_sum(data=gated_outputs,
                                                        segment_ids=self.placeholders['graph_nodes_list'],
                                                        num_segments=self.placeholders['num_graphs'])  # [g x 1]
        output = tf.squeeze(graph_representations)  # [g]
        self.output = output
        return output

    # ----- Data preprocessing and chunking into minibatches: 
Example #30
Source File: graph_builder.py    From Action_Recognition_Zoo with MIT License 5 votes vote down vote up
def EmbeddingLookupFeatures(params, sparse_features, allow_weights):
  """Computes embeddings for each entry of sparse features sparse_features.

  Args:
    params: list of 2D tensors containing vector embeddings
    sparse_features: 1D tensor of strings. Each entry is a string encoding of
      dist_belief.SparseFeatures, and represents a variable length list of
      feature ids, and optionally, corresponding weights values.
    allow_weights: boolean to control whether the weights returned from the
      SparseFeatures are used to multiply the embeddings.

  Returns:
    A tensor representing the combined embeddings for the sparse features.
    For each entry s in sparse_features, the function looks up the embeddings
    for each id and sums them into a single tensor weighing them by the
    weight of each id. It returns a tensor with each entry of sparse_features
    replaced by this combined embedding.
  """
  if not isinstance(params, list):
    params = [params]
  # Lookup embeddings.
  sparse_features = tf.convert_to_tensor(sparse_features)
  indices, ids, weights = gen_parser_ops.unpack_sparse_features(sparse_features)
  embeddings = tf.nn.embedding_lookup(params, ids)

  if allow_weights:
    # Multiply by weights, reshaping to allow broadcast.
    broadcast_weights_shape = tf.concat(0, [tf.shape(weights), [1]])
    embeddings *= tf.reshape(weights, broadcast_weights_shape)

  # Sum embeddings by index.
  return tf.unsorted_segment_sum(embeddings, indices, tf.size(sparse_features))