def prune_completely_outside_window(boxlist, window, scope=None):
"""Prunes bounding boxes that fall completely outside of the given window.
The function clip_to_window prunes bounding boxes that fall
completely outside the window, but also clips any bounding boxes that
partially overflow. This function does not clip partially overflowing boxes.
Args:
boxlist: a BoxList holding M_in boxes.
window: a float tensor of shape [4] representing [ymin, xmin, ymax, xmax]
of the window
scope: name scope.
Returns:
pruned_boxlist: a new BoxList with all bounding boxes partially or fully in
the window.
valid_indices: a tensor with shape [M_out] indexing the valid bounding boxes
in the input tensor.
"""
with tf.name_scope(scope, 'PruneCompleteleyOutsideWindow'):
y_min, x_min, y_max, x_max = tf.split(value=boxlist.get(),
num_or_size_splits=4,
axis=1)
win_y_min, win_x_min, win_y_max, win_x_max = tf.unstack(window)
coordinate_violations = tf.concat([
tf.greater_equal(y_min, win_y_max),
tf.greater_equal(x_min, win_x_max),
tf.less_equal(y_max, win_y_min),
tf.less_equal(x_max, win_x_min)
], 1)
valid_indices = tf.reshape(
tf.where(tf.logical_not(tf.reduce_any(coordinate_violations, 1))),
[-1])
return gather(boxlist, valid_indices), valid_indices
def _get_endpoint_b(i, j, q1, q0, batch_shape):
"""Determine the end of the interval, `b` as either `q0[j]` or `q1[i]`."""
# Get `i`-th element of `q1` and the `j`-th of `q0` (for batched data).
i_geq_0 = tf.where(tf.greater_equal(i, 0))
q1_i = _get_q_slice(q1, i, i_geq_0, batch_shape=batch_shape)
j_geq_0 = tf.where(tf.greater_equal(j, 0))
q0_j = _get_q_slice(q0, j, j_geq_0, batch_shape=batch_shape)
# Find `b`. If `b==q0[j]`, decrement `j`; if `b==q1[i]`, decrement `i`.
# Note: we could have just said "b = a;" at the end of this loop but then
# we'd have to handle corner cases before the loop and within. Grabbing
# it each time is minimal more work and saves added code complexity.
# if i < 0: i, j, b = i, j-1, q0[j]
i_lt_0 = tf.less(i, 0)
ind = tf.where(i_lt_0)
b = _update_batch(ind, q0_j, batch_shape=batch_shape)
j = _update_batch(ind, j - 1, b=j)
# elif j < 0: i, j, b = i - 1, j, q1[i]
j_lt_0 = tf.less(j, 0)
cond_from_previous = ~i_lt_0
ind = tf.where(j_lt_0 & cond_from_previous)
b = _update_batch(ind, q1_i, b=b)
i = _update_batch(ind, i - 1, b=i)
# elif q1[i] == q0[j]: i, j, b = i - 1, j - 1, q1[i]
q_equal = tf.equal(q1_i, q0_j)
cond_from_previous = cond_from_previous & ~j_lt_0
ind = tf.where(q_equal & cond_from_previous)
b = _update_batch(ind, q1_i, b=b)
i = _update_batch(ind, i - 1, b=i)
j = _update_batch(ind, j - 1, b=j)
# elif q1[i] > q0[j]: i, j, b = i - 1, j, q1[i]
q1_gt_q0 = tf.greater(q1_i, q0_j)
cond_from_previous = cond_from_previous & ~q_equal
ind = tf.where(q1_gt_q0 & cond_from_previous)
b = _update_batch(ind, q1_i, b=b)
i = _update_batch(ind, i - 1, b=i)
# else: i, j, b = i, j - 1, q0[j]
ind = tf.where(cond_from_previous & ~q1_gt_q0)
b = _update_batch(ind, q0_j, b=b)
j = _update_batch(ind, j - 1, b=j)
return b, i, j
def prune_small_boxes(boxlist, min_side, scope=None):
"""Prunes small boxes in the boxlist which have a side smaller than min_side.
Args:
boxlist: BoxList holding N boxes.
min_side: Minimum width AND height of box to survive pruning.
scope: name scope.
Returns:
A pruned boxlist.
"""
with tf.name_scope(scope, 'PruneSmallBoxes'):
height, width = height_width(boxlist)
is_valid = tf.logical_and(tf.greater_equal(width, min_side),
tf.greater_equal(height, min_side))
return gather(boxlist, tf.reshape(tf.where(is_valid), [-1]))
def prune_non_overlapping_boxes(boxlist1,
boxlist2,
min_overlap=0.0,
scope=None):
"""Prunes the boxes in boxlist1 that overlap less than thresh with boxlist2.
For each box in boxlist1, we want its IOA to be more than minoverlap with
at least one of the boxes in boxlist2. If it does not, we remove it.
Args:
boxlist1: BoxList holding N boxes.
boxlist2: BoxList holding M boxes.
min_overlap: Minimum required overlap between boxes, to count them as
overlapping.
scope: name scope.
Returns:
new_boxlist1: A pruned boxlist with size [N', 4].
keep_inds: A tensor with shape [N'] indexing kept bounding boxes in the
first input BoxList `boxlist1`.
"""
with tf.name_scope(scope, 'PruneNonOverlappingBoxes'):
ioa_ = ioa(boxlist2, boxlist1) # [M, N] tensor
ioa_ = tf.reduce_max(ioa_, reduction_indices=[0]) # [N] tensor
keep_bool = tf.greater_equal(ioa_, tf.constant(min_overlap))
keep_inds = tf.squeeze(tf.where(keep_bool), axis=[1])
new_boxlist1 = gather(boxlist1, keep_inds)
return new_boxlist1, keep_inds
def call(self, inputs, count_weights=None):
inputs = utils.ensure_tensor(inputs)
if count_weights is not None:
if self.output_mode != COUNT:
raise ValueError(
"`count_weights` is not used when `output_mode` is not `'count'`. "
"Received `count_weights={}`.".format(count_weights))
count_weights = utils.ensure_tensor(count_weights, self.compute_dtype)
depth = self.num_tokens
if isinstance(inputs, tf.SparseTensor):
max_value = tf.reduce_max(inputs.values)
min_value = tf.reduce_min(inputs.values)
else:
max_value = tf.reduce_max(inputs)
min_value = tf.reduce_min(inputs)
condition = tf.logical_and(
tf.greater(tf.cast(depth, max_value.dtype), max_value),
tf.greater_equal(min_value, tf.cast(0, min_value.dtype)))
assertion = tf.Assert(condition, [
"Input values must be in the range 0 <= values < num_tokens"
" with num_tokens={}".format(depth)
])
with tf.control_dependencies([assertion]):
return utils.encode_categorical_inputs(
inputs,
output_mode=self.output_mode,
depth=depth,
dtype=self.compute_dtype,
sparse=self.sparse,
count_weights=count_weights)
def get_minimal_coverage_box(boxlist, default_box=None, scope=None):
"""Creates a single bounding box which covers all boxes in the boxlist.
Args:
boxlist: A Boxlist.
default_box: A [1, 4] float32 tensor. If no boxes are present in `boxlist`,
this default box will be returned. If None, will use a default box of
[[0., 0., 1., 1.]].
scope: Name scope.
Returns:
A [1, 4] float32 tensor with a bounding box that tightly covers all the
boxes in the box list. If the boxlist does not contain any boxes, the
default box is returned.
"""
with tf.name_scope(scope, 'CreateCoverageBox'):
num_boxes = boxlist.num_boxes()
def coverage_box(bboxes):
y_min, x_min, y_max, x_max = tf.split(value=bboxes,
num_or_size_splits=4,
axis=1)
y_min_coverage = tf.reduce_min(y_min, axis=0)
x_min_coverage = tf.reduce_min(x_min, axis=0)
y_max_coverage = tf.reduce_max(y_max, axis=0)
x_max_coverage = tf.reduce_max(x_max, axis=0)
return tf.stack([
y_min_coverage, x_min_coverage, y_max_coverage, x_max_coverage
],
axis=1)
default_box = default_box or tf.constant([[0., 0., 1., 1.]])
return tf.cond(tf.greater_equal(num_boxes, 1),
true_fn=lambda: coverage_box(boxlist.get()),
false_fn=lambda: default_box)
def call(self, inputs, count_weights=None):
if isinstance(inputs, (list, np.ndarray)):
inputs = tf.convert_to_tensor(inputs)
if inputs.shape.rank == 1:
inputs = tf.compat.v1.expand_dims(inputs, 1)
if count_weights is not None and self.output_mode != COUNT:
raise ValueError("count_weights is not used in "
"`output_mode='multi_hot'`. Please pass a single input.")
out_depth = self.num_tokens
multi_hot_output = (self.output_mode == MULTI_HOT)
if isinstance(inputs, tf.SparseTensor):
max_value = tf.reduce_max(inputs.values)
min_value = tf.reduce_min(inputs.values)
else:
max_value = tf.reduce_max(inputs)
min_value = tf.reduce_min(inputs)
condition = tf.logical_and(
tf.greater(
tf.cast(out_depth, max_value.dtype), max_value),
tf.greater_equal(
min_value, tf.cast(0, min_value.dtype)))
tf.Assert(condition, [
"Input values must be in the range 0 <= values < num_tokens"
" with num_tokens={}".format(out_depth)
])
if self.sparse:
return sparse_bincount(inputs, out_depth, multi_hot_output, count_weights)
else:
return dense_bincount(inputs, out_depth, multi_hot_output, count_weights)
def matched_column_indicator(self):
"""Returns column indices that are matched.
Returns:
column_indices: int32 tensor of shape [K] with column indices.
"""
return tf.greater_equal(self._match_results, 0)
def pad_to_fixed_size(input_tensor, size, constant_values=0):
"""Pads data to a fixed length at the first dimension.
Args:
input_tensor: `Tensor` with any dimension.
size: `int` number for the first dimension of output Tensor.
constant_values: `int` value assigned to the paddings.
Returns:
`Tensor` with the first dimension padded to `size`.
"""
input_shape = input_tensor.get_shape().as_list()
padding_shape = []
# Computes the padding length on the first dimension.
padding_length = size - tf.shape(input=input_tensor)[0]
assert_length = tf.Assert(tf.greater_equal(padding_length, 0),
[padding_length])
with tf.control_dependencies([assert_length]):
padding_shape.append(padding_length)
# Copies shapes of the rest of input shape dimensions.
for i in range(1, len(input_shape)):
padding_shape.append(tf.shape(input=input_tensor)[i])
# Pads input tensor to the fixed first dimension.
paddings = tf.cast(constant_values * tf.ones(padding_shape),
input_tensor.dtype)
padded_tensor = tf.concat([input_tensor, paddings], axis=0)
output_shape = input_shape
output_shape[0] = size
padded_tensor.set_shape(output_shape)
return padded_tensor
def _lower_triangular_mask(shape):
"""Creates a lower-triangular boolean mask over the last 2 dimensions."""
row_index = tf.cumsum(
tf.ones(shape=shape, dtype=tf.int32), axis=-2)
col_index = tf.cumsum(
tf.ones(shape=shape, dtype=tf.int32), axis=-1)
return tf.greater_equal(row_index, col_index)
def _get_values_from_start_and_end(self, input_tensor, num_start_samples,
num_end_samples, total_num_samples):
"""slices num_start_samples and last num_end_samples from input_tensor.
Args:
input_tensor: An int32 tensor of shape [N] to be sliced.
num_start_samples: Number of examples to be sliced from the beginning
of the input tensor.
num_end_samples: Number of examples to be sliced from the end of the
input tensor.
total_num_samples: Sum of is num_start_samples and num_end_samples. This
should be a scalar.
Returns:
A tensor containing the first num_start_samples and last num_end_samples
from input_tensor.
"""
input_length = tf.shape(input=input_tensor)[0]
start_positions = tf.less(tf.range(input_length), num_start_samples)
end_positions = tf.greater_equal(tf.range(input_length),
input_length - num_end_samples)
selected_positions = tf.logical_or(start_positions, end_positions)
selected_positions = tf.cast(selected_positions, tf.float32)
indexed_positions = tf.multiply(tf.cumsum(selected_positions),
selected_positions)
one_hot_selector = tf.one_hot(tf.cast(indexed_positions, tf.int32) - 1,
total_num_samples,
dtype=tf.float32)
return tf.cast(
tf.tensordot(tf.cast(input_tensor, tf.float32),
one_hot_selector,
axes=[0, 0]), tf.int32)
def manual_stepping(global_step, boundaries, rates):
boundaries = [0] + boundaries
num_boundaries = len(boundaries)
rate_index = tf.reduce_max(
tf.where(tf.greater_equal(global_step, boundaries),
list(range(num_boundaries)), [0] * num_boundaries))
return tf.reduce_sum(rates * tf.one_hot(rate_index, depth=num_boundaries))
def is_greater_equal_1(label):
"""Computes whether label is greater or equal to 1.
Args:
label: A `Tensor` or anything that can be converted to a tensor using
`tf.convert_to_tensor`.
Returns:
A `Tensor` that has each input element transformed as `x` to `I(x > 1)`.
"""
return tf.greater_equal(label, 1.0)
def _get_interval_proportion(k, n, q_k, q_kp1, a, b, batch_shape):
"""Calculate the proportion `[a, b)` represents of `[q[i], q[i+1])`."""
k_valid = tf.greater_equal(k, 0) & tf.less(k + 1, n)
q_equal = tf.equal(q_kp1, q_k)
d = tf.where(q_equal, tf.ones((batch_shape), dtype=q_k.dtype),
((b - a) / (q_kp1 - q_k)))
d = tf.where(k_valid, d, tf.zeros((batch_shape), dtype=q_k.dtype))
# Turn proportions in to probabilities. This is where we assume that
# the input vectors are quantiles.
return d / tf.cast(n - 1, d.dtype)
def call(self, inputs, count_weights=None):
if isinstance(inputs, (list, np.ndarray)):
inputs = tf.convert_to_tensor(inputs)
def expand_dims(inputs, axis):
if tf_utils.is_sparse(inputs):
return tf.sparse.expand_dims(inputs, axis)
else:
return tf.compat.v1.expand_dims(inputs, axis)
original_shape = inputs.shape
# In all cases, we should uprank scalar input to a single sample.
if inputs.shape.rank == 0:
inputs = expand_dims(inputs, -1)
# One hot will unprank only if the final output dimension is not already 1.
if self.output_mode == ONE_HOT:
if inputs.shape[-1] != 1:
inputs = expand_dims(inputs, -1)
# TODO(b/190445202): remove output rank restriction.
if inputs.shape.rank > 2:
raise ValueError(
"Received input shape {}, which would result in output rank {}. "
"Currently only outputs up to rank 2 are supported.".format(
original_shape, inputs.shape.rank))
if count_weights is not None and self.output_mode != COUNT:
raise ValueError(
"`count_weights` is not used when `output_mode` is not `'count'`. "
"Received `count_weights={}`.".format(count_weights))
out_depth = self.num_tokens
binary_output = self.output_mode in (MULTI_HOT, ONE_HOT)
if isinstance(inputs, tf.SparseTensor):
max_value = tf.reduce_max(inputs.values)
min_value = tf.reduce_min(inputs.values)
else:
max_value = tf.reduce_max(inputs)
min_value = tf.reduce_min(inputs)
condition = tf.logical_and(
tf.greater(tf.cast(out_depth, max_value.dtype), max_value),
tf.greater_equal(min_value, tf.cast(0, min_value.dtype)))
assertion = tf.Assert(condition, [
"Input values must be in the range 0 <= values < num_tokens"
" with num_tokens={}".format(out_depth)
])
with tf.control_dependencies([assertion]):
if self.sparse:
return sparse_bincount(inputs, out_depth, binary_output,
count_weights)
else:
return dense_bincount(inputs, out_depth, binary_output,
count_weights)
def _match_when_rows_are_non_empty():
"""Performs matching when the rows of similarity matrix are non empty.
Returns:
matches: int32 tensor indicating the row each column matches to.
"""
# Matches for each column
matches = tf.argmax(input=similarity_matrix,
axis=0,
output_type=tf.int32)
# Deal with matched and unmatched threshold
if self._matched_threshold is not None:
# Get logical indices of ignored and unmatched columns as tf.int64
matched_vals = tf.reduce_max(input_tensor=similarity_matrix,
axis=0)
below_unmatched_threshold = tf.greater(
self._unmatched_threshold, matched_vals)
between_thresholds = tf.logical_and(
tf.greater_equal(matched_vals, self._unmatched_threshold),
tf.greater(self._matched_threshold, matched_vals))
if self._negatives_lower_than_unmatched:
matches = self._set_values_using_indicator(
matches, below_unmatched_threshold, -1)
matches = self._set_values_using_indicator(
matches, between_thresholds, -2)
else:
matches = self._set_values_using_indicator(
matches, below_unmatched_threshold, -2)
matches = self._set_values_using_indicator(
matches, between_thresholds, -1)
if self._force_match_for_each_row:
similarity_matrix_shape = shape_utils.combined_static_and_dynamic_shape(
similarity_matrix)
force_match_column_ids = tf.argmax(input=similarity_matrix,
axis=1,
output_type=tf.int32)
force_match_column_indicators = tf.one_hot(
force_match_column_ids, depth=similarity_matrix_shape[1])
force_match_row_ids = tf.argmax(
input=force_match_column_indicators,
axis=0,
output_type=tf.int32)
force_match_column_mask = tf.cast(
tf.reduce_max(input_tensor=force_match_column_indicators,
axis=0), tf.bool)
final_matches = tf.where(force_match_column_mask,
force_match_row_ids, matches)
return final_matches
else:
return matches
def construct(im, num_levels, wavelet_type):
"""Constructs a wavelet decomposition of an image.
Args:
im: A numpy or TF tensor of single or double precision floats of size
(batch_size, width, height)
num_levels: The number of levels (or scales) of the wavelet decomposition to
apply. A value of 0 returns a "wavelet decomposition" that is just the
image.
wavelet_type: The kind of wavelet to use, see generate_filters().
Returns:
A wavelet decomposition of `im` that has `num_levels` levels (not including
the coarsest residual level) and is of type `wavelet_type`. This
decomposition is represented as a tuple of 3-tuples, with the final element
being a tensor:
((band00, band01, band02), (band10, band11, band12), ..., resid)
Where band** and resid are TF tensors. Each element of these nested tuples
is of shape [batch_size, width * 2^-(level+1), height * 2^-(level+1)],
though the spatial dimensions may be off by 1 if width and height are not
factors of 2. The residual image is of the same (rough) size as the last set
of bands. The floating point precision of these tensors matches that of
`im`.
"""
if len(im.shape) != 3:
raise ValueError(
'Expected `im` to have a rank of 3, but is of size {}'.format(
im.shape))
if num_levels == 0:
return (tf.convert_to_tensor(value=im), )
max_num_levels = get_max_num_levels(tf.shape(im))
assert_ops = [
tf.Assert(tf.greater_equal(max_num_levels, num_levels),
[tf.shape(im), num_levels, max_num_levels])
]
with tf.control_dependencies(assert_ops):
filters = generate_filters(wavelet_type)
pyr = []
for _ in range(num_levels):
hi = _downsample(im, filters.analysis_hi, 0, 1)
lo = _downsample(im, filters.analysis_lo, 0, 0)
pyr.append(
(_downsample(hi, filters.analysis_hi, 1,
1), _downsample(lo, filters.analysis_hi, 1, 1),
_downsample(hi, filters.analysis_lo, 1,
0))) # pyformat: disable
im = _downsample(lo, filters.analysis_lo, 1, 0)
pyr.append(im)
pyr = tuple(pyr)
return pyr
def _sample_control_dependencies(self, x):
assertions = []
if not self.validate_args:
return assertions
loc = tf.convert_to_tensor(self.loc)
scale = tf.convert_to_tensor(self.scale)
concentration = tf.convert_to_tensor(self.concentration)
assertions.append(assert_util.assert_greater_equal(
x, loc, message='Sample must be greater than or equal to `loc`.'))
assertions.append(assert_util.assert_equal(
tf.logical_or(tf.greater_equal(concentration, 0),
tf.less_equal(x, loc - scale / concentration)),
True,
message=('If `concentration < 0`, sample must be less than or '
'equal to `loc - scale / concentration`.'),
summarize=100))
return assertions
def collater_fn(batch: Dict[str, tf.Tensor]) -> Dict[str, tf.Tensor]:
new_batch = mention_collater_fn(batch)
# Only generate text identifiers and mention hashes for
# the target (linked) mentions.
new_batch['target_text_identifiers'] = tf.gather(
new_batch['text_identifiers'],
new_batch['mention_target_batch_positions'])
new_batch[
'target_mention_hashes'] = mention_preprocess_utils.modified_cantor_pairing(
new_batch['mention_target_start_positions'],
new_batch['target_text_identifiers'])
seq_len = tf.shape(batch['text_ids'])[1]
starts_far_from_passage_boundary = tf.greater_equal(
new_batch['mention_target_start_positions'],
min_distance_from_passage_boundary)
ends_far_from_passage_boundary = tf.less(
new_batch['mention_target_end_positions'],
tf.cast(seq_len,
new_batch['mention_target_end_positions'].dtype) -
min_distance_from_passage_boundary)
far_from_passage_boundary = tf.logical_and(
starts_far_from_passage_boundary,
ends_far_from_passage_boundary)
far_from_passage_boundary = tf.cast(
far_from_passage_boundary,
dtype=new_batch['mention_target_weights'].dtype)
new_batch['mention_target_weights'] = (
new_batch['mention_target_weights'] *
far_from_passage_boundary)
# Collect unique mention IDs per sample in the batch
unique_mention_ids = []
# Mask-out not linked entities.
dense_mention_ids = batch['dense_mention_ids'] * batch[
'dense_mention_mask']
for i in range(bsz):
unique_mention_ids_per_i = tf.unique(dense_mention_ids[i]).y
unique_mention_ids_per_i = tf.cast(unique_mention_ids_per_i,
tf.int32)
unique_mention_ids_per_i = mention_preprocess_utils.dynamic_padding_1d(
unique_mention_ids_per_i, max_mentions_per_sample)
unique_mention_ids.append(unique_mention_ids_per_i)
new_batch['unique_mention_ids'] = tf.stack(unique_mention_ids)
return new_batch
def dropped_inputs(inputs=inputs, rate=self.rate): # pylint: disable=missing-docstring
alpha = 1.6732632423543772848170429916717
scale = 1.0507009873554804934193349852946
alpha_p = -alpha * scale
kept_idx = tf.greater_equal(
self._random_generator.random_uniform(noise_shape), rate)
kept_idx = tf.cast(kept_idx, inputs.dtype)
# Get affine transformation params
a = ((1 - rate) * (1 + rate * alpha_p**2))**-0.5
b = -a * alpha_p * rate
# Apply mask
x = inputs * kept_idx + alpha_p * (1 - kept_idx)
# Do affine transformation
return a * x + b
def _validate_input_shape_tensor(self, input_shape):
input_dim = tf.gather(
input_shape, [ps.rank_from_shape(input_shape) + self.axis])
if self.split_sizes is None:
return [assert_util.assert_equal(
0,
tf.math.floormod(input_dim, self.num_splits))]
else:
split_sizes = tf.convert_to_tensor(self.split_sizes)
splits_are_known = tf.reduce_all(tf.greater_equal(split_sizes, 0))
return [assert_util.assert_equal(
True,
((~splits_are_known &
(tf.reduce_sum(split_sizes) + 1 <= input_dim)) |
(splits_are_known & tf.equal(input_dim, tf.reduce_sum(split_sizes)))
),
message='The size of the input along `axis` does not match '
'`split_sizes`.')]
def meta_init(model,
loss,
x_shape,
y_shape,
n_params,
learning_rate=0.001,
momentum=0.9,
meta_steps=1000,
eps=1e-5,
mask_gradient_fn=None):
"""Run MetaInit algorithm. See `https://papers.nips.cc/paper/9427-metainit-initializing-learning-by-learning-to-initialize`"""
optimizer = ScaleSGD(learning_rate, momentum=momentum)
for _ in range(meta_steps):
x = np.random.normal(0, 1, x_shape)
y = np.random.randint(0, y_shape[1], y_shape[0])
with tf.GradientTape(persistent=True) as tape:
batch_loss = loss(y, model(x, training=True))
grad = tape.gradient(batch_loss, model.trainable_variables)
if mask_gradient_fn is not None:
grad = mask_gradient_fn(model, grad, model.trainable_variables)
prod = tape.gradient(
tf.reduce_sum([tf.reduce_sum(g**2) / 2 for g in grad]),
model.trainable_variables)
if mask_gradient_fn is not None:
prod = mask_gradient_fn(model, prod, model.trainable_variables)
meta_loss = [
tf.abs(1 - ((g - p) / (g + eps * tf.stop_gradient(
(2 * tf.cast(tf.greater_equal(g, 0), tf.float32)) - 1))))
for g, p in zip(grad, prod)
]
if mask_gradient_fn is not None:
meta_loss = mask_gradient_fn(model, meta_loss,
model.trainable_variables)
meta_loss = sum([tf.reduce_sum(m) for m in meta_loss]) / n_params
tf.summary.scalar("meta_loss", meta_loss)
gradients = tape.gradient(meta_loss, model.trainable_variables)
if mask_gradient_fn is not None:
gradients = mask_gradient_fn(model, gradients,
model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
def to_absolute_coordinates(boxlist,
height,
width,
check_range=True,
maximum_normalized_coordinate=1.1,
scope=None):
"""Converts normalized box coordinates to absolute pixel coordinates.
This function raises an assertion failed error when the maximum box coordinate
value is larger than maximum_normalized_coordinate (in which case coordinates
are already absolute).
Args:
boxlist: BoxList with coordinates in range [0, 1].
height: Maximum value for height of absolute box coordinates.
width: Maximum value for width of absolute box coordinates.
check_range: If True, checks if the coordinates are normalized or not.
maximum_normalized_coordinate: Maximum coordinate value to be considered
as normalized, default to 1.1.
scope: name scope.
Returns:
boxlist with absolute coordinates in terms of the image size.
"""
with tf.name_scope(scope, 'ToAbsoluteCoordinates'):
height = tf.cast(height, tf.float32)
width = tf.cast(width, tf.float32)
# Ensure range of input boxes is correct.
if check_range:
box_maximum = tf.reduce_max(boxlist.get())
max_assert = tf.Assert(
tf.greater_equal(maximum_normalized_coordinate, box_maximum), [
'maximum box coordinate value is larger '
'than %f: ' % maximum_normalized_coordinate, box_maximum
])
with tf.control_dependencies([max_assert]):
width = tf.identity(width)
return scale(boxlist, height, width)
def maybe_update_alpha():
"""Maybe update the alpha param.
Checks if global_step is between begin_compression_step and
end_compression_step, and if the current training step is a
compression step.
Returns:
Boolean tensor whether the training step is a compression step.
"""
is_step_within_compression_range = tf.logical_and(
tf.greater_equal(tf.cast(self._global_step, tf.int32),
self._spec.begin_compression_step),
tf.logical_or(
tf.less_equal(tf.cast(self._global_step, tf.int32),
self._spec.end_compression_step),
tf.less(self._spec.end_compression_step, 0)))
is_compression_step = tf.less_equal(
tf.add(self.last_alpha_update_step,
self._spec.compression_frequency),
tf.cast(self._global_step, tf.int32))
return tf.logical_and(is_step_within_compression_range,
is_compression_step)
def remove_above_nyquist(frequency_envelopes: tf.Tensor,
amplitude_envelopes: tf.Tensor,
sample_rate: int = 16000) -> tf.Tensor:
"""Set amplitudes for oscillators above nyquist to 0.
Args:
frequency_envelopes: Sample-wise oscillator frequencies (Hz). Shape
[batch_size, n_samples, n_sinusoids].
amplitude_envelopes: Sample-wise oscillator amplitude. Shape [batch_size,
n_samples, n_sinusoids].
sample_rate: Sample rate in samples per a second.
Returns:
amplitude_envelopes: Sample-wise filtered oscillator amplitude.
Shape [batch_size, n_samples, n_sinusoids].
"""
frequency_envelopes = tf_float32(frequency_envelopes)
amplitude_envelopes = tf_float32(amplitude_envelopes)
amplitude_envelopes = tf.where(
tf.greater_equal(frequency_envelopes, sample_rate / 2.0),
tf.zeros_like(amplitude_envelopes), amplitude_envelopes)
return amplitude_envelopes
def quantile_auc(q0, n0, q1, n1, curve='ROC', name=None):
"""Calculate ranking stats AUROC and AUPRC.
Computes AUROC and AUPRC from quantiles (one for positive trials and one for
negative trials).
We use `pi(x)` to denote a score. We assume that if `pi(x) > k` for some
threshold `k` then the event is predicted to be "1", and otherwise it is
predicted to be a "0". Its actual label is `y`, which may or may not be the
same.
Area Under Curve: Receiver Operator Characteristic (AUROC) is defined as:
/ 1
| TruePositiveRate(k) d FalsePositiveRate(k)
/ 0
where,
```none
TruePositiveRate(k) = P(pi(x) > k | y = 1)
FalsePositiveRate(k) = P(pi(x) > k | y = 0)
```
Area Under Curve: Precision-Recall (AUPRC) is defined as:
/ 1
| Precision(k) d Recall(k)
/ 0
where,
```none
Precision(k) = P(y = 1 | pi(x) > k)
Recall(k) = TruePositiveRate(k) = P(pi(x) > k | y = 1)
```
Notice that AUROC and AUPRC exchange the role of Recall in the
integration, i.e.,
Integrand Measure
+------------+-----------+
AUROC | Recall | FPR |
+------------+-----------+
AUPRC | Precision | Recall |
+------------+-----------+
To learn more about the relationship between AUROC and AUPRC see [1].
Args:
q0: `N-D` `Tensor` of `float`, Quantiles of predicted probabilities given a
negative trial. The first `N-1` dimensions are batch dimensions, and the
AUC is calculated over the final dimension.
n0: `float` or `(N-1)-D Tensor`, Number of negative trials. If `Tensor`,
dimensions must match the first `N-1` dimensions of `q0`.
q1: `N-D` `Tensor` of `float`, Quantiles of predicted probabilities given a
positive trial. The first `N-1` dimensions are batch dimensions, which
must match those of `q0`.
n1: `float` or `(N-1)-D Tensor`, Number of positive trials. If `Tensor`,
dimensions must match the first `N-1` dimensions of `q0`.
curve: `str`, Specifies the name of the curve to be computed. Must be 'ROC'
[default] or 'PR' for the Precision-Recall-curve.
name: `str`, An optional name_scope name.
Returns:
auc: `Tensor` of `float`, area under the ROC or PR curve.
#### Examples
```python
n = 1000
m = 500
predictions_given_positive = np.random.rand(n)
predictions_given_negative = np.random.rand(m)
q1 = tfp.stats.quantiles(predictions_given_positive, num_quantiles=50)
q0 = tfp.stats.quantiles(predictions_given_negative, num_quantiles=50)
auroc = tfp.stats.quantile_auc(q0, m, q1, n, curve='ROC')
```
### Mathematical Details
The algorithm proceeds by partitioning the combined quantile data into a
series of intervals `[a, b)`, approximating the probability mass of
predictions conditioned on a positive trial (`d1`) and probability mass of
predictions conditioned on a negative trial (`d0`) in each interval `[a, b)`,
and accumulating the incremental AUROC/AUPRC as functions of `d0` and `d1`.
We assume that pi(x) is uniform within a given bucket of each quantile. Thus
it will also be uniform within an interval [a, b) as long as the interval does
not cross the quantile's bucket boundaries.
A consequence of this assumption is that the cdf is piecewise linear.
That is,
P( pi(x) > k | y = 0 ), and,
P( pi(x) > k | y = 1 ),
are linear in `k`.
Standard AUROC is fairly easier to calculate. Under the conditional uniformity
assumptions we have a piece's contribution, [a, b), as:
/ b
|
| P(y = 1 | pi > k) d P(pi > k | y = 0)
|
/ a
/ b
|
= - | P(pi > k | y = 1) P(pi = k | y = 0) d k
|
/ a
/ b
-1 / (len(q0) - 1) |
= -------------------- | P(pi > k | y = 1) d k
q0[j + 1] - q0[j] |
/ a
/ 1
1 / (len(q0) - 1) |
= ------------------- (b - a) | (p1 + u d1) d u
q0[j + 1] - q0[j] |
/ 0
1 / (len(q0) - 1)
= ------------------- (b - a) (p1 + d1 / 2)
q0[j + 1] - q0[j]
AUPRC is a bit harder to calculate since the integrand,
`P(y > 0 | pi(x) > k)`, is conditional on `k` rather than a probability over
`k`.
We proceed by formulating Precision in terms of the quantiles we have
available to us.
Precision(k) = P(y = 1 | pi(x) > k)
P(pi(x) > delta | y = 1 ) P(y = 1)
= -----------------------------------------------------------------------
P(pi(x) > delta | y = 1 ) P(y = 1) + P(pi(x) > delta | y = 0 ) P(y = 0)
Since the cdf's are piecewise linear, we calculate this piece's contribution
to AUPRC by the integral:
/ b
|
| P(y = 1 | pi(x) > delta) d P(pi > delta | y = 1)
|
/ a
1 / (len(q1) - 1)
= ------------------- (b - a) *
q1[i + 1] - q1[i]
/ 1
| n1 * (u d1 + p1)
* | -------------------------------------- du
| n1 * (u d1 + p1) + n0 * (u d0 + p0)
/ 0
/ 1
1 / (len(q1) - 1) | n1 * (u d1 + p1)
= ------------------- (b - a) | ------------------------------------ du
q1[i + 1] - q1[i] | n1 * (u d1 + p1) + n0 * (u d0 + p0)
/ 0
where the equality is a consequence of the piecewise uniformity assumption.
The solution to the integral is given by Mathematica:
```
Integrate[n1 (u d1 + p1) / (n1 (u d1 + p1) + n0 (u d0 + p0)), {u, 0, 1},
Assumptions -> {p1 > 0, d1 > 0, p0 > 0, d0 > 0, n1 > 0, n0 > 0}]
```
This integral can be solved by hand by noticing that:
f(x) / (f(x) + g(x)) = 1 / (1 + g(x)/f(x))
Thus define: u = 1 + g(x)/f(x)
for which: du = [g'(x)h(x) - g(x)f'(x)] / h(x)^2 dx
and solving integral 1/u du.
#### References
[1]: Jesse Davis and Mark Goadrich. The relationship between
Precision-Recall and ROC curves. In _International Conference on
Machine Learning_, 2006. http://dl.acm.org/citation.cfm?id=1143874
"""
with tf.name_scope(name or 'quantile_auc'):
dtype = dtype_util.common_dtype([q0, q1, n1, n0],
dtype_hint=tf.float32)
q1 = tf.convert_to_tensor(q1, dtype=dtype)
q0 = tf.convert_to_tensor(q0, dtype=dtype)
n1 = tf.convert_to_tensor(n1, dtype=dtype)
n0 = tf.convert_to_tensor(n0, dtype=dtype)
# Broadcast `q0` and `q1` to at least 2D. This makes `q1[i]` and `q0[j]` at
# least 1D, allowing `tf.where` to operate on them.
q1 = q1 + tf.zeros((1, 1), dtype=dtype)
q0 = q0 + tf.zeros((1, 1), dtype=dtype)
q1_shape = ps.shape(q1)
batch_shape = q1_shape[:-1]
n_q1 = q1_shape[-1]
n_q0 = ps.shape(q0)[-1]
static_batch_shape = q1.shape[:-1]
n0 = tf.broadcast_to(n0, batch_shape)
n1 = tf.broadcast_to(n1, batch_shape)
def loop_body(auc_prev, p0_prev, p1_prev, i, j):
"""Body of the loop to integrate over the ROC/PR curves."""
# We will align intervals from q0 and q1 by moving from end to beginning,
# stopping with whatever boundary we encounter first. This boundary is
# `b`. We find `a` by continuing the search from `b` to the left.
b, i, j = _get_endpoint_b(i, j, q1, q0, batch_shape)
# Get q1[i] and q1[i+1]
ind = tf.where(tf.greater_equal(i, 0))
q1_i = _get_q_slice(q1, i, ind, batch_shape=batch_shape)
ip1 = tf.minimum(i + 1, n_q1 - 1)
q1_ip1 = _get_q_slice(q1, ip1, ind, batch_shape=batch_shape)
# Get q0[j] and q0[j+1]
ind = tf.where(tf.greater_equal(j, 0))
q0_j = _get_q_slice(q0, j, ind, batch_shape=batch_shape)
jp1 = tf.minimum(j + 1, n_q0 - 1)
q0_jp1 = _get_q_slice(q0, jp1, ind, batch_shape=batch_shape)
a = _get_endpoint_a(i, j, q1_i, q0_j, batch_shape)
# Calculate the proportion [a, b) represents of [q1[i], q1[i+1]).
d1 = _get_interval_proportion(i, n_q1, q1_i, q1_ip1, a, b,
batch_shape)
# Calculate the proportion [a, b) represents of [q1[i], q1[i+1]).
d0 = _get_interval_proportion(j, n_q0, q0_j, q0_jp1, a, b,
batch_shape)
# Notice that because we assumed within bucket values are distributed
# uniformly, we end up with something which is identical to the
# trapezoidal rule: definite_integral += (b - a) * (f(a) + f(b)) / 2.
if curve == 'ROC':
auc = auc_prev + d0 * (p1_prev + d1 / 2.)
else:
total_scaled_delta = n0 * d0 + n1 * d1
total_scaled_cdf_at_b = n0 * p0_prev + n1 * p1_prev
def get_auprc_update():
proportion = (total_scaled_cdf_at_b /
(total_scaled_cdf_at_b + total_scaled_delta))
definite_integral = (
(n1 / tf.square(total_scaled_delta)) *
(d1 * total_scaled_delta + tf.math.log(proportion) *
(d1 * total_scaled_cdf_at_b -
p1_prev * total_scaled_delta)))
return d1 * definite_integral
# Values should be non-negative and we use > 0.0 rather than != 0.0 to
# catch possible numerical imprecision.
delta_gt_0 = tf.greater(total_scaled_delta, 0.)
cdf_gt_0 = tf.greater(total_scaled_cdf_at_b, 0.)
d1_gt_0 = tf.greater(d1, 0.)
ind = tf.where(delta_gt_0 & cdf_gt_0 & d1_gt_0)
auc_update = tf.gather_nd(get_auprc_update(), ind)
auc = tf.tensor_scatter_nd_add(auc_prev, ind, auc_update)
# TODO(jvdillon): In calculating AUROC and AUPRC, we probably should
# resolve ties randomly, making the following eight states possible (where
# e = 1 means we expected a positive trial and e = 0 means we expected a
# negative trial):
#
# P(y = 1 | pi(x) > delta)
# P(y = 1 | pi(x) = delta, e = 1) 0.5
# P(y = 1 | pi(x) = delta, e = 0) 0.5
# P(y = 1 | pi(x) < delta)
#
# P(y = 0 | pi(x) > delta)
# P(y = 0 | pi(x) = delta, e = 1) 0.5
# P(y = 0 | pi(x) = delta, e = 0) 0.5
# P(y = 0 | pi(x) < delta)
#
# This makes the math a bit harder and its not clear this adds much,
# especially when we're assuming piecewise uniformity.
# Accumulate this mass (d1, d0) for the next iteration.
p1 = p1_prev + d1
p0 = p0_prev + d0
return auc, p0, p1, i, j
init_auc = tf.zeros(batch_shape, dtype=dtype)
init_p0 = tf.zeros(batch_shape, dtype=dtype)
init_p1 = tf.zeros(batch_shape, dtype=dtype)
init_i = tf.zeros(batch_shape, dtype=tf.int64) + n_q1 - 1
init_j = tf.zeros(batch_shape, dtype=tf.int64) + n_q0 - 1
loop_cond = lambda auc, p0, p1, i, j: tf.reduce_all( # pylint: disable=g-long-lambda
tf.greater_equal(i, 0) | tf.greater_equal(j, 0))
init_vars = [init_auc, init_p0, init_p1, init_i, init_j]
auc, _, _, _, _ = tf.while_loop(loop_cond,
loop_body,
init_vars,
shape_invariants=[static_batch_shape] *
5)
return tf.squeeze(auc)
def __call__(self, w):
return w * tf.cast(tf.greater_equal(w, 0.0), backend.floatx())
def call(self, inputs, training=False):
x = inputs
training = training and self.trainable
self.will_ema_freeze = self.will_ema_freeze and self.trainable
# Update the step count if the optimizer step count is unknown
self.step.assign_add(
K.switch(
tf.math.logical_and(self.is_estimating_step_count, training),
tf.constant(1, tf.int64), tf.constant(0, tf.int64)))
# Perform the quantization
if training:
# Calculate the qnoise, a scalar from 0 to 1 that represents the level of
# quantization noise to use. At training start, we want no quantization,
# so qnoise_factor = 0.0. After quantization_delay steps, we want normal
# quantization, so qnoise_factor = 1.0.
qnoise_factor = K.switch(
tf.greater_equal(self.step, self.quantization_delay),
lambda: tf.constant(1.0), lambda: tf.constant(0.0))
qx = self.quantizer(x, qnoise_factor=qnoise_factor)
else: # If not training, we always want to use full quantization
qx = self.quantizer(x, qnoise_factor=tf.constant(1.0))
# Calculate the axis along where to find the min and max EMAs
len_axis = len(x.shape)
if len_axis > 1:
if self.per_channel:
if K.image_data_format() == "channels_last":
axis = list(range(len_axis - 1))
else:
axis = list(range(1, len_axis))
else:
axis = list(range(len_axis))
else:
axis = [0]
# Determine if freezing the EMA
is_ema_training = tf.constant(training, dtype=tf.bool)
if self.will_ema_freeze:
is_ema_training = tf.cond(
tf.greater(self.step, self.ema_freeze_delay),
lambda: tf.constant(False), lambda: tf.constant(True))
# Update the moving average
if is_ema_training:
new_min = tf.squeeze(K.min(qx, axis=axis, keepdims=True))
K.moving_average_update(self.ema_min, new_min, self.ema_decay)
new_max = tf.squeeze(K.max(qx, axis=axis, keepdims=True))
K.moving_average_update(self.ema_max, new_max, self.ema_decay)
# Set the integer bits for the quantizer
integer_bits = _get_integer_bits(min_value=self.ema_min,
max_value=self.ema_max,
bits=self.total_bits,
symmetric=self.symmetric,
keep_negative=self.keep_negative,
is_clipping=self.po2_rounding)
self.quantizer.integer.assign(integer_bits)
return qx
def _at_least_x_are_equal(a, b, x):
"""At least `x` of `a` and `b` `Tensors` are equal."""
match = tf.equal(a, b)
match = tf.cast(match, tf.int32)
return tf.greater_equal(tf.reduce_sum(match), x)
def loop_body(auc_prev, p0_prev, p1_prev, i, j):
"""Body of the loop to integrate over the ROC/PR curves."""
# We will align intervals from q0 and q1 by moving from end to beginning,
# stopping with whatever boundary we encounter first. This boundary is
# `b`. We find `a` by continuing the search from `b` to the left.
b, i, j = _get_endpoint_b(i, j, q1, q0, batch_shape)
# Get q1[i] and q1[i+1]
ind = tf.where(tf.greater_equal(i, 0))
q1_i = _get_q_slice(q1, i, ind, batch_shape=batch_shape)
ip1 = tf.minimum(i + 1, n_q1 - 1)
q1_ip1 = _get_q_slice(q1, ip1, ind, batch_shape=batch_shape)
# Get q0[j] and q0[j+1]
ind = tf.where(tf.greater_equal(j, 0))
q0_j = _get_q_slice(q0, j, ind, batch_shape=batch_shape)
jp1 = tf.minimum(j + 1, n_q0 - 1)
q0_jp1 = _get_q_slice(q0, jp1, ind, batch_shape=batch_shape)
a = _get_endpoint_a(i, j, q1_i, q0_j, batch_shape)
# Calculate the proportion [a, b) represents of [q1[i], q1[i+1]).
d1 = _get_interval_proportion(i, n_q1, q1_i, q1_ip1, a, b,
batch_shape)
# Calculate the proportion [a, b) represents of [q1[i], q1[i+1]).
d0 = _get_interval_proportion(j, n_q0, q0_j, q0_jp1, a, b,
batch_shape)
# Notice that because we assumed within bucket values are distributed
# uniformly, we end up with something which is identical to the
# trapezoidal rule: definite_integral += (b - a) * (f(a) + f(b)) / 2.
if curve == 'ROC':
auc = auc_prev + d0 * (p1_prev + d1 / 2.)
else:
total_scaled_delta = n0 * d0 + n1 * d1
total_scaled_cdf_at_b = n0 * p0_prev + n1 * p1_prev
def get_auprc_update():
proportion = (total_scaled_cdf_at_b /
(total_scaled_cdf_at_b + total_scaled_delta))
definite_integral = (
(n1 / tf.square(total_scaled_delta)) *
(d1 * total_scaled_delta + tf.math.log(proportion) *
(d1 * total_scaled_cdf_at_b -
p1_prev * total_scaled_delta)))
return d1 * definite_integral
# Values should be non-negative and we use > 0.0 rather than != 0.0 to
# catch possible numerical imprecision.
delta_gt_0 = tf.greater(total_scaled_delta, 0.)
cdf_gt_0 = tf.greater(total_scaled_cdf_at_b, 0.)
d1_gt_0 = tf.greater(d1, 0.)
ind = tf.where(delta_gt_0 & cdf_gt_0 & d1_gt_0)
auc_update = tf.gather_nd(get_auprc_update(), ind)
auc = tf.tensor_scatter_nd_add(auc_prev, ind, auc_update)
# TODO(jvdillon): In calculating AUROC and AUPRC, we probably should
# resolve ties randomly, making the following eight states possible (where
# e = 1 means we expected a positive trial and e = 0 means we expected a
# negative trial):
#
# P(y = 1 | pi(x) > delta)
# P(y = 1 | pi(x) = delta, e = 1) 0.5
# P(y = 1 | pi(x) = delta, e = 0) 0.5
# P(y = 1 | pi(x) < delta)
#
# P(y = 0 | pi(x) > delta)
# P(y = 0 | pi(x) = delta, e = 1) 0.5
# P(y = 0 | pi(x) = delta, e = 0) 0.5
# P(y = 0 | pi(x) < delta)
#
# This makes the math a bit harder and its not clear this adds much,
# especially when we're assuming piecewise uniformity.
# Accumulate this mass (d1, d0) for the next iteration.
p1 = p1_prev + d1
p0 = p0_prev + d0
return auc, p0, p1, i, j
