def slice_bounds_by_doubling(x_initial,
target_log_prob,
log_slice_heights,
max_doublings,
step_size,
seed=None,
name=None):
"""Returns the bounds of the slice at each stage of doubling procedure.
Precomputes the x coordinates of the left (L) and right (R) endpoints of the
interval `I` produced in the "doubling" algorithm [Neal 2003][1] P713. Note
that we simultaneously compute all possible doubling values for each chain,
for the reason that at small-medium densities, the gains from parallel
evaluation might cause a speed-up, but this will be benchmarked against the
while loop implementation.
Args:
x_initial: `tf.Tensor` of any shape and any real dtype consumable by
`target_log_prob`. The initial points.
target_log_prob: A callable taking a `tf.Tensor` of shape and dtype as
`x_initial` and returning a tensor of the same shape. The log density of
the target distribution.
log_slice_heights: `tf.Tensor` with the same shape as `x_initial` and the
same dtype as returned by `target_log_prob`. The log of the height of the
slice for each chain. The values must be bounded above by
`target_log_prob(x_initial)`.
max_doublings: Scalar positive int32 `tf.Tensor`. The maximum number of
doublings to consider.
step_size: `tf.Tensor` with same dtype as and shape compatible with
`x_initial`. The size of the initial interval.
seed: (Optional) positive int or Tensor seed pair. The random seed.
name: Python `str` name prefixed to Ops created by this function.
Default value: `None` (i.e., 'find_slice_bounds').
Returns:
upper_bounds: A tensor of same shape and dtype as `x_initial`. Slice upper
bounds for each chain.
lower_bounds: A tensor of same shape and dtype as `x_initial`. Slice lower
bounds for each chain.
both_ok: A tensor of shape `x_initial` and boolean dtype. Indicates if both
the chosen upper and lower bound lie outside of the slice.
#### References
[1]: Radford M. Neal. Slice Sampling. The Annals of Statistics. 2003, Vol 31,
No. 3 , 705-767.
https://projecteuclid.org/download/pdf_1/euclid.aos/1056562461
"""
with tf.name_scope(name or 'slice_bounds_by_doubling'):
left_seed, increments_seed = samplers.split_seed(
seed, salt='slice_bounds_by_doubling')
x_initial = tf.convert_to_tensor(value=x_initial)
batch_shape = tf.shape(x_initial)
dtype = dtype_util.base_dtype(step_size.dtype)
left_endpoints = x_initial + step_size * samplers.uniform(
batch_shape, minval=-1.0, maxval=0.0, dtype=dtype, seed=left_seed)
# Compute the increments by which we need to step the upper and lower bounds
# part of the doubling procedure.
left_increments, widths = _left_doubling_increments(
batch_shape, max_doublings, step_size, seed=increments_seed)
# The left and right end points. Shape (max_doublings+1,) + batch_shape.
left_endpoints = left_endpoints - left_increments
right_endpoints = left_endpoints + widths
# Test if these end points lie outside of the slice.
# Checks if the end points of the slice are outside the graph of the pdf.
left_ep_values = tf.map_fn(target_log_prob, left_endpoints)
right_ep_values = tf.map_fn(target_log_prob, right_endpoints)
left_ok = left_ep_values < log_slice_heights
right_ok = right_ep_values < log_slice_heights
both_ok = left_ok & right_ok
both_ok_f = tf.reshape(both_ok, [max_doublings + 1, -1])
best_interval_idx = _find_best_interval_idx(
tf.cast(both_ok_f, dtype=tf.int32))
# Formats the above index as required to use with gather_nd.
point_index_gather = tf.stack(
[best_interval_idx,
tf.range(tf.size(best_interval_idx))],
axis=1,
name='point_index_gather')
left_ep_f = tf.reshape(left_endpoints, [max_doublings + 1, -1])
right_ep_f = tf.reshape(right_endpoints, [max_doublings + 1, -1])
# The x values of the uppper and lower bounds of the slices for each chain.
lower_bounds = tf.reshape(tf.gather_nd(left_ep_f, point_index_gather),
batch_shape)
upper_bounds = tf.reshape(tf.gather_nd(right_ep_f, point_index_gather),
batch_shape)
both_ok = tf.reduce_any(both_ok, axis=0)
return upper_bounds, lower_bounds, both_ok
def brier_decomposition(labels, logits, name=None):
r"""Decompose the Brier score into uncertainty, resolution, and reliability.
[Proper scoring rules][1] measure the quality of probabilistic predictions;
any proper scoring rule admits a [unique decomposition][2] as
`Score = Uncertainty - Resolution + Reliability`, where:
* `Uncertainty`, is a generalized entropy of the average predictive
distribution; it can both be positive or negative.
* `Resolution`, is a generalized variance of individual predictive
distributions; it is always non-negative. Difference in predictions reveal
information, that is why a larger resolution improves the predictive score.
* `Reliability`, a measure of calibration of predictions against the true
frequency of events. It is always non-negative and a lower value here
indicates better calibration.
This method estimates the above decomposition for the case of the Brier
scoring rule for discrete outcomes. For this, we need to discretize the space
of probability distributions; we choose a simple partition of the space into
`nlabels` events: given a distribution `p` over `nlabels` outcomes, the index
`k` for which `p_k > p_i` for all `i != k` determines the discretization
outcome; that is, `p in M_k`, where `M_k` is the set of all distributions for
which `p_k` is the largest value among all probabilities.
The estimation error of each component is O(k/n), where n is the number
of instances and k is the number of labels. There may be an error of this
order when compared to `brier_score`.
#### References
[1]: Tilmann Gneiting, Adrian E. Raftery.
Strictly Proper Scoring Rules, Prediction, and Estimation.
Journal of the American Statistical Association, Vol. 102, 2007.
https://www.stat.washington.edu/raftery/Research/PDF/Gneiting2007jasa.pdf
[2]: Jochen Broecker. Reliability, sufficiency, and the decomposition of
proper scores.
Quarterly Journal of the Royal Meteorological Society, Vol. 135, 2009.
https://rmets.onlinelibrary.wiley.com/doi/epdf/10.1002/qj.456
Args:
labels: Tensor, (n,), with tf.int32 or tf.int64 elements containing ground
truth class labels in the range [0,nlabels].
logits: Tensor, (n, nlabels), with logits for n instances and nlabels.
name: Python `str` name prefixed to Ops created by this function.
Returns:
uncertainty: Tensor, scalar, the uncertainty component of the
decomposition.
resolution: Tensor, scalar, the resolution component of the decomposition.
reliability: Tensor, scalar, the reliability component of the
decomposition.
"""
with tf.name_scope(name or 'brier_decomposition'):
labels = tf.convert_to_tensor(labels)
logits = tf.convert_to_tensor(logits)
num_classes = logits.shape[-1]
# Compute pbar, the average distribution
pred_class = tf.argmax(logits, axis=-1, output_type=labels.dtype)
if tensorshape_util.rank(logits.shape) > 2:
flatten, unflatten = _make_flatten_unflatten_fns(logits.shape[:-2])
def fn_to_map(args):
yhat, y = args
return tf.math.confusion_matrix(yhat,
y,
num_classes=num_classes,
dtype=logits.dtype)
confusion_matrix = tf.map_fn(
fn_to_map,
[flatten(pred_class), flatten(labels)],
fn_output_signature=logits.dtype)
confusion_matrix = unflatten(confusion_matrix)
else:
confusion_matrix = tf.math.confusion_matrix(
pred_class,
labels,
num_classes=num_classes,
dtype=logits.dtype)
dist_weights = tf.reduce_sum(confusion_matrix, axis=-1)
dist_weights /= tf.reduce_sum(dist_weights, axis=-1, keepdims=True)
pbar = tf.reduce_sum(confusion_matrix, axis=-2)
pbar /= tf.reduce_sum(pbar, axis=-1, keepdims=True)
eps = np.finfo(dtype_util.as_numpy_dtype(confusion_matrix.dtype)).eps
# dist_mean[k,:] contains the empirical distribution for the set M_k
# Some outcomes may not realize, corresponding to dist_weights[k] = 0
dist_mean = confusion_matrix / (
eps + tf.reduce_sum(confusion_matrix, axis=-1, keepdims=True))
# Uncertainty: quadratic entropy of the average label distribution
uncertainty = -tf.reduce_sum(tf.square(pbar), axis=-1)
# Resolution: expected quadratic divergence of predictive to mean
resolution = tf.square(tf.expand_dims(pbar, -1) - dist_mean)
resolution = tf.reduce_sum(dist_weights *
tf.reduce_sum(resolution, axis=-1),
axis=-1)
# Reliability: expected quadratic divergence of predictive to true
if tensorshape_util.rank(logits.shape) > 2:
# TODO(b/139094519): Avoid using tf.map_fn here.
prob_true = tf.map_fn(
lambda args: tf.gather(args[0], args[1]),
[flatten(dist_mean), flatten(pred_class)],
fn_output_signature=dist_mean.dtype)
prob_true = unflatten(prob_true)
else:
prob_true = tf.gather(dist_mean, pred_class, axis=0)
log_prob_true = tf.math.log(prob_true)
log_prob_pred = logits - tf.math.reduce_logsumexp(
logits, axis=-1, keepdims=True)
log_reliability = _reduce_log_l2_exp(log_prob_pred,
log_prob_true,
axis=-1)
log_reliability = tf.math.reduce_logsumexp(
log_reliability,
axis=-1,
)
num_samples = tf.cast(tf.shape(logits)[-2], logits.dtype)
reliability = tf.exp(log_reliability - tf.math.log(num_samples))
return uncertainty, resolution, reliability
def lu_reconstruct(lower_upper, perm, validate_args=False, name=None):
"""The inverse LU decomposition, `X == lu_reconstruct(*tf.linalg.lu(X))`.
Args:
lower_upper: `lu` as returned by `tf.linalg.lu`, i.e., if
`matmul(P, matmul(L, U)) = X` then `lower_upper = L + U - eye`.
perm: `p` as returned by `tf.linag.lu`, i.e., if
`matmul(P, matmul(L, U)) = X` then `perm = argmax(P)`.
validate_args: Python `bool` indicating whether arguments should be checked
for correctness.
Default value: `False` (i.e., don't validate arguments).
name: Python `str` name given to ops managed by this object.
Default value: `None` (i.e., 'lu_reconstruct').
Returns:
x: The original input to `tf.linalg.lu`, i.e., `x` as in,
`lu_reconstruct(*tf.linalg.lu(x))`.
#### Examples
```python
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
x = [[[3., 4], [1, 2]],
[[7., 8], [3, 4]]]
x_reconstructed = tfp.math.lu_reconstruct(*tf.linalg.lu(x))
tf.assert_near(x, x_reconstructed)
# ==> True
```
"""
with tf.name_scope(name or 'lu_reconstruct'):
lower_upper = tf.convert_to_tensor(lower_upper,
dtype_hint=tf.float32,
name='lower_upper')
perm = tf.convert_to_tensor(perm, dtype_hint=tf.int32, name='perm')
assertions = lu_reconstruct_assertions(lower_upper, perm,
validate_args)
if assertions:
with tf.control_dependencies(assertions):
lower_upper = tf.identity(lower_upper)
perm = tf.identity(perm)
shape = tf.shape(lower_upper)
lower = tf.linalg.set_diag(
tf.linalg.band_part(lower_upper, num_lower=-1, num_upper=0),
tf.ones(shape[:-1], dtype=lower_upper.dtype))
upper = tf.linalg.band_part(lower_upper, num_lower=0, num_upper=-1)
x = tf.matmul(lower, upper)
if (tensorshape_util.rank(lower_upper.shape) is None
or tensorshape_util.rank(lower_upper.shape) != 2):
# We either don't know the batch rank or there are >0 batch dims.
batch_size = tf.reduce_prod(shape[:-2])
d = shape[-1]
x = tf.reshape(x, [batch_size, d, d])
perm = tf.reshape(perm, [batch_size, d])
perm = tf.map_fn(tf.math.invert_permutation, perm)
batch_indices = tf.broadcast_to(
tf.range(batch_size)[:, tf.newaxis], [batch_size, d])
x = tf.gather_nd(x, tf.stack([batch_indices, perm], axis=-1))
x = tf.reshape(x, shape)
else:
x = tf.gather(x, tf.math.invert_permutation(perm))
tensorshape_util.set_shape(x, lower_upper.shape)
return x
def objective_func(population):
return tf.map_fn(quadratic, population)
def count_integers(arr,
weights=None,
minlength=None,
maxlength=None,
axis=None,
dtype=tf.int32,
name=None):
"""Counts the number of occurrences of each value in an integer array `arr`.
Works like `tf.math.bincount`, but provides an `axis` kwarg that specifies
dimensions to reduce over. With
`~axis = [i for i in range(arr.ndim) if i not in axis]`,
this function returns a `Tensor` of shape `[K] + arr.shape[~axis]`.
If `minlength` and `maxlength` are not given, `K = tf.reduce_max(arr) + 1`
if `arr` is non-empty, and 0 otherwise.
If `weights` are non-None, then index `i` of the output stores the sum of the
value in `weights` at each index where the corresponding value in `arr` is
`i`.
Args:
arr: An `int32` `Tensor` of non-negative values.
weights: If non-None, must be the same shape as arr. For each value in
`arr`, the bin will be incremented by the corresponding weight instead of
1.
minlength: If given, ensures the output has length at least `minlength`,
padding with zeros at the end if necessary.
maxlength: If given, skips values in `arr` that are equal or greater than
`maxlength`, ensuring that the output has length at most `maxlength`.
axis: A `0-D` or `1-D` `int32` `Tensor` (with static values) designating
dimensions in `arr` to reduce over.
`Default value:` `None`, meaning reduce over all dimensions.
dtype: If `weights` is None, determines the type of the output bins.
name: A name scope for the associated operations (optional).
Returns:
A vector with the same dtype as `weights` or the given `dtype`. The bin
values.
"""
with tf.name_scope(name or 'count_integers'):
if axis is None:
return tf.math.bincount(arr,
weights=weights,
minlength=minlength,
maxlength=maxlength,
dtype=dtype)
arr = tf.convert_to_tensor(arr, dtype=tf.int32, name='arr')
arr_ndims = _get_static_ndims(arr, expect_static=True)
axis = _make_static_axis_non_negative_list(axis, arr_ndims)
# ~axis from docstring. Dims in arr that are not in axis.
not_axis = sorted(set(range(arr_ndims)).difference(axis))
# If we're reducing over everything, just use standard bincount.
if not not_axis:
return tf.math.bincount(arr,
weights=weights,
minlength=minlength,
maxlength=maxlength,
dtype=dtype)
# Move dims in ~axis to the left, so we can tf.map_fn bincount over them,
# Producing counts for every index I in ~axis.
# Thus, flat_arr is not totally flat, it just has the dims in ~axis
# flattened.
flat_arr = _move_dims_to_flat_end(arr,
not_axis,
arr_ndims,
right_end=False)
minlength = minlength if minlength is not None else tf.reduce_max(
arr) + 1
maxlength = maxlength if maxlength is not None else tf.reduce_max(
arr) + 1
# tf.map_fn over dim 0.
if weights is None:
def one_bincount(arr_slice):
return tf.math.bincount(arr_slice,
weights=None,
minlength=minlength,
maxlength=maxlength,
dtype=dtype)
flat_counts = tf.map_fn(one_bincount,
elems=flat_arr,
fn_output_signature=dtype)
else:
weights = tf.convert_to_tensor(weights, name='weights')
_get_static_ndims(weights,
expect_static=True,
expect_ndims=arr_ndims)
flat_weights = _move_dims_to_flat_end(weights,
not_axis,
arr_ndims,
right_end=False)
def one_bincount(arr_and_weights_slices):
arr_slice, weights_slice = arr_and_weights_slices
return tf.math.bincount(arr_slice,
weights=weights_slice,
minlength=minlength,
maxlength=maxlength,
dtype=dtype)
flat_counts = tf.map_fn(one_bincount,
elems=[flat_arr, flat_weights],
fn_output_signature=weights.dtype)
# flat_counts.shape = [prod(~axis), K], because map_fn stacked on axis 0.
# bincount needs to have the K bins in axis 0, so transpose...
flat_counts_t = tf.transpose(a=flat_counts, perm=[1, 0])
# Throw in this assert, to ensure shape assumptions are correct.
_get_static_ndims(flat_counts_t, expect_ndims=2, expect_static=True)
# not_axis_shape = arr.shape[~axis]
not_axis_shape = ps.gather(ps.shape(arr), indices=not_axis)
# The first index of flat_counts_t indexes bins 0,..,K-1, the rest are ~axis
out_shape = ps.concat([[-1], not_axis_shape], axis=0)
return tf.reshape(flat_counts_t, out_shape)
def sample_and_preprocess(video,
labels,
seq_label,
seq_len,
name,
num_steps,
augment,
sample_all=False,
sample_all_stride=1,
add_shape=False):
"""Samples frames and prepares them for training."""
if sample_all:
# When dealing with very long videos we can choose to sub-sample to fit
# data in memory. But be aware this also evaluates over a subset of frames.
# Subsampling the validation set videos when reporting performance is not
# recommended.
steps = tf.range(0, seq_len, sample_all_stride)
seq_len = tf.shape(steps)[0]
chosen_steps = steps
else:
stride = CONFIG.DATA.STRIDE
sampling_strategy = CONFIG.DATA.SAMPLING_STRATEGY
# TODO(debidatta) : More flexible sampling
if sampling_strategy == 'stride':
# Offset can be set between 0 and maximum location from which we can get
# total coverage of the video without having to pad.
# This handles sampling over longer sequences.
offset = tf.random.uniform(
(),
0,
tf.maximum(tf.cast(1, tf.int64), seq_len - stride * num_steps),
dtype=tf.int64)
# This handles sampling over shorter sequences by padding the last frame
# many times. This is not ideal for the way alignment training batches are
# created.
steps = tf.minimum(
seq_len - 1,
tf.range(offset, offset + num_steps * stride + 1, stride))
steps = steps[:num_steps]
elif sampling_strategy == 'offset_uniform':
# Sample a random offset less than a provided max offset. Among all frames
# higher than the chosen offset, randomly sample num_frames
check1 = tf.debugging.assert_greater_equal(
seq_len,
tf.cast(CONFIG.DATA.RANDOM_OFFSET, tf.int64),
message='Random offset is more than sequence length.')
check2 = tf.less_equal(
tf.cast(num_steps, tf.int64),
seq_len - tf.cast(CONFIG.DATA.RANDOM_OFFSET, tf.int64),
)
def _sample_random():
with tf.control_dependencies([tf.identity(check1.outputs[0])]):
offset = CONFIG.DATA.RANDOM_OFFSET
steps = tf.random.shuffle(tf.range(offset, seq_len))
steps = tf.gather(steps, tf.range(0, num_steps))
steps = tf.gather(
steps,
tf.nn.top_k(steps, k=num_steps).indices[::-1])
return steps
def _sample_all():
return tf.range(0, num_steps, dtype=tf.int64)
steps = tf.cond(check2, _sample_random, _sample_all)
else:
raise ValueError(
'Sampling strategy %s is unknown. Supported values are '
'stride, offset_uniform .' % sampling_strategy)
if not sample_all and 'tcn' in CONFIG.TRAINING_ALGO:
pos_window = CONFIG.TCN.POSITIVE_WINDOW
# pylint: disable=g-long-lambda
pos_steps = tf.map_fn(
lambda step: tf.random.uniform(
(), minval=step - pos_window, maxval=step, dtype=tf.int64),
steps)
# pylint: enable=g-long-lambda
steps = tf.stack([pos_steps, steps])
steps = tf.reshape(tf.transpose(steps), (-1, ))
# Store chosen indices.
chosen_steps = steps
# Get multiple context steps depending on config at selected steps.
steps = tf.reshape(tf.map_fn(get_steps, steps), [-1])
steps = tf.maximum(tf.cast(0, tf.int64), steps)
steps = tf.minimum(seq_len - 1, steps)
shape_all_steps = CONFIG.DATA.NUM_STEPS * num_steps
if not sample_all and 'tcn' in CONFIG.TRAINING_ALGO:
shape_all_steps *= 2
# Select data based on steps/
video = tf.gather(video, steps)
# Decode the encoded JPEG images
video = tf.map_fn(tf.image.decode_jpeg,
video,
parallel_iterations=FLAGS.num_parallel_calls,
dtype=tf.uint8)
# Take images in range [0, 255] and normalize to [0, 1]
video = tf.map_fn(normalize_input,
video,
parallel_iterations=FLAGS.num_parallel_calls,
dtype=tf.float32)
# Perform data-augmentation and return images in range [-1, 1]
video = preprocess_input(video, augment)
if add_shape:
video.set_shape(
[shape_all_steps, CONFIG.IMAGE_SIZE, CONFIG.IMAGE_SIZE, 3])
if CONFIG.DATA.FRAME_LABELS:
labels = tf.gather(labels, steps)
if add_shape:
labels.set_shape([shape_all_steps])
return {
'frames': video,
'frame_labels': labels,
'chosen_steps': chosen_steps,
'seq_lens': seq_len,
'seq_labels': seq_label,
'name': name
}
def convert_image_sequence_dtype(tensor, dtype=tf.float32):
return tf.map_fn(lambda x: tf.image.convert_image_dtype(x, dtype),
tensor,
dtype=dtype,
back_prop=False)
def draw_bounding_boxes_on_image_tensors(images,
boxes,
classes,
scores,
category_index,
original_image_spatial_shape=None,
true_image_shape=None,
instance_masks=None,
keypoints=None,
max_boxes_to_draw=20,
min_score_thresh=0.2,
use_normalized_coordinates=True):
"""Draws bounding boxes, masks, and keypoints on batch of image tensors.
Args:
images: A 4D uint8 image tensor of shape [N, H, W, C]. If C > 3, additional
channels will be ignored. If C = 1, then we convert the images to RGB
images.
boxes: [N, max_detections, 4] float32 tensor of detection boxes.
classes: [N, max_detections] int tensor of detection classes. Note that
classes are 1-indexed.
scores: [N, max_detections] float32 tensor of detection scores.
category_index: a dict that maps integer ids to category dicts. e.g.
{1: {1: 'dog'}, 2: {2: 'cat'}, ...}
original_image_spatial_shape: [N, 2] tensor containing the spatial size of
the original image.
true_image_shape: [N, 3] tensor containing the spatial size of unpadded
original_image.
instance_masks: A 4D uint8 tensor of shape [N, max_detection, H, W] with
instance masks.
keypoints: A 4D float32 tensor of shape [N, max_detection, num_keypoints, 2]
with keypoints.
max_boxes_to_draw: Maximum number of boxes to draw on an image. Default 20.
min_score_thresh: Minimum score threshold for visualization. Default 0.2.
use_normalized_coordinates: Whether to assume boxes and kepoints are in
normalized coordinates (as opposed to absolute coordiantes).
Default is True.
Returns:
4D image tensor of type uint8, with boxes drawn on top.
"""
# Additional channels are being ignored.
if images.shape[3] > 3:
images = images[:, :, :, 0:3]
elif images.shape[3] == 1:
images = tf.image.grayscale_to_rgb(images)
visualization_keyword_args = {
'use_normalized_coordinates': use_normalized_coordinates,
'max_boxes_to_draw': max_boxes_to_draw,
'min_score_thresh': min_score_thresh,
'agnostic_mode': False,
'line_thickness': 4
}
if true_image_shape is None:
true_shapes = tf.constant(-1, shape=[images.shape.as_list()[0], 3])
else:
true_shapes = true_image_shape
if original_image_spatial_shape is None:
original_shapes = tf.constant(-1, shape=[images.shape.as_list()[0], 2])
else:
original_shapes = original_image_spatial_shape
if instance_masks is not None and keypoints is None:
visualize_boxes_fn = functools.partial(
_visualize_boxes_and_masks,
category_index=category_index,
**visualization_keyword_args)
elems = [
true_shapes, original_shapes, images, boxes, classes, scores,
instance_masks
]
elif instance_masks is None and keypoints is not None:
visualize_boxes_fn = functools.partial(
_visualize_boxes_and_keypoints,
category_index=category_index,
**visualization_keyword_args)
elems = [
true_shapes, original_shapes, images, boxes, classes, scores, keypoints
]
elif instance_masks is not None and keypoints is not None:
visualize_boxes_fn = functools.partial(
_visualize_boxes_and_masks_and_keypoints,
category_index=category_index,
**visualization_keyword_args)
elems = [
true_shapes, original_shapes, images, boxes, classes, scores,
instance_masks, keypoints
]
else:
visualize_boxes_fn = functools.partial(
_visualize_boxes,
category_index=category_index,
**visualization_keyword_args)
elems = [
true_shapes, original_shapes, images, boxes, classes, scores
]
def draw_boxes(image_and_detections):
"""Draws boxes on image."""
true_shape = image_and_detections[0]
original_shape = image_and_detections[1]
if true_image_shape is not None:
image = shape_utils.pad_or_clip_nd(
image_and_detections[2], [true_shape[0], true_shape[1], 3])
if original_image_spatial_shape is not None:
image_and_detections[2] = _resize_original_image(image, original_shape)
image_with_boxes = tf.compat.v1.py_func(visualize_boxes_fn,
image_and_detections[2:], tf.uint8)
return image_with_boxes
images = tf.map_fn(draw_boxes, elems, dtype=tf.uint8, back_prop=False)
return images
def _build_tables(self, prior):
"""Computes integer-valued probability tables used by the range coder.
These tables must not be re-generated independently on the sending and
receiving side, since small numerical discrepancies between both sides can
occur in this process. If the tables differ slightly, this in turn would
very likely cause catastrophic error propagation during range decoding. For
a more in-depth discussion of this, see:
> "Integer Networks for Data Compression with Latent-Variable Models"<br />
> J. Ballé, N. Johnston, D. Minnen<br />
> https://openreview.net/forum?id=S1zz2i0cY7
The tables are stored in `tf.Variable`s as attributes of this object. The
recommended way is to train the model, instantiate an entropy model with
`compression=True`, and then distribute the model to a sender and a
receiver.
Arguments:
prior: The `tfp.distributions.Distribution` object (see initializer).
"""
offset = helpers.quantization_offset(prior)
lower_tail = helpers.lower_tail(prior, self.tail_mass)
upper_tail = helpers.upper_tail(prior, self.tail_mass)
# Largest distance observed between lower tail and median, and between
# median and upper tail.
minima = offset - lower_tail
minima = tf.cast(tf.math.ceil(minima), tf.int32)
minima = tf.math.maximum(minima, 0)
maxima = upper_tail - offset
maxima = tf.cast(tf.math.ceil(maxima), tf.int32)
maxima = tf.math.maximum(maxima, 0)
# PMF starting positions and lengths.
pmf_start = offset - tf.cast(minima, self.dtype)
pmf_length = maxima + minima + 1
# Sample the densities in the computed ranges, possibly computing more
# samples than necessary at the upper end.
max_length = tf.math.reduce_max(pmf_length)
if tf.executing_eagerly() and max_length > 2048:
logging.warning(
"Very wide PMF with %d elements may lead to out of memory issues. "
"Consider priors with smaller dispersion or increasing `tail_mass` "
"parameter.", int(max_length))
samples = tf.range(tf.cast(max_length, self.dtype), dtype=self.dtype)
samples = tf.reshape(samples, [-1] + len(self.prior_shape) * [1])
samples += pmf_start
pmf = prior.prob(samples)
# Collapse batch dimensions of distribution.
pmf = tf.reshape(pmf, [max_length, -1])
pmf = tf.transpose(pmf)
pmf_length = tf.broadcast_to(pmf_length, self.prior_shape)
pmf_length = tf.reshape(pmf_length, [-1])
cdf_length = pmf_length + 2
cdf_offset = tf.broadcast_to(-minima, self.prior_shape)
cdf_offset = tf.reshape(cdf_offset, [-1])
# Prevent tensors from bouncing back and forth between host and GPU.
with tf.device("/cpu:0"):
def loop_body(args):
prob, length = args
prob = prob[:length]
prob = tf.concat(
[prob, 1 - tf.reduce_sum(prob, keepdims=True)], axis=0)
cdf = range_coding_ops.pmf_to_quantized_cdf(
prob, precision=self.range_coder_precision)
return tf.pad(cdf, [[0, max_length - length]],
mode="CONSTANT",
constant_values=0)
# TODO(jonycgn,ssjhv): Consider switching to Python control flow.
cdf = tf.map_fn(loop_body, (pmf, pmf_length),
dtype=tf.int32,
name="pmf_to_cdf")
self._cdf = tf.Variable(cdf, trainable=False, name="cdf")
self._cdf_offset = tf.Variable(cdf_offset,
trainable=False,
name="cdf_offset")
self._cdf_length = tf.Variable(cdf_length,
trainable=False,
name="cdf_length")
def draw_sample(num_samples, num_classes, logits, num_trials, dtype, seed):
"""Sample a multinomial.
The batch shape is given by broadcasting num_trials with
remove_last_dimension(logits).
Args:
num_samples: Python int or singleton integer Tensor: number of multinomial
samples to draw.
num_classes: Python int or singleton integer Tensor: number of classes.
logits: Floating Tensor with last dimension k, of (unnormalized) logit
probabilities per class.
num_trials: Tensor of number of categorical trials each multinomial consists
of. num_trials[..., tf.newaxis] must broadcast with logits.
dtype: dtype at which to emit samples.
seed: Random seed.
Returns:
samples: Tensor of given dtype and shape [n] + batch_shape + [k].
"""
with tf.name_scope('draw_sample'):
# broadcast the num_trials and logits to same shape
num_trials = tf.ones_like(
logits[..., 0], dtype=num_trials.dtype) * num_trials
logits = tf.ones_like(
num_trials[..., tf.newaxis], dtype=logits.dtype) * logits
# flatten the total_count and logits
# flat_logits has shape [B1B2...Bm, num_classes]
flat_logits = tf.reshape(logits, [-1, num_classes])
flat_num_trials = num_samples * tf.reshape(num_trials, [-1]) # [B1B2...Bm]
# Computes each logits and num_trials situation by map_fn.
# Using just one batch tf.random.categorical call doesn't work because that
# requires num_trials to be the same across all members of the batch of
# logits. This restriction makes sense for tf.random.categorical because
# for it, num_trials is part of the returned shape. However, the
# multinomial sampler does not need that restriction, because it sums out
# exactly that dimension.
# One possibility would be to draw a batch categorical whose sample count is
# max(num_trials) and mask out the excess ones. However, if the elements of
# num_trials vary widely, this can be wasteful of memory.
# TODO(b/123763054, b/112152209): Revisit the possibility of writing this
# with a batch categorical followed by batch unsorted_segment_sum, once both
# of those work and are memory-efficient enough.
def _sample_one_batch_member(args):
logits, num_cat_samples = args[0], args[1] # [K], []
# x has shape [1, num_cat_samples = num_samples * num_trials]
x = tf.random.categorical(
logits[tf.newaxis, ...], num_cat_samples, seed=seed)
x = tf.reshape(x, shape=[num_samples, -1]) # [num_samples, num_trials]
x = tf.one_hot(
x, depth=num_classes) # [num_samples, num_trials, num_classes]
x = tf.reduce_sum(x, axis=-2) # [num_samples, num_classes]
return tf.cast(x, dtype=dtype)
if seed is not None:
# Force parallel_iterations to 1 to ensure reproducibility
# b/139210489
x = tf.map_fn(
_sample_one_batch_member, [flat_logits, flat_num_trials],
fn_output_signature=dtype, # [B1B2...Bm, num_samples, num_classes]
parallel_iterations=1)
else:
# Invoke default parallel_iterations behavior
x = tf.map_fn(
_sample_one_batch_member, [flat_logits, flat_num_trials],
fn_output_signature=dtype) # [B1B2...Bm, num_samples, num_classes]
# reshape the results to proper shape
x = tf.transpose(a=x, perm=[1, 0, 2])
final_shape = tf.concat(
[[num_samples], tf.shape(num_trials), [num_classes]], axis=0)
x = tf.reshape(x, final_shape)
return x
def estimate_reward_ci(self,
dataset: dataset_lib.OffpolicyDataset,
target_policy: tf_policy.TFPolicy,
episode_limit: Optional[int] = None,
num_grid: Optional[int] = 100,
eps: Optional[float] = 1e-6,
num_bootstraps: Optional[int] = 10000,
num_bootstrap_samples: Optional[int] = 10000):
"""Estimate the confidence interval of reward."""
is_weighted_reward_samples = self.get_is_weighted_reward_samples(
dataset, target_policy, episode_limit)
episodes, valid_steps = dataset.get_all_episodes(limit=episode_limit)
num_episodes = tf.shape(valid_steps)[0]
max_abs_reward = tf.reduce_max(
tf.where(valid_steps, tf.abs(self._reward_fn(episodes)), 0.))
# mean estimate
center = self.estimate_average_reward(dataset, target_policy)
delta_tail_half = self._delta_tail / 2.0
num_episodes_float = tf.cast(num_episodes, tf.float32)
if self._ci_method == 'CH': # Chernoff-Hoeffding
width = max_abs_reward * tf.math.sqrt(
tf.math.log(1.0 / delta_tail_half) / num_episodes_float)
lb = center - width
ub = center + width
elif self._ci_method == 'BE': # Empirical Bernstein
constant_term = 7 * max_abs_reward * tf.math.log(
2.0 / delta_tail_half) / (3 * (num_episodes_float - 1))
variance_term = tf.reduce_sum(
tf.square(is_weighted_reward_samples - center))
variance_term *= tf.math.log(
2.0 / delta_tail_half) / (num_episodes_float - 1)
width = constant_term + tf.math.sqrt(
variance_term) / num_episodes_float
lb = center - width
ub = center + width
elif self._ci_method == 'C-BE': # Clipped empirical Bernstein
# need to learn c
def compute_center_width(c_const):
"""Compute the center and width of CI."""
c_vec = c_const * tf.ones_like(is_weighted_reward_samples)
c_is_weighted_reward_samples = tf.minimum(
is_weighted_reward_samples, c_vec) / c_vec
constant_term = 7 * num_episodes_float * tf.math.log(
2.0 / delta_tail_half) / (3 * (num_episodes_float - 1))
center = tf.reduce_sum(
c_is_weighted_reward_samples) / tf.reduce_sum(1.0 / c_vec)
variance_term = tf.reduce_sum(
tf.square(c_is_weighted_reward_samples - center))
variance_term *= tf.math.log(
2.0 / delta_tail_half) / (num_episodes_float - 1)
width = (constant_term + tf.math.sqrt(variance_term)
) / tf.reduce_sum(1.0 / c_vec)
return center, width
def compute_bdd(c_const):
center, width = compute_center_width(c_const)
return center - width, center + width
def compute_obj(c_const, obj='width'):
center, width = compute_center_width(c_const)
if obj == 'lb':
return center - width
elif obj == 'ub': # minimize ub
return -(center + width)
elif obj == 'width':
return width
elif obj == 'lb_ub':
return -2 * width
else:
ValueError('Objective is not implemented')
c_grid = tf.linspace(eps, max_abs_reward, num_grid)
objs = tf.map_fn(compute_obj, c_grid, dtype=tf.float32)
star_index = tf.argmax(objs)
c_star = tf.gather(c_grid, star_index)
lb, ub = compute_bdd(c_star)
elif self._ci_method == 'TT': # Student-t test
# Two-tailed confidence intervals
t_statistic_quantile = stats.t.ppf(1 - delta_tail_half,
num_episodes_float - 1)
std_term = tf.math.sqrt(
tf.reduce_sum(tf.square(is_weighted_reward_samples - center)) /
(num_episodes_float - 1))
width = t_statistic_quantile * std_term / tf.math.sqrt(
num_episodes_float)
lb = center - width
ub = center + width
elif self._ci_method == 'BCa': # Bootstrap
# see references
# https://faculty.washington.edu/heagerty/Courses/b572/public/GregImholte-1.pdf
# http://users.stat.umn.edu/~helwig/notes/bootci-Notes.pdf
gaussian_rv = tfp.distributions.Normal(loc=0, scale=1)
def _compute_bootstrap_lb_ub(reward_samples):
"""Compute Efron's bootstrap lb."""
sample_mean = tf.reduce_mean(reward_samples)
# Step 1, sample with replacement and compute subsampled mean
uniform_log_prob = tf.tile(
tf.expand_dims(tf.zeros(num_episodes), 0),
[num_bootstraps, 1])
ind = tf.random.categorical(uniform_log_prob,
num_bootstrap_samples)
bootstrap_subsamples = tf.gather(reward_samples, ind)
subsample_means = tf.reduce_mean(bootstrap_subsamples, axis=1)
# Step 2, sort subsample means, compute y, z_0, and a
sorted_subsample_means = tf.sort(subsample_means,
axis=0,
direction='ASCENDING')
# bias factor
z_0 = gaussian_rv.quantile(
tf.reduce_sum(
tf.cast(
tf.greater(sample_mean, sorted_subsample_means),
tf.float32)) / float(num_bootstraps))
# y is the leave-one-out, jackknife sample mean
mask_matrix = tf.ones([num_episodes, num_episodes
]) - tf.eye(num_episodes)
leave_one_out_subsample_sums = tf.einsum(
'j,jk->k', reward_samples, mask_matrix)
ys = leave_one_out_subsample_sums / (num_episodes_float - 1)
# average of jackknife estimate
y_bar = tf.reduce_mean(ys)
# acceleration factor
d_ys = y_bar - ys
a = tf.reduce_sum(tf.pow(d_ys, 3.0)) / tf.maximum(
eps, 6.0 * tf.pow(tf.reduce_sum(tf.pow(d_ys, 2.0)), 1.5))
# Step 3, compute z_scores for lb and ub
z_score_delta_tail = gaussian_rv.quantile(delta_tail_half)
z_score_1_delta_tail = gaussian_rv.quantile(1.0 -
delta_tail_half)
z_lb = z_0 + (z_score_delta_tail + z_0) / tf.maximum(
eps, 1 - a * (z_score_delta_tail + z_0))
z_ub = z_0 + (z_score_1_delta_tail + z_0) / tf.maximum(
eps, 1 - a * (z_score_1_delta_tail + z_0))
# Step 4, compute corresponding quantiles and get bootstrap intervals
lb_index = tf.cast(
tf.maximum(
tf.minimum(
tf.floor(num_bootstraps * gaussian_rv.cdf(z_lb)),
num_bootstraps - 1), 1), tf.int64)
ub_index = tf.cast(
tf.maximum(
tf.minimum(
tf.floor(num_bootstraps * gaussian_rv.cdf(z_ub)),
num_bootstraps - 1), 1), tf.int64)
lb = tf.gather(sorted_subsample_means, lb_index)
ub = tf.gather(sorted_subsample_means, ub_index)
return lb, ub
lb, ub = _compute_bootstrap_lb_ub(is_weighted_reward_samples)
else:
ValueError('Confidence interval is not implemented!')
return [lb, ub]
