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.TFTreeAggregator( new_value_fn=new_value_fn) tree_aggregator_truth = tree_aggregation.TFTreeAggregator( 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.assertAllEqual, val, expected_result)
def __init__(self, learning_rate: float, momentum: float, noise_std: float, model_weight_shape: Collection[tf.Tensor], efficient_tree: bool = True, use_nesterov: bool = False): """Initialize the momemtum DPFTRL Optimizer.""" _check_momentum(momentum) if use_nesterov and momentum == 0: raise ValueError('Use a positive momentum for Nesterov') self.lr = learning_rate self.momentum = momentum self.model_weight_shape = model_weight_shape self.use_nesterov = use_nesterov random_generator = tf.random.Generator.from_non_deterministic_state() def _noise_fn(): return tf.nest.map_structure( lambda x: random_generator.normal(x, stddev=noise_std), model_weight_shape) if efficient_tree: self.noise_generator = tree_aggregation.TFEfficientTreeAggregator( new_value_fn=_noise_fn) else: self.noise_generator = tree_aggregation.TFTreeAggregator( new_value_fn=_noise_fn)
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 __init__(self, learning_rate: float, momentum: float, noise_std: float, model_weight_shape: Collection[tf.Tensor], efficient_tree: bool = True): """Initialize the momemtum DPFTRL Optimizer.""" self.lr = learning_rate self.momentum = momentum self.model_weight_shape = model_weight_shape random_generator = tf.random.Generator.from_non_deterministic_state() def _noise_fn(): return tf.nest.map_structure( lambda x: random_generator.normal(x, stddev=noise_std), model_weight_shape) if efficient_tree: self.noise_generator = tree_aggregation.TFEfficientTreeAggregator( new_value_fn=_noise_fn) else: self.noise_generator = tree_aggregation.TFTreeAggregator( new_value_fn=_noise_fn)
def __init__(self, learning_rate: float, momentum: float, noise_std: float, model_weight_specs: Collection[tf.TensorSpec], efficient_tree: bool = True, use_nesterov: bool = False): """Initialize the momemtum DPFTRL Optimizer.""" _check_momentum(momentum) if use_nesterov and momentum == 0: raise ValueError('Use a positive momentum for Nesterov') self.lr = learning_rate self.momentum = momentum self.model_weight_specs = model_weight_specs self.use_nesterov = use_nesterov random_generator = tree_aggregation.GaussianNoiseGenerator( noise_std, model_weight_specs) if efficient_tree: self.noise_generator = tree_aggregation.TFEfficientTreeAggregator( value_generator=random_generator) else: self.noise_generator = tree_aggregation.TFTreeAggregator( value_generator=random_generator)
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_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_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_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`) 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))