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)