def apply_hard_selection(self, scores):
# <int32>[batch_size, seq_length]
# Using argsort(argsort(.)) gives the rank of the score.
sorted_indexes = tf.argsort(
tf.argsort(scores, direction="DESCENDING", axis=-1, stable=True))
# <float32>[batch_size, seq_length]
mask = tf.cast(sorted_indexes <
tf.ones_like(sorted_indexes) * self._max_num_tokens,
dtype=tf.float32)
return mask
def _face_model_map_fn(example):
vertices = example['vertices']
# Randomly shift vertices
if apply_random_shift:
vertices = random_shift(vertices)
example['num_vertices'] = tf.shape(vertices)[0]
# Optionally shuffle vertices and re-order faces to match
if shuffle_vertices:
permutation = tf.random_shuffle(tf.range(example['num_vertices']))
vertices = tf.gather(vertices, permutation)
face_permutation = tf.concat([
tf.constant([0, 1], dtype=tf.int32),
tf.argsort(permutation) + 2
],
axis=0)
example['faces'] = tf.cast(
tf.gather(face_permutation, example['faces']), tf.int64)
# Vertices are quantized. So convert to floats for input to face model
example['vertices'] = modules.dequantize_verts(vertices,
quantization_bits)
example['vertices_mask'] = tf.ones_like(example['vertices'][Ellipsis,
0],
dtype=tf.float32)
example['faces_mask'] = tf.ones_like(example['faces'],
dtype=tf.float32)
return example
def boolean_mask(tensor, mask):
"""Select elements from tensor where mask is True, and pads the rest with 0.
Example:
tensor = [1 2 3 4 5 6 7]
mask = [T T F F T T T]
output = [1 2 5 6 7 0 0]
Args:
tensor: <T> [batch_size, dim, ...]
mask: <bool> [batch_size, dim]
Returns:
masked_tensor: <T> [batch_size, dim, ...]. The first k items of row i
correspond to the k truthy values in row i of the mask tensor,
indexed into the provided tensor.
"""
batch_size = shape(mask, 0)
dim = shape(mask, 1)
# Bring the coordinates of un-masked elements to the front.
mask = tf.cast(mask, tensor.dtype)
indices = tf.argsort(mask, axis=1, direction="DESCENDING", stable=True)
batch_indices = tf.expand_dims(tf.range(tf.shape(mask)[0]), 1)
batch_indices = tf.tile(batch_indices, [1, tf.shape(indices)[1]])
coords = tf.stack([batch_indices, indices], axis=2)
# Gather and zero masked elements.
selected_tensor = tf.gather_nd(tensor, coords)
selected_mask = tf.gather_nd(mask, coords)
reshape = [batch_size, dim] + [1] * (tensor.get_shape().ndims - 2)
selected_mask = tf.reshape(selected_mask, reshape)
masked_tensor = selected_tensor * selected_mask
return masked_tensor
def _top_k_sample(logits, ignore_ids=None, num_samples=1, k=10):
"""
Does top-k sampling. if ignore_ids is on, then we will zero out those logits.
:param logits: [batch_size, vocab_size] tensor
:param ignore_ids: [vocab_size] one-hot representation of the indices we'd like to ignore and never predict,
like padding maybe
:param p: topp threshold to use, either a float or a [batch_size] vector
:return: [batch_size, num_samples] samples
# TODO FIGURE OUT HOW TO DO THIS ON TPUS. IT'S HELLA SLOW RIGHT NOW, DUE TO ARGSORT I THINK
"""
with tf.variable_scope('top_p_sample'):
batch_size, vocab_size = get_shape_list(logits, expected_rank=2)
probs = tf.nn.softmax(logits if ignore_ids is None else logits - tf.cast(ignore_ids[None], tf.float32) * 1e10,
axis=-1)
# [batch_size, vocab_perm]
indices = tf.argsort(probs, direction='DESCENDING')
# find the top pth index to cut off. careful we don't want to cutoff everything!
# result will be [batch_size, vocab_perm]
k_expanded = k if isinstance(k, int) else k[:, None]
exclude_mask = tf.range(vocab_size)[None] >= k_expanded
# OPTION A - sample in the sorted space, then unsort.
logits_to_use = tf.batch_gather(logits, indices) - tf.cast(exclude_mask, tf.float32) * 1e10
sample_perm = tf.random.categorical(logits=logits_to_use, num_samples=num_samples)
sample = tf.batch_gather(indices, sample_perm)
return {
'probs': probs,
'sample': sample,
}
def nms_tf(dets, thresh):
"""Non-maximum suppression with tf graph mode."""
x1 = dets[:, 0]
y1 = dets[:, 1]
x2 = dets[:, 2]
y2 = dets[:, 3]
scores = dets[:, 4]
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
order = tf.argsort(scores, direction='DESCENDING')
keep = tf.TensorArray(tf.int32, size=0, dynamic_size=True)
index = 0
while tf.shape(order)[0] > 0:
i = order[0]
keep = keep.write(index, i)
xx1 = tf.maximum(x1[i], tf.gather(x1, order[1:]))
yy1 = tf.maximum(y1[i], tf.gather(y1, order[1:]))
xx2 = tf.minimum(x2[i], tf.gather(x2, order[1:]))
yy2 = tf.minimum(y2[i], tf.gather(y2, order[1:]))
w = tf.maximum(0.0, xx2 - xx1 + 1)
h = tf.maximum(0.0, yy2 - yy1 + 1)
intersection = w * h
overlap = intersection / (
areas[i] + tf.gather(areas, order[1:]) - intersection)
inds = tf.where_v2(overlap <= thresh)
order = tf.concat(tf.gather(order, inds + 1), axis=1)
order = tf.squeeze(order, axis=-1)
index += 1
return keep.stack()
def build_decoder(encoded_image, caption_input, params, is_training):
decoder_output = transformer.decoder_layer(caption_input, encoded_image,
params, is_training) # n,20,512
logit = tf.layers.dense(decoder_output, params.vocab_size) # n,16,5000
probs = tf.nn.softmax(logit, name='probs')
sorted_indices = tf.identity(tf.argsort(probs,
axis=-1,
direction='DESCENDING'),
name='sorted_indices')
return probs, sorted_indices
def _top_p_sample(logits, ignore_ids=None, num_samples=1, p=0.9):
"""
Does top-p sampling. if ignore_ids is on, then we will zero out those logits.
:param logits: [batch_size, vocab_size] tensor
:param ignore_ids: [vocab_size] one-hot representation of the indices we'd like to ignore and never predict,
like padding maybe
:param p: topp threshold to use, either a float or a [batch_size] vector
:return: [batch_size, num_samples] samples
# TODO FIGURE OUT HOW TO DO THIS ON TPUS. IT'S HELLA SLOW RIGHT NOW, DUE TO ARGSORT I THINK
"""
with tf.variable_scope('top_p_sample'):
batch_size, vocab_size = get_shape_list(logits, expected_rank=2)
probs = tf.nn.softmax(logits if ignore_ids is None else logits - tf.cast(ignore_ids[None], tf.float32) * 1e10,
axis=-1)
if isinstance(p, float) and p > 0.999999:
# Don't do top-p sampling in this case
print("Top-p sampling DISABLED", flush=True)
return {
'probs': probs,
'sample': tf.random.categorical(
logits=logits if ignore_ids is None else logits - tf.cast(ignore_ids[None], tf.float32) * 1e10,
num_samples=num_samples, dtype=tf.int32),
}
# [batch_size, vocab_perm]
indices = tf.argsort(probs, direction='DESCENDING')
cumulative_probabilities = tf.math.cumsum(tf.batch_gather(probs, indices), axis=-1, exclusive=False)
# find the top pth index to cut off. careful we don't want to cutoff everything!
# result will be [batch_size, vocab_perm]
p_expanded = p if isinstance(p, float) else p[:, None]
exclude_mask = tf.logical_not(
tf.logical_or(cumulative_probabilities < p_expanded, tf.range(vocab_size)[None] < 1))
# OPTION A - sample in the sorted space, then unsort.
logits_to_use = tf.batch_gather(logits, indices) - tf.cast(exclude_mask, tf.float32) * 1e10
sample_perm = tf.random.categorical(logits=logits_to_use, num_samples=num_samples)
sample = tf.batch_gather(indices, sample_perm)
# OPTION B - unsort first - Indices need to go back to 0 -> N-1 -- then sample
# unperm_indices = tf.argsort(indices, direction='ASCENDING')
# include_mask_unperm = tf.batch_gather(include_mask, unperm_indices)
# logits_to_use = logits - (1 - tf.cast(include_mask_unperm, tf.float32)) * 1e10
# sample = tf.random.categorical(logits=logits_to_use, num_samples=num_samples, dtype=tf.int32)
return {
'probs': probs,
'sample': sample,
}
def _classification_loss(self, pred_label, gt_label, num_matched_boxes):
"""Computes the classification loss.
Computes the classification loss with hard negative mining.
Args:
pred_label: a flatten tensor that includes all predicted class. The shape
is [batch_size, num_anchors, num_classes].
gt_label: a tensor that represents the classification groundtruth targets.
The shape is [batch_size, num_anchors, 1].
num_matched_boxes: the number of anchors that are matched to a groundtruth
targets. This is used as the loss normalizater.
Returns:
box_loss: a float32 representing total box regression loss.
"""
cross_entropy = tf.losses.sparse_softmax_cross_entropy(
gt_label, pred_label, reduction=tf.losses.Reduction.NONE)
mask = tf.greater(tf.squeeze(gt_label), 0)
float_mask = tf.cast(mask, tf.float32)
# Hard example mining
neg_masked_cross_entropy = cross_entropy * (1 - float_mask)
relative_position = tf.argsort(
tf.argsort(neg_masked_cross_entropy, direction='DESCENDING'))
num_neg_boxes = tf.minimum(
tf.to_int32(num_matched_boxes) * ssd_constants.NEGS_PER_POSITIVE,
ssd_constants.NUM_SSD_BOXES)
top_k_neg_mask = tf.cast(
tf.less(
relative_position,
tf.tile(num_neg_boxes[:, tf.newaxis],
(1, ssd_constants.NUM_SSD_BOXES))), tf.float32)
class_loss = tf.reduce_sum(tf.multiply(cross_entropy,
float_mask + top_k_neg_mask),
axis=1)
return tf.reduce_mean(class_loss / num_matched_boxes)
def _compute_bucket_id(bucket_size, header_size, token_type_id, input_mask):
"""Reorder inputs based on token_type index and compute the bucket index.
Args:
bucket_size: <int32> Only attend to position that fall in consecutive
equally sized buckets, or to/from the first bucket.
header_size: <int32> The size of the first bucket.
token_type_id: <int32>[batch_size, seq_length]
input_mask: <int32>[batch_size, seq_length]
Returns:
bucket_id: <int32>[batch_size, seq_lengh]
"""
# Ensure that the padding tokens get sorted last
masked_token_type_id = token_type_id + (-input_mask + 1) * 10000
# Argsorting twice gives you the rank. Sorting once and using scatter could
# be slightly faster, but not relevant for simulation code
token_type_ranks = tf.argsort(tf.argsort(masked_token_type_id,
stable=True),
stable=True)
# Elements from 0 to `header_size` - 1 are mapped to 0, the rest are dividied
# evenly into groups of size `bucket_size` starting from 1.
return tf.math.maximum((token_type_ranks - header_size) // bucket_size + 1,
0)
def diffsort(offsets):
"""Calculate the argsort of the difference between the input offsets.
Useful for sorting row indices in sparse matrices.
Args:
offsets: Tensor, array of offsets for the sparse of each row, where
`offset[i+1] - offsets[i]` is the length of row i. Length `m+1`,
where 'm' is the number of rows.
Returns:
Tensor, array of row indices sorted by row length, from largest to
smallest.
"""
diffs = (offsets - tf.roll(offsets, shift=-1, axis=0))[:-1]
return tf.cast(tf.argsort(diffs, direction="DESCENDING"), tf.uint32)
def body(past, prev, output, addr):
next_outputs = step(hparams, prev, past=past)
logits = next_outputs['logits'][:, -1, :]
samples = tf.gather_nd(
tf.argsort(logits[0, :], direction="DESCENDING"), [
[addr[0]],
])
#print (">>>> next past=" + str (tf.concat([past, next_outputs['presents']], axis=-2)))
#print (">>>> next prev=" + str(samples))
#print (">>>> next output="+str(tf.concat([output, samples], axis=0)))
#print (">>>> next addr="+str(addr[1:]))
return [
tf.concat([past, next_outputs['presents']], axis=-2),
samples,
tf.concat([output, samples], axis=0),
addr[1:],
]
def nucleus_sampling(self, logits):
"""Nucleus sampling."""
p = self.hparams.nucleus_sampling
tf.logging.info("Nucleus sampling top_p = {}".format(p))
sort_indices = tf.argsort(logits, axis=-1, direction="DESCENDING")
probs = tf.gather(tf.nn.softmax(logits), sort_indices, batch_dims=1)
cumprobs = tf.cumsum(probs, axis=-1, exclusive=True)
# The top 1 candidate always will not be masked.
# This way ensures at least 1 indices will be selected.
sort_mask = tf.cast(tf.greater(cumprobs, p), logits.dtype)
batch_indices = tf.tile(
tf.expand_dims(tf.range(logits.shape[0]), axis=-1),
[1, logits.shape[1]])
top_p_mask = tf.scatter_nd(
tf.stack([batch_indices, sort_indices], axis=-1), sort_mask,
logits.shape)
logits -= top_p_mask * logits.dtype.max
return logits
def sort2d(tensor2d,
ref_indices2d,
first_k=float("inf"),
direction="DESCENDING"):
"""Perform sort in 2D based on tensor2d values with indices in ref_indices2d.
Args:
tensor2d: 3D Tensor of size [batch_size, height, width]
ref_indices2d: 4D Tensor of size [batch_size, height, width, 2]
with reference coordinates in 2D.
first_k: (Integer) return indices of first_k elements.
direction: "ASCENDING" or "DESCENDING".
Returns:
sorted_tensor: 2D Tensor of size (first_k, batch_size)
sorted_ref_indices: 3D Tensor of size (first_k, batch_size, 2) with
ref_indices sorted order based on the order of the input (tensor2d).
"""
_, height, width = tensor2d.shape.as_list()
batch_size = tf.shape(tensor2d)[0]
first_k = min(first_k, height*width)
tensor2d_reshaped = tf.reshape(
tensor2d, [batch_size, height*width])
ref_indices2d_reshaped = tf.reshape(
ref_indices2d, [batch_size, height*width, 2])
sort_indices = tf.argsort(
tensor2d_reshaped, axis=1, direction=direction)
sort_indices = tf.gather(sort_indices,
tf.cast(np.array(range(first_k)), tf.int32), axis=1)
sorted_ref_indices = tf.reduce_sum(
tf.eye(batch_size, batch_size)[:, :, tf.newaxis, tf.newaxis] * tf.gather(
ref_indices2d_reshaped, sort_indices, axis=1),
axis=1)
sorted_tensor = tf.reduce_sum(
tf.eye(batch_size, batch_size)[:, :, tf.newaxis] * tf.gather(
tensor2d_reshaped, sort_indices, axis=1),
axis=1)
sorted_tensor = tf.transpose(sorted_tensor, [1, 0])
sorted_ref_indices = tf.transpose(sorted_ref_indices, [1, 0, 2])
return sorted_tensor, sorted_ref_indices
def idx2csr(indices, rows, columns):
"""Convert index format meta-data to compressed sparse row format.
Args:
indices: Tensor, the linear indices for the sparse matrix.
rows: Tensor, the number of rows in the sparse matrix.
columns: Tensor, the number of columns in the sparse matrix.
Returns:
row_indices: Tensor, the row indices sorted by size.
row_offset: Tensor, the offsets for each row in the sparse matrix.
column_indices: Tensor, the column indices for each nonzero in the
matrix.
"""
# Calculate the length of each row by histogramming the indices.
row_lengths = tf.histogram_fixed_width(indices, [0, rows * columns],
nbins=rows)
row_offsets = tf.concat([[0], tf.cumsum(row_lengths)], axis=0)
row_indices = tf.argsort(row_lengths, direction="DESCENDING")
column_indices = tf.mod(indices, columns)
return row_indices, row_offsets, column_indices
def d2s(self, matrix, weight):
"""dense to sparse operation"""
# Extract the nonzero values.
mask = tf.math.not_equal(matrix, 0)
values = tf.boolean_mask(weight, mask)
# Calculate offset of each row
mask = tf.cast(mask, tf.int32)
row_offsets = tf.concat(
[[0], tf.cumsum(tf.reduce_sum(mask, axis=1))], axis=0)
# Create the row indices and sort them.
row_indices = tf.argsort(-(row_offsets[1:] - row_offsets[:-1]))
# Extract the column indices for the nonzero values.
x = mask * (tf.range(matrix.shape[1]) + 1)
column_indices = tf.boolean_mask(x, tf.math.not_equal(matrix, 0))
column_indices = column_indices - 1
row_indices = tf.cast(row_indices, tf.uint32)
row_offsets = tf.cast(row_offsets, tf.uint32)
column_indices = tf.cast(column_indices, tf.uint32)
return self._rows, self._columns, values, row_indices, row_offsets, column_indices
def _face_model_map_fn(example):
vertices = example['vertices']
# Randomly shift vertices
if apply_random_shift:
vertices = random_shift(vertices)
example['num_vertices'] = tf.shape(vertices)[0]
# Optionally shuffle vertices and re-order faces to match
if shuffle_vertices:
permutation = tf.random_shuffle(tf.range(example['num_vertices']))
vertices = tf.gather(vertices, permutation)
face_permutation = tf.concat([
tf.constant([0, 1], dtype=tf.int32),
tf.argsort(permutation) + 2
],
axis=0)
example['faces'] = tf.cast(
tf.gather(face_permutation, example['faces']), tf.int64)
def _dequantize_verts(verts, n_bits):
min_range = -0.5
max_range = 0.5
range_quantize = 2**n_bits - 1
verts = tf.cast(verts, tf.float32)
verts = verts * (max_range -
min_range) / range_quantize + min_range
return verts
# Vertices are quantized. So convert to floats for input to face model
example['vertices'] = _dequantize_verts(vertices, quantization_bits)
example['vertices_mask'] = tf.ones_like(example['vertices'][..., 0],
dtype=tf.float32)
example['faces_mask'] = tf.ones_like(example['faces'],
dtype=tf.float32)
return example
def random_crop_to_size(protein,
crop_size,
max_templates,
shape_schema,
subsample_templates=False):
"""Crop randomly to `crop_size`, or keep as is if shorter than that."""
seq_length = protein['seq_length']
if 'template_mask' in protein:
num_templates = tf.cast(
shape_helpers.shape_list(protein['template_mask'])[0], tf.int32)
else:
num_templates = tf.constant(0, dtype=tf.int32)
num_res_crop_size = tf.math.minimum(seq_length, crop_size)
# Ensures that the cropping of residues and templates happens in the same way
# across ensembling iterations.
# Do not use for randomness that should vary in ensembling.
seed_maker = utils.SeedMaker(
initial_seed=protein['random_crop_to_size_seed'])
if subsample_templates:
templates_crop_start = tf.random.stateless_uniform(
shape=(),
minval=0,
maxval=num_templates + 1,
dtype=tf.int32,
seed=seed_maker())
else:
templates_crop_start = 0
num_templates_crop_size = tf.math.minimum(
num_templates - templates_crop_start, max_templates)
num_res_crop_start = tf.random.stateless_uniform(shape=(),
minval=0,
maxval=seq_length -
num_res_crop_size + 1,
dtype=tf.int32,
seed=seed_maker())
templates_select_indices = tf.argsort(
tf.random.stateless_uniform([num_templates], seed=seed_maker()))
for k, v in protein.items():
if k not in shape_schema or ('template' not in k
and NUM_RES not in shape_schema[k]):
continue
# randomly permute the templates before cropping them.
if k.startswith('template') and subsample_templates:
v = tf.gather(v, templates_select_indices)
crop_sizes = []
crop_starts = []
for i, (dim_size, dim) in enumerate(
zip(shape_schema[k], shape_helpers.shape_list(v))):
is_num_res = (dim_size == NUM_RES)
if i == 0 and k.startswith('template'):
crop_size = num_templates_crop_size
crop_start = templates_crop_start
else:
crop_start = num_res_crop_start if is_num_res else 0
crop_size = (num_res_crop_size if is_num_res else
(-1 if dim is None else dim))
crop_sizes.append(crop_size)
crop_starts.append(crop_start)
protein[k] = tf.slice(v, crop_starts, crop_sizes)
protein['seq_length'] = num_res_crop_size
return protein
def body_cross_entropy(itr, cond_terminate, sample_mean,
sample_covariance_diag, top_k_value,
top_k_action_samples):
"""Function for cross entropy search of actions."""
del top_k_action_samples
top_k_value_prev = top_k_value
batch_sample_mean = tf.reshape(
tf.tile(sample_mean, [1, num_samples]),
[self._dynamic_batch_size * num_samples, self.action_dim])
batch_sample_covariance_diag = tf.reshape(
tf.tile(sample_covariance_diag, [1, num_samples]),
[self._dynamic_batch_size * num_samples, self.action_dim])
action_samples = self._action_projection(
batch_sample_mean +
batch_sample_covariance_diag * tf.cast(random_sampler.sample(
sample_shape=[self._dynamic_batch_size * num_samples]),
dtype=tf.float32))
state_samples = tf.reshape(
tf.tile(self._state_tensor, [1, num_samples]),
[self._dynamic_batch_size * num_samples, self.state_dim])
action_samples = tf.reshape(
action_samples,
[self._dynamic_batch_size * num_samples, self.action_dim])
values = tf.reshape(
self._build_q_function_net(state_samples, action_samples),
[self._dynamic_batch_size, num_samples])
# everything is in batch mode
top_k_index = tf.argsort(values, axis=1,
direction="DESCENDING")[:, 0:top_k_num]
top_k_index_1d = tf.reshape(
top_k_index, [self._dynamic_batch_size * top_k_num, 1])
counter_tensor_1d = tf.reshape(
tf.tile(
tf.reshape(tf.range(self._dynamic_batch_size),
[self._dynamic_batch_size, 1]), [1, top_k_num]),
[self._dynamic_batch_size * top_k_num, 1])
top_k_index_2d = tf.concat([counter_tensor_1d, top_k_index_1d],
axis=1)
action_samples = tf.reshape(
action_samples,
[self._dynamic_batch_size, num_samples, self.action_dim])
top_k_action_samples = tf.gather_nd(action_samples, top_k_index_2d)
top_k_action_samples = tf.reshape(
top_k_action_samples,
[self._dynamic_batch_size, top_k_num, self.action_dim])
top_k_values = tf.gather_nd(values, top_k_index_2d)
top_k_values = tf.reshape(top_k_values,
[self._dynamic_batch_size, top_k_num])
# it's a batch_size x 1 tensor
top_k_value = tf.reshape(tf.reduce_mean(top_k_values, axis=1),
[self._dynamic_batch_size, 1])
sample_mean = tf.reduce_mean(top_k_action_samples, axis=1)
sample_covariance_diag = tf.math.reduce_variance(
top_k_action_samples, axis=1)
itr = itr + 1
cond_terminate = tf.less_equal(
tf.reduce_mean(tf.math.abs(top_k_value - top_k_value_prev)),
self._tolerance_tensor)
return itr, cond_terminate, sample_mean, sample_covariance_diag, \
top_k_value, top_k_action_samples
def create_cm_sketch(topk_obj_ids, topk_obj_weights, all_entity_sketches,
cm_width):
"""Create cm sketches for a set of weighted entities.
Args:
topk_obj_ids: batch_size, topk
topk_obj_weights: batch_size, topk
all_entity_sketches: num_entities, depth
cm_width: width of count-min sketch
Returns:
k hot dense vectors: batch_size, depth, width
"""
topk_fact_obj_sketches = tf.gather(all_entity_sketches,
topk_obj_ids,
axis=0)
# batch_size, topk, depth
batch_size = tf.shape(topk_fact_obj_sketches)[0]
topk = tf.shape(topk_fact_obj_sketches)[1]
cm_depth = tf.shape(topk_fact_obj_sketches)[2]
# We first create a sparse matrix from the hash values. We will then
# convert it into dense matrix. This is more efficient than creating
# k one-hot vectors and then aggregating them into one k-hot vector.
# First prepare ids of non-zero values in the sparse matrix
flattened_topk_hash_ids = tf.reshape(topk_fact_obj_sketches, shape=[-1])
# batch_size * topk * depth
topk_obj_weights = tf.tile(tf.expand_dims(topk_obj_weights, axis=2),
multiples=[1, 1, cm_depth])
# batch_size, topk, depth
flattened_topk_obj_weights = tf.reshape(topk_obj_weights, shape=[-1])
# batch_size * topk * depth
batch_ids = tf.range(batch_size)
# batch_size,
batch_ids = tf.expand_dims(tf.expand_dims(batch_ids, axis=1), axis=2)
# batch_size, 1, 1
batch_ids = tf.tile(batch_ids, multiples=[1, topk, cm_depth])
# batch_size, topk, depth
flattened_batch_ids = tf.reshape(batch_ids, shape=[-1])
# batch_size * topk * depth
depth_ids = tf.range(cm_depth)
# depth,
depth_ids = tf.expand_dims(tf.expand_dims(depth_ids, axis=0), axis=1)
# 1, 1, depth
depth_ids = tf.tile(depth_ids, multiples=[batch_size, topk, 1])
# batch_size, topk, depth
flattened_depth_ids = tf.reshape(depth_ids, shape=[-1])
# batch_size * topk * depth
sparse_value_ids = tf.cast(tf.stack(
[flattened_batch_ids, flattened_depth_ids, flattened_topk_hash_ids],
axis=1),
dtype=tf.int64)
# Then prepare values of non-zero values in the sparse matrix. Values
# are sorted to ascending order. If there are duplicates, later (larger)
# values will be kept.
sorted_orders = tf.argsort(flattened_topk_obj_weights,
direction='ASCENDING',
stable=True)
# batch_size * topk * depth
sorted_flattened_topk_obj_weights = tf.gather(flattened_topk_obj_weights,
sorted_orders)
sorted_sparse_value_ids = tf.gather(sparse_value_ids, sorted_orders)
# Finally create sketch in sparse tensors and convert it to dense tensors.
# We donot validate indices here. If multiple values are about to be assigned
# to the same row and column, we will keep the last value, because the last
# value is the larger one. This behaviour is by design.
sparse_k_hot_sketch = tf.SparseTensor(
indices=sorted_sparse_value_ids,
values=sorted_flattened_topk_obj_weights,
dense_shape=[batch_size, cm_depth, cm_width])
dense_k_hot_sketch = tf.sparse.to_dense(sparse_k_hot_sketch,
validate_indices=False)
# batch_size, cm_depth, cm_width
return dense_k_hot_sketch
def _generate_detections_tf(cls_outputs, box_outputs, anchor_boxes, indices,
classes, image_id, image_scale, num_classes,
min_score_thresh=0.2, max_boxes_to_draw=50,
soft_nms_sigma=0.0, iou_threshold=0.5,
use_native_nms=False):
"""Generates detections with model outputs and anchors.
Args:
cls_outputs: a numpy array with shape [N, 1], which has the highest class
scores on all feature levels. The N is the number of selected
top-K total anchors on all levels. (k being MAX_DETECTION_POINTS)
box_outputs: a numpy array with shape [N, 4], which stacks box regression
outputs on all feature levels. The N is the number of selected top-k
total anchors on all levels. (k being MAX_DETECTION_POINTS)
anchor_boxes: a numpy array with shape [N, 4], which stacks anchors on all
feature levels. The N is the number of selected top-k total anchors on
all levels.
indices: a numpy array with shape [N], which is the indices from top-k
selection.
classes: a numpy array with shape [N], which represents the class
prediction on all selected anchors from top-k selection.
image_id: an integer number to specify the image id.
image_scale: a float tensor representing the scale between original image
and input image for the detector. It is used to rescale detections for
evaluating with the original groundtruth annotations.
num_classes: a integer that indicates the number of classes.
min_score_thresh: A float representing the threshold for deciding when to
remove boxes based on score.
max_boxes_to_draw: Max number of boxes to draw.
soft_nms_sigma: A scalar float representing the Soft NMS sigma parameter;
See Bodla et al, https://arxiv.org/abs/1704.04503). When
`soft_nms_sigma=0.0` (which is default), we fall back to standard (hard)
NMS.
iou_threshold: A float representing the threshold for deciding whether boxes
overlap too much with respect to IOU.
use_native_nms: a bool that indicates whether to use native nms.
Returns:
detections: detection results in a tensor with each row representing
[image_id, y, x, height, width, score, class]
"""
anchor_boxes = tf.gather(anchor_boxes, indices)
scores = tf.math.sigmoid(cls_outputs)
# apply bounding box regression to anchors
boxes = decode_box_outputs_tf(
tf.transpose(box_outputs, [1, 0]), tf.transpose(anchor_boxes, [1, 0]))
def _else(detections, class_id, indices):
"""Else branch for generating detections."""
boxes_cls = tf.gather(boxes, indices)
scores_cls = tf.gather(scores, indices)
# Select top-scoring boxes in each class and apply non-maximum suppression
# (nms) for boxes in the same class. The selected boxes from each class are
# then concatenated for the final detection outputs.
if use_native_nms:
top_detection_idx, scores_cls = tf.image.non_max_suppression_with_scores(
boxes_cls,
scores_cls,
max_boxes_to_draw,
iou_threshold=iou_threshold,
score_threshold=min_score_thresh,
soft_nms_sigma=soft_nms_sigma)
scores_cls = tf.expand_dims(scores_cls, axis=1)
boxes_cls = tf.gather(boxes_cls, top_detection_idx)
top_detections_cls = tf.concat([boxes_cls, scores_cls], axis=1)
else:
scores_cls = tf.expand_dims(scores_cls, axis=1)
all_detections_cls = tf.concat([boxes_cls, scores_cls], axis=1)
top_detection_idx = nms_tf(all_detections_cls, iou_threshold)
top_detections_cls = tf.gather(all_detections_cls, top_detection_idx)
height = top_detections_cls[:, 2] - top_detections_cls[:, 0]
width = top_detections_cls[:, 3] - top_detections_cls[:, 1]
top_detections_cls = tf.stack([top_detections_cls[:, 0] * image_scale,
top_detections_cls[:, 1] * image_scale,
height * image_scale, width * image_scale,
top_detections_cls[:, 4]], axis=-1)
top_detections_cls = tf.stack(
[
tf.cast(
tf.repeat(image_id, tf.size(top_detection_idx)), tf.float32),
*tf.unstack(top_detections_cls, 5, axis=1),
tf.repeat(class_id + 1.0, tf.size(top_detection_idx))
],
axis=1)
detections = tf.concat([detections, top_detections_cls], axis=0)
return detections
detections = tf.constant([], tf.float32, [0, 7])
for c in range(num_classes):
indices_cls = tf.squeeze(tf.where_v2(tf.equal(classes, c)), axis=-1)
detections = tf.cond(
tf.equal(tf.size(indices), 0), lambda: detections,
lambda id=c, id_cls=indices_cls: _else(detections, id, id_cls))
indices_final = tf.argsort(detections[:, -2], direction='DESCENDING')
detections = tf.gather(
detections, indices_final[:max_boxes_to_draw], name='detection')
return detections
def build(self,
input_list,
noisy_params=None,
noise_rate=0.05,
is_training=False,
**kwargs):
"""Create embedding RNN sequence-to-sequence model. No support for noisy parameters.
Args:
input_list: (list<tf.Tensor>) A list of tensors containing the features
for a list of documents.
is_training: (bool) A flag indicating whether the model is running in training mode.
Returns:
A list of tf.Tensor containing the ranking scores for each instance in input_list.
"""
feed_previous = self.feed_previous
embed_size = input_list[0].get_shape()[-1].value
dtype = dtypes.float32
output_projection = None
list_size = len(input_list) # len_seq
with tf.variable_scope("cell", reuse=tf.AUTO_REUSE):
# single_cell = tf.contrib.rnn.GRUCell(
single_cell = tf.keras.layers.GRUCell(embed_size +
self.expand_embed_size)
# double_cell = tf.contrib.rnn.GRUCell(
double_cell = tf.keras.layers.GRUCell(
(embed_size + self.expand_embed_size) * 2)
if self.hparams.use_lstm:
# single_cell = tf.contrib.rnn.BasicLSTMCell(
single_cell = tf.keras.layers.LSTMCell(
(embed_size + self.expand_embed_size))
# double_cell = tf.contrib.rnn.BasicLSTMCell(
double_cell = tf.keras.layers.LSTMCell(
(embed_size + self.expand_embed_size) * 2)
cell = single_cell
self.double_cell = double_cell
self.cell = cell
self.output_projection = output_projection
if self.hparams.num_layers > 1:
# cell = tf.contrib.rnn.MultiRNNCell(
cell = tf.keras.layers.StackedRNNCells([single_cell] *
self.hparams.num_layers)
# self.double_cell = tf.contrib.rnn.MultiRNNCell(
self.double_cell = tf.keras.layers.StackedRNNCells(
[double_cell] * self.hparams.num_layers)
with tf.variable_scope(tf.get_variable_scope()
or "embedding_rnn_seq2seq",
reuse=tf.AUTO_REUSE):
def abstract(input_data, index):
reuse = None if index < 1 else True
# print(reuse, "reuse or not", tf.AUTO_REUSE, "tf.AUTO_REUSE")
with tf.variable_scope(tf.get_variable_scope(),
reuse=tf.AUTO_REUSE):
output_data = input_data
output_sizes = [
int((embed_size + self.expand_embed_size) / 2),
self.expand_embed_size
]
current_size = embed_size
for i in xrange(2):
expand_W = self.get_variable(
"expand_W_%d" % i, [current_size, output_sizes[i]])
expand_b = self.get_variable("expand_b_%d" % i,
[output_sizes[i]])
output_data = tf.nn.bias_add(
tf.matmul(output_data, expand_W), expand_b)
output_data = tf.nn.elu(output_data)
current_size = output_sizes[i]
return output_data
for i in xrange(list_size):
input_list[i] = self.Layer_embedding(input_list[i])
if self.expand_embed_size > 0:
input_list[i] = tf.concat(
axis=1,
values=[
input_list[i],
abstract(input_list[i], i)
]) # [batch,feature_size+expand_embed_size]*len_seq
# input_list= [tf.reshape(e, [1, -1,self.expand_embed_size+embed_size])
# for e in input_list]
# input_list = tf.concat(axis=0, values=input_list)###[len_seq,batch,feature_size+expand_embed_size]
# [len_seq,batch,feature_size+expand_embed_size]
input_list = tf.stack(input_list, axis=0)
enc_cell = copy.deepcopy(cell)
ind = tf.range(0, list_size)
print(self.hparams.input_sequence)
if self.hparams.input_sequence == "reverse":
ind = ind
elif self.hparams.input_sequence == "initial":
ind = tf.range(list_size - 1, -1, -1)
elif self.hparams.input_sequence == "random":
ind = tf.random.shuffle(ind)
# [len_seq,batch,feature_size+expand_embed_size]
input_list_input = tf.nn.embedding_lookup(input_list, ind)
# [batch,feature_size+expand_embed_size]*len_seq
input_list_input_list = tf.unstack(input_list_input, axis=0)
# cprint('input_list_input_list: {}'.format(input_list_input_list), 'green')
'''
[<tf.Tensor 'ranking_model_1/ranking_model/unstack:0' shape=(?, 186) dtype=float32>,
<tf.Tensor 'ranking_model_1/ranking_model/unstack:1' shape=(?, 186) dtype=float32>,
<tf.Tensor 'ranking_model_1/ranking_model/unstack:2' shape=(?, 186) dtype=float32>,
<tf.Tensor 'ranking_model_1/ranking_model/unstack:3' shape=(?, 186) dtype=float32>,
<tf.Tensor 'ranking_model_1/ranking_model/unstack:4' shape=(?, 186) dtype=float32>,
<tf.Tensor 'ranking_model_1/ranking_model/unstack:5' shape=(?, 186) dtype=float32>,
<tf.Tensor 'ranking_model_1/ranking_model/unstack:6' shape=(?, 186) dtype=float32>,
<tf.Tensor 'ranking_model_1/ranking_model/unstack:7' shape=(?, 186) dtype=float32>,
<tf.Tensor 'ranking_model_1/ranking_model/unstack:8' shape=(?, 186) dtype=float32>]
'''
encoder_outputs_some_order, encoder_state = tf.nn.static_rnn( # encoder_state就是final state
enc_cell,
input_list_input_list,
dtype=dtype)
# _rnn = tf.keras.layers.RNN(enc_cell, unroll=True)
# _rnn = tf.keras.layers.RNN(enc_cell, unroll=True, return_sequences=True)
# _output = _rnn(input_list_input_list)
# cprint('_output: {}'.format(_output), 'green')
# cprint('encoder_outputs_some_order: {}'.format(encoder_outputs_some_order), 'green')
'''
下面是运行tf.nn.static_rnn版本的代码输出的结果：
[<tf.Tensor 'ranking_model/ranking_model/rnn/gru_cell/add_5:0' shape=(?, 186) dtype=float32>,
<tf.Tensor 'ranking_model/ranking_model/rnn/gru_cell_1/add_5:0' shape=(?, 186) dtype=float32>,
<tf.Tensor 'ranking_model/ranking_model/rnn/gru_cell_2/add_5:0' shape=(?, 186) dtype=float32>,
<tf.Tensor 'ranking_model/ranking_model/rnn/gru_cell_3/add_5:0' shape=(?, 186) dtype=float32>,
<tf.Tensor 'ranking_model/ranking_model/rnn/gru_cell_4/add_5:0' shape=(?, 186) dtype=float32>,
<tf.Tensor 'ranking_model/ranking_model/rnn/gru_cell_5/add_5:0' shape=(?, 186) dtype=float32>,
<tf.Tensor 'ranking_model/ranking_model/rnn/gru_cell_6/add_5:0' shape=(?, 186) dtype=float32>,
<tf.Tensor 'ranking_model/ranking_model/rnn/gru_cell_7/add_5:0' shape=(?, 186) dtype=float32>,
<tf.Tensor 'ranking_model/ranking_model/rnn/gru_cell_8/add_5:0' shape=(?, 186) dtype=float32>]
'''
# cprint('encoder_state: {}'.format(encoder_state), 'green')
'''
下面是运行tf.nn.static_rnn版本的代码输出的结果：
Tensor("ranking_model/ranking_model/rnn/gru_cell_8/add_5:0", shape=(?, 186), dtype=float32)
'''
ind_sort = tf.argsort(ind) # find the order of sequence
# ind_sort=tf.Print(tf.argsort(ind),[tf.argsort(ind),ind],"sequence")
self.ind_sort = [ind_sort, ind]
# [len_seq,batch,feature_size+expand_embed_size]
encoder_outputs_some_order = tf.stack(encoder_outputs_some_order,
axis=0)
# encoder_outputs=[None]*list_size
# for i in range(list_size):
# encoder_outputs[ind[i]]=encoder_outputs_some_order[i]## back to the order of initial list.
# [len_seq,batch,feature_size+expand_embed_size]
input_list_output = tf.nn.embedding_lookup(
encoder_outputs_some_order, ind_sort)
# [feature_size+expand_embed_size] *len_seq
input_list_output_list = tf.unstack(input_list_output, axis=0)
top_states = [
tf.reshape(self.layer_norm_hidden(e),
[-1, 1, cell.output_size])
for e in input_list_output_list
] # [batch,1,encoder_out]*len_seq
encoder_state = self.layer_norm_final(
encoder_state) # [batch,encoder_state]
# top_states = [tf.reshape(e, [-1, 1, cell.output_size])
# for e in encoder_outputs] ##[]
# encoder_state=encoder_state
# for e in input_list]
# [batch,len_seq,encoder_out]
attention_states = tf.concat(axis=1, values=top_states)
# print(attention_states.get_shape(),"attention_states.get_shape()")
if isinstance(feed_previous, bool):
outputs, state = self.embedding_rnn_decoder(
encoder_state,
cell,
attention_states,
input_list,
num_heads=self.hparams.num_heads,
output_projection=output_projection,
feed_previous=feed_previous)
print(outputs[0].get_shape(), "outputs[0].get_shape()")
return outputs[0]
# If feed_previous is a Tensor, we construct 2 graphs and use cond.
def decoder(feed_previous_bool):
reuse = None if feed_previous_bool else True
with tf.variable_scope(tf.get_variable_scope(), reuse=reuse):
outputs, state = self.embedding_rnn_decoder(
encoder_state,
cell,
attention_states,
input_list,
num_heads=self.hparams.num_heads,
output_projection=output_projection,
feed_previous=feed_previous_bool,
update_embedding_for_previous=False)
return outputs + [state]
outputs_and_state = control_flow_ops.cond(feed_previous,
lambda: decoder(True),
lambda: decoder(False))
print(outputs[0].get_shape(), "outputs[0].get_shape()")
return outputs_and_state[0]
def __init__(self, item_num, args, reuse=None):
self.args = args
self.is_training = tf.placeholder(tf.bool, shape=())
self.input_seq = tf.placeholder(tf.int32, shape=(None, args.maxlen))
self.pos = tf.placeholder(tf.int32, shape=None)
self.exemplar_logits = tf.placeholder(tf.float32, shape=(None, None))
self.exemplar_pos = tf.placeholder(tf.int32, shape=None)
self.max_item = tf.placeholder(tf.int32, shape=())
self.lr = tf.placeholder(tf.float32, shape=())
self.dropout_rate = tf.placeholder(tf.float32, shape=())
pos = self.pos
mask = tf.expand_dims(tf.to_float(tf.not_equal(self.input_seq, 0)), -1)
with tf.variable_scope("SASRec", reuse=reuse):
# sequence embedding, item embedding table
self.seq, item_emb_table = embedding(self.input_seq,
vocab_size=item_num + 1,
num_units=args.hidden_units,
zero_pad=True,
scale=True,
l2_reg=args.l2_emb,
scope="input_embeddings",
with_t=True,
reuse=reuse
)
# # Positional Encoding
t, pos_emb_table = embedding(
tf.tile(tf.expand_dims(tf.range(tf.shape(self.input_seq)[1]), 0), [tf.shape(self.input_seq)[0], 1]),
vocab_size=args.maxlen,
num_units=args.hidden_units,
zero_pad=False,
scale=False,
l2_reg=args.l2_emb,
scope="dec_pos",
reuse=reuse,
with_t=True
)
self.seq += t
# Dropout
self.seq = tf.layers.dropout(self.seq,
rate=self.dropout_rate,
training=tf.convert_to_tensor(self.is_training),
seed=args.random_seed)
self.seq *= mask
# Build blocks
for i in range(args.num_blocks):
with tf.variable_scope("num_blocks_%d" % i):
# Self-attention
self.seq = multihead_attention(queries=normalize(self.seq),
keys=self.seq,
num_units=args.hidden_units,
num_heads=args.num_heads,
dropout_rate=self.dropout_rate,
seed=args.random_seed,
is_training=self.is_training,
causality=True,
scope="self_attention")
# Feed forward
self.seq = feedforward(normalize(self.seq), num_units=[args.hidden_units, args.hidden_units],
dropout_rate=self.dropout_rate, is_training=self.is_training,
seed=args.random_seed)
self.seq *= mask
self.seq = normalize(self.seq)
# find representation
self.rep = self.seq[:, -1, :]
# define loss
seq_emb = tf.reshape(self.rep, [tf.shape(self.input_seq)[0], args.hidden_units])
indices = pos - 1
self.labels = tf.one_hot(indices, self.max_item)
item_emb = tf.nn.embedding_lookup(item_emb_table, tf.range(1, self.max_item + 1))
self.logits = tf.matmul(seq_emb, tf.transpose(item_emb))
self.loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=self.labels, logits=self.logits))
self.global_step = tf.Variable(0, name='global_step', trainable=False)
self.optimizer = tf.train.AdamOptimizer(learning_rate=self.lr)
# prediction
self.test_item = tf.placeholder(tf.int32, shape=None)
self.test_item_emb = tf.nn.embedding_lookup(item_emb_table, self.test_item)
self.test_logits = tf.matmul(seq_emb, tf.transpose(self.test_item_emb))
self.test_logits = tf.reshape(self.test_logits, [tf.shape(self.input_seq)[0], tf.shape(self.test_item)[0]])
self.pred_last = tf.argsort(tf.argsort(-self.test_logits))
def interleave_shells(shells, radii): """Interleave spherical shell tensors out-to-in by radii.""" radius_order = tf.argsort(radii, direction='DESCENDING') shells_interleaved = tf.gather(shells, radius_order, axis=3) return shells_interleaved
def create_bucketed_attention_layer(input_mask, input_header, bucket_size,
header_size, token_type_ids,
sort_after_projection):
"""Returns a drop-in replacement for attention_layer using sparsse attention.
Args:
input_mask: int32<batch_size, seq_length> The values should be 1 or 0. The
attention scores will effectively be set to -infinity for any positions in
the mask that are 0, and will be unchanged for positions that are 1.
input_header: bool<batch_size, seq_length> The values should be 1 or 0.
Attention will not be restricted to or from tokens where header is 1.
bucket_size: int. Size of sections where self attention happens.
header_size: Size of the first bucket that will attend to/from everything.
If None is passed will use the same as `bucket_size`.
token_type_ids: List<(int, bool, <int32>[batch_size, seq_length])> contains
the number of heads for each token type, whether they are sorted, and the
ids of each position. Attention is restricted between tokens with the same
type id and this field is used to sort/bucket. ids must be non-negative.
sort_after_projection: Sorting can happen on the layer input or after
applying the projection to keys, queries and values. Depending on the
accelerator, one option could be more convenient.
Returns:
Function with same signature as `attention_layer`.
See `_bucketed_attention_layer`.
"""
type_counts = [cnt for cnt, _, _ in token_type_ids]
num_heads = sum(type_counts)
# Ensure that the padding tokens get sorted last.
additive_mask = _additive_mask(input_mask)
type_sorted = [is_sorted for _, is_sorted, _ in token_type_ids]
type_ids_masked = [ids - additive_mask for _, _, ids in token_type_ids]
type_ids_argsorted = [
None if is_sorted else tf.argsort(ids, stable=True)
for is_sorted, ids in zip(type_sorted, type_ids_masked)
]
def _gather_and_repeat(tensors, name):
"""Reorder each tensor using its type_id and copies it once per head."""
if len(tensors) == 1:
tensors = tensors * len(type_ids_argsorted)
with tf.variable_scope(name):
gathered_tensors = []
for tensor, indices in zip(tensors, type_ids_argsorted):
gathered_tensors.append(tensor if indices is None else tf.
gather(tensor, indices, batch_dims=1))
return tf.repeat(tf.stack(gathered_tensors, axis=1),
repeats=type_counts,
axis=1)
# <int32>[batch_size, num_types, seq_length]
type_ids = tf.stack(type_ids_masked, axis=1)
type_order = tf.argsort(type_ids, stable=True)
type_rank = tf.argsort(type_order, stable=True)
# For each head the correct order to sort the embeddings.
# <int32>[batch_size, num_attention_heads, seq_length].
type_order_repeated = tf.repeat(type_order, repeats=type_counts, axis=1)
# For each head the inverse of the correct order to sort the embeddings.
# <int32>[batch_size, num_attention_heads, seq_length].
type_rank_repeated = tf.repeat(type_rank, repeats=type_counts, axis=1)
input_header = input_header & tf.math.equal(input_mask, 1)
if sort_after_projection:
# No need to sort in this case since this will happen in the BucketedTensor.
type_ids_repeated = tf.repeat(type_ids, repeats=type_counts, axis=1)
attend_always_repeated = tf.repeat(tf.expand_dims(input_header,
axis=1),
repeats=num_heads,
axis=1)
else:
type_ids_repeated = _gather_and_repeat(type_ids_masked,
'type_id_gather')
attend_always_repeated = _gather_and_repeat([input_header],
'attend_always_gather')
if header_size is None:
header_size = bucket_size
args = {
'order_indices': type_order_repeated,
'bucket_size': bucket_size,
'header_size': header_size,
'num_heads': num_heads,
'sort_after_projection': sort_after_projection,
}
ids = _create_bucketed_tensor(type_ids_repeated, 'type_id', **args)
attend_always = _create_bucketed_tensor(attend_always_repeated,
'attend_always', **args)
head_full_mask = _compute_bucketed_attention_mask('head', 'full', ids,
attend_always)
tail_head_mask = _compute_bucketed_attention_mask('tail', 'head', ids,
attend_always)
tail_window_mask = _compute_bucketed_attention_mask(
'tail', 'window', ids, attend_always)
return functools.partial(_bucketed_attention_layer,
sort_after_projection=sort_after_projection,
gather_and_repeat=_gather_and_repeat,
bucketed_tensor_args=args,
header_size=header_size,
type_ranks=type_rank_repeated,
head_full_mask=head_full_mask,
tail_head_mask=tail_head_mask,
tail_window_mask=tail_window_mask)
def add_single_image_info(self, image_id, eval_dict):
groundtruth_boxes = eval_dict[
standard_fields.InputDataFields.groundtruth_boxes]
groundtruth_classes = eval_dict[
standard_fields.InputDataFields.groundtruth_classes]
detection_boxes = eval_dict[
standard_fields.DetectionResultFields.detection_boxes]
detection_scores = eval_dict[
standard_fields.DetectionResultFields.detection_scores]
detection_classes = eval_dict[
standard_fields.DetectionResultFields.detection_classes]
groundtruth_has_rotation = groundtruth_classes > 1
groundtruth_boxes_with_rotation = groundtruth_boxes[
groundtruth_has_rotation]
#ensure classes are both not 'dot' class, so they have a meaningful rotation value
detection_within_score = detection_scores > self._score_threshold
detection_class_has_rotation = detection_classes > 1
detection_has_rotation_and_score = tf.logical_and(
detection_within_score, detection_class_has_rotation)
detection_boxes_within_score = detection_boxes[
detection_has_rotation_and_score]
detection_classes_within_score = detection_classes[
detection_has_rotation_and_score]
gt_boxlist = box_list.BoxList(
tf.convert_to_tensor(groundtruth_boxes_with_rotation))
det_boxlist = box_list.BoxList(
tf.convert_to_tensor(detection_boxes_within_score))
detection_y_rotation_angles = eval_dict[
additional_fields.DetectionResultFields.y_rotation_angles]
groundtruth_y_rotation_angles = eval_dict[
additional_fields.GroundtruthResultFields.y_rotation_angles]
detection_y_rotation_angles_within_score = detection_y_rotation_angles[
detection_has_rotation_and_score]
for iou_threshold, assigner in self._iou_thresholds_and_assigners:
cls_targets, cls_weights, reg_targets, reg_weights, match = assigner.assign(
det_boxlist, gt_boxlist)
fg_detections = match >= 0
fg_detection_boxes = detection_boxes_within_score[fg_detections, :]
fg_matches = match[fg_detections]
fg_matches_argsort = tf.argsort(fg_matches)
fg_matches_sorted = tf.gather(fg_matches, fg_matches_argsort)
gt_match_indices, fg_match_sorted_indices_with_repeats, fg_match_sorted_indices_counts = tf.unique_with_counts(
fg_matches_sorted)
fg_match_sorted_indices_no_repeats = tf.cumsum(
tf.pad(fg_match_sorted_indices_counts, [[1, 0]]))[:-1]
fg_match_indices_no_repeats = tf.gather(
fg_matches_argsort, fg_match_sorted_indices_no_repeats)
def get_matches_and_angle_difference(fg_match_idx_tensor,
gt_match_idx_tensor):
if debug_get_matching_boxes:
gt_matching_detection_boxes = tf.gather(
groundtruth_boxes_with_rotation,
gt_match_idx_tensor,
axis=0)
fg_matching_detection_boxes = tf.gather(
fg_detection_boxes, fg_match_idx_tensor, axis=0)
pass
fg_matching_detection_y_rot_angles = tf.gather(
detection_y_rotation_angles_within_score,
fg_match_idx_tensor,
axis=0)
groundtruth_y_rotation_angles_matches = tf.gather(
groundtruth_y_rotation_angles, gt_match_idx_tensor, axis=0)
groundtruth_has_y_rot = tf.math.logical_not(
tf.math.equal(groundtruth_y_rotation_angles_matches, 0))
groundtruth_existant_y_rot_angle = groundtruth_y_rotation_angles_matches[
groundtruth_has_y_rot]
detection_existant_y_rot_angle = fg_matching_detection_y_rot_angles[
groundtruth_has_y_rot]
angle_diff = detection_existant_y_rot_angle - groundtruth_existant_y_rot_angle
angle_diff_unwrapped = tf.math.atan2(tf.math.sin(angle_diff),
tf.math.cos(angle_diff))
angle_diff_abs = tf.math.abs(angle_diff_unwrapped)
n_angle_matches = len(angle_diff)
return n_angle_matches, angle_diff_abs
num_angle_matches, abs_angle_differences = get_matches_and_angle_difference(
fg_match_indices_no_repeats, gt_match_indices)
angle_diff_sum_square = tf.reduce_sum(
tf.math.square(abs_angle_differences * 180 / np.pi))
match_angle_diff_histogram = tf.histogram_fixed_width(
abs_angle_differences * 180 / np.pi,
self._histogram_range,
nbins=self._num_angle_bins,
dtype=tf.dtypes.int32)
self.total_num_angle_matches[iou_threshold] += num_angle_matches
self.total_angle_diff_sum_squared[
iou_threshold] += angle_diff_sum_square
self.angle_histograms[iou_threshold] += match_angle_diff_histogram