def testConditionalMaskUpdate(self):
    param_list = ['pruning_frequency=2', 'begin_pruning_step=1', 'end_pruning_step=6', 'nbins=100']
    test_spec = ','.join(param_list)
    pruning_hparams = pruning.get_pruning_hparams().parse(test_spec)
    weights = variables.VariableV1(math_ops.linspace(1.0, 100.0, 100), name='weights')
    masked_weights = pruning.apply_mask(weights)
    sparsity = variables.VariableV1(0.0, name='sparsity')
    p = pruning.Pruning(pruning_hparams, sparsity=sparsity)
    p._spec.threshold_decay = 0.0
    mask_update_op = p.conditional_mask_update_op()
    sparsity_val = math_ops.linspace(0.0, 0.9, 10)
    increment_global_step = state_ops.assign_add(self.global_step, 1)
    non_zero_count = []
    with self.cached_session() as session:
        variables.global_variables_initializer().run()
        for i in range(10):
            session.run(state_ops.assign(sparsity, sparsity_val[i]))
            session.run(mask_update_op)
            session.run(increment_global_step)
            non_zero_count.append(np.count_nonzero(masked_weights.eval()))
    expected_non_zero_count = [100, 100, 80, 80, 60, 60, 40, 40, 40, 40]
    self.assertAllEqual(expected_non_zero_count, non_zero_count)
def testConditionalMaskUpdate(self):
    param_list = ['pruning_frequency=2', 'begin_pruning_step=1', 'end_pruning_step=6', 'nbins=100']
    test_spec = ','.join(param_list)
    pruning_hparams = pruning.get_pruning_hparams().parse(test_spec)
    weights = variables.VariableV1(math_ops.linspace(1.0, 100.0, 100), name='weights')
    masked_weights = pruning.apply_mask(weights)
    sparsity = variables.VariableV1(0.0, name='sparsity')
    p = pruning.Pruning(pruning_hparams, sparsity=sparsity)
    p._spec.threshold_decay = 0.0
    mask_update_op = p.conditional_mask_update_op()
    sparsity_val = math_ops.linspace(0.0, 0.9, 10)
    increment_global_step = state_ops.assign_add(self.global_step, 1)
    non_zero_count = []
    with self.cached_session() as session:
        variables.global_variables_initializer().run()
        for i in range(10):
            session.run(state_ops.assign(sparsity, sparsity_val[i]))
            session.run(mask_update_op)
            session.run(increment_global_step)
            non_zero_count.append(np.count_nonzero(masked_weights.eval()))
    expected_non_zero_count = [100, 100, 80, 80, 60, 60, 40, 40, 40, 40]
    self.assertAllEqual(expected_non_zero_count, non_zero_count)