in neural_structured_learning/lib/utils.py [0:0]
def random_in_norm_ball(tensor_structure, radius, norm_type):
"""Generates a random sample in a norm ball, conforming the given structure.
This function expects `tensor_structure` to represent a batch of examples,
where the tensors in the structure are features. The norms are calculated in
the whole feature space (including all tensors).
For L-inf norm, each feature dimension can be sampled independently from a
uniform distribution. For L2 and L1 norms, we use the result (Corollary 4 in
Barthe et al., 2005) that the first `n` coordinates of a uniformly random
point on the `p`-norm unit sphere in `R^(n+p)` are uniformly distributed in
the `p`-norm unit ball in `R^n`. For L2 norm, we draw random points on the
unit sphere from the standard normal distribution which density is
proportional to `exp(-x^2)`(Corollary 1 in Harman and Lacko, 2010). For L1
norm, we draw random points on the 1-norm unit sphere from the standard
Laplace distribution which density is proportional to `exp(-|x|)`. More
details on the papers: https://arxiv.org/abs/math/0503650 and
https://doi.org/10.1016/j.jmva.2010.06.002
Args:
tensor_structure: A (nested) collection of batched, floating-point tensors.
Only the structure (number of tensors, shape, and dtype) is used, not the
values. The first dimension of each tensor is the batch size and must all
be equal.
radius: The radius of the norm ball to sample from.
norm_type: One of `nsl.configs.NormType` which specifies the norm of the
norm ball.
Returns:
A (nested) collection of tensors in the same structure as `tensor_structure`
but has its values drawn randomly. The norm of each example in the batch
is less than or equal to the given radius.
"""
tensors = tf.nest.flatten(tensor_structure)
if not all(t.dtype.is_floating for t in tensors):
raise ValueError('random_in_norm_ball() only supports floating tensors.')
# Reject unless the set of first dimensions is {N}, {None}, or {N, None} for
# some number N.
if len(set(t.shape.as_list()[0] for t in tensors) - set([None])) > 1:
raise ValueError('The first dimension of all tensors should be the same.')
if isinstance(norm_type, six.string_types): # Convert string to NormType.
norm_type = configs.NormType(norm_type)
if norm_type == configs.NormType.INFINITY:
return tf.nest.map_structure(
lambda t: tf.random.uniform(t.shape, -radius, radius, t.dtype),
tensor_structure)
elif norm_type == configs.NormType.L2:
# Sample each coordinate with Normal(0, 1).
samples = [tf.random.normal(t.shape, dtype=t.dtype) for t in tensors]
squared_norm = _reduce_across_tensors(tf.reduce_sum,
[tf.square(t) for t in samples])
# Squared L2 norm of two extra coordinates:
# Normal(0, 1)^2 + Normal(0, 1)^2 ~ Chi^2(2) ~ Exponential(0.5)
# Exponential(lambda) ~ Gamma(1, beta=lambda)
extra = tf.random.gamma(squared_norm.shape, 1.0, 0.5, squared_norm.dtype)
scale = radius / tf.math.sqrt(squared_norm + extra)
return tf.nest.pack_sequence_as(
tensor_structure,
[t * _expand_to_rank(scale, t.shape.rank) for t in samples])
elif norm_type == configs.NormType.L1:
# Sample each coordinate w/ Laplace(0, 1) ~ Exponential(1) - Exponential(1)
samples = []
for t in tensors:
samples.append(
tf.random.gamma(t.shape, 1., 1., t.dtype) -
tf.random.gamma(t.shape, 1., 1., t.dtype))
l1_norm = _reduce_across_tensors(tf.reduce_sum,
[tf.abs(t) for t in samples])
# L1 norm of one extra coordinate: |Laplace(0, 1)| ~ Exponential(1).
extra = tf.random.gamma(l1_norm.shape, 1.0, 1.0, l1_norm.dtype)
scale = radius / (l1_norm + extra)
return tf.nest.pack_sequence_as(
tensor_structure,
[t * _expand_to_rank(scale, t.shape.rank) for t in samples])
else:
raise ValueError(f'NormType not supported: {norm_type}')