def test_load_state_interoperability(_algos, _model_wrapper, is_resume):
config_save = get_empty_config()
config_save['compression'] = [{
'algorithm': algo
} for algo in _algos['save_algos']]
register_bn_adaptation_init_args(config_save)
compressed_model_save, _ = create_compressed_model_and_algo_for_test(
BasicConvTestModel(), config_save)
model_save = _model_wrapper['save_model'](compressed_model_save)
saved_model_state = model_save.state_dict()
ref_num_loaded = len(saved_model_state)
config_resume = get_empty_config()
config_resume['compression'] = [{
'algorithm': algo
} for algo in _algos['load_algos']]
register_bn_adaptation_init_args(config_resume)
compressed_model_resume, _ = create_compressed_model_and_algo_for_test(
BasicConvTestModel(), config_resume)
model_resume = _model_wrapper['resume_model'](compressed_model_resume)
if not is_resume or (is_resume and _algos['is_resume_ok']):
act_num_loaded = load_state(model_resume, saved_model_state, is_resume)
if ('magnitude_sparsity' in _algos['load_algos'] or 'const_sparsity' in _algos['load_algos']) \
and 'rb_sparsity' in _algos['save_algos']:
# no need to load _mask and _uniform
ref_num_loaded -= 2
assert act_num_loaded == ref_num_loaded
else:
with pytest.raises(RuntimeError):
load_state(model_resume, saved_model_state, is_resume)
def test_can_restore_binary_mask_on_magnitude_algo_resume():
config = get_empty_config()
config['compression'] = {
"algorithm": "magnitude_sparsity",
"params": {
"weight_importance": "abs",
"schedule": "multistep",
"multistep_sparsity_levels": [0.3, 0.5]
}
}
sparse_model, _ = create_compressed_model_and_algo_for_test(
MagnitudeTestModel(), config)
with torch.no_grad():
sparse_model(torch.ones([1, 1, 10, 10]))
config = get_empty_config()
config["compression"] = {"algorithm": "const_sparsity"}
const_sparse_model, _ = create_compressed_model_and_algo_for_test(
MagnitudeTestModel(), config)
load_state(const_sparse_model, sparse_model.state_dict())
op = const_sparse_model.conv1.pre_ops['0']
PTTensorListComparator.check_equal(ref_mask_1, op.operand.binary_mask)
op = const_sparse_model.conv2.pre_ops['0']
PTTensorListComparator.check_equal(ref_mask_2, op.operand.binary_mask)
def test_load_state__with_resume_checkpoint(_resume_algos, _model_wrapper,
mocker):
config_save = get_empty_config()
config_save['compression'] = [{
'algorithm': algo
} for algo in _resume_algos['save_algos'] if algo != 'EMPTY']
register_bn_adaptation_init_args(config_save)
orig_model = BasicConvTestModel()
num_model_params = len(orig_model.state_dict())
model_save, compressed_ctrl_save = create_compressed_model_and_algo_for_test(
orig_model, config_save)
saved_model_state = model_save.state_dict()
saved_checkpoint = compressed_ctrl_save.get_compression_state()
ref_num_loaded = _resume_algos[
'ref_num_compression_params'] + num_model_params + 1 # padding_value
config_resume = get_empty_config()
config_resume['compression'] = [{
'algorithm': algo
} for algo in _resume_algos['load_algos'] if algo != 'EMPTY']
register_bn_adaptation_init_args(config_resume)
from nncf.torch.checkpoint_loading import KeyMatcher
key_matcher_run_spy = mocker.spy(KeyMatcher, 'run')
model, _ = create_compressed_model_and_algo_for_test(
BasicConvTestModel(),
config_resume,
compression_state=saved_checkpoint)
load_state(model, saved_model_state, _resume_algos['is_strict'])
key_matcher_run_spy.assert_called_once()
act_num_loaded = len(key_matcher_run_spy.spy_return)
assert act_num_loaded == ref_num_loaded
def test_scheduler_can_do_epoch_step(self, algo, schedule, get_params,
ref_levels):
model = BasicConvTestModel()
config = get_empty_config()
config['compression'] = {
'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
scheduler.epoch_step()
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
for m in compression_ctrl.sparsified_module_info:
if hasattr(m.operand, "frozen"):
assert m.operand.frozen
def create_config(self):
self.with_range_init()
config = get_empty_config()
config.update(self.config_update)
for config_updater in self._config_updaters:
config = config_updater(config)
return config
def test_building_block(model_creator, input_sizes, ref_building_blocks):
model = model_creator()
nncf_config = get_empty_config(input_sample_sizes=input_sizes)
_, compressed_model = create_compressed_model(model, nncf_config)
building_blocks = get_building_blocks(compressed_model)
assert building_blocks == ref_building_blocks
def test_can_do_sparsity_freeze_epoch():
def compare_binary_mask(ref_sparse_module_info, sparse_module_info):
for ref_sparse_layer, sparse_layer in zip(ref_sparse_module_info, sparse_module_info):
if (ref_sparse_layer.operand.binary_mask != sparse_layer.operand.binary_mask).view(-1).sum() != 0:
return False
return True
model = BasicConvTestModel()
config = get_empty_config()
config['compression'] = {"algorithm": "magnitude_sparsity",
"sparsity_init": 0.1,
"params": {"sparsity_target": 0.9,
"sparsity_target_epoch": 3,
"sparsity_freeze_epoch": 3}}
_, compression_ctrl = create_compressed_model_and_algo_for_test(model, config)
sparsified_minfo_before_update = deepcopy(compression_ctrl.sparsified_module_info)
compression_ctrl.scheduler.epoch_step() # update binary_masks
compression_ctrl.scheduler.epoch_step() # update binary_masks
compression_ctrl.scheduler.epoch_step() # update binary_masks, freeze binary_masks
sparsified_minfo_after_update = deepcopy(compression_ctrl.sparsified_module_info)
assert not compare_binary_mask(sparsified_minfo_after_update, sparsified_minfo_before_update)
compression_ctrl.scheduler.epoch_step() # don't update binary_masks
sparsified_minfo_after_freeze = deepcopy(compression_ctrl.sparsified_module_info)
assert compare_binary_mask(sparsified_minfo_after_update, sparsified_minfo_after_freeze)
def test_quantization_preset_with_scope_overrides():
model = QuantizeOutputsTestModel()
config = get_empty_config(input_sample_sizes=[2, 3, 32, 32])
config['target_device'] = "TRIAL"
config['compression'] = {
'algorithm': 'quantization',
'preset': 'mixed',
'scope_overrides': {
'weights': {
'QuantizeOutputsTestModel/NNCFConv2d[conv5]/conv2d_0': {
"mode": "asymmetric",
}
}
}
}
register_bn_adaptation_init_args(config)
_, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
for wq_info in compression_ctrl.weight_quantizers.values():
if wq_info.affected_insertions[0].target_node_name !=\
'QuantizeOutputsTestModel/NNCFConv2d[conv5]/conv2d_0':
assert isinstance(wq_info.quantizer_module_ref, SymmetricQuantizer)
else:
assert isinstance(wq_info.quantizer_module_ref,
AsymmetricQuantizer)
for aq_info in compression_ctrl.non_weight_quantizers.values():
assert isinstance(aq_info.quantizer_module_ref, AsymmetricQuantizer)
def test_magnitude_algo_binary_masks_are_applied():
model = BasicConvTestModel()
config = get_empty_config()
config['compression'] = {'algorithm': "magnitude_sparsity"}
compressed_model, compression_ctrl = create_compressed_model_and_algo_for_test(model, config)
minfo_list = compression_ctrl.sparsified_module_info # type: List[SparseModuleInfo]
minfo = minfo_list[0] # type: SparseModuleInfo
minfo.operand.binary_mask = torch.ones_like(minfo.module.weight) # 1x1x2x2
input_ = torch.ones(size=(1, 1, 5, 5))
ref_output_1 = -4 * torch.ones(size=(2, 4, 4))
output_1 = compressed_model(input_)
assert torch.all(torch.eq(output_1, ref_output_1))
minfo.operand.binary_mask[0][0][0][1] = 0
minfo.operand.binary_mask[1][0][1][0] = 0
ref_output_2 = - 3 * torch.ones_like(ref_output_1)
output_2 = compressed_model(input_)
assert torch.all(torch.eq(output_2, ref_output_2))
minfo.operand.binary_mask[1][0][0][1] = 0
ref_output_3 = ref_output_2.clone()
ref_output_3[1] = -2 * torch.ones_like(ref_output_1[1])
output_3 = compressed_model(input_)
assert torch.all(torch.eq(output_3, ref_output_3))
def test_activation_quantizers_order_is_the_same__for_resnet50(
tmp_path, runs_subprocess_in_precommit):
if not torch.cuda.is_available():
pytest.skip("Skipping CUDA test cases for CPU only setups")
config = get_empty_config(input_sample_sizes=[1, 3, 224, 224])
config['compression'] = {
'algorithm': 'quantization',
"initializer": {
"range": {
"num_init_samples": 0
}
}
}
register_bn_adaptation_init_args(config)
ngpus_per_node = torch.cuda.device_count()
torch.multiprocessing.spawn(activation_quantizers_dumping_worker,
nprocs=ngpus_per_node,
args=(config, tmp_path),
join=True)
with open(get_path_to_keys(tmp_path, 0), 'r', encoding='utf8') as f:
ref_list = f.readlines()
for i in range(1, ngpus_per_node):
with open(get_path_to_keys(tmp_path, i), 'r', encoding='utf8') as f:
curr_list = f.readlines()
assert curr_list == ref_list
def test_export_stacked_bi_lstm(tmp_path):
p = LSTMTestSizes(3, 3, 3, 3)
config = get_empty_config(
input_sample_sizes=[1, p.hidden_size, p.input_size])
config['compression'] = {'algorithm': 'quantization'}
register_bn_adaptation_init_args(config)
# TODO: batch_first=True fails with building graph: ambiguous call to mul or sigmoid
test_rnn = NNCF_RNN('LSTM',
input_size=p.input_size,
hidden_size=p.hidden_size,
num_layers=2,
bidirectional=True,
batch_first=False)
model, algo = create_compressed_model_and_algo_for_test(test_rnn, config)
test_path = str(tmp_path.joinpath('test.onnx'))
algo.export_model(test_path)
assert os.path.exists(test_path)
onnx_num = 0
model = onnx.load(test_path)
# pylint: disable=no-member
for node in model.graph.node:
if node.op_type == 'FakeQuantize':
onnx_num += 1
assert onnx_num == 54
def test_can_restore_binary_mask_on_magnitude_quant_algo_resume(
tmp_path, use_data_parallel):
config = get_empty_config()
config["compression"] = [{
"algorithm": "magnitude_sparsity",
"params": {
"schedule": "multistep",
"multistep_sparsity_levels": [0.3, 0.5],
"weight_importance": "abs"
}
}, {
"algorithm": "quantization"
}]
register_bn_adaptation_init_args(config)
sparse_model, _ = create_compressed_model_and_algo_for_test(
MagnitudeTestModel(), config)
if use_data_parallel:
if not torch.cuda.is_available():
pytest.skip("Skipping CUDA test cases for CPU only setups")
sparse_model.cuda()
sparse_model = torch.nn.DataParallel(sparse_model)
with torch.no_grad():
sparse_model(torch.ones([1, 1, 10, 10]))
config = get_empty_config()
config["compression"] = [{
"algorithm": "const_sparsity"
}, {
"algorithm": "quantization"
}]
register_bn_adaptation_init_args(config)
const_sparse_model, _ = create_compressed_model_and_algo_for_test(
MagnitudeTestModel(), config)
load_state(const_sparse_model, sparse_model.state_dict())
op = const_sparse_model.get_nncf_wrapped_model().conv1.pre_ops['0']
PTTensorListComparator.check_equal(ref_mask_1, op.operand.binary_mask)
op = const_sparse_model.get_nncf_wrapped_model().conv2.pre_ops['0']
PTTensorListComparator.check_equal(ref_mask_2, op.operand.binary_mask)
def test_can_not_create_magnitude_algo__with_adaptive_scheduler():
config = get_empty_config()
config['compression'] = {
'algorithm': 'magnitude_sparsity',
"params": {
"schedule": 'adaptive'
}
}
with pytest.raises(TypeError):
_, _ = create_compressed_model_and_algo_for_test(MockModel(), config)
def get_basic_quantization_config(quantization_type='symmetric', input_sample_sizes=None, input_info: Dict = None):
config = get_empty_config(input_sample_sizes=input_sample_sizes, input_info=input_info)
config["compression"] = {"algorithm": "quantization",
"activations": {
"mode": quantization_type
},
"weights": {
"mode": quantization_type
}}
return config
def test_can_create_rb_algo__with_adaptive_scheduler():
config = get_empty_config()
config['compression'] = {
'algorithm': 'rb_sparsity',
"params": {
"schedule": 'adaptive'
}
}
_, compression_ctrl = create_compressed_model_and_algo_for_test(
MockModel(), config)
assert isinstance(compression_ctrl.scheduler, AdaptiveSparsityScheduler)
def test_can_choose_scheduler(algo, schedule_type, scheduler_class):
config = get_empty_config()
config['compression'] = {
'algorithm': algo,
"params": {
"schedule": schedule_type
}
}
_, compression_ctrl = create_compressed_model_and_algo_for_test(
MockModel(), config)
assert isinstance(compression_ctrl.scheduler, scheduler_class)
def get_quantization_config_without_range_init(model_size=4) -> NNCFConfig:
config = get_empty_config(input_sample_sizes=[1, 1, model_size, model_size])
config["compression"] = {
"algorithm": "quantization",
"initializer": {
"range": {
"num_init_samples": 0
}
}
}
return config
def test_create_rb_algo_with_local_sparsity_mode():
config = get_empty_config()
config['compression'] = {
'algorithm': 'rb_sparsity',
"params": {
"sparsity_level_setting_mode": 'local'
}
}
_, compression_ctrl = create_compressed_model_and_algo_for_test(
MockModel(), config)
assert compression_ctrl.compression_stage(
) == CompressionStage.FULLY_COMPRESSED
def create_config():
config = get_empty_config()
config['compression'] = {
'algorithm': 'quantization',
'initializer': {
'range': {
'num_init_samples': 1
}
}
}
register_bn_adaptation_init_args(config)
return config
def test_create_rb_algo_with_per_layer_loss():
config = get_empty_config()
config['compression'] = {
'algorithm': 'rb_sparsity',
"params": {
"sparsity_level_setting_mode": 'local'
}
}
_, compression_ctrl = create_compressed_model_and_algo_for_test(
MockModel(), config)
# pylint: disable=protected-access
assert isinstance(compression_ctrl._loss, SparseLossForPerLayerSparsity)
def test_can_freeze_binary_masks():
model = BasicConvTestModel()
config = get_empty_config()
config['compression'] = {'algorithm': "magnitude_sparsity"}
_, compression_ctrl = create_compressed_model_and_algo_for_test(model, config)
for sparse_layer in compression_ctrl.sparsified_module_info:
assert not sparse_layer.operand.frozen
compression_ctrl.freeze()
for sparse_layer in compression_ctrl.sparsified_module_info:
assert sparse_layer.operand.frozen
def test_can_create_sparse_scheduler__with_defaults(self, algo):
config = get_empty_config()
config['compression'] = {
'algorithm': algo,
"params": {
"schedule": 'polynomial'
}
}
_, compression_ctrl = create_compressed_model_and_algo_for_test(
MockModel(), config)
scheduler = compression_ctrl.scheduler
assert scheduler.initial_level == 0
assert scheduler.target_level == 0.5
assert scheduler.target_epoch == 90
assert scheduler.freeze_epoch == 100
def test_sparse_network(self, desc: ModelDesc, algo):
model = desc.model_builder()
config = get_empty_config(input_sample_sizes=desc.input_sample_sizes)
config["compression"] = {"algorithm": algo}
compressed_model, compression_ctrl = \
create_compressed_model_and_algo_for_test(model, config, dummy_forward_fn=desc.dummy_forward_fn,
wrap_inputs_fn=desc.wrap_inputs_fn)
sparsifiable_modules = self.get_sparsifiable_modules(algo)
ref_num_sparsed = len(get_all_modules_by_type(model, sparsifiable_modules))
assert ref_num_sparsed == len(compression_ctrl.sparsified_module_info)
check_model_graph(compressed_model, desc.dot_filename, algo)
def test_ordinary_load(algo, _model_wrapper, is_resume):
config = get_empty_config()
if algo:
config['compression'] = {'algorithm': algo}
register_bn_adaptation_init_args(config)
compressed_model_save, _ = create_compressed_model_and_algo_for_test(
BasicConvTestModel(), config)
model_save = _model_wrapper['save_model'](compressed_model_save)
compressed_model_resume, _ = create_compressed_model_and_algo_for_test(
BasicConvTestModel(), config)
model_resume = _model_wrapper['resume_model'](compressed_model_resume)
num_loaded = load_state(model_resume, model_save.state_dict(), is_resume)
assert num_loaded == len(model_save.state_dict())
def test_sparsity_statistics_add_module(algo):
config = get_empty_config()
sparsity_init = 0.5
config['compression'] = {'algorithm': algo, 'sparsity_init': sparsity_init}
model = TwoConvTestModel()
fill_params_of_model_by_normal(model)
submodule = TwoConvTestModel()
fill_params_of_model_by_normal(submodule)
model, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
statistics_before = compression_ctrl.statistics()
model.add_module('submodule', submodule)
statistics_after = compression_ctrl.statistics()
assert getattr(statistics_before, algo).model_statistics.sparsity_level == \
getattr(statistics_after, algo).model_statistics.sparsity_level
assert getattr(statistics_before, algo).model_statistics.sparsity_level_for_layers == \
getattr(statistics_after, algo).model_statistics.sparsity_level_for_layers
def test_export_lstm_cell(tmp_path):
config = get_empty_config(model_size=1, input_sample_sizes=[1, 1])
config['compression'] = {'algorithm': 'quantization'}
register_bn_adaptation_init_args(config)
model, algo = create_compressed_model_and_algo_for_test(
LSTMCellNNCF(1, 1), config)
test_path = str(tmp_path.joinpath('test.onnx'))
algo.export_model(test_path)
assert os.path.exists(test_path)
onnx_num = 0
model = onnx.load(test_path)
# pylint: disable=no-member
for node in model.graph.node:
if node.op_type == 'FakeQuantize':
onnx_num += 1
assert onnx_num == 14
def test_sparse_quantize_network(self, desc: ModelDesc):
model = desc.model_builder()
config = get_empty_config(input_sample_sizes=desc.input_sample_sizes)
config["compression"] = [
{"algorithm": "rb_sparsity"},
{"algorithm": "quantization"}
]
register_bn_adaptation_init_args(config)
compressed_model, compression_ctrl = \
create_compressed_model_and_algo_for_test(model, config, dummy_forward_fn=desc.dummy_forward_fn,
wrap_inputs_fn=desc.wrap_inputs_fn)
sparsifiable_modules = self.get_sparsifiable_modules('rb_sparsity')
ref_num_sparsed = len(get_all_modules_by_type(compressed_model, sparsifiable_modules))
assert ref_num_sparsed == len(compression_ctrl.child_ctrls[0].sparsified_module_info)
check_model_graph(compressed_model, desc.dot_filename, "quantized_rb_sparsity")
def test_rb_sparsity__can_set_sparsity_level_for_module():
config = get_empty_config()
config['compression'] = {
'algorithm': 'rb_sparsity',
"params": {
"sparsity_level_setting_mode": 'local'
}
}
_, compression_ctrl = create_compressed_model_and_algo_for_test(
BasicConvTestModel(), config)
# pylint: disable=protected-access
assert list(compression_ctrl._loss.per_layer_target.values())[0] == 1
compression_ctrl.set_sparsity_level(
0.7, compression_ctrl.sparsified_module_info[0])
assert list(
compression_ctrl._loss.per_layer_target.values())[0] == pytest.approx(
0.3)
def test_kd_sparsity_statistics(algo: str):
model = TwoConvTestModel()
fill_params_of_model_by_normal(model)
model_with_kd = deepcopy(model)
config = get_empty_config()
sparsity_init = 0.5
config['compression'] = {'algorithm': algo, 'sparsity_init': sparsity_init}
config_with_kd = deepcopy(config)
config_with_kd = get_kd_config(config_with_kd)
model, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
model_with_kd, compression_ctrl_with_kd = create_compressed_model_and_algo_for_test(
model_with_kd, config_with_kd)
statistics = compression_ctrl.statistics()
statistics_with_kd = compression_ctrl_with_kd.statistics()
assert getattr(statistics, algo).model_statistics.sparsity_level ==\
getattr(statistics_with_kd, algo).model_statistics.sparsity_level
assert getattr(statistics, algo).model_statistics.sparsity_level_for_layers ==\
getattr(statistics_with_kd, algo).model_statistics.sparsity_level_for_layers
def test_quantization_preset(data):
model = BasicConvTestModel()
config = get_empty_config(input_sample_sizes=[1, 1, 4, 4])
config['target_device'] = data['target_device']
config['compression'] = {
'algorithm': 'quantization',
'preset': data['preset']
}
config['compression'].update(data['overrided_param'])
register_bn_adaptation_init_args(config)
_, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
for wq_info in compression_ctrl.weight_quantizers.values():
assert isinstance(wq_info.quantizer_module_ref,
data['expected_weights_q'])
for aq_info in compression_ctrl.non_weight_quantizers.values():
assert isinstance(aq_info.quantizer_module_ref,
data['expected_activations_q'])
