def loop_and_compute():
p13n_metrics = collections.OrderedDict()
for name, personalize_fn in personalize_fns.items():
tf_computation_utils.assign(model_weights, initial_model_weights)
py_typecheck.check_callable(personalize_fn)
p13n_metrics[name] = personalize_fn(model, train_data, test_data,
context)
return p13n_metrics
def _tf_client_eval(incoming_model_weights, dataset):
"""Evaluation TF work."""
tf_computation_utils.assign(model.weights, incoming_model_weights)
def reduce_fn(prev_loss, batch):
model_output = model.forward_pass(batch, training=False)
return prev_loss + tf.cast(model_output.loss, tf.float64)
dataset.reduce(tf.constant(0.0, dtype=tf.float64), reduce_fn)
return collections.OrderedDict([('local_outputs',
model.report_local_outputs())])
def _mnist_batch_train(model, batch):
optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.01)
model_vars = tf_computation_utils.create_variables('v', _mnist_model_type)
assign_vars_op = tf_computation_utils.assign(model_vars, model)
with tf.control_dependencies([assign_vars_op]):
train_op = optimizer.minimize(_mnist_batch_loss(model_vars, batch))
with tf.control_dependencies([train_op]):
return tf_computation_utils.identity(model_vars)
def _tf_client_eval(incoming_model_weights, dataset):
"""Evaluation TF work."""
tf_computation_utils.assign(model.weights, incoming_model_weights)
def reduce_fn(prev_loss, batch):
model_output = model.forward_pass(batch, training=False)
return prev_loss + tf.cast(model_output.loss, tf.float64)
dataset_reduce_fn = dataset_reduce.build_dataset_reduce_fn(
use_experimental_simulation_loop)
dataset_reduce_fn(
reduce_fn=reduce_fn,
dataset=dataset,
initial_state_fn=lambda: tf.constant(0, dtype=tf.float64))
return collections.OrderedDict([('local_outputs',
model.report_local_outputs())])
def test_assign_with_no_nesting(self):
with tf.Graph().as_default() as graph:
v = tf.Variable(0, name='foo', dtype=tf.int32, shape=[])
c = tf.constant(10, dtype=tf.int32, shape=[])
op = tf_computation_utils.assign(v, c)
with tf.compat.v1.Session(graph=graph) as sess:
sess.run(v.initializer)
sess.run(op)
result = sess.run(v)
self.assertEqual(result, 10)
def test_assign_with_anonymous_tuple(self):
with tf.Graph().as_default() as graph:
v = tf.Variable(0, name='foo', dtype=tf.int32, shape=[])
c = tf.constant(10, dtype=tf.int32, shape=[])
v_tuple = anonymous_tuple.AnonymousTuple([('bar', v)])
c_tuple = anonymous_tuple.AnonymousTuple([('bar', c)])
op = tf_computation_utils.assign(v_tuple, c_tuple)
with tf.compat.v1.Session(graph=graph) as sess:
sess.run(v.initializer)
sess.run(op)
result = sess.run(v)
self.assertEqual(result, 10)
def test_assign_with_unordered_dict(self):
with tf.Graph().as_default() as graph:
v = tf.Variable(0, name='foo', dtype=tf.int32, shape=[])
c = tf.constant(10, dtype=tf.int32, shape=[])
v_dict = {'bar': v}
c_dict = {'bar': c}
op = tf_computation_utils.assign(v_dict, c_dict)
with tf.compat.v1.Session(graph=graph) as sess:
sess.run(v.initializer)
sess.run(op)
result = sess.run(v)
self.assertEqual(result, 10)
def assign_and_compute():
tf_computation_utils.assign(model_weights, initial_model_weights)
py_typecheck.check_callable(baseline_evaluate_fn)
return baseline_evaluate_fn(model, test_data)
