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))