def _replicate_sources(self, sources, targets):
"""Replicates `sources` to match the shape of `targets`.
`targets` should either have an additional neighborhood size dimension at
axis -2 or be of the same rank as `sources`. If `targets` has an additional
dimension and `sources` has rank k, the first k - 1 dimensions and last
dimension of `sources` and `targets` should match. If `sources` and
`targets` have the same rank, the last k - 1 dimensions should match and the
first dimension of `targets` should be a multiple of the first dimension of
`sources`. This multiple represents the fixed neighborhood size of each
sample.
Args:
sources: Tensor with shape [..., feature_size] from which distance will be
calculated.
targets: Either a tensor with shape [..., neighborhood_size, feature_size]
or [sources.shape[0] * neighborhood_size] + sources.shape[1:].
Returns:
`sources` replicated to be shape-compatible with `targets`.
"""
# Depending on the rank of `sources` and `targets`, decide to broadcast
# first, or replicate directly.
if (sources.shape.ndims is not None and targets.shape.ndims is not None
and sources.shape.ndims + 1 == targets.shape.ndims):
return tf.broadcast_to(tf.expand_dims(sources, axis=-2),
tf.shape(targets))
return utils.replicate_embeddings(
sources,
tf.shape(targets)[0] // tf.shape(sources)[0])
def testInvalidRepeatTimes(self):
"""Test the replicate_embeddings function with invalid repeat_times."""
input_embeddings = tf.constant([
[[1., 2., 4.], [3., 5., 8.]],
[[2., 10., 3.], [1., 1., 1.]],
[[4., 8., 1.], [8., 4., 1.]],
])
replicate_times = tf.constant([-1, 0, 1])
with self.assertRaises(tf.errors.InvalidArgumentError):
self.evaluate(
utils.replicate_embeddings(input_embeddings, replicate_times))
def testInvalidRepeatTimes(self):
"""Test the replicate_embeddings function with invalid repeat_times."""
input_embeddings = tf.constant(
[[[1, 2, 4], [3, 5, 8]], [[2, 10, 3], [1, 1, 1]],
[[4, 8, 1], [8, 4, 1]]],
dtype='float32')
replicate_times = tf.constant([-1, 0, 1])
with self.cached_session():
with self.assertRaises(tf.errors.InvalidArgumentError):
output_embeddings = utils.replicate_embeddings(
input_embeddings, replicate_times)
output_embeddings.eval()
def testReplicateEmbeddingsWithIndexArray(self):
"""Test the replicate_embeddings function with 1-D replicate_times."""
input_embeddings = tf.constant(
[[[1, 2, 4], [3, 5, 8]], [[2, 10, 3], [1, 1, 1]],
[[4, 8, 1], [8, 4, 1]]],
dtype='float32')
replicate_times = tf.constant([2, 0, 1])
output_embeddings = utils.replicate_embeddings(input_embeddings,
replicate_times)
with self.cached_session() as sess:
self.assertAllEqual(
[[[1, 2, 4], [3, 5, 8]], [[1, 2, 4], [3, 5, 8]],
[[4, 8, 1], [8, 4, 1]]], sess.run(output_embeddings))
def testReplicateEmbeddingsWithIndexArray(self):
"""Test the replicate_embeddings function with 1-D replicate_times."""
input_embeddings = tf.constant([
[[1., 2., 4.], [3., 5., 8.]],
[[2., 10., 3.], [1., 1., 1.]],
[[4., 8., 1.], [8., 4., 1.]],
])
replicate_times = tf.constant([2, 0, 1])
output_embeddings = self.evaluate(
utils.replicate_embeddings(input_embeddings, replicate_times))
expected_embeddings = [
[[1., 2., 4.], [3., 5., 8.]],
[[1., 2., 4.], [3., 5., 8.]],
[[4., 8., 1.], [8., 4., 1.]],
]
self.assertAllEqual(expected_embeddings, output_embeddings)
def _replicate_with_dynamic_batch_size(embeddings, replicate_times):
return utils.replicate_embeddings(embeddings, replicate_times)
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)
