def testWeightedDistance(self):
source_tensor = tf.constant([[1, 1], [2, 2], [0, 2], [5, 5]],
dtype='float32')
target_tensor = tf.constant([[1, 1], [0, 2], [4, 4], [1, 4]],
dtype='float32')
weights = tf.constant([[1], [0], [0.5], [0.5]], dtype='float32')
l1_distance_config = configs.DistanceConfig('l1', sum_over_axis=-1)
l1_distance_tensor = distances.pairwise_distance_wrapper(
source_tensor, target_tensor, weights, l1_distance_config)
l2_distance_config = configs.DistanceConfig('l2', sum_over_axis=-1)
l2_distance_tensor = distances.pairwise_distance_wrapper(
source_tensor, target_tensor, weights, l2_distance_config)
with self.cached_session() as sess:
l1_distance_value = sess.run(l1_distance_tensor)
self.assertAllClose(l1_distance_value, 5.5 / 3)
l2_distance_value = sess.run(l2_distance_tensor)
self.assertAllClose(l2_distance_value, 18.5 / 3)
def testDistanceInvalidAxis(self):
source_tensor = tf.constant(1.0, dtype='float32', shape=[4, 2])
target_tensor = tf.constant(1.0, dtype='float32', shape=[4, 2])
weights = tf.constant(1.0, dtype='float32', shape=[4, 2])
distance_config = configs.DistanceConfig(sum_over_axis=2)
with self.assertRaises(ValueError):
distance_tensor = distances.pairwise_distance_wrapper(
source_tensor, target_tensor, weights, distance_config)
distance_tensor.eval()
def testL2Distance(self):
source_tensor = tf.constant([[1, 1], [2, 2], [0, 2], [5, 5]],
dtype='float32')
target_tensor = tf.constant([[1, 1], [0, 2], [4, 4], [1, 4]],
dtype='float32')
distance_config = configs.DistanceConfig('l2', sum_over_axis=-1)
distance_tensor = distances.pairwise_distance_wrapper(
source_tensor, target_tensor, distance_config=distance_config)
with self.cached_session() as sess:
distance_value = sess.run(distance_tensor)
self.assertAllClose(distance_value, 10.25)
def testCosineDistance(self):
source_tensor = tf.constant([[1, 1], [1, 1], [3, 4], [-1, -1]],
dtype='float32')
target_tensor = tf.constant([[1, 1], [5, 5], [4, 3], [1, 1]],
dtype='float32')
distance_config = configs.DistanceConfig('cosine', sum_over_axis=-1)
distance_tensor = distances.pairwise_distance_wrapper(
source_tensor, target_tensor, distance_config=distance_config)
with self.cached_session() as sess:
distance_value = sess.run(distance_tensor)
self.assertAllClose(distance_value,
0.51) # sum([0.0, 1.0, 0.04, 2.0]) / 4
def testDistanceWithoutSumOverAxis(self):
source_tensor = tf.constant([[1, 1], [2, 2], [0, 2], [5, 5]],
dtype='float32')
target_tensor = tf.constant([[1, 1], [0, 2], [4, 4], [1, 4]],
dtype='float32')
weights = tf.constant([[1], [0], [0.5], [0.5]], dtype='float32')
distance_config = configs.DistanceConfig('l1')
distance_tensor = distances.pairwise_distance_wrapper(
source_tensor, target_tensor, weights, distance_config)
with self.cached_session() as sess:
distance_value = sess.run(distance_tensor)
self.assertAllClose(distance_value, 5.5 / 6)
def testDistanceReductionMean(self):
source_tensor = tf.constant([[1, 1], [2, 2], [0, 2], [5, 5]],
dtype='float32')
target_tensor = tf.constant([[1, 1], [0, 2], [4, 4], [1, 4]],
dtype='float32')
weights = tf.constant([[1], [0], [0.5], [0.5]], dtype='float32')
distance_mean_config = configs.DistanceConfig(
'l1', tf.compat.v1.losses.Reduction.MEAN, sum_over_axis=-1)
distance_mean_tensor = distances.pairwise_distance_wrapper(
source_tensor, target_tensor, weights, distance_mean_config)
with self.cached_session() as sess:
distance_mean_value = sess.run(distance_mean_tensor)
self.assertAllClose(distance_mean_value, 5.5 / 2.0)
def testJensenShannonDistance(self):
source_tensor = np.array([[1, 0, 0], [0.1, 0.2, 0.7]], dtype='float32')
target_tensor = np.array([[1, 0, 0], [0.1, 0.9, 0]], dtype='float32')
expected_tensor = np.sum(self._jsd_func(source_tensor, target_tensor), -1)
expected_value = np.mean(expected_tensor)
distance_config = configs.DistanceConfig(
'jensen_shannon_divergence', sum_over_axis=-1)
distance_tensor = distances.pairwise_distance_wrapper(
tf.constant(source_tensor),
tf.constant(target_tensor),
distance_config=distance_config)
with self.cached_session() as sess:
distance_value = sess.run(distance_tensor)
self.assertAllClose(distance_value, expected_value)
def testDistanceWithTransformButNoSumOverAxis(self):
source = np.array([[1, 1], [2, 2], [0, 2], [10, -10]], dtype='float32')
target = np.array([[0, 0], [0, 2], [1, 3], [3, 3]], dtype='float32')
distance_config = configs.DistanceConfig(
distance_type='l1',
reduction=tf.compat.v1.losses.Reduction.NONE,
transform_fn='softmax')
distance_tensor = distances.pairwise_distance_wrapper(
tf.constant(source),
tf.constant(target),
distance_config=distance_config)
expected_distance = np.abs(
self._softmax_func(source) - self._softmax_func(target))
with self.cached_session() as sess:
distance = sess.run(distance_tensor)
self.assertAllClose(distance, expected_distance)
def testKLDistanceFromLogit(self):
source = np.array([[1, 2, 3], [1, -1, 2]], dtype='float32')
target = np.array([[1, 2, 3], [1, 0, -1]], dtype='float32')
expected_value = np.mean(
np.sum(
self._kl_func(
self._softmax_func(source), self._softmax_func(target)), -1))
distance_config = configs.DistanceConfig(
'kl_divergence', transform_fn='softmax', sum_over_axis=-1)
distance_tensor = distances.pairwise_distance_wrapper(
tf.constant(source),
tf.constant(target),
distance_config=distance_config)
with self.cached_session() as sess:
distance_value = sess.run(distance_tensor)
self.assertAllClose(distance_value, expected_value)
def call(self, inputs, weights=None):
"""Replicates sources and computes pairwise distance.
Args:
inputs: Symbolic inputs. Should be (sources, targets) if `weights` is
non-symbolic. Otherwise, should be (sources, targets, weights).
weights: If target weights are not symbolic, `weights` should be passed as
a separate argument. In this case, `inputs` should have length 2.
Returns:
Pairwise distance tensor.
"""
if weights is None:
sources, targets, weights = inputs
else:
sources, targets = inputs
return distances.pairwise_distance_wrapper(
sources=self._replicate_sources(sources, targets),
targets=targets,
weights=weights,
distance_config=self._distance_config)
def loss_fn(embedding, perturbed_embedding):
return distances.pairwise_distance_wrapper(
sources=embedding,
targets=perturbed_embedding,
distance_config=virtual_adv_config.distance_config)
def graph_reg_model_fn(features, labels, mode, params=None, config=None):
"""The graph-regularized model function.
Args:
features: This is the first item returned from the `input_fn` passed to
`train`, `evaluate`, and `predict`. This should be a dictionary
containing sample features as well as corresponding neighbor features
and neighbor weights.
labels: This is the second item returned from the `input_fn` passed to
`train`, `evaluate`, and `predict`. This should be a single `Tensor` or
`dict` of same (for multi-head models). If mode is
`tf.estimator.ModeKeys.PREDICT`, `labels=None` will be passed. If the
`model_fn`'s signature does not accept `mode`, the `model_fn` must still
be able to handle `labels=None`.
mode: Optional. Specifies if this is training, evaluation, or prediction.
See `tf.estimator.ModeKeys`.
params: Optional `dict` of hyperparameters. Will receive what is passed to
Estimator in the `params` parameter. This allows users to configure
Estimators from hyper parameter tuning.
config: Optional `tf.estimator.RunConfig` object. Will receive what is
passed to Estimator as its `config` parameter, or a default value.
Allows setting up things in the `model_fn` based on configuration such
as `num_ps_replicas`, or `model_dir`. Unused currently.
Returns:
A `tf.estimator.EstimatorSpec` with graph regularization.
"""
# Parameters 'params' and 'config' are optional. If they are not passed,
# then it is possible for base_model_fn not to accept these arguments.
# See documentation for tf.estimator.Estimator for additional context.
kwargs = {'mode': mode}
embedding_fn_kwargs = dict()
if 'params' in base_model_fn_args:
kwargs['params'] = params
embedding_fn_kwargs['params'] = params
if 'config' in base_model_fn_args:
kwargs['config'] = config
# Uses the same variable scope for calculating the original objective and
# the graph regularization loss term.
with tf.compat.v1.variable_scope(tf.compat.v1.get_variable_scope(),
reuse=tf.compat.v1.AUTO_REUSE,
auxiliary_name_scope=False):
nbr_features = dict()
nbr_weights = None
if mode == tf.estimator.ModeKeys.TRAIN:
# Extract sample features, neighbor features, and neighbor weights if we
# are in training mode.
sample_features, nbr_features, nbr_weights = (
utils.unpack_neighbor_features(
features, graph_reg_config.neighbor_config))
else:
# Otherwise, we strip out all neighbor features and use just the
# sample's features.
sample_features = utils.strip_neighbor_features(
features, graph_reg_config.neighbor_config)
base_spec = base_model_fn(sample_features, labels, **kwargs)
has_nbr_inputs = nbr_weights is not None and nbr_features
# Graph regularization happens only if all the following conditions are
# satisfied:
# - the mode is training
# - neighbor inputs exist
# - the graph regularization multiplier is greater than zero.
# So, return early if any of these conditions is false.
if (not has_nbr_inputs or mode != tf.estimator.ModeKeys.TRAIN
or graph_reg_config.multiplier <= 0):
return base_spec
# Compute sample embeddings.
sample_embeddings = embedding_fn(sample_features, mode,
**embedding_fn_kwargs)
# Compute the embeddings of the neighbors.
nbr_embeddings = embedding_fn(nbr_features, mode,
**embedding_fn_kwargs)
replicated_sample_embeddings = utils.replicate_embeddings(
sample_embeddings,
graph_reg_config.neighbor_config.max_neighbors)
# Compute the distance between the sample embeddings and each of their
# corresponding neighbor embeddings.
graph_loss = distances.pairwise_distance_wrapper(
replicated_sample_embeddings,
nbr_embeddings,
weights=nbr_weights,
distance_config=graph_reg_config.distance_config)
scaled_graph_loss = graph_reg_config.multiplier * graph_loss
tf.compat.v1.summary.scalar('loss/scaled_graph_loss',
scaled_graph_loss)
supervised_loss = base_spec.loss
tf.compat.v1.summary.scalar('loss/supervised_loss',
supervised_loss)
total_loss = supervised_loss + scaled_graph_loss
if not optimizer_fn:
# Default to Adagrad optimizer, the same as the canned DNNEstimator.
optimizer = tf.compat.v1.train.AdagradOptimizer(
learning_rate=0.05)
else:
optimizer = optimizer_fn()
train_op = optimizer.minimize(
loss=total_loss,
global_step=tf.compat.v1.train.get_global_step())
update_ops = tf.compat.v1.get_collection(
tf.compat.v1.GraphKeys.UPDATE_OPS)
if update_ops:
train_op = tf.group(train_op, *update_ops)
return base_spec._replace(loss=total_loss, train_op=train_op)
def graph_reg_model_fn(features, labels, mode, params=None, config=None):
"""The graph-regularized model function.
Args:
features: This is the first item returned from the `input_fn` passed to
`train`, `evaluate`, and `predict`. This should be a dictionary
containing sample features as well as corresponding neighbor features
and neighbor weights.
labels: This is the second item returned from the `input_fn` passed to
`train`, `evaluate`, and `predict`. This should be a single `Tensor` or
`dict` of same (for multi-head models). If mode is
`tf.estimator.ModeKeys.PREDICT`, `labels=None` will be passed. If the
`model_fn`'s signature does not accept `mode`, the `model_fn` must still
be able to handle `labels=None`.
mode: Optional. Specifies if this is training, evaluation, or prediction.
See `tf.estimator.ModeKeys`.
params: Optional `dict` of hyperparameters. Will receive what is passed to
Estimator in the `params` parameter. This allows users to configure
Estimators from hyper parameter tuning.
config: Optional `tf.estimator.RunConfig` object. Will receive what is
passed to Estimator as its `config` parameter, or a default value.
Allows setting up things in the `model_fn` based on configuration such
as `num_ps_replicas`, or `model_dir`. Unused currently.
Returns:
A `tf.EstimatorSpec` whose loss incorporates graph-based regularization.
"""
# Uses the same variable scope for calculating the original objective and
# the graph regularization loss term.
with tf.compat.v1.variable_scope(
tf.compat.v1.get_variable_scope(),
reuse=tf.compat.v1.AUTO_REUSE,
auxiliary_name_scope=False):
# Extract sample features, neighbor features, and neighbor weights.
sample_features, nbr_features, nbr_weights = (
utils.unpack_neighbor_features(features,
graph_reg_config.neighbor_config))
# If no 'params' is passed, then it is possible for base_model_fn not to
# accept a 'params' argument. See documentation for tf.estimator.Estimator
# for additional context.
if params:
base_spec = base_model_fn(sample_features, labels, mode, params, config)
else:
base_spec = base_model_fn(sample_features, labels, mode, config)
has_nbr_inputs = nbr_weights is not None and nbr_features
# Graph regularization happens only if all the following conditions are
# satisfied:
# - the mode is training
# - neighbor inputs exist
# - the graph regularization multiplier is greater than zero.
# So, return early if any of these conditions is false.
if (not has_nbr_inputs or mode != tf.estimator.ModeKeys.TRAIN or
graph_reg_config.multiplier <= 0):
return base_spec
# Compute sample embeddings.
sample_embeddings = embedding_fn(sample_features, mode)
# Compute the embeddings of the neighbors.
nbr_embeddings = embedding_fn(nbr_features, mode)
replicated_sample_embeddings = utils.replicate_embeddings(
sample_embeddings, graph_reg_config.neighbor_config.max_neighbors)
# Compute the distance between the sample embeddings and each of their
# corresponding neighbor embeddings.
graph_loss = distances.pairwise_distance_wrapper(
replicated_sample_embeddings,
nbr_embeddings,
weights=nbr_weights,
distance_config=graph_reg_config.distance_config)
total_loss = base_spec.loss + graph_reg_config.multiplier * graph_loss
if not optimizer_fn:
# Default to Adagrad optimizer, the same as the canned DNNEstimator.
optimizer = tf.train.AdagradOptimizer(learning_rate=0.05)
else:
optimizer = optimizer_fn()
final_train_op = optimizer.minimize(
loss=total_loss, global_step=tf.compat.v1.train.get_global_step())
return base_spec._replace(loss=total_loss, train_op=final_train_op)
