def test_variable_v2(self):
with tf.variable_scope('', use_resource=False):
relu = self.build_model()
gamma_tensor, _ = self.get_gamma(relu)
# Check that maybe_convert_to_variable ignores VariableV2 (i.e., is no op).
self.assertEqual(tpu_util.maybe_convert_to_variable(gamma_tensor),
gamma_tensor)
def __init__(self, weight_tensor, reduce_dims, threshold, l1_fraction=0.0):
"""Creates an instance.
Args:
weight_tensor: A tensor with the weights of the op (potentially sliced).
reduce_dims: A tuple indictaing the dimensions of `weight_tensor`
to reduce over. Most often it will include all dimensions except
the output size.
threshold: A float. When the norm of the group associated with an
activation is below the threshold, it will be considered dead.
l1_fraction: A float, controls the balance between L1 and L2 grouping (see
above).
"""
weight_tensor = tpu_util.maybe_convert_to_variable(weight_tensor)
if l1_fraction < 0.0 or l1_fraction > 1.0:
raise ValueError('l1_fraction should be in [0.0, 1.0], not %e.' %
l1_fraction)
self._threshold = threshold
l2_norm = tf.sqrt(
tf.reduce_mean(tf.square(weight_tensor), axis=reduce_dims))
if l1_fraction > 0.0:
l1_norm = tf.reduce_mean(tf.abs(weight_tensor), axis=reduce_dims)
norm = l1_fraction * l1_norm + (1.0 - l1_fraction) * l2_norm
else:
norm = l2_norm
self._regularization_vector = norm
self._alive_vector = norm > threshold
def create_regularizer(self, op_slice):
"""Create a regularizer for this logistic-sigmoid gating OpSlice.
Args:
op_slice: op_regularizer_manager.OpSlice that is a batch norm OpSlice.
Returns:
OpRegularizer for this batch norm op.
"""
start_index = op_slice.slice.start_index
size = op_slice.slice.size
logits = op_slice.op.inputs[0] # Input 0 are the logits.
mask = op_slice.op.outputs[0]
# If OpSlice size matches tensor size, use the entire tensor. Otherwise,
# slice the tensor accordingly.
mask = tpu_util.read_from_variable(mask)
if start_index == 0 and size == logits.shape.as_list()[-1]:
return prob_gating_regularizer.ProbGatingRegularizer(
logits,
mask,
regularize_on_mask=self._regularize_on_mask,
alive_threshold=self._alive_threshold,
mask_as_alive_vector=self._mask_as_alive_vector)
else:
logits = tpu_util.maybe_convert_to_variable(logits)
return prob_gating_regularizer.ProbGatingRegularizer(
logits[start_index:start_index + size],
mask[start_index:start_index + size],
regularize_on_mask=self._regularize_on_mask,
alive_threshold=self._alive_threshold,
mask_as_alive_vector=self._mask_as_alive_vector)
def __init__(self, logits, mask, regularize_on_mask=True,
alive_threshold=0.1, mask_as_alive_vector=True):
"""Creates an instance.
Args:
logits: A tf.Tensor of shape (n_channels,) with the
log odds ratio of the channel being on: log(p / 1-p).
mask: A tf.Tensor of the same shape as `logits`.
The sampled mask/gating vector.
regularize_on_mask: Bool. If True uses the mask as the
regularization vector. Else uses probabilities. Default True.
alive_threshold: Float. Threshold below which values are considered dead.
This can be used both when mask_as_alive_vector is True and then the
threshold is used to binarize the sampled values and
when mask_as_alive_vector is False, and then the threshold is on the
channel probability.
mask_as_alive_vector: Bool. If True use the thresholded sampled mask
as the alive vector. Else, use thresholded probabilities from the
logits.
"""
if len(logits.shape.as_list()) != 1:
raise ValueError('logits tensor should be 1D.')
if len(mask.shape.as_list()) != 1:
raise ValueError('mask tensor should be 1D.')
self._logits = tpu_util.maybe_convert_to_variable(logits)
self._mask = mask
self._probs = tf.sigmoid(self._logits)
alive_vector = self._mask if mask_as_alive_vector else self._probs
self._alive_vector = alive_vector >= alive_threshold
self._regularization_vector = (
self._mask if regularize_on_mask else self._probs)
def create_regularizer(self, op_slice):
"""Create a regularizer for this conv2d OpSlice.
Args:
op_slice: op_regularizer_manager.OpSlice that is a conv2d OpSlice.
Returns:
OpRegularizer for this conv2d op.
"""
start_index = op_slice.slice.start_index
size = op_slice.slice.size
weights = op_slice.op.inputs[1] # Input 1 are the weights.
weights = tpu_util.maybe_convert_to_variable(weights)
reduce_dims = self._reduce_dims(op_slice.op)
rank = len(weights.shape.as_list())
if rank != len(reduce_dims) + 1:
raise ValueError('Rank %d incompatible with reduce_dims %s for op %s' %
(rank, reduce_dims, op_slice.op.name))
def _slice_weights():
"""Slices the weight tensor according to op_slice information."""
if rank == 2:
if reduce_dims[0] == 0:
return weights[:, start_index:start_index + size]
else:
return weights[start_index:start_index + size, :]
if rank == 3:
if 2 not in reduce_dims:
return weights[:, :, start_index:start_index + size]
if 1 not in reduce_dims:
return weights[:, start_index:start_index + size, :]
if 0 not in reduce_dims:
return weights[start_index:start_index + size, :, :]
if rank == 4:
if 3 not in reduce_dims:
return weights[:, :, :, start_index:start_index + size]
if 2 not in reduce_dims:
return weights[:, :, start_index:start_index + size, :]
if 1 not in reduce_dims:
return weights[:, start_index:start_index + size, :, :]
if 0 not in reduce_dims:
return weights[start_index:start_index + size, :, :, :]
if rank == 5:
if 4 not in reduce_dims:
return weights[:, :, :, :, start_index:start_index + size]
raise ValueError('Unsupported reduce_dim for rank 5 tensors (Conv3D)')
raise ValueError('Unsupported rank or bad reduce_dim')
weight_tensor = _slice_weights()
# If OpSlice size matches tensor size, use the entire tensor. Otherwise,
# slice the tensor accordingly.
return group_lasso_regularizer.GroupLassoRegularizer(
weight_tensor=weight_tensor,
reduce_dims=self._reduce_dims(op_slice.op),
threshold=self._threshold,
l1_fraction=self._l1_fraction)
def test_resource_variable(self):
with tf.variable_scope('', use_resource=True):
relu = self.build_model()
gamma_tensor, gamma_source_op = self.get_gamma(relu)
variable = tpu_util.maybe_convert_to_variable(gamma_tensor)
# First assert that we didn't return the original tensor
self.assertNotEqual(variable, gamma_tensor)
# Now check that the variable created by maybe_convert_to_variable is
# driven by the same op as the tensor passed as input.
self.assertEqual(variable.op, gamma_source_op)
# If input tensor is separated from a variable by an extra hop of Identity,
# maybe_read_variable pretends the Identity op isn't there.
identity_tensor = tf.identity(gamma_tensor)
self.assertEqual(tpu_util.maybe_convert_to_variable(identity_tensor),
variable)
def __init__(self, gamma, gamma_threshold):
"""Creates an instance.
Args:
gamma: A tf.Tensor of shape (n_channels,) with the gammas.
gamma_threshold: A float scalar, the threshold above which a gamma is
considered 'alive'.
"""
self._gamma = tpu_util.maybe_convert_to_variable(gamma)
self._gamma_threshold = gamma_threshold
abs_gamma = tf.abs(self._gamma)
self._alive_vector = abs_gamma > gamma_threshold
self._regularization_vector = abs_gamma
def create_regularizer(self, op_slice):
"""Create a regularizer for this batch norm OpSlice.
Args:
op_slice: op_regularizer_manager.OpSlice that is a batch norm OpSlice.
Returns:
OpRegularizer for this batch norm op.
"""
start_index = op_slice.slice.start_index
size = op_slice.slice.size
gamma = op_slice.op.inputs[1] # Input 1 is gamma.
# If OpSlice size matches tensor size, use the entire tensor. Otherwise,
# slice the tensor accordingly.
if start_index == 0 and size == gamma.shape.as_list()[-1]:
return gamma_l1_regularizer.GammaL1Regularizer(
gamma, self._gamma_threshold)
else:
# Note: this conversion is also attempted inside GammaL1Regularizer
# because it may be invoked from another call site.
gamma = tpu_util.maybe_convert_to_variable(gamma)
return gamma_l1_regularizer.GammaL1Regularizer(
gamma[start_index:start_index + size], self._gamma_threshold)
def test_noop(self):
with tf.variable_scope('', use_resource=True):
relu = self.build_model()
# Check tensors that are not variable reads are ignored.
self.assertEqual(tpu_util.maybe_convert_to_variable(relu), relu)