Python tensorflow.unsorted_segment_max() Examples

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 

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 

def get_pixel_value(img, x, y):

    """Cantor pairing for removing non-unique updates and indices. At the time of implementation, unfixed issue with scatter_nd causes problems with int32 update values. Till resolution, implemented on cpu """

    with tf.device('/cpu:0'):
        indices = tf.stack([y, x], 2)
        indices = tf.reshape(indices, (375*1242, 2))
        values = tf.reshape(img, [-1])

        Y = indices[:,0]
        X = indices[:,1]
        Z = (X + Y)*(X + Y + 1)/2 + Y

        filtered, idx = tf.unique(tf.squeeze(Z))
        updated_values  = tf.unsorted_segment_max(values, idx, tf.shape(filtered)[0])

        # updated_indices = tf.map_fn(fn=lambda i: reverse(i), elems=filtered, dtype=tf.float32)
        updated_indices = reverse_all(filtered)
        updated_indices = tf.cast(updated_indices, 'int32')
        resolved_map = tf.scatter_nd(updated_indices, updated_values, img.shape)

        return resolved_map 

def aggregate(data, agg_idx, new_size, method="sum"):
  """ Aggregate data

  Args:
    data: tf tensor, see "unsorted_segment_x" in tf documents for more detail
    agg_idx: tf tensor of int, index for aggregation
    new_size: tf tensor of int, size of the data after aggregation
    method: aggregation method

  Returns:
    agg_data: tf tensor, aggregated data
  """

  if method == "sum":
    agg_data = tf.unsorted_segment_sum(data, agg_idx, new_size)
  elif method == "avg":
    agg_data = tf.unsorted_segment_sum(data, agg_idx, new_size)
    denom_const = tf.unsorted_segment_sum(tf.ones_like(data), agg_idx, new_size)
    agg_data = tf.div(agg_data, (denom_const + tf.constant(1.0e-10)))
  elif method == "max":
    agg_data = tf.unsorted_segment_max(data, agg_idx, new_size)
  elif method == "min":
    agg_data = tf.unsorted_segment_max(-data, agg_idx, new_size)
  else:
    raise ValueError("Unsupported aggregation method!")

  return agg_data 

def xqa_crossentropy_loss(start_scores, end_scores, answer_span, answer2support, support2question, use_sum=True):
    """Very common XQA loss function."""
    num_questions = tf.reduce_max(support2question) + 1

    start, end = answer_span[:, 0], answer_span[:, 1]

    start_probs = segment_softmax(start_scores, support2question)
    start_probs = tf.gather_nd(start_probs, tf.stack([answer2support, start], 1))

    # only start probs are normalized on multi-paragraph, end probs conditioned on start only on per support level
    num_answers = tf.shape(answer_span)[0]
    is_aligned = tf.equal(tf.shape(end_scores)[0], num_answers)
    end_probs = tf.cond(
        is_aligned,
        lambda: tf.gather_nd(tf.nn.softmax(end_scores), tf.stack([tf.range(num_answers, dtype=tf.int32), end], 1)),
        lambda: tf.gather_nd(segment_softmax(end_scores, support2question), tf.stack([answer2support, end], 1))
    )

    answer2question = tf.gather(support2question, answer2support)
    # compute losses individually
    if use_sum:
        span_probs = tf.unsorted_segment_sum(
            start_probs, answer2question, num_questions) * tf.unsorted_segment_sum(
            end_probs, answer2question, num_questions)
    else:
        span_probs = tf.unsorted_segment_max(
            start_probs, answer2question, num_questions) * tf.unsorted_segment_max(
            end_probs, answer2question, num_questions)

    return -tf.reduce_mean(tf.log(tf.maximum(1e-6, span_probs + 1e-6))) 

def segment_is_max(inputs, segment_ids):
    num_segments = tf.reduce_max(segment_ids) + 1
    if len(inputs.get_shape()) > 1:
        inputs_max = tf.reduce_max(inputs, reduction_indices=list(range(1, len(inputs.get_shape()))))
    else:
        inputs_max = inputs
    max_per_partition = tf.unsorted_segment_max(inputs_max, segment_ids, num_segments)
    return tf.equal(inputs, tf.gather(max_per_partition, segment_ids)) 

def get_component_bounds(image):
  """Returns the bounding box of each connected component in `image`.

  Connected components are segments of adjacent True pixels in the image.

  Args:
    image: A 2D boolean image tensor.

  Returns:
    A tensor of shape (num_components, 5), where each row represents a connected
        component of the image as `(x0, y0, x1, y1, size)`. `size` is the count
        of True pixels in the component, and the coordinates are the top left
        and bottom right corners of the bounding box.
  """
  with tf.name_scope('get_component_bounds'):
    components = tf.contrib.image.connected_components(image)
    num_components = tf.reduce_max(components) + 1
    width = tf.shape(image)[1]
    height = tf.shape(image)[0]
    xs, ys = tf.meshgrid(tf.range(width), tf.range(height))
    component_x0 = _unsorted_segment_min(xs, components, num_components)[1:]
    component_x1 = tf.unsorted_segment_max(xs, components, num_components)[1:]
    component_y0 = _unsorted_segment_min(ys, components, num_components)[1:]
    component_y1 = tf.unsorted_segment_max(ys, components, num_components)[1:]
    component_size = tf.bincount(components)[1:]
    return tf.stack([
        component_x0, component_y0, component_x1, component_y1, component_size
    ],
                    axis=1) 

def _unsorted_segment_min(data, segment_ids, num_segments):
  return -tf.unsorted_segment_max(-data, segment_ids, num_segments) 

def get_aggregation_function(aggregation_fun: Optional[str]):
    if aggregation_fun in ['sum', 'unsorted_segment_sum']:
        return tf.unsorted_segment_sum
    if aggregation_fun in ['max', 'unsorted_segment_max']:
        return tf.unsorted_segment_max
    if aggregation_fun in ['mean', 'unsorted_segment_mean']:
        return tf.unsorted_segment_mean
    if aggregation_fun in ['sqrt_n', 'unsorted_segment_sqrt_n']:
        return tf.unsorted_segment_sqrt_n
    else:
        raise ValueError("Unknown aggregation function '%s'!" % aggregation_fun) 

def unsorted_segment_logsumexp(scores, segment_ids, num_segments):
    """Perform an unsorted segment safe logsumexp."""
    # 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=scores,
                                              segment_ids=segment_ids,
                                              num_segments=num_segments)
    scattered_log_maxes = tf.gather(params=max_per_segment,
                                    indices=segment_ids)
    recentered_scores = scores - scattered_log_maxes
    exped_recentered_scores = tf.exp(recentered_scores)

    per_segment_sums = tf.unsorted_segment_sum(exped_recentered_scores, segment_ids, num_segments)
    per_segment_logs = tf.log(per_segment_sums)
    return per_segment_logs + max_per_segment 

def unsorted_segment_softmax(logits, segment_ids, num_segments):
    """Perform a safe unsorted segment softmax."""
    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)

    probs = exped_recentered_scores / (tf.gather(params=per_segment_sums, indices=segment_ids) + SMALL_NUMBER)
    return probs 

def vertical_run_length_encoding(image):
  """Returns the runs in each column of the image.

  A run is a subsequence of consecutive pixels that all have the same value.
  Internally, we treat the image as batches of single-column images in order to
  use connected component analysis.

  Args:
    image: A 2D image.

  Returns:
    The column index of each vertical run.
    The value in the image for each vertical run.
    The length of each vertical run.
  """
  with tf.name_scope('run_length_encoding'):
    image = tf.convert_to_tensor(image, name='image', dtype=tf.bool)
    # Set arbitrary, distinct, nonzero values for True and False pixels.
    # True pixels map to 2, and False pixels map to 1.
    # Transpose the image, and insert an extra dimension. This creates a batch
    # of "images" for connected component analysis, where each image is a single
    # column of the original image. Therefore, the connected components are
    # actually runs from a single column.
    components = tf.contrib.image.connected_components(
        tf.to_int32(tf.expand_dims(tf.transpose(image), axis=1)) + 1)
    # Flatten in order to use with unsorted segment ops.
    flat_components = tf.reshape(components, [-1])

    num_components = tf.maximum(0, tf.reduce_max(components) + 1)
    # Get the column index corresponding to each pixel present in
    # flat_components.
    column_indices = tf.reshape(
        tf.tile(
            # Count 0 through `width - 1` on axis 0, then repeat each element
            # `height` times.
            tf.expand_dims(tf.range(tf.shape(image)[1]), axis=1),
            multiples=[1, tf.shape(image)[0]]),
        # pyformat: disable
        [-1])
    # Take the column index for each component. For each component index k,
    # we want any entry of column_indices where the corresponding entry in
    # flat_components is k. column_indices should be the same for all pixels in
    # the same component, so we can just take the max of all of them. Disregard
    # component 0, which just represents all of the zero pixels across the
    # entire array (should be empty, because we pass in a nonzero image).
    component_columns = tf.unsorted_segment_max(column_indices, flat_components,
                                                num_components)[1:]
    # Take the original value of each component. Again, the value should be the
    # same for all pixels in a single component, so we can just take the max.
    component_values = tf.unsorted_segment_max(
        tf.to_int32(tf.reshape(tf.transpose(image), [-1])), flat_components,
        num_components)[1:]
    # Take the length of each component (run), by counting the number of pixels
    # in the component.
    component_lengths = tf.to_int32(tf.bincount(flat_components)[1:])
    return component_columns, component_values, component_lengths