def test_checkpoint_callback_make_checkpoints(mocker, tmp_path): save_freq = 2 config = get_basic_quantization_config() gen_setup_spy = mocker.spy(QuantizationBuilder, '_get_quantizer_setup') model, compression_ctrl = create_compressed_model_and_algo_for_test( get_basic_conv_test_model(), config, force_no_init=True) assert isinstance(compression_ctrl, QuantizationController) quantizer_setup = gen_setup_spy.spy_return compression_callbacks = create_compression_callbacks(compression_ctrl, log_tensorboard=False) dataset_len = 8 dummy_x = tf.random.normal((dataset_len, ) + model.input_shape[1:]) dummy_y = tf.random.normal((dataset_len, ) + model.output_shape[1:]) model.compile(loss=tf.losses.CategoricalCrossentropy()) ckpt_path = tmp_path / 'checkpoint' ckpt = tf.train.Checkpoint( model=model, compression_state=TFCompressionState(compression_ctrl)) model.fit(dummy_x, dummy_y, epochs=5, batch_size=2, callbacks=[ CheckpointManagerCallback(ckpt, str(ckpt_path), save_freq), *compression_callbacks ]) assert sorted(os.listdir(ckpt_path)) == REF_CKPT_DIR[save_freq] new_compression_state = load_compression_state(ckpt_path) new_model, new_compression_ctrl = create_compressed_model_and_algo_for_test( get_basic_conv_test_model(), config, new_compression_state) new_model.compile(loss=tf.losses.CategoricalCrossentropy()) new_ckpt = tf.train.Checkpoint( model=new_model, compression_state=TFCompressionState(new_compression_ctrl)) load_checkpoint(new_ckpt, ckpt_path) builder = QuantizationBuilder(config) builder.load_state(new_compression_state['builder_state']) # pylint:disable=protected-access new_quantizer_setup = builder._quantizer_setup assert _quantization_setup_cmp(quantizer_setup, new_quantizer_setup) assert new_compression_ctrl.get_state() == compression_ctrl.get_state() assert tf.reduce_all([ tf.reduce_all(w_new == w) for w_new, w in zip(new_model.weights, model.weights) ])
def test_magnitude_algo_binary_masks_are_applied(): input_shape = (1, 5, 5, 1) model = get_basic_conv_test_model(input_shape=input_shape[1:]) config = get_empty_config(input_sample_sizes=input_shape) config.update(Dict({'compression': {'algorithm': "magnitude_sparsity"}})) compressed_model, _ = create_compressed_model_and_algo_for_test( model, config) conv = compressed_model.layers[1] op_name = list(conv.ops_weights.keys())[0] conv.ops_weights[op_name] = {'mask': tf.ones_like(conv.weights[0])} input_ = tf.ones(input_shape) ref_output_1 = -4 * tf.ones((1, 4, 4, 2)) output_1 = compressed_model(input_) tf.assert_equal(output_1, ref_output_1) np_mask = conv.ops_weights[op_name]['mask'].numpy() np_mask[0, 1, 0, 0] = 0 np_mask[1, 0, 0, 1] = 0 conv.ops_weights[op_name] = {'mask': tf.constant(np_mask)} ref_output_2 = -3 * tf.ones_like(ref_output_1) output_2 = compressed_model(input_) tf.assert_equal(output_2, ref_output_2) np_mask[0, 1, 0, 1] = 0 conv.ops_weights[op_name] = {'mask': tf.constant(np_mask)} ref_output_3 = ref_output_2.numpy() ref_output_3[..., 1] = -2 * np.ones_like(ref_output_1[..., 1]) ref_output_3 = tf.constant(ref_output_3) output_3 = compressed_model(input_) tf.assert_equal(output_3, ref_output_3)
def test_scheduler_can_do_epoch_step(self, algo, schedule, get_params, ref_levels): model = get_basic_conv_test_model() config = get_empty_config() config['compression'] = Dict({ 'algorithm': algo, 'sparsity_init': 0.2, "params": { **get_params(), "schedule": schedule } }) _, compression_ctrl = create_compressed_model_and_algo_for_test( model, config) scheduler = compression_ctrl.scheduler assert pytest.approx(scheduler.current_sparsity_level) == ref_levels[0] for ref_level in ref_levels[1:]: scheduler.epoch_step() assert pytest.approx(scheduler.current_sparsity_level) == ref_level _, compression_ctrl = create_compressed_model_and_algo_for_test( model, config) scheduler = compression_ctrl.scheduler assert pytest.approx(scheduler.current_sparsity_level) == ref_levels[0] for i, ref_level in enumerate(ref_levels[1:]): scheduler.epoch_step(i) assert pytest.approx(scheduler.current_sparsity_level) == ref_level
def test_quantization_configs__with_defaults(): model = get_basic_conv_test_model() config = get_basic_quantization_config() compression_model, compression_ctrl = create_compressed_model_and_algo_for_test( model, config, force_no_init=True) assert isinstance(compression_ctrl, QuantizationController) check_default_qspecs(compression_model)
def test_can_set_sparse_layers_to_loss(): model = get_basic_conv_test_model() config = get_basic_sparsity_config() _, compression_ctrl = create_compressed_model_and_algo_for_test( model, config) loss = compression_ctrl.loss assert isinstance(loss, SparseLoss) # pylint: disable=protected-access for op, _ in loss._target_ops: assert isinstance(op, RBSparsifyingWeight)
def test_quantization_configs__disable_overflow_fix(): model = get_basic_conv_test_model() config = get_basic_quantization_config() config['compression'].update({'overflow_fix': 'disable'}) compression_model, compression_ctrl = create_compressed_model_and_algo_for_test( model, config, force_no_init=True) assert isinstance(compression_ctrl, QuantizationController) check_specs_for_disabled_overflow_fix(compression_model)
def test_basic_model_has_expected_params(): default_weight = tf.constant([[[[0., -1.], [-1., 0.]]], [[[0., -1.], [-1., 0.]]]]) default_weight = tf.transpose(default_weight, (2, 3, 1, 0)) default_bias = tf.constant([-2., -2.]) model = get_basic_conv_test_model() act_weights = model.layers[1].weights[0] ref_weights = default_weight act_bias = model.layers[1].weights[1] ref_bias = default_bias TFTensorListComparator.check_equal(act_bias, ref_bias) TFTensorListComparator.check_equal(act_weights, ref_weights)
def test_scheduler_can_do_epoch_step__with_rb_algo(): config = NNCFConfig() config['input_info'] = [{"sample_size": [1, 4, 4, 1]}] config['compression'] = { 'algorithm': 'rb_sparsity', 'sparsity_init': 0.2, "params": { 'schedule': 'polynomial', 'power': 1, 'sparsity_target_epoch': 2, 'sparsity_target': 0.6, 'sparsity_freeze_epoch': 3 } } _, compression_ctrl = create_compressed_model_and_algo_for_test( get_basic_conv_test_model(), config) scheduler = compression_ctrl.scheduler loss = compression_ctrl.loss assert not loss.disabled # pylint: disable=protected-access for op, op_weights in loss._target_ops: assert op.get_trainable_weight(op_weights) scheduler.epoch_step() assert pytest.approx(loss.target_sparsity_rate, abs=1e-3) == 0.2 assert pytest.approx(loss(), abs=1e-3) == 16 assert not loss.disabled scheduler.epoch_step() assert pytest.approx(loss.target_sparsity_rate, abs=1e-3) == 0.4 assert pytest.approx(loss(), abs=1e-3) == 64 assert not loss.disabled scheduler.epoch_step() assert pytest.approx(loss.target_sparsity_rate, abs=1e-3) == 0.6 assert pytest.approx(loss(), abs=1e-3) == 144 assert not loss.disabled scheduler.epoch_step() assert loss.disabled assert pytest.approx(loss.target_sparsity_rate, abs=1e-3) == 0.6 assert loss() == 0 for op, op_weights in loss._target_ops: assert not op.get_trainable_weight(op_weights)
def test_struct_auxiliary_nodes_nncf_graph(): model = get_basic_conv_test_model() config = get_basic_quantization_config() compressed_model, _ = create_compressed_model_and_algo_for_test( model, config, force_no_init=True) nncf_graph = convert_keras_model_to_nncf_graph(compressed_model) input_nodes = nncf_graph.get_input_nodes() output_nodes = nncf_graph.get_output_nodes() assert len(input_nodes) == 1 assert len(output_nodes) == 1 assert input_nodes[0].metatype in INPUT_NOOP_METATYPES assert output_nodes[0].metatype in OUTPUT_NOOP_METATYPES
def test_quantization_preset(data): model = get_basic_conv_test_model() config = get_basic_quantization_config() config['target_device'] = data['target_device'] config['compression'] = { 'algorithm': 'quantization', 'preset': data['preset'] } config['compression'].update(data['overrided_param']) compression_model, _ = create_compressed_model_and_algo_for_test( model, config, force_no_init=True) activation_quantizers, weight_quantizers = get_quantizers( compression_model) for aq in activation_quantizers: assert aq.mode == data['expected_activations_q'] for wq in weight_quantizers: assert wq.mode == data['expected_weights_q']
def test_can_create_sparse_loss_and_scheduler(): model = get_basic_conv_test_model() config = get_basic_sparsity_config() _, compression_ctrl = create_compressed_model_and_algo_for_test( model, config) scheduler = compression_ctrl.scheduler scheduler.epoch_step() loss = compression_ctrl.loss assert isinstance(loss, SparseLoss) assert not loss.disabled assert loss.target_sparsity_rate == approx(0.02) assert loss.p == approx(0.05) assert isinstance(scheduler, PolynomialSparsityScheduler) assert scheduler.current_sparsity_level == approx(0.02) assert scheduler.target_level == approx(0.5) assert scheduler.target_epoch == 2 assert scheduler.freeze_epoch == 3
def test_quantization_configs__custom(): model = get_basic_conv_test_model() config = get_basic_quantization_config() config['target_device'] = 'TRIAL' config['compression'].update({ "weights": { "mode": "asymmetric", "per_channel": True, "bits": 4 }, "activations": { "mode": "asymmetric", "bits": 4, "signed": True, }, }) compression_model, compression_ctrl = create_compressed_model_and_algo_for_test( model, config, force_no_init=True) assert isinstance(compression_ctrl, QuantizationController) activation_quantizers, weight_quantizers = get_quantizers( compression_model) ref_weight_qspec = TFQuantizerSpec(mode=QuantizationMode.ASYMMETRIC, num_bits=4, signedness_to_force=None, per_channel=True, narrow_range=True, half_range=False) for wq in weight_quantizers: compare_qspecs(ref_weight_qspec, wq) ref_activation_qspec = TFQuantizerSpec(mode=QuantizationMode.ASYMMETRIC, num_bits=4, signedness_to_force=True, per_channel=False, narrow_range=False, half_range=False) for wq in activation_quantizers: compare_qspecs(ref_activation_qspec, wq)
def test_quantization_configs__on_resume_with_compression_state( tmp_path, mocker): model = get_basic_conv_test_model() config = get_basic_quantization_config() init_spy = mocker.spy(QuantizationBuilder, 'initialize') gen_setup_spy = mocker.spy(QuantizationBuilder, '_get_quantizer_setup') dataset = get_dataset_for_test(shape=[4, 4, 1]) config = register_default_init_args(config, dataset, 10) _, compression_ctrl = create_compressed_model_and_algo_for_test( model, config) assert isinstance(compression_ctrl, QuantizationController) init_spy.assert_called() gen_setup_spy.assert_called() saved_quantizer_setup = gen_setup_spy.spy_return check_serialization(saved_quantizer_setup, _quantization_setup_cmp) compression_state_to_load = _save_and_load_compression_state( compression_ctrl, tmp_path) init_spy.reset_mock() gen_setup_spy.reset_mock() compression_model, compression_ctrl = create_compressed_model_and_algo_for_test( model, config, compression_state_to_load) assert isinstance(compression_ctrl, QuantizationController) init_spy.assert_not_called() gen_setup_spy.assert_not_called() check_default_qspecs(compression_model) builder = QuantizationBuilder(config) builder.load_state(compression_state_to_load['builder_state']) # pylint:disable=protected-access loaded_quantizer_setup = builder._quantizer_setup assert _quantization_setup_cmp(loaded_quantizer_setup, saved_quantizer_setup)
from tests.tensorflow.helpers import get_empty_config from tests.tensorflow.helpers import get_op_by_cls from tests.tensorflow.helpers import get_weight_by_name CONF = Path( __file__ ).parent.parent.parent / 'data' / 'configs' / 'sequential_model_cifar10_rb_sparsity.json' MASKS_SEEDS_PATH = Path(__file__).parent / 'output_seeds.txt' TEST_MODELS = { 'Dense': lambda: get_basic_fc_test_model(input_shape=(4, ), out_shape=10), 'Conv2D': lambda: get_basic_conv_test_model(input_shape=(4, 4, 1), out_channels=2, kernel_size=2, weight_init=-1., bias_init=-2., transpose=False), 'Conv2DTranspose': lambda: get_basic_conv_test_model(input_shape=(4, 4, 1), out_channels=1, kernel_size=2, weight_init=-1., bias_init=-2., transpose=True), } def get_basic_rb_sparse_model(model_name, local=False, config=CONF,
def test_basic_model_is_valid(): model = get_basic_conv_test_model() input_ = tf.ones([1, 4, 4, 1]) ref_output = tf.ones((1, 3, 3, 2)) * (-4) act_output = model(input_) TFTensorListComparator.check_equal(ref_output, act_output)
def get_simple_compressed_model(compression_state=None): model = get_basic_conv_test_model() config = get_empty_config() config.update({'compression': {'algorithm': 'magnitude_sparsity'}}) compression_ctrl, model = create_compressed_model(model, config, compression_state=compression_state) return compression_ctrl, model