def test_compression_controller_state():
from nncf.common.compression import BaseControllerStateNames as CtrlStateNames
model = get_magnitude_test_model()
config = get_basic_magnitude_sparsity_config()
algo_name = config['compression']['algorithm']
config['compression']['params'] = {'schedule': 'multistep'}
_, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
# Test get state
compression_ctrl.scheduler.current_step = 100
compression_ctrl.scheduler.current_epoch = 5
state_content = compression_ctrl.get_state()[algo_name]
assert state_content[CtrlStateNames.SCHEDULER] == {
'current_step': 100,
'current_epoch': 5
}
# Test load state
new_state = {
algo_name: {
CtrlStateNames.SCHEDULER: {
'current_step': 500,
'current_epoch': 10
},
CtrlStateNames.LOSS: {},
CtrlStateNames.COMPRESSION_STAGE: None,
}
}
compression_ctrl.load_state(new_state)
assert compression_ctrl.scheduler.current_step == 500
assert compression_ctrl.scheduler.current_epoch == 10
assert compression_ctrl.get_state() == new_state
def test_magnitude_algo_set_binary_mask_on_forward():
config = get_basic_magnitude_sparsity_config()
config['compression']['params'] = {'weight_importance': 'abs'}
sparse_model, compression_ctrl = create_compressed_model_and_algo_for_test(
get_magnitude_test_model(), config)
compression_ctrl.set_sparsity_level(0.3)
TFTensorListComparator.check_equal(ref_mask_1,
sparse_model.layers[1].weights[-1])
TFTensorListComparator.check_equal(ref_mask_2,
sparse_model.layers[2].weights[-1])
def test_magnitude_sparse_algo_sets_threshold(weight_importance,
sparsity_level, threshold):
model = get_magnitude_test_model()
config = get_basic_magnitude_sparsity_config()
config['compression']['params'] = {
'schedule': 'multistep',
'weight_importance': weight_importance
}
_, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
if sparsity_level:
compression_ctrl.set_sparsity_level(sparsity_level)
assert compression_ctrl._threshold == pytest.approx(threshold, 0.01) # pylint: disable=protected-access
def test_magnitude_scheduler_can_do_epoch_step__with_multistep():
config = get_empty_config()
config["compression"] = Dict({"algorithm": "magnitude_sparsity",
"params": {"schedule": "multistep", 'multistep_steps': [1]}})
model = get_magnitude_test_model(config['input_info'][0]['sample_size'][1:])
_, compression_ctrl = create_compressed_model_and_algo_for_test(model, config)
scheduler = compression_ctrl.scheduler
assert isinstance(scheduler, MultiStepSparsityScheduler)
assert compression_ctrl.scheduler.current_sparsity_level == 0.1
assert scheduler.schedule.values == [0.1, 0.5]
expected_levels = [0.1, 0.5, 0.5]
for expected_level in expected_levels:
scheduler.epoch_step()
assert compression_ctrl.scheduler.current_sparsity_level == expected_level
def test_magnitude_scheduler_can_do_epoch_step__with_last():
config = get_multistep_normed_abs_config()
model = get_magnitude_test_model(config['input_info']['sample_size'][1:])
_, compression_ctrl = create_compressed_model_and_algo_for_test(model, config)
scheduler = compression_ctrl.scheduler
scheduler.epoch_step(3)
assert compression_ctrl.scheduler.current_sparsity_level == 0.9
nncf_stats = compression_ctrl.statistics()
for layer_info in nncf_stats.magnitude_sparsity.thresholds:
assert layer_info.threshold == pytest.approx(0.371, 0.01)
scheduler.epoch_step()
assert compression_ctrl.scheduler.current_sparsity_level == 0.9
nncf_stats = compression_ctrl.statistics()
for layer_info in nncf_stats.magnitude_sparsity.thresholds:
assert layer_info.threshold == pytest.approx(0.371, 0.01)
def test_magnitude_scheduler_can_do_epoch_step__with_norm():
config = get_multistep_normed_abs_config()
model = get_magnitude_test_model(config['input_info']['sample_size'][1:])
_, compression_ctrl = create_compressed_model_and_algo_for_test(model, config)
scheduler = compression_ctrl.scheduler
assert isinstance(scheduler, MultiStepSparsityScheduler)
assert compression_ctrl.scheduler.current_sparsity_level == 0.1
expected_levels = [0.1, 0.5, 0.5, 0.9]
expected_thresholds = [0.219, 0.243, 0.243, 0.371]
for expected_level, expected_threshold in zip(expected_levels, expected_thresholds):
scheduler.epoch_step()
assert compression_ctrl.scheduler.current_sparsity_level == expected_level
nncf_stats = compression_ctrl.statistics()
for layer_info in nncf_stats.magnitude_sparsity.thresholds:
assert layer_info.threshold == pytest.approx(expected_threshold, 0.01)
def test_can_create_magnitude_sparse_algo__with_defaults():
model = get_magnitude_test_model()
config = get_basic_magnitude_sparsity_config()
config['compression']['params'] = \
{'schedule': 'multistep'}
sparse_model, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
assert isinstance(compression_ctrl, MagnitudeSparsityController)
assert compression_ctrl.scheduler.current_sparsity_level == approx(0.1)
conv_names = [
layer.name for layer in model.layers
if isinstance(layer, tf.keras.layers.Conv2D)
]
wrappers = [
layer for layer in sparse_model.layers
if isinstance(layer, NNCFWrapper)
]
correct_wrappers = [
wrapper for wrapper in wrappers if wrapper.name in conv_names
]
assert len(conv_names) == len(wrappers)
assert len(conv_names) == len(correct_wrappers)
assert compression_ctrl._threshold == approx(0.24, 0.1) # pylint: disable=protected-access
# pylint: disable=protected-access
assert isinstance(compression_ctrl._weight_importance_fn,
type(normed_magnitude))
for i, wrapper in enumerate(wrappers):
ref_mask = tf.ones_like(wrapper.weights[-1]) if i == 0 else ref_mask_2
mask = list(wrapper.ops_weights.values())[0]['mask']
op = list(wrapper.weights_attr_ops['kernel'].values())[0]
tf.assert_equal(mask, ref_mask)
assert isinstance(op, BinaryMask)