def flatten_list_of_vars(var_list):
flat_vars = [tf.reshape(v, -1) for v in var_list]
return tf.concat(flat_vars, axis=-1)
def _sample_shape(self): return tf.concat([self.batch_shape_tensor(), self.event_shape_tensor()], -1)
def crop_mask_in_target_box(masks, boxes, target_boxes, output_size):
"""Crop masks in target boxes.
Args:
masks: A tensor with a shape of [batch_size, num_masks, height, width].
boxes: a float tensor representing box cooridnates that tightly enclose
masks with a shape of [batch_size, num_masks, 4] in un-normalized
coordinates. A box is represented by [ymin, xmin, ymax, xmax].
target_boxes: a float tensor representing target box cooridnates for
masks with a shape of [batch_size, num_masks, 4] in un-normalized
coordinates. A box is represented by [ymin, xmin, ymax, xmax].
output_size: A scalar to indicate the output crop size. It currently only
supports to output a square shape outputs.
Returns:
A 4-D tensor representing feature crop of shape
[batch_size, num_boxes, output_size, output_size].
"""
with tf.name_scope('crop_mask_in_target_box'):
batch_size, num_masks, height, width = masks.get_shape().as_list()
masks = tf.reshape(masks, [batch_size * num_masks, height, width, 1])
# Pad zeros on the boundary of masks.
masks = tf.image.pad_to_bounding_box(masks, 2, 2, height + 4,
width + 4)
masks = tf.reshape(masks,
[batch_size, num_masks, height + 4, width + 4, 1])
# Projects target box locations and sizes to corresponding cropped
# mask coordinates.
gt_y_min, gt_x_min, gt_y_max, gt_x_max = tf.split(value=boxes,
num_or_size_splits=4,
axis=2)
bb_y_min, bb_x_min, bb_y_max, bb_x_max = tf.split(value=target_boxes,
num_or_size_splits=4,
axis=2)
y_transform = (bb_y_min - gt_y_min) * height / (gt_y_max - gt_y_min +
_EPSILON) + 2
x_transform = (bb_x_min - gt_x_min) * height / (gt_x_max - gt_x_min +
_EPSILON) + 2
h_transform = (bb_y_max - bb_y_min) * width / (gt_y_max - gt_y_min +
_EPSILON)
w_transform = (bb_x_max - bb_x_min) * width / (gt_x_max - gt_x_min +
_EPSILON)
boundaries = tf.concat([
tf.cast(tf.ones_like(y_transform) * ((height + 4) - 1),
dtype=tf.float32),
tf.cast(tf.ones_like(x_transform) * ((width + 4) - 1),
dtype=tf.float32)
],
axis=-1)
# Reshape tensors to have the right shape for selective_crop_and_resize.
trasnformed_boxes = tf.concat(
[y_transform, x_transform, h_transform, w_transform], -1)
levels = tf.tile(tf.reshape(tf.range(num_masks), [1, num_masks]),
[batch_size, 1])
cropped_masks = selective_crop_and_resize(masks,
trasnformed_boxes,
levels,
boundaries,
output_size,
sample_offset=0)
cropped_masks = tf.squeeze(cropped_masks, axis=-1)
return cropped_masks
def concat_vectors(*args):
"""Concatenates input vectors, statically if possible."""
args_ = [tf.get_static_value(x) for x in args]
if any(vec is None for vec in args_):
return tf.concat(args, axis=0)
return [val for vec in args_ for val in vec]
def _queue_push(queue, should_update, new_vecs):
"""Conditionally push new vectors into a batch of first-in-first-out queues.
The `queue` of shape `[k, ..., n]` can be thought of as a batch of queues,
each holding `k` n-D vectors; while `new_vecs` of shape `[..., n]` is a
fresh new batch of n-D vectors. The `should_update` batch of Boolean scalars,
i.e. shape `[...]`, indicates batch members whose corresponding n-D vector in
`new_vecs` should be added at the back of its queue, pushing out the
corresponding n-D vector from the front. Batch members in `new_vecs` for
which `should_update` is False are ignored.
Note: the choice of placing `k` at the dimension 0 of the queue is
constrained by the L-BFGS two-loop algorithm above. The algorithm uses
tf.scan to iterate over the `k` correction pairs simulatneously across all
batches, and tf.scan itself can only iterate over dimension 0.
For example:
```python
k, b, n = (3, 2, 5)
queue = tf.reshape(tf.range(30), (k, b, n))
# => [[[ 0, 1, 2, 3, 4],
# [ 5, 6, 7, 8, 9]],
#
# [[10, 11, 12, 13, 14],
# [15, 16, 17, 18, 19]],
#
# [[20, 21, 22, 23, 24],
# [25, 26, 27, 28, 29]]]
element = tf.reshape(tf.range(30, 40), (b, n))
# => [[30, 31, 32, 33, 34],
[35, 36, 37, 38, 39]]
should_update = tf.constant([True, False]) # Shape: (b,)
_queue_add(should_update, queue, element)
# => [[[10, 11, 12, 13, 14],
# [ 5, 6, 7, 8, 9]],
#
# [[20, 21, 22, 23, 24],
# [15, 16, 17, 18, 19]],
#
# [[30, 31, 32, 33, 34],
# [25, 26, 27, 28, 29]]]
```
Args:
queue: A `tf.Tensor` of shape `[k, ..., n]`; a batch of queues each with
`k` n-D vectors.
should_update: A Boolean `tf.Tensor` of shape `[...]` indicating batch
members where new vectors should be added to their queues.
new_vecs: A `tf.Tensor` of shape `[..., n]`; a batch of n-D vectors to add
at the end of their respective queues, pushing out the first element from
each.
Returns:
A new `tf.Tensor` of shape `[k, ..., n]`.
"""
new_queue = tf.concat([queue[1:], [new_vecs]], axis=0)
return tf.where(should_update[tf.newaxis, ..., tf.newaxis], new_queue,
queue)
def fill_triangular(x, upper=False, name=None):
"""Creates a (batch of) triangular matrix from a vector of inputs.
Created matrix can be lower- or upper-triangular. (It is more efficient to
create the matrix as upper or lower, rather than transpose.)
Triangular matrix elements are filled in a clockwise spiral. See example,
below.
If `x.shape` is `[b1, b2, ..., bB, d]` then the output shape is
`[b1, b2, ..., bB, n, n]` where `n` is such that `d = n(n+1)/2`, i.e.,
`n = int(np.sqrt(0.25 + 2. * m) - 0.5)`.
Example:
```python
fill_triangular([1, 2, 3, 4, 5, 6])
# ==> [[4, 0, 0],
# [6, 5, 0],
# [3, 2, 1]]
fill_triangular([1, 2, 3, 4, 5, 6], upper=True)
# ==> [[1, 2, 3],
# [0, 5, 6],
# [0, 0, 4]]
```
The key trick is to create an upper triangular matrix by concatenating `x`
and a tail of itself, then reshaping.
Suppose that we are filling the upper triangle of an `n`-by-`n` matrix `M`
from a vector `x`. The matrix `M` contains n**2 entries total. The vector `x`
contains `n * (n+1) / 2` entries. For concreteness, we'll consider `n = 5`
(so `x` has `15` entries and `M` has `25`). We'll concatenate `x` and `x` with
the first (`n = 5`) elements removed and reversed:
```python
x = np.arange(15) + 1
xc = np.concatenate([x, x[5:][::-1]])
# ==> array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 15, 14, 13,
# 12, 11, 10, 9, 8, 7, 6])
# (We add one to the arange result to disambiguate the zeros below the
# diagonal of our upper-triangular matrix from the first entry in `x`.)
# Now, when reshapedlay this out as a matrix:
y = np.reshape(xc, [5, 5])
# ==> array([[ 1, 2, 3, 4, 5],
# [ 6, 7, 8, 9, 10],
# [11, 12, 13, 14, 15],
# [15, 14, 13, 12, 11],
# [10, 9, 8, 7, 6]])
# Finally, zero the elements below the diagonal:
y = np.triu(y, k=0)
# ==> array([[ 1, 2, 3, 4, 5],
# [ 0, 7, 8, 9, 10],
# [ 0, 0, 13, 14, 15],
# [ 0, 0, 0, 12, 11],
# [ 0, 0, 0, 0, 6]])
```
From this example we see that the resuting matrix is upper-triangular, and
contains all the entries of x, as desired. The rest is details:
- If `n` is even, `x` doesn't exactly fill an even number of rows (it fills
`n / 2` rows and half of an additional row), but the whole scheme still
works.
- If we want a lower triangular matrix instead of an upper triangular,
we remove the first `n` elements from `x` rather than from the reversed
`x`.
For additional comparisons, a pure numpy version of this function can be found
in `distribution_util_test.py`, function `_fill_triangular`.
Args:
x: `Tensor` representing lower (or upper) triangular elements.
upper: Python `bool` representing whether output matrix should be upper
triangular (`True`) or lower triangular (`False`, default).
name: Python `str`. The name to give this op.
Returns:
tril: `Tensor` with lower (or upper) triangular elements filled from `x`.
Raises:
ValueError: if `x` cannot be mapped to a triangular matrix.
"""
with tf.name_scope(name or 'fill_triangular'):
x = tf.convert_to_tensor(x, name='x')
m = tf.compat.dimension_value(
tensorshape_util.with_rank_at_least(x.shape, 1)[-1])
if m is not None:
# Formula derived by solving for n: m = n(n+1)/2.
m = np.int32(m)
n = np.sqrt(0.25 + 2. * m) - 0.5
if n != np.floor(n):
raise ValueError(
'Input right-most shape ({}) does not '
'correspond to a triangular matrix.'.format(m))
n = np.int32(n)
static_final_shape = tensorshape_util.concatenate(
x.shape[:-1], [n, n])
else:
m = tf.shape(x)[-1]
# For derivation, see above. Casting automatically lops off the 0.5, so we
# omit it. We don't validate n is an integer because this has
# graph-execution cost; an error will be thrown from the reshape, below.
n = tf.cast(tf.sqrt(0.25 + tf.cast(2 * m, dtype=tf.float32)),
dtype=tf.int32)
static_final_shape = tensorshape_util.concatenate(
tensorshape_util.with_rank_at_least(x.shape, 1)[:-1],
[None, None])
# Try it out in numpy:
# n = 3
# x = np.arange(n * (n + 1) / 2)
# m = x.shape[0]
# n = np.int32(np.sqrt(.25 + 2 * m) - .5)
# x_tail = x[(m - (n**2 - m)):]
# np.concatenate([x_tail, x[::-1]], 0).reshape(n, n) # lower
# # ==> array([[3, 4, 5],
# [5, 4, 3],
# [2, 1, 0]])
# np.concatenate([x, x_tail[::-1]], 0).reshape(n, n) # upper
# # ==> array([[0, 1, 2],
# [3, 4, 5],
# [5, 4, 3]])
#
# Note that we can't simply do `x[..., -(n**2 - m):]` because this doesn't
# correctly handle `m == n == 1`. Hence, we do nonnegative indexing.
# Furthermore observe that:
# m - (n**2 - m)
# = n**2 / 2 + n / 2 - (n**2 - n**2 / 2 + n / 2)
# = 2 (n**2 / 2 + n / 2) - n**2
# = n**2 + n - n**2
# = n
ndims = prefer_static.rank(x)
if upper:
x_list = [x, tf.reverse(x[..., n:], axis=[ndims - 1])]
else:
x_list = [x[..., n:], tf.reverse(x, axis=[ndims - 1])]
new_shape = (tensorshape_util.as_list(static_final_shape)
if tensorshape_util.is_fully_defined(static_final_shape)
else tf.concat([tf.shape(x)[:-1], [n, n]], axis=0))
x = tf.reshape(tf.concat(x_list, axis=-1), new_shape)
x = tf.linalg.band_part(x,
num_lower=(0 if upper else -1),
num_upper=(-1 if upper else 0))
tensorshape_util.set_shape(x, static_final_shape)
return x
def _test_finding_sparse_solution(self, batch_shape=None):
# Test that Proximal Hessian descent prefers sparse solutions when
# l1_regularizer is large enough.
#
# Define
#
# Loss(x) := (x[0] - a[0])**2 + epsilon * sum(
# (x[i] - a[i])**2 for i in range(1, n))
#
# where `a` is a constant and epsilon is small. Set
# l2_regularizer = 0 and set l1_regularizer such that
#
# epsilon << l1_regularizer << 1.
#
# L1 regularization should cause the computed optimum to have zeros in all
# but the 0th coordinate: optimal_x ~= [a[0], 0, ..., 0].
n = 10
epsilon = 1e-6
if batch_shape is None:
batch_shape = []
# Set a[0] explicitly to make sure it's not very close to zero
a0 = 6.
a_ = np.concatenate([
np.full(batch_shape + [1], a0),
np.random.random(size=batch_shape + [n - 1])
],
axis=-1)
a = self._adjust_dtype_and_shape_hints(a_)
def _grad_and_hessian_unregularized_loss_fn(x):
diff = x - a
grad = 2. * tf.concat([diff[..., :1], epsilon * diff[..., 1:]],
axis=-1)
hessian_outer = tf.SparseTensor(
indices=[
b + (i, i) for i in range(n)
for b in np.ndindex(*batch_shape)
],
values=tf.ones(shape=[np.prod(batch_shape, dtype=int) * n],
dtype=self.dtype),
dense_shape=batch_shape + [n, n])
hessian_middle_per_batch = 2 * tf.concat(
[[1.], epsilon * tf.ones([n - 1], dtype=self.dtype)], axis=0)
hessian_middle = tf.zeros(
batch_shape + [n], dtype=self.dtype) + hessian_middle_per_batch
return grad, hessian_outer, hessian_middle
w, is_converged, num_iter = tfp.optimizer.proximal_hessian_sparse_minimize(
_grad_and_hessian_unregularized_loss_fn,
x_start=tf.zeros(batch_shape + [n], dtype=self.dtype),
l1_regularizer=1e-2,
l2_regularizer=None,
maximum_iterations=10,
maximum_full_sweeps_per_iteration=10,
tolerance=1e-5,
learning_rate=1.)
w_, is_converged_, _ = self.evaluate([w, is_converged, num_iter])
expected_w = tf.concat(
[a[..., :1],
tf.zeros(batch_shape + [n - 1], self.dtype)], axis=-1)
# Using atol=0 ensures that w must be exactly zero in all coordinates
# where expected_w is exactly zero.
self.assertAllEqual(is_converged_, True)
self.assertAllClose(w_, expected_w, atol=0., rtol=1e-3)
def _sample_n(self, n, seed=None): probs = self._probs_parameter_no_checks() new_shape = tf.concat([[n], tf.shape(probs)], 0) uniform = samplers.uniform(new_shape, seed=seed, dtype=probs.dtype) sample = tf.less(uniform, probs) return tf.cast(sample, self.dtype)
def _forward(self, x):
with tf.control_dependencies(self._assertions(x)):
y0 = x[..., :1]
yk = tf.math.log(x[..., 1:] - x[..., :-1])
y = tf.concat([y0, yk], axis=-1)
return y
def l1norm(x):
x = tf.concat([x, tf.ones_like(x[..., :1]) * 1e-6], axis=-1)
x = x / tf.linalg.norm(x, ord=1, axis=-1, keepdims=True)
return x
def _validate_sample_arg(self, x):
"""Helper which validates sample arg, e.g., input to `log_prob`."""
with tf.name_scope('validate_sample_arg'):
x_ndims = (
tf.rank(x) if tensorshape_util.rank(x.shape) is None else
tensorshape_util.rank(x.shape))
event_ndims = (
tf.size(self.event_shape_tensor())
if tensorshape_util.rank(self.event_shape) is None else
tensorshape_util.rank(self.event_shape))
batch_ndims = (
tf.size(self._batch_shape_unexpanded)
if tensorshape_util.rank(self.batch_shape) is None else
tensorshape_util.rank(self.batch_shape))
expected_batch_event_ndims = batch_ndims + event_ndims
if (isinstance(x_ndims, int) and
isinstance(expected_batch_event_ndims, int)):
if x_ndims < expected_batch_event_ndims:
raise NotImplementedError(
'Broadcasting is not supported; too few batch and event dims '
'(expected at least {}, saw {}).'.format(
expected_batch_event_ndims, x_ndims))
ndims_assertion = []
elif self.validate_args:
ndims_assertion = [
assert_util.assert_greater_equal(
x_ndims,
expected_batch_event_ndims,
message=('Broadcasting is not supported; too few '
'batch and event dims.'),
name='assert_batch_and_event_ndims_large_enough'),
]
if (tensorshape_util.is_fully_defined(self.batch_shape) and
tensorshape_util.is_fully_defined(self.event_shape)):
expected_batch_event_shape = np.int32(
tensorshape_util.concatenate(self.batch_shape, self.event_shape))
else:
expected_batch_event_shape = tf.concat(
[
self.batch_shape_tensor(),
self.event_shape_tensor(),
], axis=0)
sample_ndims = x_ndims - expected_batch_event_ndims
if isinstance(sample_ndims, int):
sample_ndims = max(sample_ndims, 0)
if (isinstance(sample_ndims, int) and
tensorshape_util.is_fully_defined(x.shape[sample_ndims:])):
actual_batch_event_shape = np.int32(x.shape[sample_ndims:])
else:
sample_ndims = tf.maximum(sample_ndims, 0)
actual_batch_event_shape = tf.shape(x)[sample_ndims:]
if (isinstance(expected_batch_event_shape, np.ndarray) and
isinstance(actual_batch_event_shape, np.ndarray)):
if any(expected_batch_event_shape != actual_batch_event_shape):
raise NotImplementedError('Broadcasting is not supported; '
'unexpected batch and event shape '
'(expected {}, saw {}).'.format(
expected_batch_event_shape,
actual_batch_event_shape))
# We need to set the final runtime-assertions to `ndims_assertion` since
# its possible this assertion was created. We could add a condition to
# only do so if `self.validate_args == True`, however this is redundant
# as `ndims_assertion` already encodes this information.
runtime_assertions = ndims_assertion
elif self.validate_args:
# We need to make the `ndims_assertion` a control dep because otherwise
# TF itself might raise an exception owing to this assertion being
# ill-defined, ie, one cannot even compare different rank Tensors.
with tf.control_dependencies(ndims_assertion):
shape_assertion = assert_util.assert_equal(
expected_batch_event_shape,
actual_batch_event_shape,
message=('Broadcasting is not supported; '
'unexpected batch and event shape.'),
name='assert_batch_and_event_shape_same')
runtime_assertions = [shape_assertion]
else:
runtime_assertions = []
return runtime_assertions
def testLSTMCellRank1BatchEnsemble(self, alpha_initializer, gamma_initializer,
recurrent_alpha_initializer,
recurrent_gamma_initializer,
bias_initializer, use_bias, implementation,
use_additive_perturbation):
tf.keras.backend.set_learning_phase(1) # training time
ensemble_size = 4
examples_per_model = 4
input_dim = 5
output_dim = 5
inputs = tf.random.normal([examples_per_model, input_dim])
batched_inputs = tf.tile(inputs, [ensemble_size, 1])
layer = rank1_bnn_layers.LSTMCellRank1(
output_dim,
use_bias=use_bias,
alpha_initializer=alpha_initializer,
gamma_initializer=gamma_initializer,
recurrent_alpha_initializer=recurrent_alpha_initializer,
recurrent_gamma_initializer=recurrent_gamma_initializer,
bias_initializer=bias_initializer,
alpha_regularizer=None,
gamma_regularizer=None,
recurrent_alpha_regularizer=None,
recurrent_gamma_regularizer=None,
implementation=implementation,
use_additive_perturbation=use_additive_perturbation,
ensemble_size=ensemble_size)
h0 = tf.random.normal([examples_per_model, output_dim])
c0 = tf.random.normal([examples_per_model, output_dim])
def compute_rank1_lstm_cell(i):
if use_additive_perturbation:
ifgo = tf.linalg.matmul(
inputs + layer.alpha[i], layer.kernel) + layer.gamma[i]
ifgo += tf.linalg.matmul(
h0 + layer.recurrent_alpha[i],
layer.recurrent_kernel) + layer.recurrent_gamma[i]
else:
ifgo = tf.linalg.matmul(
inputs * layer.alpha[i], layer.kernel) * layer.gamma[i]
ifgo += tf.linalg.matmul(
h0 * layer.recurrent_alpha[i],
layer.recurrent_kernel) * layer.recurrent_gamma[i]
if use_bias:
ifgo += layer.bias[i]
i, f, g, o = tf.split(ifgo, num_or_size_splits=4, axis=1)
i = tf.nn.sigmoid(i)
f = tf.nn.sigmoid(f)
g = tf.nn.tanh(g)
o = tf.nn.sigmoid(o)
c = f*c0 + i*g
h = o * tf.nn.tanh(c)
return h
h0_batched = tf.tile(h0, [ensemble_size, 1])
c0_batched = tf.tile(c0, [ensemble_size, 1])
outputs, _ = layer(batched_inputs, (h0_batched, c0_batched))
manual_outputs = tf.concat(
[compute_rank1_lstm_cell(i) for i in range(ensemble_size)], axis=0)
expected_shape = (ensemble_size*examples_per_model, output_dim)
self.assertEqual(outputs.shape, expected_shape)
self.assertAllClose(outputs, manual_outputs)
layer2 = rank1_bnn_layers.LSTMCellRank1.from_config(layer.get_config())
layer2(batched_inputs, (h0_batched, c0_batched)) # force initialization
layer2.set_weights(layer.get_weights())
outputs2, _ = layer2(batched_inputs, (h0_batched, c0_batched))
self.assertAllClose(outputs, outputs2)
def _batch_of_zeros_with_rightmost_singletons(n_singletons):
"""Return Tensor of zeros with some singletons on the rightmost dims."""
ones = tf.ones(shape=[n_singletons], dtype=tf.int32)
return tf.zeros(shape=tf.concat([batch_shape, ones], axis=0),
dtype=dtype)
def step_fn(inputs):
"""Per-Replica StepFn."""
images, labels = inputs
if FLAGS.ensemble_size > 1:
images = tf.tile(images, [FLAGS.ensemble_size, 1, 1, 1])
labels = tf.tile(labels, [FLAGS.ensemble_size])
with tf.GradientTape() as tape:
logits = model(images, training=True)
if FLAGS.use_bfloat16:
logits = tf.cast(logits, tf.float32)
negative_log_likelihood = tf.reduce_mean(
tf.keras.losses.sparse_categorical_crossentropy(
labels, logits, from_logits=True))
filtered_variables = []
for var in model.trainable_variables:
# Apply l2 on the BN parameters and bias terms. This
# excludes only fast weight approximate posterior/prior parameters,
# but pay caution to their naming scheme.
if ('kernel' in var.name or 'batch_norm' in var.name
or 'bias' in var.name):
filtered_variables.append(tf.reshape(var, (-1, )))
l2_loss = FLAGS.l2 * 2 * tf.nn.l2_loss(
tf.concat(filtered_variables, axis=0))
kl = sum(model.losses) / train_dataset_size
kl_scale = tf.cast(optimizer.iterations + 1, kl.dtype)
kl_scale /= steps_per_epoch * FLAGS.kl_annealing_epochs
kl_scale = tf.minimum(1., kl_scale)
kl_loss = kl_scale * kl
# Scale the loss given the TPUStrategy will reduce sum all gradients.
loss = negative_log_likelihood + l2_loss + kl_loss
scaled_loss = loss / strategy.num_replicas_in_sync
grads = tape.gradient(scaled_loss, model.trainable_variables)
# Separate learning rate implementation.
if FLAGS.fast_weight_lr_multiplier != 1.0:
grads_and_vars = []
for grad, var in zip(grads, model.trainable_variables):
# Apply different learning rate on the fast weight approximate
# posterior/prior parameters. This is excludes BN and slow weights,
# but pay caution to the naming scheme.
if ('kernel' not in var.name
and 'batch_norm' not in var.name
and 'bias' not in var.name):
grads_and_vars.append(
(grad * FLAGS.fast_weight_lr_multiplier, var))
else:
grads_and_vars.append((grad, var))
optimizer.apply_gradients(grads_and_vars)
else:
optimizer.apply_gradients(zip(grads,
model.trainable_variables))
probs = tf.nn.softmax(logits)
metrics['train/ece'].update_state(labels, probs)
metrics['train/loss'].update_state(loss)
metrics['train/negative_log_likelihood'].update_state(
negative_log_likelihood)
metrics['train/kl'].update_state(kl)
metrics['train/kl_scale'].update_state(kl_scale)
metrics['train/accuracy'].update_state(labels, logits)
def fill_triangular_inverse(x, upper=False, name=None):
"""Creates a vector from a (batch of) triangular matrix.
The vector is created from the lower-triangular or upper-triangular portion
depending on the value of the parameter `upper`.
If `x.shape` is `[b1, b2, ..., bB, n, n]` then the output shape is
`[b1, b2, ..., bB, d]` where `d = n (n + 1) / 2`.
Example:
```python
fill_triangular_inverse(
[[4, 0, 0],
[6, 5, 0],
[3, 2, 1]])
# ==> [1, 2, 3, 4, 5, 6]
fill_triangular_inverse(
[[1, 2, 3],
[0, 5, 6],
[0, 0, 4]], upper=True)
# ==> [1, 2, 3, 4, 5, 6]
```
Args:
x: `Tensor` representing lower (or upper) triangular elements.
upper: Python `bool` representing whether output matrix should be upper
triangular (`True`) or lower triangular (`False`, default).
name: Python `str`. The name to give this op.
Returns:
flat_tril: (Batch of) vector-shaped `Tensor` representing vectorized lower
(or upper) triangular elements from `x`.
"""
with tf.name_scope(name or 'fill_triangular_inverse'):
x = tf.convert_to_tensor(x, name='x')
n = tf.compat.dimension_value(
tensorshape_util.with_rank_at_least(x.shape, 2)[-1])
if n is not None:
n = np.int32(n)
m = np.int32((n * (n + 1)) // 2)
static_final_shape = tensorshape_util.concatenate(
x.shape[:-2], [m])
else:
n = tf.shape(x)[-1]
m = (n * (n + 1)) // 2
static_final_shape = tensorshape_util.concatenate(
tensorshape_util.with_rank_at_least(x.shape, 2)[:-2], [None])
ndims = prefer_static.rank(x)
if upper:
initial_elements = x[..., 0, :]
triangular_portion = x[..., 1:, :]
else:
initial_elements = tf.reverse(x[..., -1, :], axis=[ndims - 2])
triangular_portion = x[..., :-1, :]
rotated_triangular_portion = tf.reverse(tf.reverse(triangular_portion,
axis=[ndims - 1]),
axis=[ndims - 2])
consolidated_matrix = triangular_portion + rotated_triangular_portion
end_sequence = tf.reshape(
consolidated_matrix,
tf.concat([tf.shape(x)[:-2], [n * (n - 1)]], axis=0))
y = tf.concat([initial_elements, end_sequence[..., :m - n]], axis=-1)
tensorshape_util.set_shape(y, static_final_shape)
return y
def _inverse(self, y):
x0 = y[..., :1]
xk = tf.exp(y[..., 1:])
x = tf.concat([x0, xk], axis=-1)
return tf.cumsum(x, axis=-1)
def lu_solve(lower_upper, perm, rhs, validate_args=False, name=None):
"""Solves systems of linear eqns `A X = RHS`, given LU factorizations.
Note: this function does not verify the implied matrix is actually invertible
nor is this condition checked even when `validate_args=True`.
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)`.
rhs: Matrix-shaped float `Tensor` representing targets for which to solve;
`A X = RHS`. To handle vector cases, use:
`lu_solve(..., rhs[..., tf.newaxis])[..., 0]`.
validate_args: Python `bool` indicating whether arguments should be checked
for correctness. Note: this function does not verify the implied matrix is
actually invertible, even when `validate_args=True`.
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_solve').
Returns:
x: The `X` in `A @ X = RHS`.
#### Examples
```python
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
x = [[[1., 2],
[3, 4]],
[[7, 8],
[3, 4]]]
inv_x = tfp.math.lu_solve(*tf.linalg.lu(x), rhs=tf.eye(2))
tf.assert_near(tf.matrix_inverse(x), inv_x)
# ==> True
```
"""
with tf.name_scope(name or 'lu_solve'):
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')
rhs = tf.convert_to_tensor(rhs,
dtype_hint=lower_upper.dtype,
name='rhs')
assertions = _lu_solve_assertions(lower_upper, perm, rhs,
validate_args)
if assertions:
with tf.control_dependencies(assertions):
lower_upper = tf.identity(lower_upper)
perm = tf.identity(perm)
rhs = tf.identity(rhs)
if (tensorshape_util.rank(rhs.shape) == 2
and tensorshape_util.rank(perm.shape) == 1):
# Both rhs and perm have scalar batch_shape.
permuted_rhs = tf.gather(rhs, perm, axis=-2)
else:
# Either rhs or perm have non-scalar batch_shape or we can't determine
# this information statically.
rhs_shape = tf.shape(rhs)
broadcast_batch_shape = tf.broadcast_dynamic_shape(
rhs_shape[:-2],
tf.shape(perm)[:-1])
d, m = rhs_shape[-2], rhs_shape[-1]
rhs_broadcast_shape = tf.concat([broadcast_batch_shape, [d, m]],
axis=0)
# Tile out rhs.
broadcast_rhs = tf.broadcast_to(rhs, rhs_broadcast_shape)
broadcast_rhs = tf.reshape(broadcast_rhs, [-1, d, m])
# Tile out perm and add batch indices.
broadcast_perm = tf.broadcast_to(perm, rhs_broadcast_shape[:-1])
broadcast_perm = tf.reshape(broadcast_perm, [-1, d])
broadcast_batch_size = tf.reduce_prod(broadcast_batch_shape)
broadcast_batch_indices = tf.broadcast_to(
tf.range(broadcast_batch_size)[:, tf.newaxis],
[broadcast_batch_size, d])
broadcast_perm = tf.stack(
[broadcast_batch_indices, broadcast_perm], axis=-1)
permuted_rhs = tf.gather_nd(broadcast_rhs, broadcast_perm)
permuted_rhs = tf.reshape(permuted_rhs, rhs_broadcast_shape)
lower = tf.linalg.set_diag(
tf.linalg.band_part(lower_upper, num_lower=-1, num_upper=0),
tf.ones(tf.shape(lower_upper)[:-1], dtype=lower_upper.dtype))
return linear_operator_util.matrix_triangular_solve_with_broadcast(
lower_upper, # Only upper is accessed.
linear_operator_util.matrix_triangular_solve_with_broadcast(
lower, permuted_rhs),
lower=False)
def compute_deterministic_alignment_loss(embs, steps, seq_lens, num_steps,
batch_size, loss_type,
similarity_type, temperature,
label_smoothing, variance_lambda,
huber_delta, normalize_indices):
"""Compute cycle-consistency loss for all steps in each sequence.
This aligns each pair of videos in the batch except with itself.
When aligning it also matters which video is the starting video. So for N
videos in the batch, we have N * (N-1) alignments happening.
For example, a batch of size 3 has 6 pairs of sequence alignments.
Args:
embs: Tensor, sequential embeddings of the shape [N, T, D] where N is the
batch size, T is the number of timesteps in the sequence, D is the size
of the embeddings.
steps: Tensor, step indices/frame indices of the embeddings of the shape
[N, T] where N is the batch size, T is the number of the timesteps.
seq_lens: Tensor, Lengths of the sequences from which the sampling was
done. This can provide additional information to the alignment loss.
num_steps: Integer/Tensor, Number of timesteps in the embeddings.
batch_size: Integer, Size of the batch.
loss_type: String, This specifies the kind of loss function to use.
Currently supported loss functions: 'classification', 'regression_mse',
'regression_mse_var', 'regression_huber'.
similarity_type: String, Currently supported similarity metrics: 'l2' ,
'cosine' .
temperature: Float, temperature scaling used to scale the similarity
distributions calculated using the softmax function.
label_smoothing: Float, Label smoothing argument used in
tf.keras.losses.categorical_crossentropy function and described in this
paper https://arxiv.org/pdf/1701.06548.pdf.
variance_lambda: Float, Weight of the variance of the similarity
predictions while cycling back. If this is high then the low variance
similarities are preferred by the loss while making this term low
results in high variance of the similarities (more uniform/random
matching).
huber_delta: float, Huber delta described in tf.keras.losses.huber_loss.
normalize_indices: Boolean, If True, normalizes indices by sequence
lengths. Useful for ensuring numerical instabilities doesn't arise as
sequence indices can be large numbers.
Returns:
loss: Tensor, Scalar loss tensor that imposes the chosen variant of the
cycle-consistency loss.
"""
labels_list = []
logits_list = []
steps_list = []
seq_lens_list = []
for i in range(batch_size):
for j in range(batch_size):
# We do not align the sequence with itself.
if i != j:
logits, labels = align_pair_of_sequences(
embs[i], embs[j], similarity_type, temperature)
logits_list.append(logits)
labels_list.append(labels)
steps_list.append(tf.tile(steps[i:i + 1], [num_steps, 1]))
seq_lens_list.append(tf.tile(seq_lens[i:i + 1], [num_steps]))
logits = tf.concat(logits_list, axis=0)
# print('----------------', logits)
labels = tf.concat(labels_list, axis=0)
steps = tf.concat(steps_list, axis=0)
seq_lens = tf.concat(seq_lens_list, axis=0)
if loss_type == 'classification':
loss = classification_loss(logits, labels, label_smoothing)
elif 'regression' in loss_type:
loss = regression_loss(logits, labels, num_steps, steps, seq_lens,
loss_type, normalize_indices, variance_lambda,
huber_delta)
else:
raise ValueError('Unidentified loss_type %s. Currently supported loss '
'types are: regression_mse, regression_huber, '
'classification.' % loss_type)
return loss
def _joint_sample_n(self, n, seed=None):
"""Draw a joint sample from the prior over latents and observations.
This sampler is specific to LocalLinearTrend models and is faster than the
generic LinearGaussianStateSpaceModel implementation.
Args:
n: `int` `Tensor` number of samples to draw.
seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
Returns:
latents: `float` `Tensor` of shape `concat([[n], self.batch_shape,
[self.num_timesteps, self.latent_size]], axis=0)` representing samples
of latent trajectories.
observations: `float` `Tensor` of shape `concat([[n], self.batch_shape,
[self.num_timesteps, self.observation_size]], axis=0)` representing
samples of observed series generated from the sampled `latents`.
"""
with tf.name_scope('joint_sample_n'):
(initial_state_seed,
level_jumps_seed,
slope_jumps_seed,
prior_observation_seed) = samplers.split_seed(
seed, n=4, salt='LocalLinearTrendStateSpaceModel_joint_sample_n')
if self.batch_shape.is_fully_defined():
batch_shape = self.batch_shape
else:
batch_shape = self.batch_shape_tensor()
sample_and_batch_shape = ps.cast(
ps.concat([[n], batch_shape], axis=0), tf.int32)
# Sample the initial timestep from the prior. Since we want
# this sample to have full batch shape (not just the batch shape
# of the self.initial_state_prior object which might in general be
# smaller), we augment the sample shape to include whatever
# extra batch dimensions are required.
initial_level_and_slope = self.initial_state_prior.sample(
linear_gaussian_ssm._augment_sample_shape( # pylint: disable=protected-access
self.initial_state_prior,
sample_and_batch_shape,
self.validate_args),
seed=initial_state_seed)
# Sample the latent random walk on slopes.
jumps_shape = ps.concat([sample_and_batch_shape,
[self.num_timesteps - 1]], axis=0)
slope_jumps = samplers.normal(
jumps_shape,
dtype=self.dtype,
seed=slope_jumps_seed) * self.slope_scale[..., tf.newaxis]
prior_slope_sample = tf.cumsum(
tf.concat([initial_level_and_slope[..., 1:],
slope_jumps],
axis=-1),
axis=-1)
# Sample latent levels, given latent slopes.
level_jumps = samplers.normal(
jumps_shape,
dtype=self.dtype,
seed=level_jumps_seed) * self.level_scale[..., tf.newaxis]
prior_level_sample = tf.cumsum(
tf.concat([initial_level_and_slope[..., :1],
level_jumps + prior_slope_sample[..., :-1]],
axis=-1),
axis=-1)
# Sample noisy observations, given latent levels.
prior_observation_sample = prior_level_sample + (
samplers.normal(ps.shape(prior_level_sample),
dtype=self.dtype,
seed=prior_observation_seed) *
self.observation_noise_scale[..., tf.newaxis])
return (tf.stack([prior_level_sample,
prior_slope_sample],
axis=-1),
prior_observation_sample[..., tf.newaxis])
def _sample_n(self, n, seed=None): shape = tf.concat([[n], self.batch_shape_tensor()], 0) samples = tf.random.uniform(shape=shape, dtype=self.dtype, seed=seed) return self.low + self.range() * samples
def true_fn():
old_shape = tf.shape(result_t)
new_shape = tf.concat([tf.ones(ndmin - ndims, tf.int32), old_shape],
axis=0)
return tf.reshape(result_t, new_shape)
def _sample_n(self, n, seed=None):
dim0_seed, otherdims_seed = samplers.split_seed(
seed, salt='von_mises_fisher')
# The sampling strategy relies on the fact that vMF variates are symmetric
# about the mean direction. Accordingly, if we have a sampling strategy for
# the away-from-mean angle, then we can uniformly sample the remaining
# dimensions on the S^{dim-2} sphere for , and rotate these samples from a
# (1, 0, 0, ..., 0)-mode distribution into the target orientation.
#
# This is easy to imagine on the 1-sphere (S^1; in 2-D space): sample a
# von-Mises distributed `x` value in [-1, 1], then uniformly select what
# amounts to a "up" or "down" additional degree of freedom after unit
# normalizing, followed by a final rotation to the desired mean direction
# from a basis of (1, 0).
#
# On S^2 (in 3-D), selecting a vMF `x` identifies a circle in `yz` on the
# unit sphere over which the distribution is uniform, in particular the
# circle where x = \hat{x} intersects the unit sphere. We pick a point on
# that circle, then rotate to the desired mean direction from a basis of
# (1, 0, 0).
mean_direction = tf.convert_to_tensor(self.mean_direction)
concentration = tf.convert_to_tensor(self.concentration)
event_dim = (
tf.compat.dimension_value(self.event_shape[0])
or self._event_shape_tensor(mean_direction=mean_direction)[0])
sample_batch_shape = ps.concat(
[[n],
self._batch_shape_tensor(mean_direction=mean_direction,
concentration=concentration)],
axis=0)
dim = tf.cast(event_dim - 1, self.dtype)
if event_dim == 3:
samples_dim0 = self._sample_3d(n,
mean_direction=mean_direction,
concentration=concentration,
seed=dim0_seed)
else:
# Wood'94 provides a rejection algorithm to sample the x coordinate.
# Wood'94 definition of b:
# b = (-2 * kappa + tf.sqrt(4 * kappa**2 + dim**2)) / dim
# https://stats.stackexchange.com/questions/156729 suggests:
b = dim / (2 * concentration +
tf.sqrt(4 * concentration**2 + dim**2))
# TODO(bjp): Integrate any useful numerical tricks from hyperspherical VAE
# https://github.com/nicola-decao/s-vae-tf/
x = (1 - b) / (1 + b)
c = concentration * x + dim * tf.math.log1p(-x**2)
beta = beta_lib.Beta(dim / 2, dim / 2)
def cond_fn(w, should_continue, seed):
del w, seed
return tf.reduce_any(should_continue)
def body_fn(w, should_continue, seed):
"""While loop body for sampling the angle `w`."""
beta_seed, unif_seed, next_seed = samplers.split_seed(seed,
n=3)
z = beta.sample(sample_shape=sample_batch_shape,
seed=beta_seed)
# set_shape needed here because of b/139013403
tensorshape_util.set_shape(z, w.shape)
w = tf.where(should_continue,
(1. - (1. + b) * z) / (1. - (1. - b) * z), w)
if not self.allow_nan_stats:
w = tf.debugging.check_numerics(w, 'w')
unif = samplers.uniform(sample_batch_shape,
seed=unif_seed,
dtype=self.dtype)
# set_shape needed here because of b/139013403
tensorshape_util.set_shape(unif, w.shape)
should_continue = should_continue & (
concentration * w + dim * tf.math.log1p(-x * w) - c <
# Use log1p(-unif) to prevent log(0) and ensure that log(1) is
# possible.
tf.math.log1p(-unif))
return w, should_continue, next_seed
w = tf.zeros(sample_batch_shape, dtype=self.dtype)
should_continue = tf.ones(sample_batch_shape, dtype=tf.bool)
samples_dim0, _, _ = tf.while_loop(cond=cond_fn,
body=body_fn,
loop_vars=(w, should_continue,
dim0_seed))
samples_dim0 = samples_dim0[..., tf.newaxis]
if not self._allow_nan_stats:
# Verify samples are w/in -1, 1, with useful error output tensors (top
# value rather than all values).
with tf.control_dependencies([
assert_util.assert_less_equal(
samples_dim0,
dtype_util.as_numpy_dtype(self.dtype)(1.01)),
assert_util.assert_greater_equal(
samples_dim0,
dtype_util.as_numpy_dtype(self.dtype)(-1.01)),
]):
samples_dim0 = tf.identity(samples_dim0)
samples_otherdims_shape = ps.concat(
[sample_batch_shape, [event_dim - 1]], axis=0)
unit_otherdims = tf.math.l2_normalize(samplers.normal(
samples_otherdims_shape, seed=otherdims_seed, dtype=self.dtype),
axis=-1)
samples = tf.concat(
[
samples_dim0, # we must avoid sqrt(1 - (>1)**2)
tf.sqrt(tf.maximum(1 - samples_dim0**2, 0.)) * unit_otherdims
],
axis=-1)
samples = tf.math.l2_normalize(samples, axis=-1)
if not self.allow_nan_stats:
samples = tf.debugging.check_numerics(samples, 'samples')
# Runtime assert that samples are unit length.
if not self.allow_nan_stats:
worst, _ = tf.math.top_k(
tf.reshape(tf.abs(1 - tf.linalg.norm(samples, axis=-1)), [-1]))
with tf.control_dependencies([
assert_util.assert_near(dtype_util.as_numpy_dtype(
self.dtype)(0),
worst,
atol=1e-4,
summarize=100)
]):
samples = tf.identity(samples)
# The samples generated are symmetric around a mode at (1, 0, 0, ...., 0).
# Now, we move the mode to `self.mean_direction` using a rotation matrix.
if not self.allow_nan_stats:
# Assert that the basis vector rotates to the mean direction, as expected.
basis = tf.cast(
tf.concat([[1.], tf.zeros([event_dim - 1])], axis=0),
self.dtype)
with tf.control_dependencies([
assert_util.assert_less(
tf.linalg.norm(self._rotate(
basis, mean_direction=mean_direction) -
mean_direction,
axis=-1),
dtype_util.as_numpy_dtype(self.dtype)(1e-5))
]):
return self._rotate(samples, mean_direction=mean_direction)
return self._rotate(samples, mean_direction=mean_direction)
def _mean(self):
shape = tf.concat([
self.batch_shape_tensor(),
self.event_shape_tensor(),
], 0)
return tf.broadcast_to(self.loc, shape)
def add_points_bboxes(points_xyz, bboxes, bboxes_label, bboxes_mask,
is_2d=False):
"""Assign bboxes to points."""
one_bbox = bboxes[-1:, :]
one_bbox_mask = tf.zeros([1])
bboxes = bboxes[bboxes_mask == 1]
bboxes_mask = bboxes_mask[bboxes_mask == 1]
bboxes = tf.concat([bboxes, one_bbox], axis=0)
bboxes_mask = tf.concat([bboxes_mask, one_bbox_mask], axis=0)
theta = bboxes[:, 6]
rz = tf.stack(
[tf.cos(theta), -tf.sin(theta), tf.zeros_like(theta),
tf.sin(theta), tf.cos(theta), tf.zeros_like(theta),
tf.zeros_like(theta), tf.zeros_like(theta), tf.ones_like(theta)],
axis=-1)
rz = tf.reshape(rz, [-1, 3, 3])
points_xyz = tf.reshape(points_xyz, [-1, 1, 3])
points_xyz_in_bbox_frame = (tf.reshape(points_xyz, [-1, 1, 3])-
tf.reshape(bboxes[:, 0:3], [1, -1, 3]))
# points_xyz -> (n, m, 1, 3) * (1, m, 3, 3) -> (n, m, 1, 3)
points_xyz_in_bbox_frame = tf.expand_dims(points_xyz_in_bbox_frame, axis=2)
rz = tf.expand_dims(rz, axis=0)
# (n, m, 3)
points_xyz_in_bbox_frame = tf.squeeze(tf.matmul(points_xyz_in_bbox_frame, rz),
axis=2)
# (1, m, 3)
bboxes_size_min = tf.expand_dims(-bboxes[:, 3:6] / 2, axis=0)
bboxes_size_max = tf.expand_dims(bboxes[:, 3:6] / 2, axis=0)
# (n, m)
if is_2d:
points_if_in_bboxes = tf.reduce_all(
(points_xyz_in_bbox_frame[:, :, 0:2] >= bboxes_size_min[:, :, 0:2]) &
(points_xyz_in_bbox_frame[:, :, 0:2] <= bboxes_size_max[:, :, 0:2]),
axis=-1)
points_centerness_left = tf.math.abs(
points_xyz_in_bbox_frame[:, :, 0:2] - bboxes_size_min[:, :, 0:2])
points_centerness_right = tf.math.abs(
bboxes_size_max[:, :, 0:2] - points_xyz_in_bbox_frame[:, :, 0:2])
points_centerness_min = tf.math.minimum(points_centerness_left,
points_centerness_right)
points_centerness_max = tf.math.maximum(points_centerness_left,
points_centerness_right)
points_centerness = tf.math.sqrt(
tf.math.reduce_prod(points_centerness_min / points_centerness_max,
axis=-1))
else:
points_if_in_bboxes = tf.reduce_all(
(points_xyz_in_bbox_frame >= bboxes_size_min) &
(points_xyz_in_bbox_frame <= bboxes_size_max), axis=-1)
points_centerness_left = tf.math.abs(
points_xyz_in_bbox_frame - bboxes_size_min)
points_centerness_right = tf.math.abs(
bboxes_size_max - points_xyz_in_bbox_frame)
points_centerness_min = tf.math.minimum(points_centerness_left,
points_centerness_right)
points_centerness_max = tf.math.maximum(points_centerness_left,
points_centerness_right)
# (N, M)
points_centerness = tf.math.sqrt(
tf.math.reduce_prod(points_centerness_min / points_centerness_max,
axis=-1))
points_if_in_bboxes = tf.cast(points_if_in_bboxes, tf.dtypes.float32)
# n, m
points_if_in_bboxes = points_if_in_bboxes * tf.reshape(bboxes_mask, [1, -1])
# n
points_bboxes_index = tf.argmax(points_if_in_bboxes, axis=1)
points_if_in_bboxes = tf.gather(points_if_in_bboxes,
points_bboxes_index, batch_dims=1)
points_centerness = tf.gather(points_centerness,
points_bboxes_index, batch_dims=1)
points_bboxes = tf.gather(bboxes, points_bboxes_index, axis=0)
points_bboxes_label = tf.gather(bboxes_label, points_bboxes_index, axis=0)
return (points_bboxes, points_bboxes_label,
points_if_in_bboxes, points_centerness, points_bboxes_index)
def _sample_n(self, n, seed=None):
return self.loc + self.scale * tf.random.normal(tf.concat(
[[n], self.batch_shape_tensor()], axis=0),
seed=seed)
def _replicate(n, tensor):
"""Replicate the input tensor n times along a new (major) dimension."""
# TODO(axch) Does this already exist somewhere? Should it get contributed?
multiples = tf.concat([[n], tf.ones([tf.rank(tensor)], dtype=n.dtype)],
axis=0)
return tf.tile(tensor[tf.newaxis], multiples)
def multilevel_crop_and_resize(features, boxes, output_size=7):
"""Crop and resize on multilevel feature pyramid.
Generate the (output_size, output_size) set of pixels for each input box
by first locating the box into the correct feature level, and then cropping
and resizing it using the correspoding feature map of that level.
Args:
features: A dictionary with key as pyramid level and value as features.
The features are in shape of [batch_size, height_l, width_l, num_filters].
boxes: A 3-D Tensor of shape [batch_size, num_boxes, 4]. Each row
represents a box with [y1, x1, y2, x2] in un-normalized coordinates.
output_size: A scalar to indicate the output crop size.
Returns:
A 5-D tensor representing feature crop of shape
[batch_size, num_boxes, output_size, output_size, num_filters].
"""
with tf.name_scope('multilevel_crop_and_resize'):
levels = features.keys()
min_level = min(levels)
max_level = max(levels)
_, max_feature_height, max_feature_width, _ = (
features[min_level].get_shape().as_list())
# Stacks feature pyramid into a features_all of shape
# [batch_size, levels, height, width, num_filters].
features_all = []
for level in range(min_level, max_level + 1):
features_all.append(
tf.image.pad_to_bounding_box(features[level], 0, 0,
max_feature_height,
max_feature_width))
features_all = tf.stack(features_all, axis=1)
# Assigns boxes to the right level.
box_width = boxes[:, :, 3] - boxes[:, :, 1]
box_height = boxes[:, :, 2] - boxes[:, :, 0]
areas_sqrt = tf.sqrt(box_height * box_width)
levels = tf.cast(tf.math.floordiv(
tf.math.log(tf.divide(areas_sqrt, 224.0)), tf.math.log(2.0)) + 4.0,
dtype=tf.int32)
# Maps levels between [min_level, max_level].
levels = tf.minimum(max_level, tf.maximum(levels, min_level))
# Projects box location and sizes to corresponding feature levels.
scale_to_level = tf.cast(tf.pow(tf.constant(2.0),
tf.cast(levels, tf.float32)),
dtype=boxes.dtype)
boxes /= tf.expand_dims(scale_to_level, axis=2)
box_width /= scale_to_level
box_height /= scale_to_level
boxes = tf.concat([
boxes[:, :, 0:2],
tf.expand_dims(box_height, -1),
tf.expand_dims(box_width, -1)
],
axis=-1)
# Maps levels to [0, max_level-min_level].
levels -= min_level
level_strides = tf.pow([[2.0]], tf.cast(levels, tf.float32))
boundary = tf.cast(
tf.concat([
tf.expand_dims([[tf.cast(max_feature_height, tf.float32)]] /
level_strides - 1,
axis=-1),
tf.expand_dims([[tf.cast(max_feature_width, tf.float32)]] /
level_strides - 1,
axis=-1),
],
axis=-1), boxes.dtype)
return selective_crop_and_resize(features_all, boxes, levels, boundary,
output_size)
def sample_lkj(num_samples,
dimension,
concentration,
cholesky_space=False,
seed=None,
name=None):
"""Returns a Tensor of samples from an LKJ distribution.
Args:
num_samples: Python `int`. The number of samples to draw.
dimension: Python `int`. The dimension of correlation matrices.
concentration: `Tensor` representing the concentration of the LKJ
distribution.
cholesky_space: Python `bool`. Whether to take samples from LKJ or
Chol(LKJ).
seed: Python integer seed for RNG
name: Python `str` name prefixed to Ops created by this function.
Returns:
samples: A Tensor of correlation matrices (or Cholesky factors of
correlation matrices if `cholesky_space = True`) with shape
`[n] + B + [D, D]`, where `B` is the shape of the `concentration`
parameter, and `D` is the `dimension`.
Raises:
ValueError: If `dimension` is negative.
"""
if dimension < 0:
raise ValueError(
'Cannot sample negative-dimension correlation matrices.')
# Notation below: B is the batch shape, i.e., tf.shape(concentration)
seed = SeedStream(seed, 'sample_lkj')
with tf.name_scope('sample_lkj' or name):
concentration = tf.convert_to_tensor(concentration)
if not dtype_util.is_floating(concentration.dtype):
raise TypeError(
'The concentration argument should have floating type, not '
'{}'.format(dtype_util.name(concentration.dtype)))
concentration = _replicate(num_samples, concentration)
concentration_shape = tf.shape(concentration)
if dimension <= 1:
# For any dimension <= 1, there is only one possible correlation matrix.
shape = tf.concat([concentration_shape, [dimension, dimension]],
axis=0)
return tf.ones(shape=shape, dtype=concentration.dtype)
beta_conc = concentration + (dimension - 2.) / 2.
beta_dist = beta.Beta(concentration1=beta_conc,
concentration0=beta_conc)
# Note that the sampler below deviates from [1], by doing the sampling in
# cholesky space. This does not change the fundamental logic of the
# sampler, but does speed up the sampling.
# This is the correlation coefficient between the first two dimensions.
# This is also `r` in reference [1].
corr12 = 2. * beta_dist.sample(seed=seed()) - 1.
# Below we construct the Cholesky of the initial 2x2 correlation matrix,
# which is of the form:
# [[1, 0], [r, sqrt(1 - r**2)]], where r is the correlation between the
# first two dimensions.
# This is the top-left corner of the cholesky of the final sample.
first_row = tf.concat([
tf.ones_like(corr12)[..., tf.newaxis],
tf.zeros_like(corr12)[..., tf.newaxis]
],
axis=-1)
second_row = tf.concat(
[corr12[..., tf.newaxis],
tf.sqrt(1 - corr12**2)[..., tf.newaxis]],
axis=-1)
chol_result = tf.concat(
[first_row[..., tf.newaxis, :], second_row[..., tf.newaxis, :]],
axis=-2)
for n in range(2, dimension):
# Loop invariant: on entry, result has shape B + [n, n]
beta_conc = beta_conc - 0.5
# norm is y in reference [1].
norm = beta.Beta(concentration1=n / 2.,
concentration0=beta_conc).sample(seed=seed())
# distance shape: B + [1] for broadcast
distance = tf.sqrt(norm)[..., tf.newaxis]
# direction is u in reference [1].
# direction shape: B + [n]
direction = _uniform_unit_norm(n, concentration_shape,
concentration.dtype, seed)
# raw_correlation is w in reference [1].
raw_correlation = distance * direction # shape: B + [n]
# This is the next row in the cholesky of the result,
# which differs from the construction in reference [1].
# In the reference, the new row `z` = chol_result @ raw_correlation^T
# = C @ raw_correlation^T (where as short hand we use C = chol_result).
# We prove that the below equation is the right row to add to the
# cholesky, by showing equality with reference [1].
# Let S be the sample constructed so far, and let `z` be as in
# reference [1]. Then at this iteration, the new sample S' will be
# [[S z^T]
# [z 1]]
# In our case we have the cholesky decomposition factor C, so
# we want our new row x (same size as z) to satisfy:
# [[S z^T] [[C 0] [[C^T x^T] [[CC^T Cx^T]
# [z 1]] = [x k]] [0 k]] = [xC^t xx^T + k**2]]
# Since C @ raw_correlation^T = z = C @ x^T, and C is invertible,
# we have that x = raw_correlation. Also 1 = xx^T + k**2, so k
# = sqrt(1 - xx^T) = sqrt(1 - |raw_correlation|**2) = sqrt(1 -
# distance**2).
new_row = tf.concat(
[raw_correlation,
tf.sqrt(1. - norm[..., tf.newaxis])],
axis=-1)
# Finally add this new row, by growing the cholesky of the result.
chol_result = tf.concat([
chol_result,
tf.zeros_like(chol_result[..., 0][..., tf.newaxis])
],
axis=-1)
chol_result = tf.concat([chol_result, new_row[..., tf.newaxis, :]],
axis=-2)
if cholesky_space:
return chol_result
result = tf.matmul(chol_result, chol_result, transpose_b=True)
# The diagonal for a correlation matrix should always be ones. Due to
# numerical instability the matmul might not achieve that, so manually set
# these to ones.
result = tf.linalg.set_diag(
result, tf.ones(shape=tf.shape(result)[:-1], dtype=result.dtype))
# This sampling algorithm can produce near-PSD matrices on which standard
# algorithms such as `tf.cholesky` or `tf.linalg.self_adjoint_eigvals`
# fail. Specifically, as documented in b/116828694, around 2% of trials
# of 900,000 5x5 matrices (distributed according to 9 different
# concentration parameter values) contained at least one matrix on which
# the Cholesky decomposition failed.
return result
def _testConvFlipout(self, layer_class): # pylint: disable=invalid-name
batch_size, depth, height, width, channels, filters = 2, 4, 4, 4, 3, 5
with self.cached_session() as sess:
(kernel_posterior, kernel_prior, kernel_divergence, bias_posterior,
bias_prior, bias_divergence, layer, inputs, outputs, kl_penalty,
kernel_shape) = self._testConvSetUp(layer_class,
batch_size,
depth=depth,
height=height,
width=width,
channels=channels,
filters=filters,
seed=44)
tf.random.set_seed(5995)
convolution_op = nn_ops.Convolution(
tf.TensorShape(inputs.shape),
filter_shape=tf.TensorShape(kernel_shape),
padding='SAME',
data_format=tf_layers_util.convert_data_format(
self.data_format, inputs.shape.rank))
expected_kernel_posterior_affine = tfd.Normal(
loc=tf.zeros_like(kernel_posterior.result_loc),
scale=kernel_posterior.result_scale)
expected_kernel_posterior_affine_tensor = (
expected_kernel_posterior_affine.sample(seed=42))
expected_outputs = convolution_op(
inputs, kernel_posterior.distribution.loc)
input_shape = tf.shape(inputs)
batch_shape = tf.expand_dims(input_shape[0], 0)
if self.data_format == 'channels_first':
channels = input_shape[1]
else:
channels = input_shape[-1]
rank = len(inputs.shape) - 2
seed_stream = tfp.util.SeedStream(layer.seed, salt='ConvFlipout')
sign_input = tfp.random.rademacher(tf.concat(
[batch_shape, tf.expand_dims(channels, 0)], 0),
dtype=inputs.dtype,
seed=seed_stream())
sign_output = tfp.random.rademacher(tf.concat(
[batch_shape, tf.expand_dims(filters, 0)], 0),
dtype=inputs.dtype,
seed=seed_stream())
if self.data_format == 'channels_first':
for _ in range(rank):
sign_input = tf.expand_dims(sign_input,
-1) # 2D ex: (B, C, 1, 1)
sign_output = tf.expand_dims(sign_output, -1)
else:
for _ in range(rank):
sign_input = tf.expand_dims(sign_input,
1) # 2D ex: (B, 1, 1, C)
sign_output = tf.expand_dims(sign_output, 1)
perturbed_inputs = convolution_op(
inputs * sign_input, expected_kernel_posterior_affine_tensor)
perturbed_inputs *= sign_output
expected_outputs += perturbed_inputs
expected_outputs = tf.nn.bias_add(
expected_outputs,
bias_posterior.result_sample,
data_format=tf_layers_util.convert_data_format(
self.data_format, 4))
[
expected_outputs_,
actual_outputs_,
expected_kernel_divergence_,
actual_kernel_divergence_,
expected_bias_,
actual_bias_,
expected_bias_divergence_,
actual_bias_divergence_,
] = sess.run([
expected_outputs,
outputs,
kernel_divergence.result,
kl_penalty[0],
bias_posterior.result_sample,
layer.bias_posterior_tensor,
bias_divergence.result,
kl_penalty[1],
])
self.assertAllClose(expected_bias_, actual_bias_, rtol=1e-6)
self.assertAllClose(expected_outputs_, actual_outputs_, rtol=1e-6)
self.assertAllClose(expected_kernel_divergence_,
actual_kernel_divergence_,
rtol=1e-6)
self.assertAllClose(expected_bias_divergence_,
actual_bias_divergence_,
rtol=1e-6)
expected_args = [kernel_posterior, kernel_prior, None]
# We expect that there was one call to kernel_divergence, with the above
# args; MockKLDivergence appends the list of args to a list, so the above
# args should be in the 0th position of that list.
actual_args = kernel_divergence.args[0]
# Test for identity with 'is'. TensorFlowTestCase.assertAllEqual actually
# coerces the inputs to numpy arrays, so we can't use that to assert that
# the arguments (which are a mixture of Distributions and Tensors) are
# equal.
for a, b in zip(expected_args, actual_args):
self.assertIs(a, b)
# Same story as above.
expected_args = [
bias_posterior, bias_prior, bias_posterior.result_sample
]
actual_args = bias_divergence.args[0]
for a, b in zip(expected_args, actual_args):
self.assertIs(a, b)
def _generate_schedule(self, cpn_frequency, roll_convention):
"""Method to generate coupon dates.
Args:
cpn_frequency: A `PeriodTensor` specifying the frequency of coupon
payments.
roll_convention: Scalar of type `BusinessDayConvention` specifying how
dates are rolled if they fall on holidays.
Returns:
A tuple containing the generated date schedule and a boolean `Tensor`
of the same shape as the schedule specifying whether the coupons are
regular coupons.
"""
if (self._first_coupon_date is None) and (self._penultimate_coupon_date
is None):
cpn_dates = dates.PeriodicSchedule(
start_date=self._start_date,
end_date=self._end_date,
tenor=cpn_frequency,
roll_convention=roll_convention).dates()
is_regular_cpn = tf.constant(True,
dtype=bool,
shape=cpn_dates[:, :-1].shape)
elif self._first_coupon_date is not None:
cpn_dates = dates.PeriodicSchedule(
start_date=self._first_coupon_date,
end_date=self._end_date,
tenor=cpn_frequency,
roll_convention=roll_convention).dates()
cpn_dates = dates.DateTensor.concat(
[self._start_date.expand_dims(-1), cpn_dates], axis=1)
is_irregular_cpn = tf.constant(False,
dtype=bool,
shape=self._start_date.shape)
is_regular_cpn = tf.concat([
tf.expand_dims(is_irregular_cpn, axis=-1),
tf.constant(True, dtype=bool, shape=cpn_dates[:, :-2].shape)
],
axis=1)
else:
cpn_dates = dates.PeriodicSchedule(
start_date=self._start_date,
end_date=self._penultimate_coupon_date,
backward=True,
tenor=cpn_frequency,
roll_convention=roll_convention).dates()
cpn_dates = dates.DateTensor.concat(
[cpn_dates, self._end_date.expand_dims(-1)], axis=1)
is_irregular_cpn = tf.constant(False,
dtype=bool,
shape=self._end_date.shape)
is_regular_cpn = tf.concat([
tf.constant(True, dtype=bool, shape=cpn_dates[:, :-2].shape),
tf.expand_dims(is_irregular_cpn, axis=-1)
],
axis=1)
return cpn_dates, is_regular_cpn
