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)