def testAllArgumentsSet(self):
"""Tests that no errors are raised when all arguments are set."""
with ops.Graph().as_default(), self.cached_session():
loss = constant_op.constant(1.)
predictions = {'loss': loss}
classes = constant_op.constant('hello')
metric_obj = metrics.Mean()
metric_obj.update_state(loss)
model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.TRAIN,
predictions=predictions,
loss=loss,
train_op=control_flow_ops.no_op(),
eval_metric_ops={
'loss': (control_flow_ops.no_op(), loss),
'mean': metric_obj,
},
export_outputs={
'head_name': export_output.ClassificationOutput(classes=classes)
},
training_chief_hooks=[_FakeHook()],
training_hooks=[_FakeHook()],
scaffold=monitored_session.Scaffold(),
evaluation_hooks=[_FakeHook()],
prediction_hooks=[_FakeHook()])
def per_example_quantile_regression_loss(labels, weights, predictions,
quantile):
"""Smoothed loss for quantile regression.
The standard quantile regression loss is quantile*(y-y') when y>y' and
(quantile-1)*(y-y') otherwise, y' is a prediction, y is a label. The impl
below is this loss but squared in the region where the loss value < 1.
Args:
labels: Rank 2 (N, D) tensor of per-example labels.
weights: Rank 2 (N, 1) tensor of per-example weights.
predictions: Rank 2 (N, D) tensor of per-example predictions.
quantile: The quantile to use.
Returns:
loss: A Rank 2 (N, 1) tensor of per-example quantile loss.
update_op: An update operation to update the loss's internal state.
"""
labels = math_ops.to_float(labels)
error = labels - predictions
square_loss_right = array_ops.where(error * quantile < 1.0,
math_ops.square(quantile * error),
quantile * error)
square_loss_left = array_ops.where(error * (quantile - 1) < 1,
math_ops.square((quantile - 1) * error),
(quantile - 1) * error)
unweighted_loss = array_ops.where(error > 0, square_loss_right,
square_loss_left)
if weights is None:
return unweighted_loss, control_flow_ops.no_op()
else:
return unweighted_loss * weights, control_flow_ops.no_op()
def testQueueRunnerSerializationRoundTrip(self):
graph = ops.Graph()
with graph.as_default():
queue = data_flow_ops.FIFOQueue(10, dtypes.float32, name="queue")
enqueue_op = control_flow_ops.no_op(name="enqueue")
close_op = control_flow_ops.no_op(name="close")
cancel_op = control_flow_ops.no_op(name="cancel")
qr0 = queue_runner_impl.QueueRunner(
queue, [enqueue_op],
close_op,
cancel_op,
queue_closed_exception_types=(errors_impl.OutOfRangeError,
errors_impl.CancelledError))
qr0_proto = queue_runner_impl.QueueRunner.to_proto(qr0)
qr0_recon = queue_runner_impl.QueueRunner.from_proto(qr0_proto)
self.assertEqual("queue", qr0_recon.queue.name)
self.assertEqual(1, len(qr0_recon.enqueue_ops))
self.assertEqual(enqueue_op, qr0_recon.enqueue_ops[0])
self.assertEqual(close_op, qr0_recon.close_op)
self.assertEqual(cancel_op, qr0_recon.cancel_op)
self.assertEqual(
(errors_impl.OutOfRangeError, errors_impl.CancelledError),
qr0_recon.queue_closed_exception_types)
# Assert we reconstruct an OutOfRangeError for QueueRunners
# created before QueueRunnerDef had a queue_closed_exception_types field.
del qr0_proto.queue_closed_exception_types[:]
qr0_legacy_recon = queue_runner_impl.QueueRunner.from_proto(qr0_proto)
self.assertEqual("queue", qr0_legacy_recon.queue.name)
self.assertEqual(1, len(qr0_legacy_recon.enqueue_ops))
self.assertEqual(enqueue_op, qr0_legacy_recon.enqueue_ops[0])
self.assertEqual(close_op, qr0_legacy_recon.close_op)
self.assertEqual(cancel_op, qr0_legacy_recon.cancel_op)
self.assertEqual((errors_impl.OutOfRangeError,),
qr0_legacy_recon.queue_closed_exception_types)
def _model_fn(features, labels, mode):
_ = labels
x = features['x']
y = features['y']
with ops.name_scope('outputs'):
predictions = {'sum': math_ops.add(x, y, name='sum'),
'product': math_ops.multiply(x, y, name='product'),
'difference': math_ops.subtract(x, y, name='difference')}
if core:
export_outputs = {k: export_output.PredictOutput({k: v})
for k, v in predictions.items()}
export_outputs[signature_constants.
DEFAULT_SERVING_SIGNATURE_DEF_KEY] = export_outputs['sum']
return model_fn.EstimatorSpec(mode=mode,
predictions=predictions,
export_outputs=export_outputs,
loss=constant_op.constant(0),
train_op=control_flow_ops.no_op())
else:
output_alternatives = {k: (constants.ProblemType.UNSPECIFIED, {k: v})
for k, v in predictions.items()}
return contrib_model_fn.ModelFnOps(
mode=mode,
predictions=predictions,
output_alternatives=output_alternatives,
loss=constant_op.constant(0),
train_op=control_flow_ops.no_op())
def per_example_maxent_loss(labels, weights, logits, num_classes, eps=1e-15):
"""Maximum entropy loss for multiclass problems.
Maximum entropy is a generalization of logistic loss for the case when more
than 2 classes are present.
Args:
labels: Rank 2 (N, 1) or Rank 1 (N) tensor of per-example labels.
weights: Rank 2 (N, 1) tensor of per-example weights.
logits: Rank 2 (N, K) tensor of per-example predictions, K - num of
classes.
num_classes: number of classes in classification task. Used to expand label
indices into one-hot encodings.
eps: tolerance, used as a minimum possible value.
Returns:
loss: A Rank 2 (N, 1) tensor of per-example maxent loss
update_op: An update operation to update the loss's internal state.
"""
labels = math_ops.to_int64(labels)
# If labels are of rank 1, make them rank 2.
labels_shape = labels.get_shape()
if len(labels_shape) != 2:
labels = array_ops.expand_dims(labels, 1)
# Labels are indices of classes, convert them to one hot encodings.
target_one_hot = array_ops.one_hot(indices=labels, depth=num_classes)
labels = math_ops.reduce_sum(
input_tensor=target_one_hot, reduction_indices=[1])
labels = math_ops.to_float(labels)
# Calculate softmax probabilities for each class.
unnormalized_probs = math_ops.exp(logits)
normalizers = math_ops.reduce_sum(unnormalized_probs, 1, keepdims=True)
softmax_predictions = math_ops.divide(unnormalized_probs,
math_ops.add(normalizers, eps))
# Pull out the probabilities for real label.
probs_for_real_class = math_ops.reduce_sum(labels * softmax_predictions, 1)
# Add handling for values near 0 and 1.
zeros = array_ops.zeros_like(probs_for_real_class, dtype=logits.dtype) + eps
one_minus_eps = array_ops.ones_like(
probs_for_real_class, dtype=logits.dtype) - eps
# Take maximum(eps, pred)
cond = (probs_for_real_class >= eps)
probs_for_real_class = array_ops.where(cond, probs_for_real_class, zeros)
# Take minimum(1-eps, pred)
cond = (probs_for_real_class <= 1 - eps)
probs_for_real_class = array_ops.where(cond, probs_for_real_class,
one_minus_eps)
unweighted_loss = array_ops.expand_dims(-math_ops.log(probs_for_real_class),
1)
if weights is None:
return unweighted_loss, control_flow_ops.no_op()
else:
return unweighted_loss * weights, control_flow_ops.no_op()
def metrics_fn(predictions, features):
# checking that the inputs are properly passed.
predict = predictions["mean"]
target = features[feature_keys.TrainEvalFeatures.VALUES][:, -1, 0]
return {
"plain_boring_metric386":
(math_ops.reduce_mean(math_ops.abs(predict - target)),
control_flow_ops.no_op()),
"fun_metric101": (math_ops.reduce_sum(predict + target),
control_flow_ops.no_op()),
}
def _cached_copy(self, var, name, pass_through=False):
"""Helper function to create a worker cached copy of a Variable.
This assigns the var (either a single Variable or a list of Variables) to
local transient cache Variable(s). Note that if var is a list of Variables,
the assignment is done sequentially to minimize the memory overheads.
Also note that if pass_through is set to True, this does not create new
Variables but simply return the input back.
Args:
var: A Variable or a list of Variables to cache.
name: name of cached Variable.
pass_through: when set to True, this simply pass through the var back
through identity operator and does not actually creates a cache.
Returns:
Tuple consisting of following three entries:
cache: the new transient Variable or list of transient Variables
corresponding one-to-one with var.
cache_init: op to initialize the Variable or the list of Variables.
cache_reset: op to reset the Variable or the list of Variables to some
default value.
"""
if var is None:
return None, None, None
elif pass_through:
cache = var
cache_init = control_flow_ops.no_op()
cache_reset = control_flow_ops.no_op()
elif isinstance(var, variables.Variable):
cache = WALSModel._transient_var(name=name)
with ops.colocate_with(cache):
cache_init = state_ops.assign(cache, var, validate_shape=False)
cache_reset = state_ops.assign(cache, 1.0, validate_shape=False)
else:
assert isinstance(var, list)
assert var
cache = [
WALSModel._transient_var(name="%s_shard_%d" % (name, i))
for i in xrange(len(var))
]
reset_ops = []
for i, c in enumerate(cache):
with ops.colocate_with(c):
if i == 0:
cache_init = state_ops.assign(c, var[i], validate_shape=False)
else:
with ops.control_dependencies([cache_init]):
cache_init = state_ops.assign(c, var[i], validate_shape=False)
reset_ops.append(state_ops.assign(c, 1.0, validate_shape=False))
cache_reset = control_flow_ops.group(*reset_ops)
return cache, cache_init, cache_reset
def _model_fn(features, labels, mode):
del features # unused
del labels
return model_fn_lib.EstimatorSpec(
mode,
train_op=control_flow_ops.no_op(),
loss=constant_op.constant(1.),
eval_metric_ops={
'nested_metric': (
((constant_op.constant(2.), constant_op.constant(1)),
constant_op.constant(3., dtype=dtypes.float64)),
control_flow_ops.no_op())})
def _define_maximization_operation(self, num_batches):
"""Maximization operations."""
# TODO(xavigonzalvo): some of these operations could be moved to C++.
# Compute the effective number of data points assigned to component k.
with ops.control_dependencies(self._w):
points_in_k = array_ops.squeeze(
math_ops.add_n(self._points_in_k), axis=[0])
# Update alpha.
if 'w' in self._params:
final_points_in_k = points_in_k / num_batches
num_examples = math_ops.cast(math_ops.reduce_sum(final_points_in_k),
dtypes.float32)
self._alpha_op = self._alpha.assign(final_points_in_k /
(num_examples + MEPS))
else:
self._alpha_op = control_flow_ops.no_op()
self._train_ops = [self._alpha_op]
# Update means.
points_in_k_expanded = array_ops.reshape(points_in_k,
[self._num_classes, 1, 1])
if 'm' in self._params:
self._means_op = self._means.assign(
math_ops.div(
math_ops.add_n(self._w_mul_x), points_in_k_expanded + MEPS))
else:
self._means_op = control_flow_ops.no_op()
# means are (num_classes x 1 x dims)
# Update covariances.
with ops.control_dependencies([self._means_op]):
b = math_ops.add_n(self._w_mul_x2) / (points_in_k_expanded + MEPS)
new_covs = []
for k in range(self._num_classes):
mean = self._means.value()[k, :, :]
square_mean = math_ops.matmul(mean, mean, transpose_a=True)
new_cov = b[k, :, :] - square_mean + self._min_var
if self._covariance_type == FULL_COVARIANCE:
new_covs.append(array_ops.expand_dims(new_cov, 0))
elif self._covariance_type == DIAG_COVARIANCE:
new_covs.append(
array_ops.expand_dims(array_ops.diag_part(new_cov), 0))
new_covs = array_ops.concat(new_covs, 0)
if 'c' in self._params:
# Train operations don't need to take care of the means
# because covariances already depend on it.
with ops.control_dependencies([self._means_op, new_covs]):
self._train_ops.append(
state_ops.assign(
self._covs, new_covs, validate_shape=False))
def test_stop_based_on_num_step(self):
h = basic_session_run_hooks.StopAtStepHook(num_steps=10)
with ops.Graph().as_default():
global_step = variables.get_or_create_global_step()
no_op = control_flow_ops.no_op()
h.begin()
with session_lib.Session() as sess:
mon_sess = monitored_session._HookedSession(sess, [h])
sess.run(state_ops.assign(global_step, 5))
h.after_create_session(sess, None)
mon_sess.run(no_op)
self.assertFalse(mon_sess.should_stop())
sess.run(state_ops.assign(global_step, 13))
mon_sess.run(no_op)
self.assertFalse(mon_sess.should_stop())
sess.run(state_ops.assign(global_step, 14))
mon_sess.run(no_op)
self.assertFalse(mon_sess.should_stop())
sess.run(state_ops.assign(global_step, 15))
mon_sess.run(no_op)
self.assertTrue(mon_sess.should_stop())
sess.run(state_ops.assign(global_step, 16))
mon_sess._should_stop = False
mon_sess.run(no_op)
self.assertTrue(mon_sess.should_stop())
def create_model_fn_ops(self, predictions, output_alternatives,
mode=model_fn.ModeKeys.INFER):
return model_fn.ModelFnOps(
model_fn.ModeKeys.INFER,
predictions=predictions,
loss=constant_op.constant([1]),
train_op=control_flow_ops.no_op(),
eval_metric_ops={
"metric_key": (constant_op.constant(1.), control_flow_ops.no_op()),
"loss": (constant_op.constant(1.), control_flow_ops.no_op()),
},
training_chief_hooks=[basic_session_run_hooks.StepCounterHook()],
training_hooks=[basic_session_run_hooks.StepCounterHook()],
output_alternatives=output_alternatives,
scaffold=monitored_session.Scaffold())
def model_fn(features, labels): # dummy variable: _ = variables.Variable([0.]) _ = labels predictions = features["x"] loss = constant_op.constant([2.]) return predictions, loss, control_flow_ops.no_op()
def model_fn(self, mode, features, labels, params):
c = variable_scope.get_variable(
'c',
initializer=constant_op.constant(10, dtype=dtypes.float64),
dtype=dtypes.float64)
predictions = math_ops.multiply(features, c)
loss = None
if mode is not model_fn_lib.ModeKeys.PREDICT:
loss = losses.absolute_difference(
labels=labels,
predictions=predictions,
reduction=losses.Reduction.SUM)
loss = math_ops.reduce_sum(loss)
metrics = {
'accuracy': metrics_lib.accuracy(labels, predictions),
'auc': metrics_lib.auc(labels, predictions)
}
return model_fn_lib.EstimatorSpec(
mode=mode,
loss=loss,
eval_metric_ops=metrics,
predictions={'probabilities': predictions},
train_op=control_flow_ops.no_op()) # This train_op isn't actually used.
def testIgnoredArguments(self):
"""Tests that JIT computations can ignore formal parameters."""
with self.session(config=NoRewriteSessionConfig()) as sess:
x = array_ops.placeholder(dtypes.int32)
y = array_ops.placeholder(dtypes.int32)
with jit_scope():
z = math_ops.add(x, x)
w = math_ops.add(y, y)
# Pulls 'w' into the same compilation via control dependencies.
with ops.control_dependencies([w]):
n = control_flow_ops.no_op()
with ops.control_dependencies([n]):
t = math_ops.add(z, z)
run_metadata = config_pb2.RunMetadata()
out = test_utils.RunWithWarmup(
sess,
t, {
x: np.int32(7),
y: np.int32(404)
},
run_metadata=run_metadata,
options=config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE))
self.assert_(MetadataHasXlaRunOp(run_metadata))
self.assertAllClose(28, out)
def assert_integer_form(
x, data=None, summarize=None, message=None,
int_dtype=None, name="assert_integer_form"):
"""Assert that x has integer components (or floats equal to integers).
Args:
x: Floating-point `Tensor`
data: The tensors to print out if the condition is `False`. Defaults to
error message and first few entries of `x` and `y`.
summarize: Print this many entries of each tensor.
message: A string to prefix to the default message.
int_dtype: A `tf.dtype` used to cast the float to. The default (`None`)
implies the smallest possible signed int will be used for casting.
name: A name for this operation (optional).
Returns:
Op raising `InvalidArgumentError` if `cast(x, int_dtype) != x`.
"""
with ops.name_scope(name, values=[x, data]):
x = ops.convert_to_tensor(x, name="x")
if x.dtype.is_integer:
return control_flow_ops.no_op()
message = message or "{} has non-integer components".format(x.op.name)
if int_dtype is None:
try:
int_dtype = {
dtypes.float16: dtypes.int16,
dtypes.float32: dtypes.int32,
dtypes.float64: dtypes.int64,
}[x.dtype.base_dtype]
except KeyError:
raise TypeError("Unrecognized type {}".format(x.dtype.name))
return check_ops.assert_equal(
x, math_ops.cast(math_ops.cast(x, int_dtype), x.dtype),
data=data, summarize=summarize, message=message, name=name)
def summary_writer_function(name, tensor, function, family=None):
"""Helper function to write summaries.
Args:
name: name of the summary
tensor: main tensor to form the summary
function: function taking a tag and a scope which writes the summary
family: optional, the summary's family
Returns:
The result of writing the summary.
"""
def record():
with summary_op_util.summary_scope(
name, family, values=[tensor]) as (tag, scope):
with ops.control_dependencies([function(tag, scope)]):
return constant_op.constant(True)
if context.context().summary_writer_resource is None:
return control_flow_ops.no_op()
with ops.device("cpu:0"):
op = utils.smart_cond(
should_record_summaries(), record, _nothing, name="")
ops.add_to_collection(ops.GraphKeys._SUMMARY_COLLECTION, op) # pylint: disable=protected-access
return op
def model_fn(features, labels, mode, params):
del features, labels, params
return tpu_estimator.TPUEstimatorSpec(
mode=mode,
loss=constant_op.constant(_EXPECTED_LOSS),
train_op=control_flow_ops.no_op(),
scaffold_fn=self._make_scaffold_fn(mode))
def testLossNotScalar(self):
with ops.Graph().as_default(), self.test_session():
with self.assertRaisesRegexp(ValueError, 'Loss must be scalar'):
model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.TRAIN,
loss=constant_op.constant([1., 2.]),
train_op=control_flow_ops.no_op())
def flush(writer=None, name=None):
"""Forces summary writer to send any buffered data to storage.
This operation blocks until that finishes.
Args:
writer: The `tf.summary.SummaryWriter` resource to flush.
The thread default will be used if this parameter is None.
Otherwise a `tf.no_op` is returned.
name: A name for the operation (optional).
Returns:
The created `tf.Operation`.
"""
if writer is None:
writer = context.context().summary_writer
if writer is None:
return control_flow_ops.no_op()
if isinstance(writer, ResourceSummaryWriter):
resource = writer._resource # pylint: disable=protected-access
else:
# Assume we were passed a raw resource tensor.
resource = writer
with ops.device("cpu:0"):
return gen_summary_ops.flush_summary_writer(resource, name=name)
def testLoss1DTensor(self):
"""Tests that no errors are raised when loss is 1D tensor."""
with ops.Graph().as_default(), self.test_session():
model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.TRAIN,
loss=constant_op.constant([1.]),
train_op=control_flow_ops.no_op())
def _TestInsertQuantOpForAddAfterConv2d(self, is_training):
graph = ops.Graph()
with graph.as_default():
batch_size, height, width, depth = 5, 128, 128, 3
input1 = array_ops.zeros((batch_size, height, width, depth))
input2 = array_ops.zeros((batch_size, height / 2, width / 2, 32))
conv = conv2d(input1, 32, [5, 5], stride=2, padding='SAME',
weights_initializer=self._WeightInit(0.09),
activation_fn=None, scope='test/test')
node = math_ops.add(conv, input2, name='test/add')
node = nn_ops.relu6(node, name='test/relu6')
update_barrier = control_flow_ops.no_op(name='update_barrier')
with ops.control_dependencies([update_barrier]):
array_ops.identity(node, name='control_dependency')
quantize.Quantize(graph, is_training, weight_bits=8, activation_bits=8)
quantization_node_name = 'FakeQuantWithMinMaxVars'
conv_quant = graph.get_operation_by_name('test/test/conv_quant/' +
quantization_node_name)
self.assertEqual(conv_quant.type, quantization_node_name)
# Scan through all FakeQuant operations, ensuring that the activation
# isn't in the consumers of the operation. Since activations are folded
# the preceding operation during inference, the FakeQuant operation after
# the activation is all that is needed.
for op in graph.get_operations():
if op.type == quantization_node_name:
quant_op = graph.get_operation_by_name(op.name)
consumers = []
for output in quant_op.outputs:
consumers.extend(output.consumers())
self.assertNotIn('test/relu6', [c.name for c in consumers])
def assert_type(tensor, tf_type, message=None, name=None):
"""Statically asserts that the given `Tensor` is of the specified type.
Args:
tensor: A tensorflow `Tensor`.
tf_type: A tensorflow type (`dtypes.float32`, `tf.int64`, `dtypes.bool`,
etc).
message: A string to prefix to the default message.
name: A name to give this `Op`. Defaults to "assert_type"
Raises:
TypeError: If the tensors data type doesn't match `tf_type`.
Returns:
A `no_op` that does nothing. Type can be determined statically.
"""
message = message or ''
with ops.name_scope(name, 'assert_type', [tensor]):
tensor = ops.convert_to_tensor(tensor, name='tensor')
if tensor.dtype != tf_type:
if context.executing_eagerly():
raise TypeError('%s tensor must be of type %s' % (message, tf_type))
else:
raise TypeError('%s %s must be of type %s' % (message, tensor.name,
tf_type))
return control_flow_ops.no_op('statically_determined_correct_type')
def testWatchingOutputSlotWithoutOutgoingEdge(self):
"""Test watching output slots not attached to any outgoing edges."""
with session.Session() as sess:
u_init_val = np.array([[5.0, 3.0], [-1.0, 0.0]])
u = constant_op.constant(u_init_val, shape=[2, 2], name="u")
# Create a control edge from a node with an output: From u to z.
# Node u will get executed only because of the control edge. The output
# tensor u:0 is not attached to any outgoing edge in the graph. This test
# checks that the debugger can watch such a tensor.
with ops.control_dependencies([u]):
z = control_flow_ops.no_op(name="z")
run_options = config_pb2.RunOptions(output_partition_graphs=True)
debug_utils.watch_graph(
run_options,
sess.graph,
debug_ops=["DebugIdentity"],
debug_urls=self._debug_urls())
run_metadata = config_pb2.RunMetadata()
sess.run(z, options=run_options, run_metadata=run_metadata)
dump = debug_data.DebugDumpDir(
self._dump_root, partition_graphs=run_metadata.partition_graphs)
# Assert that the DebugIdentity watch on u works properly.
self.assertEqual(1, len(dump.dumped_tensor_data))
datum = dump.dumped_tensor_data[0]
self.assertEqual("u", datum.node_name)
self.assertEqual(0, datum.output_slot)
self.assertEqual("DebugIdentity", datum.debug_op)
self.assertAllClose([[5.0, 3.0], [-1.0, 0.0]], datum.get_tensor())
def assert_integer(x, message=None, name=None):
"""Assert that `x` is of integer dtype.
Example of adding a dependency to an operation:
```python
with tf.control_dependencies([tf.assert_integer(x)]):
output = tf.reduce_sum(x)
```
Args:
x: `Tensor` whose basetype is integer and is not quantized.
message: A string to prefix to the default message.
name: A name for this operation (optional). Defaults to "assert_integer".
Raises:
TypeError: If `x.dtype` is anything other than non-quantized integer.
Returns:
A `no_op` that does nothing. Type can be determined statically.
"""
message = message or ''
with ops.name_scope(name, 'assert_integer', [x]):
x = ops.convert_to_tensor(x, name='x')
if not x.dtype.is_integer:
err_msg = (
'%s Expected "x" to be integer type. Found: %s of dtype %s'
% (message, x.name, x.dtype))
raise TypeError(err_msg)
return control_flow_ops.no_op('statically_determined_was_integer')
def testLegacyInitOpWithNonEmptyCollection(self):
export_dir = self._get_export_dir(
"test_legacy_init_op_with_non_empty_collection")
builder = saved_model_builder.SavedModelBuilder(export_dir)
with self.test_session(graph=ops.Graph()) as sess:
# Initialize variable `v1` to 1.
v1 = variables.Variable(1, name="v1")
ops.add_to_collection("v", v1)
# Initialize another variable `v2` to 42.
v2 = variables.Variable(42, name="v2", trainable=False, collections=[])
ops.add_to_collection("v", v2)
# Set up an assignment op to be run as part of the legacy_init_op.
assign_v2 = state_ops.assign(v2, v1)
legacy_init_op = control_flow_ops.group(assign_v2, name="legacy_init_op")
sess.run(variables.global_variables_initializer())
ops.add_to_collection(constants.LEGACY_INIT_OP_KEY,
control_flow_ops.no_op())
# AssertionError should be raised since the LEGACY_INIT_OP_KEY collection
# is not empty and we don't support multiple init ops.
with self.assertRaises(AssertionError):
builder.add_meta_graph_and_variables(
sess, ["foo"], legacy_init_op=legacy_init_op)
def testReturnDefaultScaffold(self):
with ops.Graph().as_default(), self.test_session():
estimator_spec = model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.TRAIN,
loss=constant_op.constant(1.),
train_op=control_flow_ops.no_op())
self.assertIsNotNone(estimator_spec.scaffold)
def testRequiredArgumentsSet(self):
"""Tests that no errors are raised when all required arguments are set."""
with ops.Graph().as_default(), self.test_session():
model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.TRAIN,
loss=constant_op.constant(1.),
train_op=control_flow_ops.no_op())
def test_metric_op_is_tensor(self):
"""Tests that ops.Operation is wrapped by a tensor for metric_ops."""
with context.graph_mode():
loss = {'my_loss': constant_op.constant([0])}
predictions = {u'output1': constant_op.constant(['foo'])}
metric_obj = metrics_module.Mean()
metric_obj.update_state(constant_op.constant([0]))
metrics = {
'metrics_1': metric_obj,
'metrics_2': (constant_op.constant([0]), control_flow_ops.no_op())
}
outputter = MockSupervisedOutput(loss, predictions, metrics)
self.assertTrue(outputter.metrics['metrics_1/update_op'].name.startswith(
'metric_op_wrapper'))
self.assertTrue(
isinstance(outputter.metrics['metrics_1/update_op'], ops.Tensor))
self.assertTrue(
isinstance(outputter.metrics['metrics_1/value'], ops.Tensor))
self.assertEqual(outputter.metrics['metrics_2/value'],
metrics['metrics_2'][0])
self.assertTrue(outputter.metrics['metrics_2/update_op'].name.startswith(
'metric_op_wrapper'))
self.assertTrue(
isinstance(outputter.metrics['metrics_2/update_op'], ops.Tensor))
def testLossNumber(self):
"""Tests that error is raised when loss is a number (not Tensor)."""
with ops.Graph().as_default(), self.test_session():
with self.assertRaisesRegexp(TypeError, 'loss must be Tensor'):
model_fn.EstimatorSpec(
mode=model_fn.ModeKeys.TRAIN,
loss=1.,
train_op=control_flow_ops.no_op())
def testName(self):
with ops.name_scope("scope"):
queue = data_flow_ops.FIFOQueue(10, dtypes.float32, name="queue")
qr = queue_runner_impl.QueueRunner(queue, [control_flow_ops.no_op()])
self.assertEqual("scope/queue", qr.name)
queue_runner_impl.add_queue_runner(qr)
self.assertEqual(
1, len(ops.get_collection(ops.GraphKeys.QUEUE_RUNNERS, "scope")))
def _decay_weights_op(self, var):
if not self._decay_var_list or var in self._decay_var_list:
return var.assign_sub(self._weight_decay * var, self._use_locking)
return control_flow_ops.no_op()
def testLossNotScalar(self):
with ops.Graph().as_default(), self.cached_session():
with self.assertRaisesRegexp(ValueError, 'Loss must be scalar'):
model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN,
loss=constant_op.constant([1., 2.]),
train_op=control_flow_ops.no_op())
def _decay_weights_sparse_op(self, var, indices, scatter_add):
if not self._decay_var_list or var in self._decay_var_list:
update = -self._weight_decay * array_ops.gather(var, indices)
return scatter_add(var, indices, update, self._use_locking)
return control_flow_ops.no_op()
def result():
return control_flow_ops.no_op()
def testRequiredArgumentsSet(self):
"""Tests that no errors are raised when all required arguments are set."""
with ops.Graph().as_default(), self.cached_session():
model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN,
loss=constant_op.constant(1.),
train_op=control_flow_ops.no_op())
def restore(self, restored_tensors, restored_shapes):
return control_flow_ops.no_op()
def testLoss1DTensor(self):
"""Tests that no errors are raised when loss is 1D tensor."""
with ops.Graph().as_default(), self.cached_session():
model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN,
loss=constant_op.constant([1.]),
train_op=control_flow_ops.no_op())
def _wals_factorization_model_function(features, labels, mode, params):
"""Model function for the WALSFactorization estimator.
Args:
features: Dictionary of features. See WALSMatrixFactorization.
labels: Must be None.
mode: A model_fn.ModeKeys object.
params: Dictionary of parameters containing arguments passed to the
WALSMatrixFactorization constructor.
Returns:
A ModelFnOps object.
Raises:
ValueError: If `mode` is not recognized.
"""
assert labels is None
use_factors_weights_cache = (
params["use_factors_weights_cache_for_training"]
and mode == model_fn.ModeKeys.TRAIN)
use_gramian_cache = (params["use_gramian_cache_for_training"]
and mode == model_fn.ModeKeys.TRAIN)
max_sweeps = params["max_sweeps"]
model = factorization_ops.WALSModel(
params["num_rows"],
params["num_cols"],
params["embedding_dimension"],
unobserved_weight=params["unobserved_weight"],
regularization=params["regularization_coeff"],
row_init=params["row_init"],
col_init=params["col_init"],
num_row_shards=params["num_row_shards"],
num_col_shards=params["num_col_shards"],
row_weights=params["row_weights"],
col_weights=params["col_weights"],
use_factors_weights_cache=use_factors_weights_cache,
use_gramian_cache=use_gramian_cache)
# Get input rows and cols. We either update rows or columns depending on
# the value of row_sweep, which is maintained using a session hook.
input_rows = features[WALSMatrixFactorization.INPUT_ROWS]
input_cols = features[WALSMatrixFactorization.INPUT_COLS]
# TRAIN mode:
if mode == model_fn.ModeKeys.TRAIN:
# Training consists of the following ops (controlled using a SweepHook).
# Before a row sweep:
# row_update_prep_gramian_op
# initialize_row_update_op
# During a row sweep:
# update_row_factors_op
# Before a col sweep:
# col_update_prep_gramian_op
# initialize_col_update_op
# During a col sweep:
# update_col_factors_op
is_row_sweep_var = variable_scope.variable(
True,
trainable=False,
name="is_row_sweep",
collections=[ops.GraphKeys.GLOBAL_VARIABLES])
is_sweep_done_var = variable_scope.variable(
False,
trainable=False,
name="is_sweep_done",
collections=[ops.GraphKeys.GLOBAL_VARIABLES])
completed_sweeps_var = variable_scope.variable(
0,
trainable=False,
name=WALSMatrixFactorization.COMPLETED_SWEEPS,
collections=[ops.GraphKeys.GLOBAL_VARIABLES])
loss_var = variable_scope.variable(
0.,
trainable=False,
name=WALSMatrixFactorization.LOSS,
collections=[ops.GraphKeys.GLOBAL_VARIABLES])
# The root weighted squared error =
# \sqrt( \sum_{i,j} w_ij * (a_ij - r_ij)^2 / \sum_{i,j} w_ij )
rwse_var = variable_scope.variable(
0.,
trainable=False,
name=WALSMatrixFactorization.RWSE,
collections=[ops.GraphKeys.GLOBAL_VARIABLES])
summary.scalar("loss", loss_var)
summary.scalar("root_weighted_squared_error", rwse_var)
summary.scalar("completed_sweeps", completed_sweeps_var)
def create_axis_ops(sp_input, num_items, update_fn, axis_name):
"""Creates book-keeping and training ops for a given axis.
Args:
sp_input: A SparseTensor corresponding to the row or column batch.
num_items: An integer, the total number of items of this axis.
update_fn: A function that takes one argument (`sp_input`), and that
returns a tuple of
* new_factors: A float Tensor of the factor values after update.
* update_op: a TensorFlow op which updates the factors.
* loss: A float Tensor, the unregularized loss.
* reg_loss: A float Tensor, the regularization loss.
* sum_weights: A float Tensor, the sum of factor weights.
axis_name: A string that specifies the name of the axis.
Returns:
A tuple consisting of:
* reset_processed_items_op: A TensorFlow op, to be run before the
beginning of any sweep. It marks all items as not-processed.
* axis_train_op: A Tensorflow op, to be run during this axis' sweeps.
"""
processed_items_init = array_ops.fill(dims=[num_items],
value=False)
with ops.colocate_with(processed_items_init):
processed_items = variable_scope.variable(
processed_items_init,
collections=[ops.GraphKeys.GLOBAL_VARIABLES],
trainable=False,
name="processed_" + axis_name)
_, update_op, loss, reg, sum_weights = update_fn(sp_input)
input_indices = sp_input.indices[:, 0]
with ops.control_dependencies([
update_op,
state_ops.assign(loss_var, loss + reg),
state_ops.assign(rwse_var,
math_ops.sqrt(loss / sum_weights))
]):
with ops.colocate_with(processed_items):
update_processed_items = state_ops.scatter_update(
processed_items,
input_indices,
array_ops.ones_like(input_indices, dtype=dtypes.bool),
name="update_processed_{}_indices".format(axis_name))
with ops.control_dependencies([update_processed_items]):
is_sweep_done = math_ops.reduce_all(processed_items)
axis_train_op = control_flow_ops.group(
state_ops.assign(is_sweep_done_var, is_sweep_done),
state_ops.assign_add(
completed_sweeps_var,
math_ops.cast(is_sweep_done, dtypes.int32)),
name="{}_sweep_train_op".format(axis_name))
return processed_items.initializer, axis_train_op
reset_processed_rows_op, row_train_op = create_axis_ops(
input_rows, params["num_rows"],
lambda x: model.update_row_factors(sp_input=x,
transpose_input=False), "rows")
reset_processed_cols_op, col_train_op = create_axis_ops(
input_cols, params["num_cols"],
lambda x: model.update_col_factors(sp_input=x,
transpose_input=True), "cols")
switch_op = control_flow_ops.group(state_ops.assign(
is_row_sweep_var, math_ops.logical_not(is_row_sweep_var)),
reset_processed_rows_op,
reset_processed_cols_op,
name="sweep_switch_op")
row_prep_ops = [
model.row_update_prep_gramian_op, model.initialize_row_update_op
]
col_prep_ops = [
model.col_update_prep_gramian_op, model.initialize_col_update_op
]
init_op = model.worker_init
sweep_hook = _SweepHook(is_row_sweep_var, is_sweep_done_var, init_op,
row_prep_ops, col_prep_ops, row_train_op,
col_train_op, switch_op)
global_step_hook = _IncrementGlobalStepHook()
training_hooks = [sweep_hook, global_step_hook]
if max_sweeps is not None:
training_hooks.append(_StopAtSweepHook(max_sweeps))
return model_fn.ModelFnOps(mode=model_fn.ModeKeys.TRAIN,
predictions={},
loss=loss_var,
eval_metric_ops={},
train_op=control_flow_ops.no_op(),
training_hooks=training_hooks)
# INFER mode
elif mode == model_fn.ModeKeys.INFER:
projection_weights = features.get(
WALSMatrixFactorization.PROJECTION_WEIGHTS)
def get_row_projection():
return model.project_row_factors(
sp_input=input_rows,
projection_weights=projection_weights,
transpose_input=False)
def get_col_projection():
return model.project_col_factors(
sp_input=input_cols,
projection_weights=projection_weights,
transpose_input=True)
predictions = {
WALSMatrixFactorization.PROJECTION_RESULT:
control_flow_ops.cond(
features[WALSMatrixFactorization.PROJECT_ROW],
get_row_projection, get_col_projection)
}
return model_fn.ModelFnOps(mode=model_fn.ModeKeys.INFER,
predictions=predictions,
loss=None,
eval_metric_ops={},
train_op=control_flow_ops.no_op(),
training_hooks=[])
# EVAL mode
elif mode == model_fn.ModeKeys.EVAL:
def get_row_loss():
_, _, loss, reg, _ = model.update_row_factors(
sp_input=input_rows, transpose_input=False)
return loss + reg
def get_col_loss():
_, _, loss, reg, _ = model.update_col_factors(sp_input=input_cols,
transpose_input=True)
return loss + reg
loss = control_flow_ops.cond(
features[WALSMatrixFactorization.PROJECT_ROW], get_row_loss,
get_col_loss)
return model_fn.ModelFnOps(mode=model_fn.ModeKeys.EVAL,
predictions={},
loss=loss,
eval_metric_ops={},
train_op=control_flow_ops.no_op(),
training_hooks=[])
else:
raise ValueError("mode=%s is not recognized." % str(mode))
def testLossMissing(self):
with ops.Graph().as_default(), self.cached_session():
with self.assertRaisesRegexp(ValueError, 'Missing loss'):
model_fn.EstimatorSpec(mode=model_fn.ModeKeys.TRAIN,
train_op=control_flow_ops.no_op())
def _no_op_train_fn(loss):
del loss
return control_flow_ops.no_op()
def _run_graph(self,
a_shape,
b_shape,
nnz,
num_iters,
sort=False,
transpose_a=False,
transpose_b=False):
"""Run the graph and return its average execution time.
Args:
a_shape: int list, the shape of the a matrix.
b_shape: int list, the shape of the b matrix.
nnz: int, the number of non-zero elements in the mask.
num_iters: int, the number of iterations to run (the output is the average
execution time, over num_iters).
sort: Boolean, whether to sort the indices in the mask.
transpose_a: boolean, whether to transpose the a matrix.
transpose_b: boolean, whether to transpose the b matrix.
Returns:
The average duration of the masked_matmul op in seconds.
"""
graph = ops.Graph()
with graph.as_default(), session_lib.Session(graph=graph) as session:
mask_shape = [a_shape[0], b_shape[1]]
a_shape = a_shape if not transpose_a else [a_shape[1], a_shape[0]]
b_shape = b_shape if not transpose_b else [b_shape[1], b_shape[0]]
a_var = variables.Variable(random_ops.random_normal(a_shape))
b_var = variables.Variable(random_ops.random_normal(b_shape))
mask_indices_ph = array_ops.placeholder(dtypes.int64,
shape=[nnz, 2])
a_ph = array_ops.placeholder(dtypes.float32, shape=a_shape)
b_ph = array_ops.placeholder(dtypes.float32, shape=b_shape)
mask = self._make_sparse_mask(mask_shape, nnz, sort)
masked_prod = gen_factorization_ops.masked_matmul(
a_ph, b_ph, mask_indices_ph, transpose_a, transpose_b)
with ops.control_dependencies([masked_prod]):
result = control_flow_ops.no_op()
variables.global_variables_initializer().run()
avg_wall_time = 0
for _ in range(num_iters):
a, b, mask_indices = session.run([a_var, b_var, mask.indices])
feed_dict = {mask_indices_ph: mask_indices, a_ph: a, b_ph: b}
start_time = time.time()
session.run(result, feed_dict=feed_dict)
avg_wall_time += (time.time() - start_time) / num_iters
bench_name = (
"cpu nnz:{nnz} a_shape:{a_shape} b_shape:{b_shape} tr_a:{tr_a} "
"tr_b:{tr_b} sort:{sort}").format(nnz=nnz,
a_shape=a_shape,
b_shape=b_shape,
tr_a=int(transpose_a),
tr_b=int(transpose_b),
sort=int(sort))
print(bench_name + " - %f secs" % avg_wall_time)
name = bench_name.replace(", ",
"_").replace(":", "_").replace(" ", "_")
self.report_benchmark(name=name,
iters=num_iters,
wall_time=avg_wall_time)
return avg_wall_time
def initializer(self):
if context.executing_eagerly():
return control_flow_ops.no_op()
return self._initializer
def begin(self):
self._global_step_tensor = training_util.get_global_step()
self._current_train_op = control_flow_ops.no_op()
if self._global_step_tensor is None:
raise RuntimeError(
"Global step should be created to use SwitchTrainOp.")
def Foo(x):
y = logging_ops.Print(x, [x], "Hello")
with ops.control_dependencies([y]):
z = control_flow_ops.no_op()
with ops.control_dependencies([z]):
return x * 2
def create_train_op(total_loss,
optimizer,
global_step=_USE_GLOBAL_STEP,
update_ops=None,
variables_to_train=None,
transform_grads_fn=None,
summarize_gradients=False,
gate_gradients=tf_optimizer.Optimizer.GATE_OP,
aggregation_method=None,
colocate_gradients_with_ops=False,
check_numerics=True):
"""Creates an `Operation` that evaluates the gradients and returns the loss.
Args:
total_loss: A `Tensor` representing the total loss.
optimizer: A tf.Optimizer to use for computing the gradients.
global_step: A `Tensor` representing the global step variable. If left as
`_USE_GLOBAL_STEP`, then tf.contrib.framework.global_step() is used.
update_ops: An optional list of updates to execute. If `update_ops` is
`None`, then the update ops are set to the contents of the
`tf.GraphKeys.UPDATE_OPS` collection. If `update_ops` is not `None`, but
it doesn't contain all of the update ops in `tf.GraphKeys.UPDATE_OPS`, a
warning will be displayed.
variables_to_train: an optional list of variables to train. If None, it will
default to all tf.compat.v1.trainable_variables().
transform_grads_fn: A function which takes a single argument, a list of
gradient to variable pairs (tuples), performs any requested gradient
updates, such as gradient clipping or multipliers, and returns the updated
list.
summarize_gradients: Whether or not add summaries for each gradient.
gate_gradients: How to gate the computation of gradients. See tf.Optimizer.
aggregation_method: Specifies the method used to combine gradient terms.
Valid values are defined in the class `AggregationMethod`.
colocate_gradients_with_ops: Whether or not to try colocating the gradients
with the ops that generated them.
check_numerics: Whether or not we apply check_numerics.
Returns:
A `Tensor` that when evaluated, computes the gradients and returns the total
loss value.
"""
if global_step is _USE_GLOBAL_STEP:
global_step = training_util.get_or_create_global_step()
# Update ops use GraphKeys.UPDATE_OPS collection if update_ops is None.
global_update_ops = set(ops.get_collection(ops.GraphKeys.UPDATE_OPS))
if update_ops is None:
update_ops = global_update_ops
else:
update_ops = set(update_ops)
if not global_update_ops.issubset(update_ops):
logging.warning(
'update_ops in create_train_op does not contain all the '
'update_ops in GraphKeys.UPDATE_OPS')
# Make sure update_ops are computed before total_loss.
if update_ops:
with ops.control_dependencies(update_ops):
barrier = control_flow_ops.no_op(name='update_barrier')
total_loss = control_flow_ops.with_dependencies([barrier], total_loss)
if variables_to_train is None:
# Default to tf.compat.v1.trainable_variables()
variables_to_train = tf_variables.trainable_variables()
else:
# Make sure that variables_to_train are in
# tf.compat.v1.trainable_variables()
for v in variables_to_train:
assert v.trainable or v in tf_variables.trainable_variables()
assert variables_to_train
# Create the gradients. Note that apply_gradients adds the gradient
# computation to the current graph.
grads = optimizer.compute_gradients(
total_loss,
variables_to_train,
gate_gradients=gate_gradients,
aggregation_method=aggregation_method,
colocate_gradients_with_ops=colocate_gradients_with_ops)
if transform_grads_fn:
grads = transform_grads_fn(grads)
# Summarize gradients.
if summarize_gradients:
with ops.name_scope('summarize_grads'):
add_gradients_summaries(grads)
# Create gradient updates.
grad_updates = optimizer.apply_gradients(grads, global_step=global_step)
with ops.name_scope('train_op'):
# Make sure total_loss is valid.
if check_numerics:
total_loss = array_ops.check_numerics(total_loss,
'LossTensor is inf or nan')
# Ensure the train_tensor computes grad_updates.
train_op = control_flow_ops.with_dependencies([grad_updates],
total_loss)
# Add the operation used for training to the 'train_op' collection
train_ops = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
if train_op not in train_ops:
train_ops.append(train_op)
return train_op
def _fill_meta_graph_def(meta_graph_def, saveable_view, signature_functions,
namespace_whitelist):
"""Generates a MetaGraph which calls `signature_functions`.
Args:
meta_graph_def: The MetaGraphDef proto to fill.
saveable_view: The _SaveableView being exported.
signature_functions: A dictionary mapping signature keys to concrete
functions containing signatures to add to the MetaGraph.
namespace_whitelist: List of strings containing whitelisted op namespaces.
Returns:
A tuple of (_AssetInfo, Graph) containing the captured assets and
exported Graph generated from tracing the saveable_view.
"""
# List objects from the eager context to make sure Optimizers give us the
# right Graph-dependent variables.
accessible_objects = saveable_view.nodes
resource_initializer_functions = _trace_resource_initializers(
accessible_objects)
exported_graph = ops.Graph()
resource_initializer_ops = []
with exported_graph.as_default():
object_map, resource_map, asset_info = saveable_view.map_resources()
for resource_initializer_function in resource_initializer_functions:
asset_dependencies = []
for capture in resource_initializer_function.graph.external_captures:
asset_initializer = asset_info.asset_initializers_by_resource.get(
capture, None)
if asset_initializer is not None:
asset_dependencies.append(asset_initializer)
with ops.control_dependencies(asset_dependencies):
resource_initializer_ops.append(
_call_function_with_mapped_captures(
resource_initializer_function, [], resource_map))
resource_initializer_ops.extend(
asset_info.asset_initializers_by_resource.values())
with ops.control_dependencies(resource_initializer_ops):
init_op = control_flow_ops.no_op()
# Add the same op to the main_op collection and to the init_op
# signature. The collection is for compatibility with older loader APIs;
# only one will be executed.
meta_graph_def.collection_def[constants.MAIN_OP_KEY].node_list.value.append(
init_op.name)
meta_graph_def.signature_def[constants.INIT_OP_SIGNATURE_KEY].CopyFrom(
signature_def_utils.op_signature_def(
init_op, constants.INIT_OP_SIGNATURE_KEY))
# Saving an object-based checkpoint again gathers variables. We need to do the
# gathering from the eager context so Optimizers save the right set of
# variables, but want any operations associated with the save/restore to be in
# the exported graph (thus the `to_graph` argument).
saver = functional_saver.MultiDeviceSaver(
saveable_view.checkpoint_view.frozen_saveable_objects(
object_map=object_map, to_graph=exported_graph))
with exported_graph.as_default():
signatures = _generate_signatures(signature_functions, resource_map)
for concrete_function in saveable_view.concrete_functions:
concrete_function.add_to_graph()
saver_def = saver.to_proto()
meta_graph_def.saver_def.CopyFrom(saver_def)
graph_def = exported_graph.as_graph_def(add_shapes=True)
_verify_ops(graph_def, namespace_whitelist)
meta_graph_def.graph_def.CopyFrom(graph_def)
meta_graph_def.meta_info_def.tags.append(tag_constants.SERVING)
meta_graph_def.meta_info_def.tensorflow_version = versions.__version__
meta_graph_def.meta_info_def.tensorflow_git_version = (
versions.__git_version__)
# We currently always strip default attributes.
meta_graph_def.meta_info_def.stripped_default_attrs = True
meta_graph_def.meta_info_def.stripped_op_list.MergeFrom(
meta_graph.stripped_op_list_for_graph(meta_graph_def.graph_def))
meta_graph_def.asset_file_def.extend(asset_info.asset_defs)
for signature_key, signature in signatures.items():
meta_graph_def.signature_def[signature_key].CopyFrom(signature)
meta_graph.strip_graph_default_valued_attrs(meta_graph_def)
return asset_info, exported_graph
def result():
with ops.control_dependencies([thunk()]):
return control_flow_ops.no_op()
def create_mirrored_variable(strategy, real_mirrored_creator, class_mapping,
policy_mapping, **kwargs):
"""Create distributed variables with given synchronization and aggregation."""
# Figure out what collections this variable should be added to.
# We'll add the MirroredVariable to those collections instead.
var_collections = kwargs.pop("collections", None)
if var_collections is None:
var_collections = [ops.GraphKeys.GLOBAL_VARIABLES]
kwargs["collections"] = []
synchronization = _validate_synchronization(kwargs)
# Update synchronization in kwargs in case it's AUTO, which is converted to
# ON_WRITE.
kwargs["synchronization"] = synchronization
aggregation = _validate_aggregation(kwargs)
use_var_policy = getattr(strategy.extended, "_use_var_policy", False)
# Ignore user-specified caching device, not needed for mirrored variables.
kwargs.pop("caching_device", None)
# TODO(josh11b,apassos): It would be better if variable initialization
# was never recorded on the tape instead of having to do this manually
# here.
with tape.stop_recording():
value_list = real_mirrored_creator(**kwargs)
# MirroredVariable is recreated during saved_model loading, and its
# component variables (value_list) will have None initializer. We
# set their initializers to no_op so that consumer like
# `global_variables_initializer` wouldn't complain, as it groups all
# variables' initializers thus all variables have to have initializers.
for v in value_list:
# pylint:disable=protected-access
if hasattr(v, "_initializer_op") and v._initializer_op is None:
v._initializer_op = control_flow_ops.no_op()
# pylint:enable=protected-access
if use_var_policy:
var_policy_cls = policy_mapping.get(synchronization)
var_policy = var_policy_cls(aggregation=aggregation)
var_cls = class_mapping.get("VariableClass")
result = var_cls(strategy,
value_list,
aggregation,
var_policy=var_policy)
else:
var_cls = class_mapping.get(synchronization)
result = var_cls(strategy, value_list, aggregation)
# Add the wrapped variable to the requested collections.
# The handling of eager mode and the global step matches
# ResourceVariable._init_from_args().
if not context.executing_eagerly():
g = ops.get_default_graph()
# If "trainable" is True, next_creator() will add the member variables
# to the TRAINABLE_VARIABLES collection, so we manually remove
# them and replace with the MirroredVariable. We can't set
# "trainable" to False for next_creator() since that causes functions
# like implicit_gradients to skip those variables.
if kwargs.get("trainable", True):
var_collections.append(ops.GraphKeys.TRAINABLE_VARIABLES)
l = g.get_collection_ref(ops.GraphKeys.TRAINABLE_VARIABLES)
for value in value_list:
for i, trainable_variable in enumerate(l):
if value is trainable_variable:
del l[i]
break
g.add_to_collections(var_collections, result)
elif ops.GraphKeys.GLOBAL_STEP in var_collections:
ops.add_to_collections(ops.GraphKeys.GLOBAL_STEP, result)
return result
def _decay_weights_sparse_op(self, var, indices, scatter_add):
if not self._decay_var_list or var in self._decay_var_list:
return scatter_add(var, indices, -self._weight_decay * var,
self._use_locking)
return control_flow_ops.no_op()
def _fill_meta_graph_def(meta_graph_def, obj, signature_functions,
object_saver):
"""Generates a MetaGraph which calls `signature_functions`.
Args:
meta_graph_def: The MetaGraphDef proto to fill.
obj: The checkpointable object being exported.
signature_functions: A dictionary mapping signature keys to concrete
functions containing signatures to add to the MetaGraph.
object_saver: A CheckpointableSaver to add to the MetaGraph.
Returns:
An _AssetInfo, which contains information to help creating the SavedModel.
"""
signatures = {}
# List objects from the eager context to make sure Optimizers give us the
# right Graph-dependent variables.
accessible_objects = util.list_objects(obj)
resource_initializer_functions = _trace_resource_initializers(
accessible_objects)
exported_graph = ops.Graph()
resource_initializer_ops = []
with exported_graph.as_default():
object_map, resource_map, asset_info = _map_resources(
accessible_objects)
for resource_initializer_function in resource_initializer_functions:
asset_dependencies = []
for capture in resource_initializer_function.graph.external_captures:
asset_initializer = asset_info.asset_initializers_by_resource.get(
capture, None)
if asset_initializer is not None:
asset_dependencies.append(asset_initializer)
with ops.control_dependencies(asset_dependencies):
resource_initializer_ops.append(
_call_function_with_mapped_captures(
resource_initializer_function, [], resource_map))
with ops.control_dependencies(resource_initializer_ops):
init_op = control_flow_ops.no_op()
# Add the same op to the main_op collection and to the init_op
# signature. The collection is for compatibility with older loader APIs;
# only one will be executed.
meta_graph_def.collection_def[
constants.MAIN_OP_KEY].node_list.value.append(init_op.name)
meta_graph_def.signature_def[constants.INIT_OP_SIGNATURE_KEY].CopyFrom(
signature_def_utils.op_signature_def(
init_op, constants.INIT_OP_SIGNATURE_KEY))
# Saving an object-based checkpoint again gathers variables. We need to do the
# gathering from the eager context so Optimizers save the right set of
# variables, but want any operations associated with the save/restore to be in
# the exported graph (thus the `to_graph` argument).
saver = object_saver.freeze(object_map=object_map, to_graph=exported_graph)
# We must resolve the concrete function to add to MetaGraph while in eager
# mode.
concrete_functions = []
for accessible_object in accessible_objects:
for function in function_serialization.list_all_polymorphic_functions(
accessible_object).values():
concrete_functions.extend(
function_serialization.list_all_concrete_functions(function))
with exported_graph.as_default():
signatures = _generate_signatures(signature_functions, resource_map)
for concrete_function in concrete_functions:
concrete_function.add_to_graph()
saver_def = saver.to_proto()
meta_graph_def.saver_def.CopyFrom(saver_def)
graph_def = exported_graph.as_graph_def(add_shapes=True)
# Clean reference cycles so repeated export()s don't make work for the garbage
# collector.
ops.dismantle_graph(exported_graph)
meta_graph_def.graph_def.CopyFrom(graph_def)
meta_graph_def.meta_info_def.tags.append(tag_constants.SERVING)
meta_graph_def.asset_file_def.extend(asset_info.asset_defs)
for signature_key, signature in signatures.items():
meta_graph_def.signature_def[signature_key].CopyFrom(signature)
meta_graph.strip_graph_default_valued_attrs(meta_graph_def)
return asset_info
def restore(self, restored_tensors, restored_shapes):
"""Called to restore TensorFlow state (nothing to do)."""
return control_flow_ops.no_op()
def init(self):
"""The table initialization op."""
if self._table:
return self._table.init
with ops.name_scope(None, "init"):
return control_flow_ops.no_op()
def _assert_positive_definite(self):
return control_flow_ops.no_op("assert_positive_definite")
def _assert_self_adjoint(self):
return control_flow_ops.no_op("assert_self_adjoint")
def _init_q():
with ops.control_dependencies(
[q.enqueue_many(math_ops.range(num_points))]):
return control_flow_ops.no_op()
def _assert_non_singular(self):
return control_flow_ops.no_op("assert_non_singular")
def no_op():
return control_flow_ops.no_op('no_update')
def testNoOpIsNone(self):
self.assertTrue(control_flow_ops.no_op() is None)
def apply_gradients(self, grads_and_vars, global_step=None, name=None):
"""Apply gradients to variables.
This contains most of the synchronization implementation and also wraps the
apply_gradients() from the real optimizer.
Args:
grads_and_vars: List of (gradient, variable) pairs as returned by
compute_gradients().
global_step: Optional Variable to increment by one after the
variables have been updated.
name: Optional name for the returned operation. Default to the
name passed to the Optimizer constructor.
Returns:
train_op: The op to dequeue a token so the replicas can exit this batch
and start the next one. This is executed by each replica.
Raises:
ValueError: If the grads_and_vars is empty.
ValueError: If global step is not provided, the staleness cannot be
checked.
"""
if not grads_and_vars:
raise ValueError("Must supply at least one variable")
if global_step is None:
raise ValueError("Global step is required to check staleness")
self._global_step = global_step
train_ops = []
aggregated_grad = []
var_list = []
# local_anchor op will be placed on this worker task by default.
local_anchor = control_flow_ops.no_op()
# Colocating local_step variable prevents it being placed on the PS.
with ops.colocate_with(local_anchor):
self._local_step = variable_scope.variable(
initial_value=0,
trainable=False,
collections=[ops.GraphKeys.LOCAL_VARIABLES],
dtype=global_step.dtype.base_dtype,
name="sync_rep_local_step")
self.local_step_init_op = state_ops.assign(self._local_step,
global_step)
chief_init_ops = [self.local_step_init_op]
self.ready_for_local_init_op = variables.report_uninitialized_variables(
variables.global_variables())
with ops.name_scope(None, self._name):
for grad, var in grads_and_vars:
var_list.append(var)
with ops.device(var.device):
# Dense gradients.
if grad is None:
aggregated_grad.append(None) # pass-through.
continue
elif isinstance(grad, ops.Tensor):
grad_accum = data_flow_ops.ConditionalAccumulator(
grad.dtype,
shape=var.get_shape(),
shared_name=var.name + "/grad_accum")
train_ops.append(
grad_accum.apply_grad(grad,
local_step=self._local_step))
aggregated_grad.append(
grad_accum.take_grad(self._replicas_to_aggregate))
else:
if not isinstance(grad, ops.IndexedSlices):
raise ValueError("Unknown grad type!")
grad_accum = data_flow_ops.SparseConditionalAccumulator(
grad.dtype,
shape=(),
shared_name=var.name + "/grad_accum")
train_ops.append(
grad_accum.apply_indexed_slices_grad(
grad, local_step=self._local_step))
aggregated_grad.append(
grad_accum.take_indexed_slices_grad(
self._replicas_to_aggregate))
self._accumulator_list.append((grad_accum, var.device))
aggregated_grads_and_vars = zip(aggregated_grad, var_list)
# sync_op will be assigned to the same device as the global step.
with ops.device(global_step.device), ops.name_scope(""):
update_op = self._opt.apply_gradients(
aggregated_grads_and_vars, global_step)
# Create token queue.
with ops.device(global_step.device), ops.name_scope(""):
sync_token_queue = (data_flow_ops.FIFOQueue(
-1,
global_step.dtype.base_dtype,
shapes=(),
name="sync_token_q",
shared_name="sync_token_q"))
self._sync_token_queue = sync_token_queue
# dummy_queue is passed to the queue runner. Don't use the real queues
# because the queue runner doesn't automatically reopen it once it
# closed queues in PS devices.
dummy_queue = (data_flow_ops.FIFOQueue(
1,
types_pb2.DT_INT32,
shapes=(),
name="dummy_queue",
shared_name="dummy_queue"))
with ops.device(global_step.device), ops.name_scope(""):
# Replicas have to wait until they can get a token from the token queue.
with ops.control_dependencies(train_ops):
token = sync_token_queue.dequeue()
train_op = state_ops.assign(self._local_step, token)
with ops.control_dependencies([update_op]):
# Sync_op needs to insert tokens to the token queue at the end of the
# step so the replicas can fetch them to start the next step.
tokens = array_ops.fill([self._tokens_per_step],
global_step)
sync_op = sync_token_queue.enqueue_many((tokens, ))
if self._variable_averages is not None:
with ops.control_dependencies([sync_op
]), ops.name_scope(""):
sync_op = self._variable_averages.apply(
self._variables_to_average)
self._chief_queue_runner = queue_runner.QueueRunner(
dummy_queue, [sync_op])
for accum, dev in self._accumulator_list:
with ops.device(dev):
chief_init_ops.append(
accum.set_global_step(global_step,
name="SetGlobalStep"))
self.chief_init_op = control_flow_ops.group(*(chief_init_ops))
self._gradients_applied = True
return train_op
def update(self, grads): del grads return control_flow_ops.no_op(), True
