def server_optimzier_fn(model_weights):
model_weight_shape = tf.nest.map_structure(tf.shape, model_weights)
return optimizer_utils.DPFTRLMServerOptimizer(
learning_rate=0.1,
momentum=0.9,
noise_std=1e-5,
model_weight_shape=model_weight_shape)
def server_optimizer_fn(model_weights):
model_weight_specs = tf.nest.map_structure(
lambda v: tf.TensorSpec(v.shape, v.dtype), model_weights)
return optimizer_utils.DPFTRLMServerOptimizer(
learning_rate=0.1,
momentum=0.9,
noise_std=1e-5,
model_weight_specs=model_weight_specs,
efficient_tree=True,
use_nesterov=True)
def _server_optimizer_fn(model_weights, name, learning_rate, noise_std):
"""Returns server optimizer."""
model_weight_specs = tf.nest.map_structure(
lambda v: tf.TensorSpec(v.shape, v.dtype), model_weights)
if name == 'sgd':
return optimizer_utils.SGDServerOptimizer(learning_rate)
elif name == 'sgdm':
return optimizer_utils.DPSGDMServerOptimizer(
learning_rate,
momentum=FLAGS.server_momentum,
noise_std=0,
model_weight_specs=model_weight_specs)
elif name == 'dpftrl':
return optimizer_utils.DPFTRLMServerOptimizer(
learning_rate,
momentum=0,
noise_std=noise_std,
model_weight_specs=model_weight_specs)
elif name == 'dpsgd':
return optimizer_utils.DPSGDMServerOptimizer(
learning_rate,
momentum=0,
noise_std=noise_std,
model_weight_specs=model_weight_specs)
elif name == 'dpsgdm':
return optimizer_utils.DPSGDMServerOptimizer(
learning_rate,
momentum=FLAGS.server_momentum,
noise_std=noise_std,
model_weight_specs=model_weight_specs)
elif name == 'dpftrlm':
return optimizer_utils.DPFTRLMServerOptimizer(
learning_rate,
momentum=FLAGS.server_momentum,
noise_std=noise_std,
model_weight_specs=model_weight_specs)
else:
raise ValueError('Unknown server optimizer name {}'.format(name))
def _run_ftrl():
model_variables = _create_model_variables()
model_weight_shape = tf.nest.map_structure(tf.shape, model_variables)
grad = tf.nest.map_structure(tf.ones_like, model_variables)
optimizer = optimizer_utils.DPFTRLMServerOptimizer(
learning_rate=lr,
momentum=momentum,
noise_std=0.0,
model_weight_shape=model_weight_shape,
use_nesterov=nesterov)
state = optimizer.init_state()
for i in range(steps):
state = optimizer.model_update(state, model_variables, grad, i)
self.assertLen(model_variables, 2)
return tf.nest.flatten(model_variables)