def test_server_eager_mode(self, optimizer_fn, updated_val, num_optimizer_vars): model_fn = lambda: model_examples.TrainableLinearRegression(feature_dim =2) server_state = optimizer_utils.server_init(model_fn, optimizer_fn, (), ()) train_vars = server_state.model.trainable self.assertAllClose(train_vars['a'].numpy(), np.array([[0.0], [0.0]])) self.assertEqual(train_vars['b'].numpy(), 0.0) self.assertEqual(server_state.model.non_trainable['c'].numpy(), 0.0) self.assertLen(server_state.optimizer_state, num_optimizer_vars) weights_delta = tensor_utils.to_odict({ 'a': tf.constant([[1.0], [0.0]]), 'b': tf.constant(1.0) }) server_state = optimizer_utils.server_update_model( server_state, weights_delta, model_fn, optimizer_fn) train_vars = server_state.model.trainable # For SGD: learning_Rate=0.1, update=[1.0, 0.0], initial model=[0.0, 0.0], # so updated_val=0.1 self.assertAllClose(train_vars['a'].numpy(), [[updated_val], [0.0]]) self.assertAllClose(train_vars['b'].numpy(), updated_val) self.assertEqual(server_state.model.non_trainable['c'].numpy(), 0.0)
def test_server_eager_mode(self, optimizer_fn, updated_val, num_optimizer_vars): model_fn = lambda: model_examples.TrainableLinearRegression(feature_dim =2) server_state = optimizer_utils.server_init(model_fn, optimizer_fn, (), ()) model_vars = self.evaluate(server_state.model) train_vars = model_vars.trainable self.assertLen(train_vars, 2) self.assertAllClose(train_vars['a'], [[0.0], [0.0]]) self.assertEqual(train_vars['b'], 0.0) self.assertEqual(model_vars.non_trainable, {'c': 0.0}) self.assertLen(server_state.optimizer_state, num_optimizer_vars) weights_delta = collections.OrderedDict([ ('a', tf.constant([[1.0], [0.0]])), ('b', tf.constant(1.0)), ]) server_state = optimizer_utils.server_update_model( server_state, weights_delta, model_fn, optimizer_fn) model_vars = self.evaluate(server_state.model) train_vars = model_vars.trainable # For SGD: learning_Rate=0.1, update=[1.0, 0.0], initial model=[0.0, 0.0], # so updated_val=0.1 self.assertLen(train_vars, 2) self.assertAllClose(train_vars['a'], [[updated_val], [0.0]]) self.assertAllClose(train_vars['b'], updated_val) self.assertEqual(model_vars.non_trainable, {'c': 0.0})
def test_server_graph_mode(self): optimizer_fn = lambda: gradient_descent.SGD(learning_rate=0.1) model_fn = lambda: model_examples.TrainableLinearRegression(feature_dim =2) # Explicitly entering a graph as a default enables graph-mode. with tf.Graph().as_default() as g: server_state_op = optimizer_utils.server_init( model_fn, optimizer_fn, (), ()) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) g.finalize() with self.session() as sess: sess.run(init_op) server_state = sess.run(server_state_op) train_vars = server_state.model.trainable self.assertAllClose(train_vars['a'], [[0.0], [0.0]]) self.assertEqual(train_vars['b'], 0.0) self.assertEqual(server_state.model.non_trainable['c'], 0.0) self.assertEqual(server_state.optimizer_state, [0.0]) with tf.Graph().as_default() as g: # N.B. Must use a fresh graph so variable names are the same. weights_delta = tensor_utils.to_odict({ 'a': tf.constant([[1.0], [0.0]]), 'b': tf.constant(2.0) }) update_op = optimizer_utils.server_update_model( server_state, weights_delta, model_fn, optimizer_fn) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) g.finalize() with self.session() as sess: sess.run(init_op) server_state = sess.run(update_op) train_vars = server_state.model.trainable # learning_Rate=0.1, update is [1.0, 0.0], initial model is [0.0, 0.0]. self.assertAllClose(train_vars['a'], [[0.1], [0.0]]) self.assertAllClose(train_vars['b'], 0.2) self.assertEqual(server_state.model.non_trainable['c'], 0.0)
def test_server_eager_mode(self, optimizer_fn, updated_val, num_optimizer_vars): model_fn = lambda: model_examples.LinearRegression(feature_dim=2) server_state = optimizer_utils.server_init(model_fn, optimizer_fn, (), ()) model_vars = self.evaluate(server_state.model) train_vars = model_vars.trainable self.assertLen(train_vars, 2) self.assertAllClose(train_vars, [np.zeros((2, 1)), 0.0]) self.assertAllClose(model_vars.non_trainable, [0.0]) self.assertLen(server_state.optimizer_state, num_optimizer_vars) weights_delta = [tf.constant([[1.0], [0.0]]), tf.constant(1.0)] server_state = optimizer_utils.server_update_model( server_state, weights_delta, model_fn, optimizer_fn) model_vars = self.evaluate(server_state.model) train_vars = model_vars.trainable # For SGD: learning_Rate=0.1, update=[1.0, 0.0], initial model=[0.0, 0.0], # so updated_val=0.1 self.assertLen(train_vars, 2) self.assertAllClose(train_vars, [[[updated_val], [0.0]], updated_val]) self.assertAllClose(model_vars.non_trainable, [0.0])