def test_to_odict(self):
d1 = {'b': 2, 'a': 1}
odict1 = tensor_utils.to_odict(d1)
self.assertIsInstance(odict1, collections.OrderedDict)
self.assertCountEqual(d1, odict1)
odict2 = tensor_utils.to_odict(odict1)
self.assertEqual(odict1, odict2)
with self.assertRaises(TypeError):
tensor_utils.to_odict({1: 'a', 2: 'b'})
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 client_update(model, dataset, initial_weights):
"""Updates client model.
Args:
model: A `tff.learning.Model`.
dataset: A 'tf.data.Dataset'.
initial_weights: A `tff.learning.Model.weights` from server.
Returns:
A 'ClientOutput`.
"""
tf.nest.map_structure(tf.assign, _get_weights(model), initial_weights)
@tf.function
def reduce_fn(num_examples_sum, batch):
"""Runs `tff.learning.Model.train_on_batch` on local client batch."""
output = model.train_on_batch(batch)
return num_examples_sum + tf.shape(output.predictions)[0]
num_examples_sum = dataset.reduce(initial_state=tf.constant(0),
reduce_func=reduce_fn)
weights_delta = tf.nest.map_structure(tf.subtract,
_get_weights(model).trainable,
initial_weights.trainable)
aggregated_outputs = model.report_local_outputs()
weights_delta_weight = tf.cast(num_examples_sum, tf.float32)
return ClientOutput(
weights_delta, weights_delta_weight, aggregated_outputs,
tensor_utils.to_odict({
'num_examples': num_examples_sum,
}))
def __call__(self, dataset, initial_weights):
# TODO(b/123898430): The control dependencies below have been inserted as a
# temporary workaround. These control dependencies need to be removed, and
# defuns and datasets supported together fully.
model = self._model
# TODO(b/113112108): Remove this temporary workaround and restore check for
# `tf.data.Dataset` after subclassing the currently used custom data set
# representation from it.
if 'Dataset' not in str(type(dataset)):
raise TypeError('Expected a data set, found {}.'.format(
py_typecheck.type_string(type(dataset))))
# TODO(b/120801384): We should initialize model.local_variables here.
# Or, we may just need a convention that TFF initializes all variables
# before invoking the TF function.
# We must assign to a variable here in order to use control_dependencies.
dummy_weights = nest.map_structure(tf.assign, model.weights,
initial_weights)
with tf.control_dependencies(list(dummy_weights.trainable.values())):
def reduce_fn(dummy_state, batch):
"""Runs `tff.learning.Model.train_on_batch` on local client batch."""
output = model.train_on_batch(batch)
tf.assign_add(self._num_examples,
tf.shape(output.predictions)[0])
return dummy_state
# TODO(b/124477598): Remove dummy_output when b/121400757 fixed.
dummy_output = dataset.reduce(initial_state=tf.constant(0.0),
reduce_func=reduce_fn)
with tf.control_dependencies([dummy_output]):
weights_delta = nest.map_structure(tf.subtract,
model.weights.trainable,
initial_weights.trainable)
aggregated_outputs = model.report_local_outputs()
weights_delta_weight = self._client_weight_fn(aggregated_outputs) # pylint:disable=not-callable
# TODO(b/122071074): Consider moving this functionality into
# tff.federated_average?
weights_delta, has_non_finite_delta = (
tensor_utils.zero_all_if_any_non_finite(weights_delta))
weights_delta_weight = tf.cond(tf.equal(has_non_finite_delta, 0),
lambda: weights_delta_weight,
lambda: tf.constant(0))
return optimizer_utils.ClientOutput(
weights_delta, weights_delta_weight, aggregated_outputs,
tensor_utils.to_odict({
'num_examples':
self._num_examples.value(),
'has_non_finite_delta':
has_non_finite_delta,
'workaround for b/121400757':
dummy_output,
}))
def __call__(self, dataset, initial_weights):
# TODO(b/113112108): Remove this temporary workaround and restore check for
# `tf.data.Dataset` after subclassing the currently used custom data set
# representation from it.
if 'Dataset' not in str(type(dataset)):
raise TypeError('Expected a data set, found {}.'.format(
py_typecheck.type_string(type(dataset))))
model = self._model
dummy_weights = nest.map_structure(tf.assign, model.weights,
initial_weights)
def reduce_fn(accumulated_grads, batch):
"""Runs forward_pass on batch."""
with tf.contrib.eager.GradientTape() as tape:
output = model.forward_pass(batch)
with tf.control_dependencies(list(output)):
flat_vars = nest.flatten(model.weights.trainable)
grads = nest.pack_sequence_as(
accumulated_grads, tape.gradient(output.loss, flat_vars))
if self._batch_weight_fn is not None:
batch_weight = self._batch_weight_fn(batch)
else:
batch_weight = tf.cast(
tf.shape(output.predictions)[0], tf.float32)
tf.assign_add(self._batch_weight_sum, batch_weight)
return nest.map_structure(
lambda accumulator, grad: accumulator + batch_weight * grad,
accumulated_grads, grads)
with tf.control_dependencies(list(dummy_weights.trainable.values())):
self._grad_sum_vars = dataset.reduce(
initial_state=self._grad_sum_vars, reduce_func=reduce_fn)
with tf.control_dependencies(
[tf.identity(v) for v in self._grad_sum_vars.values()]):
# For SGD, the delta is just the negative of the average gradient:
weights_delta = nest.map_structure(
lambda gradient: -1.0 * gradient / self._batch_weight_sum,
self._grad_sum_vars)
weights_delta, has_non_finite_delta = (
tensor_utils.zero_all_if_any_non_finite(weights_delta))
weights_delta_weight = tf.cond(tf.equal(has_non_finite_delta, 0),
lambda: self._batch_weight_sum,
lambda: tf.constant(0.0))
return optimizer_utils.ClientOutput(
weights_delta, weights_delta_weight,
model.report_local_outputs(),
tensor_utils.to_odict({
'client_weight':
weights_delta_weight,
'has_non_finite_delta':
has_non_finite_delta,
}))
def __call__(self, dataset, initial_weights):
# N.B. When not in eager mode, this code must be wrapped as a defun
# as it uses program-order semantics to avoid adding many explicit
# control dependencies.
model = self._model
py_typecheck.check_type(dataset, tf.data.Dataset)
nest.map_structure(tf.assign, model.weights, initial_weights)
@tf.contrib.eager.function(autograph=False)
def reduce_fn(dummy_state, batch):
"""Runs forward_pass on batch."""
with tf.contrib.eager.GradientTape() as tape:
output = model.forward_pass(batch)
flat_vars = nest.flatten(model.weights.trainable)
grads = nest.pack_sequence_as(
self._grad_sum_vars, tape.gradient(output.loss, flat_vars))
if self._batch_weight_fn is not None:
batch_weight = self._batch_weight_fn(batch)
else:
batch_weight = tf.cast(
tf.shape(output.predictions)[0], tf.float32)
tf.assign_add(self._batch_weight_sum, batch_weight)
nest.map_structure(
lambda v, g: # pylint:disable=g-long-lambda
tf.assign_add(v, batch_weight * g),
self._grad_sum_vars,
grads)
return dummy_state
# TODO(b/121400757): Remove dummy_output when bug fixed.
dummy_output = dataset.reduce(initial_state=tf.constant(0.0),
reduce_func=reduce_fn)
# For SGD, the delta is just the negative of the average gradient:
# TODO(b/109733734): Might be better to send the weighted grad sums
# and the denominator separately?
weights_delta = nest.map_structure(
lambda g: -1.0 * g / self._batch_weight_sum, self._grad_sum_vars)
weights_delta, has_non_finite_delta = (
tensor_utils.zero_all_if_any_non_finite(weights_delta))
weights_delta_weight = tf.cond(tf.equal(has_non_finite_delta, 0),
lambda: self._batch_weight_sum,
lambda: tf.constant(0.0))
return optimizer_utils.ClientOutput(
weights_delta, weights_delta_weight, model.report_local_outputs(),
tensor_utils.to_odict({
'client_weight': weights_delta_weight,
'has_non_finite_delta': has_non_finite_delta,
'workaround for b/121400757': dummy_output,
}))
def __call__(self, dataset, initial_weights):
# TODO(b/113112108): Remove this temporary workaround and restore check for
# `tf.data.Dataset` after subclassing the currently used custom data set
# representation from it.
if 'Dataset' not in str(type(dataset)):
raise TypeError('Expected a data set, found {}.'.format(
py_typecheck.type_string(type(dataset))))
model = self._model
tf.nest.map_structure(lambda a, b: a.assign(b), model.weights,
initial_weights)
@tf.function
def reduce_fn(num_examples_sum, batch):
"""Runs `tff.learning.Model.train_on_batch` on local client batch."""
output = model.train_on_batch(batch)
if output.num_examples is None:
return num_examples_sum + tf.shape(output.predictions)[0]
else:
return num_examples_sum + output.num_examples
num_examples_sum = dataset.reduce(initial_state=tf.constant(0),
reduce_func=reduce_fn)
weights_delta = tf.nest.map_structure(tf.subtract,
model.weights.trainable,
initial_weights.trainable)
aggregated_outputs = model.report_local_outputs()
# TODO(b/122071074): Consider moving this functionality into
# tff.federated_mean?
weights_delta, has_non_finite_delta = (
tensor_utils.zero_all_if_any_non_finite(weights_delta))
if self._client_weight_fn is None:
weights_delta_weight = tf.cast(num_examples_sum, tf.float32)
else:
weights_delta_weight = self._client_weight_fn(aggregated_outputs)
# Zero out the weight if there are any non-finite values.
if has_non_finite_delta > 0:
weights_delta_weight = tf.constant(0.0)
return optimizer_utils.ClientOutput(
weights_delta, weights_delta_weight, aggregated_outputs,
tensor_utils.to_odict({
'num_examples': num_examples_sum,
'has_non_finite_delta': has_non_finite_delta,
}))
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 __call__(self, dataset, initial_weights):
model = self._model
# TODO(b/113112108): Remove this temporary workaround and restore check for
# `tf.data.Dataset` after subclassing the currently used custom data set
# representation from it.
if 'Dataset' not in str(type(dataset)):
raise TypeError('Expected a data set, found {}.'.format(
py_typecheck.type_string(type(dataset))))
tf.nest.map_structure(lambda a, b: a.assign(b), model.weights,
initial_weights)
flat_trainable_weights = tuple(tf.nest.flatten(
model.weights.trainable))
@tf.function
def reduce_fn(state, batch):
"""Runs forward_pass on batch and sums the weighted gradients."""
flat_accumulated_grads, batch_weight_sum = state
with tf.GradientTape() as tape:
output = model.forward_pass(batch)
flat_grads = tape.gradient(output.loss, flat_trainable_weights)
if self._batch_weight_fn is not None:
batch_weight = self._batch_weight_fn(batch)
else:
batch_weight = tf.cast(
tf.shape(output.predictions)[0], tf.float32)
flat_accumulated_grads = tuple(
accumulator + batch_weight * grad for accumulator, grad in zip(
flat_accumulated_grads, flat_grads))
# The TF team is aware of an optimization in the reduce state to avoid
# doubling the number of required variables here (e.g. keeping two copies
# of all gradients). If you're looking to optimize memory usage this might
# be a place to look.
return (flat_accumulated_grads, batch_weight_sum + batch_weight)
def _zero_initial_state():
"""Create a tuple of (tuple of gradient accumulators, batch weight sum)."""
return (tuple(tf.zeros_like(w)
for w in flat_trainable_weights), tf.constant(0.0))
flat_grad_sums, batch_weight_sum = self._dataset_reduce_fn(
reduce_fn=reduce_fn,
dataset=dataset,
initial_state_fn=_zero_initial_state)
grad_sums = tf.nest.pack_sequence_as(model.weights.trainable,
flat_grad_sums)
# For SGD, the delta is just the negative of the average gradient:
weights_delta = tf.nest.map_structure(
lambda gradient: -1.0 * gradient / batch_weight_sum, grad_sums)
weights_delta, has_non_finite_delta = (
tensor_utils.zero_all_if_any_non_finite(weights_delta))
if has_non_finite_delta > 0:
weights_delta_weight = tf.constant(0.0)
else:
weights_delta_weight = batch_weight_sum
return optimizer_utils.ClientOutput(
weights_delta, weights_delta_weight, model.report_local_outputs(),
tensor_utils.to_odict({
'client_weight': weights_delta_weight,
'has_non_finite_delta': has_non_finite_delta,
}))
def __new__(cls, trainable, non_trainable):
return super(ModelWeights,
cls).__new__(cls, tensor_utils.to_odict(trainable),
tensor_utils.to_odict(non_trainable))
def __call__(self, model, optimizer, benign_dataset, malicious_dataset,
client_is_malicious, initial_weights):
"""Updates client model with client potentially being malicious.
Args:
model: A `tff.learning.Model`.
optimizer: A 'tf.keras.optimizers.Optimizer'.
benign_dataset: A 'tf.data.Dataset' consisting of benign dataset.
malicious_dataset: A 'tf.data.Dataset' consisting of malicious dataset.
client_is_malicious: A 'tf.bool' showing whether the client is malicious.
initial_weights: A `tff.learning.Model.weights` from server.
Returns:
A 'ClientOutput`.
"""
model_weights = _get_weights(model)
@tf.function
def clip_by_norm(gradient, norm):
"""Clip the gradient by its l2 norm."""
norm = tf.cast(norm, tf.float32)
delta_norm = _get_norm(gradient)
if delta_norm < norm:
return gradient
else:
delta_mul_factor = tf.math.divide_no_nan(norm, delta_norm)
return tf.nest.map_structure(lambda g: g * delta_mul_factor, gradient)
@tf.function
def project_weights(weights, initial_weights, norm):
"""Project the weight onto l2 ball around initial_weights with radius norm."""
weights_delta = tf.nest.map_structure(lambda a, b: a - b, weights,
initial_weights)
return tf.nest.map_structure(tf.add, clip_by_norm(weights_delta, norm),
initial_weights)
@tf.function
def reduce_fn(num_examples_sum, batch):
"""Runs `tff.learning.Model.train_on_batch` on local client batch."""
with tf.GradientTape() as tape:
output = model.forward_pass(batch)
gradients = tape.gradient(output.loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return num_examples_sum + tf.shape(output.predictions)[0]
@tf.function
def compute_benign_update():
"""compute benign update sent back to the server."""
tf.nest.map_structure(lambda a, b: a.assign(b), model_weights,
initial_weights)
num_examples_sum = benign_dataset.reduce(
initial_state=tf.constant(0), reduce_func=reduce_fn)
weights_delta_benign = tf.nest.map_structure(lambda a, b: a - b,
model_weights.trainable,
initial_weights.trainable)
aggregated_outputs = model.report_local_outputs()
return weights_delta_benign, aggregated_outputs, num_examples_sum
@tf.function
def compute_malicious_update():
"""compute malicious update sent back to the server."""
_, aggregated_outputs, num_examples_sum = compute_benign_update()
tf.nest.map_structure(lambda a, b: a.assign(b), model_weights,
initial_weights)
for _ in range(self.round_num):
benign_dataset.reduce(
initial_state=tf.constant(0), reduce_func=reduce_fn)
malicious_dataset.reduce(
initial_state=tf.constant(0), reduce_func=reduce_fn)
tf.nest.map_structure(
lambda a, b: a.assign(b), model_weights.trainable,
project_weights(model_weights.trainable, initial_weights.trainable,
tf.cast(self.norm_bound, tf.float32)))
weights_delta_malicious = tf.nest.map_structure(lambda a, b: a - b,
model_weights.trainable,
initial_weights.trainable)
weights_delta = tf.nest.map_structure(
lambda update: self.boost_factor * update, weights_delta_malicious)
return weights_delta, aggregated_outputs, num_examples_sum
if client_is_malicious:
result = compute_malicious_update()
else:
result = compute_benign_update()
weights_delta, aggregated_outputs, num_examples_sum = result
weights_delta_weight = tf.cast(num_examples_sum, tf.float32)
weight_norm = _get_norm(weights_delta)
return ClientOutput(
weights_delta, weights_delta_weight, aggregated_outputs,
tensor_utils.to_odict({
'num_examples': num_examples_sum,
'weight_norm': weight_norm,
}))
def __call__(self, model, optimizer, benign_dataset, malicious_dataset,
client_type, initial_weights):
"""Updates client model with client potentially being malicious.
Args:
model: A `tff.learning.Model`.
optimizer: A 'tf.keras.optimizers.Optimizer'.
benign_dataset: A 'tf.data.Dataset' consisting of benign dataset.
malicious_dataset: A 'tf.data.Dataset' consisting of malicious dataset.
client_type: A 'tf.bool' indicating whether the client is malicious; iff
`True` the client will construct its update using `malicious_dataset`,
otherwise will construct the update using `benign_dataset`.
initial_weights: A `tff.learning.Model.weights` from server.
Returns:
A 'ClientOutput`.
"""
model_weights = _get_weights(model)
@tf.function
def reduce_fn(num_examples_sum, batch):
"""Runs `tff.learning.Model.train_on_batch` on local client batch."""
with tf.GradientTape() as tape:
output = model.forward_pass(batch)
gradients = tape.gradient(output.loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return num_examples_sum + tf.shape(output.predictions)[0]
@tf.function
def compute_benign_update():
"""compute benign update sent back to the server."""
tf.nest.map_structure(lambda a, b: a.assign(b), model_weights,
initial_weights)
num_examples_sum = benign_dataset.reduce(
initial_state=tf.constant(0), reduce_func=reduce_fn)
weights_delta_benign = tf.nest.map_structure(lambda a, b: a - b,
model_weights.trainable,
initial_weights.trainable)
aggregated_outputs = model.report_local_outputs()
return weights_delta_benign, aggregated_outputs, num_examples_sum
@tf.function
def compute_malicious_update():
"""compute malicious update sent back to the server."""
weights_delta_benign, aggregated_outputs, num_examples_sum \
= compute_benign_update()
tf.nest.map_structure(lambda a, b: a.assign(b), model_weights,
initial_weights)
malicious_dataset.reduce(
initial_state=tf.constant(0), reduce_func=reduce_fn)
weights_delta_malicious = tf.nest.map_structure(lambda a, b: a - b,
model_weights.trainable,
initial_weights.trainable)
weights_delta = tf.nest.map_structure(
tf.add, weights_delta_benign,
tf.nest.map_structure(lambda delta: delta * self.boost_factor,
weights_delta_malicious))
return weights_delta, aggregated_outputs, num_examples_sum
result = tf.cond(
tf.equal(client_type, True), compute_malicious_update,
compute_benign_update)
weights_delta, aggregated_outputs, num_examples_sum = result
weights_delta_weight = tf.cast(num_examples_sum, tf.float32)
weight_norm = _get_norm(weights_delta)
return ClientOutput(
weights_delta, weights_delta_weight, aggregated_outputs,
tensor_utils.to_odict({
'num_examples': num_examples_sum,
'weight_norm': weight_norm,
}))
