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 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 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 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 = 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
def marginalize_table(self, table, pool_mode, axis=None, keep_dims=False):
vals = tf.reshape(table['values'], shape=[-1, self.units_in])
eps = 1e-10
out = None
if axis == 0:
inds = table['indices'][:, 1]
num_segments = table['shape'][1]
if pool_mode == 'mean':
count = self.count_entries(table, axis,
keep_dims=keep_dims) + eps
out = tf.unsorted_segment_sum(vals,
inds,
num_segments=num_segments)
if keep_dims:
out = tf.expand_dims(out, axis=0)
out = out / count
elif pool_mode == 'max':
out = tf.unsorted_segment_max(vals,
inds,
num_segments=num_segments)
if keep_dims:
out = tf.expand_dims(out, axis=0)
elif axis == 1:
inds = table['indices'][:, 0]
num_segments = table['shape'][0]
if pool_mode == 'mean':
count = self.count_entries(table, axis,
keep_dims=keep_dims) + eps
out = tf.unsorted_segment_sum(vals,
inds,
num_segments=num_segments)
if keep_dims:
out = tf.expand_dims(out, axis=1)
out = out / count
elif pool_mode == 'max':
out = tf.unsorted_segment_max(vals,
inds,
num_segments=num_segments)
if keep_dims:
out = tf.expand_dims(out, axis=1)
elif axis == None:
if 'mean' in pool_mode:
count = self.count_entries(table, axis, keep_dims=keep_dims)
out = tf.reduce_sum(vals, axis=0, keep_dims=keep_dims)
out = out / count
elif 'max' in pool_mode:
out = tf.reduce_max(vals, axis=0, keep_dims=keep_dims)
if keep_dims:
out = tf.expand_dims(out, axis=1)
else:
raise ValueError("Unknown axis: %s" % axis)
if pool_mode == 'max': ## Hack to avoid large negative number. How to deal with max over no entries?
out = tf.where(tf.greater(out, -200000000), out,
tf.zeros_like(out))
return out
def compute_candidate_softmax_and_loss(self, cand_logits):
# The following is softmax-ing over the candidates per graph.
# As the number of candidates varies, we can't just use tf.softmax.
# We implement it with the logsumexp trick:
# Step (1): Obtain shift constant as max of the logits
max_per_graph = tf.unsorted_segment_max(
data=cand_logits,
segment_ids=self.placeholders['cand_graph_nodes_list'],
num_segments=self.placeholders['num_graphs_in_batch']
) # Shape [G]
# # Step (2): Distribute max out to the corresponding logits again, and shift scores:
max_per_cand = tf.gather(params=max_per_graph,
indices=self.placeholders['cand_graph_nodes_list'])
cand_logits_shifted = cand_logits - max_per_cand
# # Step (3): Exp, sum up per target, compute exp(score) / exp(sum) as softmax:
scores_exped = tf.exp(cand_logits_shifted)
scores_sum_per_graph = tf.unsorted_segment_sum(
data=scores_exped,
segment_ids=self.placeholders['cand_graph_nodes_list'],
num_segments=self.placeholders['num_graphs_in_batch']
) # Shape [G]
scores_sum_per_cand = tf.gather(
params=scores_sum_per_graph,
indices=self.placeholders['cand_graph_nodes_list']
)
self.ops['softmax_values'] = scores_exped / (scores_sum_per_cand + utils.SMALL_NUMBER)
self.ops['log_softmax_values'] = cand_logits_shifted - tf.log(scores_sum_per_cand + utils.SMALL_NUMBER)
labels = self.placeholders['select_targets']
flat_loss_values = -tf.cast(labels, "float32") * self.ops['log_softmax_values']
losses = tf.unsorted_segment_sum(
data=flat_loss_values,
segment_ids=self.placeholders['cand_graph_nodes_list'],
num_segments=self.placeholders['num_graphs_in_batch']
)
self.ops['loss'] = tf.reduce_mean(losses)
flat_correct_prediction = tf.cast(tf.equal(cand_logits, max_per_cand), "int64") * self.placeholders[
'select_targets']
correct_prediction = tf.unsorted_segment_max(
data=tf.cast(flat_correct_prediction, "float32"),
segment_ids=self.placeholders['cand_graph_nodes_list'],
num_segments=self.placeholders['num_graphs_in_batch']
)
self.ops['accuracy_task'] = tf.reduce_mean(correct_prediction)
def _scale_and_center_cloud(self, points, ids):
num_scenes = tf.to_int32((tf.reduce_max(ids) //
self.num_objects_per_scene) + 1)
scene_ids = tf.to_int32(ids // self.num_objects_per_scene)
# convert to int, as unsorted_segment_{min, max} is faster for int
points_int = tf.to_int32(points * 10000.)
axis_min = tf.to_float(
tf.unsorted_segment_min(points_int, scene_ids,
num_scenes)) / 10000.
axis_max = tf.to_float(
tf.unsorted_segment_max(points_int, scene_ids,
num_scenes)) / 10000.
axis_min.set_shape(axis_min.shape.with_rank(1))
axis_max.set_shape(axis_max.shape.with_rank(1))
axis_span = axis_max - axis_min
with tf.name_scope("scale_to_image_cube"):
max_span = tf.reduce_max(axis_span, axis=-1, keepdims=True)
sfactor = tf.divide((self.img_size - 2.0),
max_span,
name="scale_factor")
points = points * tf.gather(sfactor, scene_ids)
# update min/max/span to save computation
with tf.name_scope("center_objects"):
axis_shift = (-axis_min * sfactor +
(((self.img_size - 2.0) - axis_span * sfactor) / 2))
gathered_shift = tf.gather(axis_shift, scene_ids)
points = points + gathered_shift
return points
def compute_attention_normalization(self, message_attention_scores, message_targets, num_nodes):
if self.params['attention_normalization'] == 'softmax':
# The following is softmax-ing over the incoming messages per node.
# As the number of incoming varies, we can't just use tf.softmax.
# Reimplement with logsumexp trick:
# Step (1): Obtain shift constant as max of messages going into a node
message_attention_score_max_per_target = tf.unsorted_segment_max(
data=message_attention_scores,
segment_ids=message_targets,
num_segments=num_nodes) # Shape [V]
# Step (2): Distribute max out to the corresponding messages again, and shift scores:
message_attention_score_max_per_message = tf.gather(
params=message_attention_score_max_per_target,
indices=message_targets) # Shape [M]
message_attention_scores -= message_attention_score_max_per_message
# Step (3): Exp, sum up per target, compute exp(score) / exp(sum) as attention prob:
message_attention_scores_exped = tf.exp(message_attention_scores) # Shape [M]
message_attention_score_sum_per_target = tf.unsorted_segment_sum(
data=message_attention_scores_exped,
segment_ids=message_targets,
num_segments=num_nodes) # Shape [V]
message_attention_normalisation_sum_per_message = tf.gather(
params=message_attention_score_sum_per_target,
indices=message_targets) # Shape [M]
message_attention = message_attention_scores_exped / (
message_attention_normalisation_sum_per_message + utils.SMALL_NUMBER) # Shape [M]
return message_attention
def bidaf_layer(seq1,
seq1_length,
seq2,
seq2_length,
seq1_to_seq2=None,
seq2_to_seq1=None):
"""Encodes seq1 conditioned on seq2, e.g., using word-by-word attention."""
attn_scores, attn_probs, seq2_weighted = attention.diagonal_bilinear_attention(
seq1, seq2, seq2_length, False, seq2_to_seq1=seq2_to_seq1)
attn_scores += tf.expand_dims(
mask_for_lengths(seq1_length,
tf.shape(attn_scores)[1]), 2)
max_seq1 = tf.reduce_max(attn_scores, 2)
if seq1_to_seq2 is None:
seq1_attention = tf.nn.softmax(max_seq1, 1)
else:
segm_max_seq1 = tf.unsorted_segment_max(
max_seq1, seq1_to_seq2,
tf.reduce_max(seq1_to_seq2) + 1)
seq1_attention = tf.nn.softmax(segm_max_seq1, 1)
seq1_attention = tf.gather(seq1_attention, seq1_to_seq2)
seq1_attention.set_shape(max_seq1.get_shape())
seq1_weighted = tf.einsum('ij,ijk->ik', seq1_attention, seq1)
seq1_weighted = tf.expand_dims(seq1_weighted, 1)
seq1_weighted = tf.tile(seq1_weighted, [1, tf.shape(seq1)[1], 1])
return tf.concat(
[seq2_weighted, seq1 * seq2_weighted, seq1 * seq1_weighted], 2)
def find_and_solve_collided_indices(indices, values, shape):
"""Sometimes ground truth indices will be collided,
shape: - a TensorShape
We attempted to solved it by :
* Find all a set of collided indices
* For each set, update to the next possible index (based on iou_score)
* If not possible, remove ground truths :(
"""
width = shape[0].value
height = shape[1].value
depth = shape[2].value
def reverse(index):
# @TODO: linearize N-d dimension
x = index / (height * depth)
y = (index - x * height * depth) / depth
z = index - x * height * depth - y * depth
return tf.stack([x, y, z], -1)
flatten = tf.matmul(tf.cast(indices, tf.int32),
[[height * depth], [depth], [1]])
filtered_idx, idx = tf.unique(tf.squeeze(flatten))
# @TODO: filter based on IOU
updated_indices = tf.cast(
tf.map_fn(fn=lambda i: reverse(i), elems=filtered_idx), tf.int64)
updated_values = tf.unsorted_segment_max(
values, idx, num_segments=tf.shape(filtered_idx)[0])
return [updated_indices, updated_values]
def pick_max_features(feature, flatten, num_voxel):
'''
Args:
feature : n x f
flatten : n x 1 -> 1-d voxel id for each point
'''
# filtered: k -> unique voxel id
# idx : n -> indices of each point in filtered
filtered, idx = tf.unique(tf.squeeze(flatten))
# k x f -> max pooling in each unique voxel
updated_features = tf.unsorted_segment_max(feature, idx, tf.shape(filtered)[0])
#updated_features = tf.unsorted_segment_sum(feature, idx, tf.shape(filtered)[0])
print(updated_features)
# k x 3 -> 3-d voxel id for each unique voxel
updated_indices = tf.map_fn(fn = lambda i: reverse(i, num_voxel), elems = filtered)
print(updated_indices)
num_features = updated_features.shape[-1]
# D x H x W x f
voxels = tf.scatter_nd(updated_indices,
updated_features,
tf.TensorShape([num_voxel, num_voxel, num_voxel, num_features]))
return voxels
def _unsorted_segment_softmax(data,
segment_ids,
num_segments,
name="unsorted_segment_softmax"):
"""Performs an elementwise softmax operation along segments of a tensor.
The input parameters are analogous to `tf.unsorted_segment_sum`. It produces
an output of the same shape as the input data, after performing an
elementwise sofmax operation between all of the rows with common segment id.
Args:
data: A tensor with at least one dimension.
segment_ids: A tensor of indices segmenting `data` across the first
dimension.
num_segments: A scalar tensor indicating the number of segments. It should
be at least `max(segment_ids) + 1`.
name: A name for the operation (optional).
Returns:
A tensor with the same shape as `data` after applying the softmax operation.
"""
with tf.name_scope(name):
segment_maxes = tf.unsorted_segment_max(data, segment_ids,
num_segments)
maxes = tf.gather(segment_maxes, segment_ids)
# Possibly refactor to `tf.stop_gradient(maxes)` for better performance.
data -= maxes
exp_data = tf.exp(data)
segment_sum_exp_data = tf.unsorted_segment_sum(exp_data, segment_ids,
num_segments)
sum_exp_data = tf.gather(segment_sum_exp_data, segment_ids)
return exp_data / sum_exp_data
def chain_pool(self,tensor,chains,num_chain_segments,empty,nonempty,mode):
bsize = self.max_board_size
assert(len(tensor.shape) == 4)
assert(len(chains.shape) == 3)
assert(len(num_chain_segments.shape) == 1)
assert(tensor.shape[1].value == bsize)
assert(tensor.shape[2].value == bsize)
assert(chains.shape[1].value == bsize)
assert(chains.shape[2].value == bsize)
assert(mode == "sum" or mode == "max")
num_channels = tensor.shape[3].value
#Since tf.unsorted_segment* doesn't operate by batches or channels, we need to manually construct
#a different shift to add to each batch and each channel so that they pool into disjoint buckets.
#Each one needs max_chain_idxs different buckets.
num_segments_by_batch_and_channel = tf.fill([1,num_channels],1) * tf.expand_dims(num_chain_segments,axis=1)
shift = tf.cumsum(tf.reshape(num_segments_by_batch_and_channel,[-1]),exclusive=True)
num_segments = tf.reduce_sum(num_chain_segments) * num_channels
shift = tf.reshape(shift,[-1,1,1,num_channels])
segments = tf.expand_dims(chains,3) + shift
if mode == "sum":
pools = tf.unsorted_segment_sum(tensor,segments,num_segments=num_segments)
elif mode == "max":
pools = tf.unsorted_segment_max(tensor,segments,num_segments=num_segments)
else:
assert False
gathered = tf.gather(pools,indices=segments)
return gathered * tf.expand_dims(nonempty,axis=3) # + tensor * empty
def unique_with_inverse(x):
y, idx = tf.unique(x)
num_segments = tf.shape(y)[0]
num_elems = tf.shape(x)[0]
return (y, idx,
tf.unsorted_segment_max(tf.range(num_elems), idx,
num_segments))
def non_zero_batchsize_op():
max_length = tf.shape(seq)[1]
encoded = tf.nn.embedding_lookup(ctxt_word_embeddings, seq)
one_hot = [0.0] * num_sequences
one_hot[i] = 1.0
mode_feature = tf.constant([[one_hot]], tf.float32)
mode_feature = tf.tile(mode_feature,
tf.stack([num_seq, max_length, 1]))
encoded = tf.concat([encoded, mode_feature], 2)
encoded = modular_encoder.modular_encoder(
sequence_module, {'text': encoded}, {'text': length},
{'text': None}, size, 1.0 - keep_prob, is_eval)[0]['text']
mask = misc.mask_for_lengths(length,
max_length,
mask_right=False,
value=1.0)
encoded = encoded * tf.expand_dims(mask, 2)
seq_lemmas = tf.gather(word2lemma_off, tf.reshape(seq, [-1]))
new_lemma_embeddings = tf.unsorted_segment_max(
tf.reshape(encoded, [-1, size]), seq_lemmas,
tf.reduce_max(word2lemma_off) + 1)
new_lemma_embeddings = tf.nn.relu(new_lemma_embeddings)
return tf.gather(new_lemma_embeddings, word2lemma_off)
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]
Y = tf.cast(Y, tf.float64)
X = tf.cast(X, tf.float64)
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 segment_pool(self, point_cloud, current, grid_size, scope=None):
with tf.variable_scope(scope):
point_cloud = (point_cloud + 1) / 2 * grid_size
point_cloud_4d = tf.concat(
(self.batch_idx, tf.reshape(point_cloud, (-1, 3))), axis=1)
centers = tf.round(point_cloud_4d)
point_cloud_idx = tf.cast(centers, tf.int64)
grid_shape = tf.constant(
[self.batch_size, grid_size, grid_size, grid_size],
dtype=tf.int64)
point_cloud_lin_idx = tf.squeeze(
ravel_index(point_cloud_idx, grid_shape))
vox_idx, idx, count = tf.unique_with_counts(point_cloud_lin_idx,
out_idx=tf.int32)
vox_idx = tf.cast(vox_idx, dtype=tf.int32)
vox_mult_idx = tf.transpose(unravel_index(vox_idx, grid_shape),
(1, 0))
# batch_idx = tf.squeeze(tf.slice(vox_mult_idx,(0,0),(-1,1)))
sh = current.get_shape().as_list()
features_vox = tf.unsorted_segment_max(
tf.reshape(current, (-1, sh[-1])), idx,
tf.reduce_max(idx) + 1)
features = tf.gather(features_vox, idx, axis=0)
features = tf.reshape(features, sh)
# max_point_per_vol = tf.segment_max(count,batch_idx)
# max_point_per_vol = tf.gather(max_point_per_vol,batch_idx)
# count_normalized = tf.divide(count,max_point_per_vol)
return features, features_vox, vox_mult_idx, count
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, (IMG_HT * IMG_WDT, 2))
values = tf.reshape(img, [-1])
Y = indices[:, 0]
X = indices[:, 1]
# The below raises an Error because
# Y is int32 but Z1 results in float64 type on the laptop CPU
# Z = (X + Y)*(X + Y + 1)/2 + Y
Z1 = (X + Y) * (X + Y + 1) / 2
Z1 = tf.cast(Z1, 'int32')
Z = Z1 + 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 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 _unsorted_segment_reduction_or_zero(reducer, values, indices, num_groups):
"""Common code for unsorted_segment_{min,max}_or_zero (below)."""
reduced = reducer(values, indices, num_groups)
present_indices = tf.unsorted_segment_max(
tf.ones_like(indices, dtype=reduced.dtype), indices, num_groups)
present_indices = tf.clip_by_value(present_indices, 0, 1)
present_indices = tf.reshape(present_indices, [num_groups] + [1] *
(reduced.shape.ndims - 1))
reduced *= present_indices
return reduced
def batch_unsrt_segment_max(data, segment_ids, num_segments):
""" Performas the `tf.unsorted_segment_max` operation batch-wise"""
# create distinct segments per batch
num_batches = tf.shape(segment_ids, out_type=tf.int64)[0]
batch_indices = tf.range(num_batches)
segment_ids_per_batch = segment_ids + num_segments * tf.expand_dims(batch_indices, axis=1)
# do the normal unsegment sum and reshape to original shape
seg_maxs = tf.unsorted_segment_max(data, segment_ids_per_batch, num_segments * num_batches)
return tf.reshape(seg_maxs, tf.stack((-1, num_segments)))
def rasterize(self,mesh):
samples_tri = mesh.faces_up
vertices_t = mesh.vertices_t
faces_values = tf.gather(vertices_t,samples_tri,axis=1) # [batch, samples, edge, xyz ]
# projection of normals on Z axis
V1 = tf.gather(faces_values,1,axis=2) - tf.gather(faces_values,0,axis=2)
V2 = tf.gather(faces_values,2,axis=2) - tf.gather(faces_values,0,axis=2)
V1x,V1y,V1z = tf.split(V1,[1,1,1],axis=2)
V2x,V2y,V2z = tf.split(V2,[1,1,1],axis=2)
normal_z = tf.reshape(V1x*V2y - V1y*V2x,(-1,))
# Up sampling the mesh
num_faces = samples_tri.get_shape().as_list()[0]
r1 = tf.random_uniform((self.batch_size,num_faces,1,1),dtype=tf.float32)
r2 = tf.random_uniform((self.batch_size,num_faces,1,1),dtype=tf.float32)
rand_samps = tf.concat(((1-tf.sqrt(r1)),(tf.sqrt(r1)*(1-r2)),(r2*tf.sqrt(r1))),axis=2)
vertices_MC = tf.reduce_sum(faces_values*rand_samps,axis=2)
num_vertices = num_faces
# Re-project
batch_idx_np = np.expand_dims(np.meshgrid(np.arange(0, self.batch_size, 1),np.arange(0, num_vertices, 1))[0].transpose(),axis=2)
batch_idx = tf.reshape(tf.convert_to_tensor(batch_idx_np.astype(np.float32)),(self.batch_size*num_vertices,1))
vertices_trans_values = tf.reshape(vertices_MC,(-1,3))
vertices_trans_values = (vertices_trans_values/2+0.5)
image_mask = tf.logical_or(tf.reduce_any(tf.less_equal(vertices_trans_values,self.minValue),axis=1),
tf.reduce_any(tf.greater_equal(vertices_trans_values,self.maxValue),axis=1))
image_mask = tf.logical_or(image_mask,tf.greater_equal(normal_z,0))
image_mask = tf.tile(tf.expand_dims(image_mask,axis=-1),(1,3))
vertices_trans_values = tf.where(image_mask,tf.zeros(image_mask.get_shape().as_list(),dtype=tf.float32),vertices_trans_values)
# vertices_up = tf.reshape(vertices_trans_values,(self.batch_size,-1,3))*2-1
vertices_trans_indices_4d = tf.concat((batch_idx,vertices_trans_values*(self.canvas_size-1)),axis=1)
vertices_trans_indices_4d = tf.cast(vertices_trans_indices_4d,tf.int64)
point_cloud_lin_idx = tf.squeeze(ravel_index_bxy(tf.slice(vertices_trans_indices_4d,[0,0],[-1,3]),(self.batch_size,self.canvas_size,self.canvas_size)))
vox_idx, idx, count = tf.unique_with_counts(point_cloud_lin_idx,out_idx=tf.int32)
# indices = tf.transpose(unravel_index(vox_idx,(self.batch_size,self.canvas_size,self.canvas_size),Type=tf.int64),(1,0))
values_z = -1*tf.unsorted_segment_max(tf.slice(-1*vertices_trans_values,[0,2],[-1,1]),idx,tf.reduce_max(idx)+1)
# get XYZ values corresponding to values_z
max_values = tf.squeeze(tf.gather(values_z,idx,axis=0),axis=1)
max_xyz_idx = tf.squeeze(tf.where(tf.equal(vertices_trans_values[:,2],max_values)),axis=1)
values_xyz = tf.gather(vertices_trans_values ,max_xyz_idx,axis=0)
values_xyz_idx = tf.gather(vertices_trans_indices_4d,max_xyz_idx,axis=0)
image = tf.scatter_nd(values_xyz_idx[:,0:3], values_xyz, (self.batch_size,self.canvas_size,self.canvas_size,3))
# Semantic labeling image
semantics = tf.tile(tf.expand_dims(mesh.semantics_up,axis=0),(self.batch_size,1,1))
semantics = tf.reshape(semantics,(-1,3))
semantics = tf.where(image_mask,tf.zeros(image_mask.get_shape().as_list(),dtype=tf.int64),semantics)
semantics_visible = tf.gather(semantics, max_xyz_idx,axis=0)
image_semantics = tf.scatter_nd(values_xyz_idx[:,0:3], semantics_visible, (self.batch_size,self.canvas_size,self.canvas_size,3))
return image, image_semantics
def define_prediction_with_loss(self, node_embeddings):
node_embeddings = tf.identity(node_embeddings)
graph_embeddings = self.define_pooling(node_embeddings)
domain_nodes_embeddings = tf.gather(
params=node_embeddings, indices=self.placeholders['domain'])
graph_embeddings_copied = tf.gather(
params=graph_embeddings,
indices=self.placeholders['domain_node_graph_ids_list'])
domain_node_score_calculator = MLP(
in_size=domain_nodes_embeddings.get_shape()[1] +
graph_embeddings_copied.get_shape()[1],
out_size=1,
hid_sizes=self.classifier_hidden_dims)
domain_node_logits = domain_node_score_calculator(
tf.concat([domain_nodes_embeddings, graph_embeddings_copied],
axis=-1))
domain_node_logits = tf.reshape(domain_node_logits, [-1])
probs, log_probs = SegmentBasedSoftmax(
data=domain_node_logits,
segment_ids=self.placeholders['domain_node_graph_ids_list'],
num_segments=self.placeholders['num_graphs'],
return_log=True)
self.ops['probabilities'] = probs
loss_per_domain_node = -tf.cast(self.placeholders['domain_labels'],
tf.float32) * log_probs
loss_per_graph = tf.unsorted_segment_sum(
data=loss_per_domain_node,
segment_ids=self.placeholders['domain_node_graph_ids_list'],
num_segments=self.placeholders['num_graphs'])
self.ops['loss'] = tf.reduce_mean(loss_per_graph)
domain_node_max_scores = tf.unsorted_segment_max(
data=domain_node_logits,
segment_ids=self.placeholders['domain_node_graph_ids_list'],
num_segments=self.placeholders['num_graphs'])
copied_domain_node_max_scores = tf.gather(
params=domain_node_max_scores,
indices=self.placeholders['domain_node_graph_ids_list'])
selected_domain_nodes = tf.cast(tf.equal(copied_domain_node_max_scores,
domain_node_logits),
dtype=tf.int32)
correct_prediction_per_node = tf.cast(
tf.equal(selected_domain_nodes,
self.placeholders['domain_labels']), tf.float32)
correct_prediction = tf.unsorted_segment_prod(
data=correct_prediction_per_node,
segment_ids=self.placeholders['domain_node_graph_ids_list'],
num_segments=self.placeholders['num_graphs'])
self.ops['accuracy'] = tf.reduce_mean(correct_prediction)
def edge_normalize(edge_states, sender_node_ids, n_nodes=None, sorted=True):
"""
Args:
edge_states: batch_size x n_edges x n_edge_dims
sender_node_ids: n_edges
sorted: the list sender_node_ids is sorted or not
Returns:
edge_states_norm: batch_size x n_nodes x n_edge_dims
"""
edge_states = tf.transpose(edge_states,
perm=[1, 0,
2]) # n_edges x batch_size x n_edge_dims
if sorted:
edge_states_max = tf.segment_max(
edge_states, sender_node_ids) # n_nodes x batch_size x n_edge_dims
edge_states_max = tf.gather(
edge_states_max,
sender_node_ids) # n_edges x batch_size x n_edge_dims
edge_states_exp = tf.exp(
edge_states -
edge_states_max) # n_edges x batch_size x n_edge_dims
edge_states_sumexp = tf.segment_sum(
edge_states_exp,
sender_node_ids) # n_nodes x batch_size x n_edge_dims
edge_states_sumexp = tf.gather(
edge_states_sumexp,
sender_node_ids) # n_edges x batch_size x n_edge_dims
edge_states_norm = edge_states_exp / edge_states_sumexp # n_edges x batch_size x n_edge_dims
else:
if n_nodes is None:
raise ValueError('`n_nodes` should not be None')
edge_states_max = tf.unsorted_segment_max(
tf.transpose(edge_states, perm=[1, 0, 2]), sender_node_ids,
n_nodes) # n_nodes x batch_size x n_edge_dims
edge_states_max = tf.gather(
edge_states_max,
sender_node_ids) # n_edges x batch_size x n_edge_dims
edge_states_exp = tf.exp(
edge_states -
edge_states_max) # n_edges x batch_size x n_edge_dims
edge_states_sumexp = tf.unsorted_segment_sum(
edge_states_exp, sender_node_ids,
n_nodes) # n_nodes x batch_size x n_edge_dims
edge_states_sumexp = tf.gather(
edge_states_sumexp,
sender_node_ids) # n_edges x batch_size x n_edge_dims
edge_states_norm = edge_states_exp / edge_states_sumexp # n_edges x batch_size x n_edge_dims
edge_states_norm = tf.transpose(
edge_states_norm, perm=[1, 0, 2]) # batch_size x n_edges x n_edge_dims
return edge_states_norm
def get_batch_max_per_key(tensor, key_uniques, dtype): # pylint: disable=missing-docstring
if tensor.get_shape().ndims < 2:
row_maxes = tensor
else:
row_maxes = tf.reduce_max(tensor,
axis=tf.range(1,
tensor.get_shape().ndims))
batch_max = tf.unsorted_segment_max(row_maxes, key_uniques.idx,
tf.size(input=key_uniques.y))
# TODO(b/112309021): Remove workaround once tf.reduce_max of a tensor of all
# NaNs produces -inf.
return _inf_to_nan(batch_max, dtype)
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 apply(self, x, segments, num_sentences_mask):
num_segments = tf.shape(num_sentences_mask)[0] * tf.reduce_max(
num_sentences_mask)
segmented = tf.unsorted_segment_max(x,
segments,
num_segments=num_segments)
reshaped = tf.reshape(segmented,
(tf.shape(x)[0], -1, x.shape.as_list()[2]))
mask = tf.expand_dims(
tf.cast(
tf.sequence_mask(num_sentences_mask,
tf.shape(reshaped)[1]), tf.float32), 2)
reshaped *= mask
return reshaped
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 message_passing(state, read_ids, write_ids, node_count, name):
values = tf.nn.embedding_lookup(
state, read_ids, name="lookup_" + name)
s_sum = tf.unsorted_segment_sum(
values, write_ids, node_count, name="segment_sum_" + name)
s_max = tf.unsorted_segment_max(
values, write_ids, node_count, name="segment_max_" + name)
s_max = tf.maximum(s_max, -1.0)
s_min = tf.unsorted_segment_min(
values, write_ids, node_count, name="segment_min_" + name)
s_min = tf.minimum(s_min, 1.0)
return [s_sum, s_max, s_min]
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 log_partition_function(num_nodes,
scores,
forest=False,
max_dynamic_range=None):
r"""Returns the log of the sum-of-product of spanning trees or forests.
Computing the sum-of-product in the log domain reduces the chance of overflow
or underflow, and ML techniques (e.g., CRF loss functions) typically require
the log partition function anyways. For similar reasons, the scores input is
assumed to be specified in the log domain.
The partition function is caluclated via application of the Matrix-Tree
theorem; see the following for details:
https://en.wikipedia.org/wiki/Kirchhoff%27s_theorem
http://www.aclweb.org/anthology/D/D07/D07-1015.pdf
Computing the gradient of the log partition function requires inverting the
Laplacian matrix. Numerical issues may occur if the Laplacian is singular or
nearly-so. (Intuitively, the Laplacian will be close to singular when the
input scores strongly favor invalid structures such as cycles). In the EMNLP
paper, we alleviated the numerical issues by clipping the difference between
the minimum and maximum score for each node to 20 (in the log domain). The
|max_dynamic_range| argument can be used for this purpose.
TODO(googleuser): Try improving the condition number of the Laplacian matrix
directly, instead of using the indirect approach above. For example, one
could add c*I to the Laplacian (i.e., Tikhonov regularization).
Args:
num_nodes: [B] vector of graph sizes per batch item.
scores: [B,M,M] tensor of padded batched arc and root scores, in the format
used by the maximum_spanning_tree() op. Padding values must be finite.
forest: If true, sum over spanning forests instead of trees.
max_dynamic_range: If specified, incoming scores for each node are clipped
to at most this far from the maximum such score (in the log domain).
Returns:
[B] vector Z of log partition function values, where
Z[b] = log(
\sum_{tree spanning batch item b}
score(root_of(tree)) \prod_{arc in tree} score(arc))
"""
orig_dtype = scores.dtype.base_dtype
scores_bxmxm = tf.to_double(scores) # use doubles to reduce under/overflow
shape_bxmxm = tf.shape(scores_bxmxm)
batch_size = shape_bxmxm[0]
max_nodes = shape_bxmxm[1]
total_nodes = batch_size * max_nodes
# To eliminate overflow, we locally normalize the scores. Specifically, for
# each node we divide its incoming arc scores and root selection score by the
# maximum such score. Since each node in a tree must select exactly one of
# these scores (i.e., it is either a root or has exactly one incoming arc),
# the local normalization factors are identical for all trees and can thus be
# factored out of the sum over trees.
#
# More concretely, we find the maximum per node, divide all scores for that
# node by the maximum, and then find the partition function of the normalized
# scores. Then we recover the un-normalized partition function by multiplying
# the per-node maxima back in. This final step is performed in the log domain
# to avoid overflow.
#
# Note that underflow is still possible, but unlikely as long as the scores
# are close to feasible (i.e., there is not too much mass on non-trees). The
# |max_dynamic_range| argument can be used to mitigate this.
# Finding the maximum incoming score is difficult, because the batch padding
# may contain arbitrary values. We restrict the maximization to valid arcs
# using tf.unsorted_segment_max() with a specially-constructed set of IDs.
_, valid_tokens_bxm = digraph_ops.ValidArcAndTokenMasks(
num_nodes, max_nodes, dtype=tf.int32)
# Create a tensor of "target IDs". In each row of each sub-matrix, the
# positions of valid source tokens are filled with the 1-origin index of that
# row in the entire batch, and zero elsewhere. For example, given a batch
# with num_nodes=[2, 3] we might have
# [[[1, 1, 0],
# [2, 2, 0],
# [3, 3, 0]],
# [[4, 4, 4],
# [5, 5, 5],
# [6, 6, 6]]]
#
# TODO(googleuser): The dynamic masking is pretty awkward. Find an op that does
# this (I looked, but maybe not hard enough), or write a custom op for this.
valid_tokens_bx1xm = tf.expand_dims(valid_tokens_bxm, 1)
valid_sources_bxmxm = tf.tile(valid_tokens_bx1xm, [1, max_nodes, 1])
sequence_bm = 1 + tf.range(total_nodes, dtype=tf.int32)
sequence_bxmx1 = tf.reshape(sequence_bm, [batch_size, max_nodes, 1])
target_ids_bxmxm = valid_sources_bxmxm * sequence_bxmx1
max_scores_bm1 = tf.unsorted_segment_max(scores_bxmxm, target_ids_bxmxm,
total_nodes + 1)
max_scores_bm = max_scores_bm1[1:] # ID 0 corresponds to padding
# Similar to above, we need to sum over the valid tokens. We analogously use
# tf.unsorted_segment_sum() with a specially-constructed set of "batch IDs".
sequence_b = 1 + tf.range(batch_size, dtype=tf.int32)
sequence_bx1 = tf.expand_dims(sequence_b, 1)
batch_ids_bxm = valid_tokens_bxm * sequence_bx1
batch_ids_bm = tf.reshape(batch_ids_bxm, [-1])
log_normalization_factor_b1 = tf.unsorted_segment_sum(
max_scores_bm, batch_ids_bm, batch_size + 1)
log_normalization_factor_b = log_normalization_factor_b1[1:]
# Locally-normalize and optionally clip the scores.
max_scores_bxmx1 = tf.reshape(max_scores_bm, [batch_size, max_nodes, 1])
scores_bxmxm -= max_scores_bxmx1
if max_dynamic_range is not None:
# After normalization, the scores are non-positive with max=0, so the
# |max_dynamic_range| can be applied directly.
#
# PyLint thinks "-max_dynamic_range" is invalid because it defaults to None.
scores_bxmxm = tf.maximum(scores_bxmxm, -max_dynamic_range)
scores_bxmxm = tf.exp(scores_bxmxm)
# Apply the Matrix-Tree theorem.
exp_normalized_laplacian_bxmxm = digraph_ops.LaplacianMatrix(
num_nodes, scores_bxmxm, forest=forest)
log_normalized_partition_function_b = tf.log(
tf.matrix_determinant(exp_normalized_laplacian_bxmxm))
# Reapply the normalization factor that was divided out.
log_partition_function_b = (
log_normalized_partition_function_b + log_normalization_factor_b)
return tf.cast(log_partition_function_b, orig_dtype)
