def test_cumsum_vector(self, total_steps=15):
def new_value_fn():
return [
tf.ones([2, 2], dtype=tf.float32),
tf.constant([2], dtype=tf.float32)
]
tree_aggregator = tree_aggregation.TFEfficientTreeAggregator(
new_value_fn=new_value_fn)
tree_aggregator_truth = tree_aggregation.TFEfficientTreeAggregator(
new_value_fn=lambda: 1.)
state = tree_aggregator.init_state()
truth_state = tree_aggregator_truth.init_state()
for leaf_node_idx in range(total_steps):
self.assertEqual(leaf_node_idx,
tree_aggregation.get_step_idx(state))
val, state = tree_aggregator.get_cumsum_and_update(state)
expected_val, truth_state = tree_aggregator_truth.get_cumsum_and_update(
truth_state)
self.assertEqual(tree_aggregation.get_step_idx(state),
tree_aggregation.get_step_idx(truth_state))
expected_result = [
expected_val * tf.ones([2, 2], dtype=tf.float32),
expected_val * tf.constant([2], dtype=tf.float32),
]
tf.nest.map_structure(self.assertAllClose, val, expected_result)
def test_dpftal_training(self, total_rounds=5):
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)
it_process = dp_fedavg.build_federated_averaging_process(
_rnn_model_fn, server_optimizer_fn=server_optimzier_fn)
server_state = it_process.initialize()
def deterministic_batch():
return collections.OrderedDict(x=np.array([[0, 1, 2, 3, 4]],
dtype=np.int32),
y=np.array([[1, 2, 3, 4, 0]],
dtype=np.int32))
batch = tff.tf_computation(deterministic_batch)()
federated_data = [[batch]]
loss_list = []
for i in range(total_rounds):
server_state, loss = it_process.next(server_state, federated_data)
loss_list.append(loss)
self.assertEqual(i + 1, server_state.round_num)
self.assertEqual(
i + 1,
tree_aggregation.get_step_idx(
server_state.optimizer_state['dp_tree_state'].level_state))
self.assertLess(np.mean(loss_list[1:]), loss_list[0])
def test_dp_momentum_training(self, model_fn, optimzer_fn, total_rounds=3):
def server_optimzier_fn(model_weights):
model_weight_shape = tf.nest.map_structure(tf.shape, model_weights)
return optimzer_fn(learning_rate=1.0,
momentum=0.9,
noise_std=1e-5,
model_weight_shape=model_weight_shape)
print('defining it process')
it_process = dp_fedavg.build_federated_averaging_process(
model_fn, server_optimizer_fn=server_optimzier_fn)
print('next type', it_process.next.type_signature.parameter[0])
server_state = it_process.initialize()
def deterministic_batch():
return collections.OrderedDict(x=np.ones([1, 28, 28, 1],
dtype=np.float32),
y=np.ones([1], dtype=np.int32))
batch = tff.tf_computation(deterministic_batch)()
federated_data = [[batch]]
loss_list = []
for i in range(total_rounds):
print('round', i)
server_state, loss = it_process.next(server_state, federated_data)
loss_list.append(loss)
self.assertEqual(i + 1, server_state.round_num)
if 'server_state_type' in server_state.optimizer_state:
self.assertEqual(
i + 1,
tree_aggregation.get_step_idx(
server_state.optimizer_state['dp_tree_state']))
self.assertLess(np.mean(loss_list[1:]), loss_list[0])
def test_dp_momentum_training(self, model_fn, optimzer_fn, total_rounds=3):
def server_optimzier_fn(model_weights):
model_weight_specs = tf.nest.map_structure(
lambda v: tf.TensorSpec(v.shape, v.dtype), model_weights)
return optimzer_fn(
learning_rate=1.0,
momentum=0.9,
noise_std=1e-5,
model_weight_specs=model_weight_specs)
it_process = dp_fedavg.build_federated_averaging_process(
model_fn, server_optimizer_fn=server_optimzier_fn)
server_state = it_process.initialize()
def deterministic_batch():
return collections.OrderedDict(
x=np.ones([1, 28, 28, 1], dtype=np.float32),
y=np.ones([1], dtype=np.int32))
batch = tff.tf_computation(deterministic_batch)()
federated_data = [[batch]]
loss_list = []
for i in range(total_rounds):
server_state, loss = it_process.next(server_state, federated_data)
loss_list.append(loss)
self.assertEqual(i + 1, server_state.round_num)
if server_state.optimizer_state is optimizer_utils.FTRLState:
self.assertEqual(
i + 1,
tree_aggregation.get_step_idx(
server_state.optimizer_state.dp_tree_state))
self.assertLess(np.mean(loss_list[1:]), loss_list[0])
def test_tree_sum_noise_efficient(self, total_steps, noise_std,
variable_shape, tolerance):
# Test the variance returned by `TFEfficientTreeAggregator` is smaller than
# `TFTreeAggregator` (within a relative `tolerance`) after `total_steps` of
# leaf nodes are traversed. Each tree node is a `variable_shape` tensor of
# Gaussian noise with `noise_std`. A small `tolerance` is used for numerical
# stability, `tolerance==0` means `TFEfficientTreeAggregator` is strictly
# better than `TFTreeAggregator` for reducing variance.
random_generator = tree_aggregation.GaussianNoiseGenerator(
noise_std, tf.TensorSpec(variable_shape))
tree_aggregator = tree_aggregation.TFEfficientTreeAggregator(
value_generator=random_generator)
tree_aggregator_baseline = tree_aggregation.TFTreeAggregator(
value_generator=random_generator)
state = tree_aggregator.init_state()
state_baseline = tree_aggregator_baseline.init_state()
for leaf_node_idx in range(total_steps):
self.assertEqual(leaf_node_idx,
tree_aggregation.get_step_idx(state))
val, state = tree_aggregator.get_cumsum_and_update(state)
val_baseline, state_baseline = tree_aggregator_baseline.get_cumsum_and_update(
state_baseline)
self.assertLess(tf.math.reduce_variance(val), (1 + tolerance) *
tf.math.reduce_variance(val_baseline))
def test_tree_sum_steps_max(self, total_steps, node_value):
tree_aggregator = tree_aggregation.TFTreeAggregator(
value_generator=ConstantValueGenerator(node_value))
max_val = node_value * math.ceil(math.log2(total_steps))
state = tree_aggregator.init_state()
for leaf_node_idx in range(total_steps):
self.assertEqual(leaf_node_idx,
tree_aggregation.get_step_idx(state))
val, state = tree_aggregator.get_cumsum_and_update(state)
self.assertLessEqual(val, max_val)
def test_tree_sum_steps_max(self, total_steps, step_value):
tree_aggregator = tree_aggregation.TFTreeAggregator(
new_value_fn=lambda: step_value)
max_val = step_value * math.ceil(math.log2(total_steps))
state = tree_aggregator.init_state()
for leaf_node_idx in range(total_steps):
val, state = tree_aggregator.get_cumsum_and_update(state)
self.assertEqual(leaf_node_idx + 1,
tree_aggregation.get_step_idx(state.level_state))
self.assertLessEqual(val, max_val)
def test_tree_sum_last_step_expected(self, total_steps, expected_value,
step_value):
tree_aggregator = tree_aggregation.TFTreeAggregator(
new_value_fn=lambda: step_value)
state = tree_aggregator.init_state()
for leaf_node_idx in range(total_steps):
val, state = tree_aggregator.get_cumsum_and_update(state)
self.assertEqual(leaf_node_idx + 1,
tree_aggregation.get_step_idx(state.level_state))
self.assertEqual(expected_value, val)
def test_tree_sum_last_step_expected_value_fn(self, total_steps,
expected_value, node_value):
# Test no-arg function as stateless value generator.
tree_aggregator = tree_aggregation.TFTreeAggregator(
value_generator=lambda: node_value)
state = tree_aggregator.init_state()
for leaf_node_idx in range(total_steps):
self.assertEqual(leaf_node_idx,
tree_aggregation.get_step_idx(state))
val, state = tree_aggregator.get_cumsum_and_update(state)
self.assertEqual(expected_value, val)
def test_tree_sum_steps_expected(self, total_steps, expected_values,
node_value):
# Test whether `tree_aggregator` will output `expected_value` in each step
# when `total_steps` of leaf nodes are traversed. The value of each tree
# node is a constant `node_value` for test purpose. Note that `node_value`
# denotes the "noise" without private values in private algorithms.
tree_aggregator = tree_aggregation.TFTreeAggregator(
new_value_fn=lambda: node_value)
state = tree_aggregator.init_state()
for leaf_node_idx in range(total_steps):
self.assertEqual(leaf_node_idx,
tree_aggregation.get_step_idx(state))
val, state = tree_aggregator.get_cumsum_and_update(state)
self.assertEqual(expected_values[leaf_node_idx], val)
def test_tree_sum_last_step_expected(self, total_steps, expected_value,
node_value):
# Test whether `tree_aggregator` will output `expected_value` after
# `total_steps` of leaf nodes are traversed. The value of each tree node
# is a constant `node_value` for test purpose. Note that `node_value`
# denotes the "noise" without private values in private algorithms.
tree_aggregator = tree_aggregation.TFTreeAggregator(
value_generator=ConstantValueGenerator(node_value))
state = tree_aggregator.init_state()
for leaf_node_idx in range(total_steps):
self.assertEqual(leaf_node_idx,
tree_aggregation.get_step_idx(state))
val, state = tree_aggregator.get_cumsum_and_update(state)
self.assertEqual(expected_value, val)
def test_cumsum_vector(self, total_steps=15):
tree_aggregator = tree_aggregation.TFTreeAggregator(
value_generator=ConstantValueGenerator([
tf.ones([2, 2], dtype=tf.float32),
tf.constant([2], dtype=tf.float32)
]))
tree_aggregator_truth = tree_aggregation.TFTreeAggregator(
value_generator=ConstantValueGenerator(1.))
state = tree_aggregator.init_state()
truth_state = tree_aggregator_truth.init_state()
for leaf_node_idx in range(total_steps):
self.assertEqual(leaf_node_idx,
tree_aggregation.get_step_idx(state))
val, state = tree_aggregator.get_cumsum_and_update(state)
expected_val, truth_state = tree_aggregator_truth.get_cumsum_and_update(
truth_state)
self.assertEqual(tree_aggregation.get_step_idx(state),
tree_aggregation.get_step_idx(truth_state))
expected_result = [
expected_val * tf.ones([2, 2], dtype=tf.float32),
expected_val * tf.constant([2], dtype=tf.float32),
]
tf.nest.map_structure(self.assertAllEqual, val, expected_result)
def test_dpftal_restart(self, total_rounds=3):
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)
it_process = dp_fedavg.build_federated_averaging_process(
_rnn_model_fn,
server_optimizer_fn=server_optimizer_fn,
use_simulation_loop=True)
server_state = it_process.initialize()
model = _rnn_model_fn()
optimizer = server_optimizer_fn(model.weights.trainable)
def server_state_update(state):
return tff.structure.update_struct(
state,
model=state.model,
optimizer_state=optimizer.restart_dp_tree(state.model.trainable),
round_num=state.round_num)
def deterministic_batch():
return collections.OrderedDict(
x=np.array([[0, 1, 2, 3, 4]], dtype=np.int32),
y=np.array([[1, 2, 3, 4, 0]], dtype=np.int32))
batch = tff.tf_computation(deterministic_batch)()
federated_data = [[batch]]
loss_list = []
for i in range(total_rounds):
server_state, loss = it_process.next(server_state, federated_data)
server_state = server_state_update(server_state)
loss_list.append(loss)
self.assertEqual(i + 1, server_state.round_num)
self.assertEqual(
0,
tree_aggregation.get_step_idx(
server_state.optimizer_state.dp_tree_state))
self.assertLess(np.mean(loss_list[1:]), loss_list[0])
def test_tree_sum_last_step_expected(self, total_steps, expected_value,
step_value):
# Test whether `tree_aggregator` will output `expected_value` after
# `total_steps` of leaf nodes are traversed. The value of each tree node
# is a constant `node_value` for test purpose. Note that `node_value`
# denotes the "noise" without private values in private algorithms. The
# `expected_value` is based on a weighting schema strongly depends on the
# depth of the binary tree.
tree_aggregator = tree_aggregation.TFEfficientTreeAggregator(
value_generator=ConstantValueGenerator(step_value))
state = tree_aggregator.init_state()
for leaf_node_idx in range(total_steps):
self.assertEqual(leaf_node_idx,
tree_aggregation.get_step_idx(state))
val, state = tree_aggregator.get_cumsum_and_update(state)
self.assertAllClose(expected_value, val)
def test_tree_sum_noise_expected(self, total_steps, expected_variance,
noise_std, variable_shape, tolerance):
# Test whether `tree_aggregator` will output `expected_variance` (within a
# relative `tolerance`) in each step when `total_steps` of leaf nodes are
# traversed. Each tree node is a `variable_shape` tensor of Gaussian noise
# with `noise_std`.
random_generator = tree_aggregation.GaussianNoiseGenerator(
noise_std, tf.TensorSpec(variable_shape), seed=2020)
tree_aggregator = tree_aggregation.TFTreeAggregator(
value_generator=random_generator)
state = tree_aggregator.init_state()
for leaf_node_idx in range(total_steps):
self.assertEqual(leaf_node_idx,
tree_aggregation.get_step_idx(state))
val, state = tree_aggregator.get_cumsum_and_update(state)
self.assertAllClose(math.sqrt(expected_variance[leaf_node_idx]),
tf.math.reduce_std(val),
rtol=tolerance)
def test_tree_sum_noise_expected(self, total_steps, expected_variance,
noise_std, variable_shape, tolerance):
random_generator = tf.random.Generator.from_seed(0)
def get_noise():
return random_generator.normal(shape=variable_shape,
stddev=noise_std)
tree_aggregator = tree_aggregation.TFTreeAggregator(
new_value_fn=get_noise)
state = tree_aggregator.init_state()
for leaf_node_idx in range(total_steps):
val, state = tree_aggregator.get_cumsum_and_update(state)
self.assertEqual(leaf_node_idx + 1,
tree_aggregation.get_step_idx(state.level_state))
self.assertAllClose(expected_variance[leaf_node_idx],
tf.math.reduce_variance(val),
rtol=tolerance)
def test_tree_sum_noise_expected(self, total_steps, expected_variance,
noise_std, variable_shape, tolerance):
# Test whether `tree_aggregator` will output `expected_variance` (within a
# relative `tolerance`) after `total_steps` of leaf nodes are traversed.
# Each tree node is a `variable_shape` tensor of Gaussian noise with
# `noise_std`. Note that the variance of a tree node is smaller than
# the given vanilla node `noise_std` because of the update rule of
# `TFEfficientTreeAggregator`.
random_generator = tree_aggregation.GaussianNoiseGenerator(
noise_std, tf.TensorSpec(variable_shape), seed=2020)
tree_aggregator = tree_aggregation.TFEfficientTreeAggregator(
value_generator=random_generator)
state = tree_aggregator.init_state()
for leaf_node_idx in range(total_steps):
self.assertEqual(leaf_node_idx,
tree_aggregation.get_step_idx(state))
val, state = tree_aggregator.get_cumsum_and_update(state)
self.assertAllClose(math.sqrt(expected_variance),
tf.math.reduce_std(val),
rtol=tolerance)
def test_tree_sum_noise_expected(self, total_steps, expected_variance,
noise_std, variable_shape, tolerance):
# Test whether `tree_aggregator` will output `expected_variance` (within a
# relative `tolerance`) in each step when `total_steps` of leaf nodes are
# traversed. Each tree node is a `variable_shape` tensor of Gaussian noise
# with `noise_std`.
random_generator = tf.random.Generator.from_seed(0)
def get_noise():
return random_generator.normal(shape=variable_shape,
stddev=noise_std)
tree_aggregator = tree_aggregation.TFTreeAggregator(
new_value_fn=get_noise)
state = tree_aggregator.init_state()
for leaf_node_idx in range(total_steps):
self.assertEqual(leaf_node_idx,
tree_aggregation.get_step_idx(state))
val, state = tree_aggregator.get_cumsum_and_update(state)
self.assertAllClose(expected_variance[leaf_node_idx],
tf.math.reduce_variance(val),
rtol=tolerance)
def test_update_level_state(self, val, state_in, expected_state_out):
tree_aggregator = tree_aggregation.TFTreeAggregator(new_value_fn=None)
state_out = tree_aggregator._update_level_state(
tf.constant(state_in, dtype=tf.int8))
self.assertAllEqual(expected_state_out, state_out)
self.assertEqual(val + 1, tree_aggregation.get_step_idx(state_out))
