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']
check_equal(ref_mask_1, op.operand.binary_mask)
op = const_sparse_model.conv2.pre_ops['0']
check_equal(ref_mask_2, op.operand.binary_mask)
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']]
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']]
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_context_independence(model_name, model_builder, input_size, _case_config):
config = get_basic_quantization_config(_case_config.quant_type, input_sample_sizes=input_size[0])
compressed_models = [create_compressed_model_and_algo_for_test(model_builder[0](), config)[0],
create_compressed_model_and_algo_for_test(model_builder[1](), config)[0]]
for i, compressed_model in enumerate(compressed_models):
check_model_graph(compressed_model, model_name[i], _case_config.graph_dir)
def test_ordinary_load(algo, _model_wrapper, is_resume):
config = get_empty_config()
if algo:
config['compression'] = {'algorithm': algo}
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_check_default_algo_params():
"""
Test for default algorithm params. Creating empty config and check for valid default
parameters.
"""
# Creating algorithm with empty config
config = get_basic_pruning_config()
config['compression']['algorithm'] = 'filter_pruning'
model = PruningTestModel()
_, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
assert isinstance(compression_ctrl, FilterPruningController)
scheduler = compression_ctrl.scheduler
# Check default algo params
assert compression_ctrl.prune_first is False
assert compression_ctrl.prune_last is False
assert compression_ctrl.prune_batch_norms is False
assert compression_ctrl.filter_importance is l2_filter_norm
assert compression_ctrl.all_weights is False
assert compression_ctrl.zero_grad is True
# Check default scheduler params
assert isinstance(scheduler, BaselinePruningScheduler)
def test_can_quantize_free_operators(mocker):
class Model(nn.Module):
def __init__(self):
super().__init__()
self.weight = nn.Parameter(torch.ones([1]))
self.bias = nn.Parameter(torch.ones([1]))
def forward(self, x):
return F.linear(x, self.weight, self.bias)
mod = Model()
config = get_basic_quantization_config(model_size=1)
config["compression"].update({"quantize_inputs": False})
quant_model, _ = create_compressed_model_and_algo_for_test(mod, config)
quantizer_list = quant_model.get_compression_modules_by_type(
CompressionModuleType.FUNCTION_QUANTIZER).values()
assert len(quantizer_list) == 2
for quantizer in quantizer_list:
mocker.spy(quantizer, 'quantize')
quant_model.do_dummy_forward()
for quantizer in quantizer_list:
assert quantizer.quantize.call_count == 1
def test_baseline_scheduler():
"""
Test baseline scheduler parameters and changes of params during epochs.
"""
config = get_pruning_baseline_config()
config['compression']['algorithm'] = 'filter_pruning'
model = PruningTestModel()
_, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
scheduler = compression_ctrl.scheduler
# Check default params
assert isinstance(scheduler, BaselinePruningScheduler)
assert pytest.approx(scheduler.pruning_target) == 0.5
assert pytest.approx(scheduler.initial_pruning) == 0.0
assert scheduler.num_init_steps == 1
# Check pruning params on epoch 0
assert pytest.approx(scheduler.current_pruning_level) == 0.0
assert pytest.approx(compression_ctrl.pruning_rate) == 0.0
assert scheduler.last_epoch == 0
assert compression_ctrl.frozen is False
# Check pruning params on epoch 1
scheduler.epoch_step()
assert pytest.approx(scheduler.current_pruning_level) == 0.5
assert pytest.approx(compression_ctrl.pruning_rate) == 0.5
assert scheduler.last_epoch == 1
assert compression_ctrl.frozen is True
def test_hawq_precision_init(_seed, dataset_dir, tmp_path, mocker):
num_data_points = 100
batch_size = 10
config = create_hawq_test_config(batch_size, num_data_points)
model = squeezenet1_1_custom(num_classes=10, pretrained=False, dropout=0)
model, compression_ctrl = create_compressed_model_and_algo_for_test(model, config)
load_state(model, model_zoo.load_url(model_urls['squeezenet1_1']))
model = model.cuda()
device = next(model.parameters()).device
criterion = nn.CrossEntropyLoss().cuda()
if not dataset_dir:
dataset_dir = str(tmp_path)
train_loader, _ = create_test_dataloaders(config.model_size, dataset_dir, batch_size)
mocked_trace = mocker.patch('nncf.quantization.hessian_trace.HessianTraceEstimator.get_average_traces')
num_traces = len(get_all_modules_by_type(model, 'NNCFConv2d'))
mock_avg_traces = [torch.Tensor([num_traces - i]).to(device) for i in range(num_traces)]
mocked_trace.return_value = mock_avg_traces
compression_ctrl.initialize(criterion=criterion, data_loader=train_loader)
act_bitwidth_per_scope = get_bitwidth_per_scope(model)
path_to_ref = str(TEST_ROOT / 'data/hawq_reference/squeezenet1_1_mixed_bitwidth_per_scope.json')
compare_with_ref_if_exists(act_bitwidth_per_scope, path_to_ref)
def test_export_stacked_bi_lstm(tmp_path):
p = LSTMTestSizes(3, 3, 3, 3)
config = get_empty_config(input_sample_size=(1, p.hidden_size,
p.input_size))
config['compression'] = {'algorithm': 'quantization'}
# 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 == 50
def test_model_can_be_loaded_with_resume(_params):
p = _params
sample_config_path = p['sample_config_path']
checkpoint_path = p['checkpoint_path']
config = SampleConfig.from_json(str(sample_config_path))
nncf_config = NNCFConfig.from_json(str(sample_config_path))
config.execution_mode = p['execution_mode']
config.current_gpu = 0
config.device = get_device(config)
config.distributed = config.execution_mode in (ExecutionMode.DISTRIBUTED, ExecutionMode.MULTIPROCESSING_DISTRIBUTED)
if config.distributed:
config.dist_url = "tcp://127.0.0.1:9898"
config.dist_backend = "nccl"
config.rank = 0
config.world_size = 1
configure_distributed(config)
model_name = config['model']
model = load_model(model_name,
pretrained=False,
num_classes=config.get('num_classes', 1000),
model_params=config.get('model_params'))
model.to(config.device)
model, compression_ctrl = create_compressed_model_and_algo_for_test(model, nncf_config)
model, _ = prepare_model_for_execution(model, config)
if config.distributed:
compression_ctrl.distributed()
checkpoint = torch.load(checkpoint_path, map_location='cpu')
load_state(model, checkpoint['state_dict'], is_resume=True)
def test_hawq_precision_init(_seed, dataset_dir, tmp_path, mocker,
config_creator: Callable, filename_suffix: str):
num_data_points = 100
batch_size = 10
config = config_creator(batch_size, num_data_points)
model = MobileNetV2(num_classes=10)
model.eval()
criterion = nn.CrossEntropyLoss().cuda()
if not dataset_dir:
dataset_dir = str(tmp_path)
train_loader, _ = create_test_dataloaders(config.get("model_size"),
dataset_dir, batch_size)
config = register_default_init_args(config, criterion, train_loader)
mocked_trace = mocker.patch(
'nncf.quantization.hessian_trace.HessianTraceEstimator.get_average_traces'
)
mock_avg_traces = get_mock_avg_traces(model)
mocked_trace.return_value = mock_avg_traces
from torchvision.models.mobilenet import model_urls
load_state(model, model_zoo.load_url(model_urls['mobilenet_v2']))
model, algo_ctrl = create_compressed_model_and_algo_for_test(model, config)
model = model.cuda()
all_quantizers_per_full_scope = get_all_quantizers_per_full_scope(model)
graph = get_bitwidth_graph(algo_ctrl, model, all_quantizers_per_full_scope)
path_to_dot = 'mobilenet_v2_mixed_bitwidth_graph_{}.dot'.format(
filename_suffix)
check_graph(graph,
path_to_dot,
os.path.join('quantized', 'hawq'),
sort_dot_graph=False)
def test_magnitude_scheduler_can_do_epoch_step__with_norm():
_ = MagnitudeTestModel()
config = get_multistep_normed_abs_config()
_, compression_ctrl = create_compressed_model_and_algo_for_test(MagnitudeTestModel(), config)
scheduler = compression_ctrl.scheduler
assert isinstance(scheduler, MultiStepSparsityScheduler)
assert pytest.approx(compression_ctrl.sparsity_level) == 0.1
assert compression_ctrl.threshold == pytest.approx(0.219, 0.01)
assert scheduler.prev_ind == 0
scheduler.epoch_step()
assert compression_ctrl.sparsity_level == 0.5
assert compression_ctrl.threshold == pytest.approx(0.243, 0.01)
assert scheduler.prev_ind == 1
scheduler.epoch_step()
assert compression_ctrl.sparsity_level == 0.5
assert compression_ctrl.threshold == pytest.approx(0.243, 0.01)
assert scheduler.prev_ind == 1
scheduler.epoch_step()
assert compression_ctrl.sparsity_level == 0.9
assert compression_ctrl.threshold == pytest.approx(0.371, 0.01)
assert scheduler.prev_ind == 2
def test_can_load_quant_algo__with_defaults():
model = BasicConvTestModel()
config = get_basic_quantization_config()
compression_algo_builder_list = create_compression_algorithm_builders(config)
assert len(compression_algo_builder_list) == 1
assert isinstance(compression_algo_builder_list[0], QuantizationBuilder)
quant_model, _ = create_compressed_model_and_algo_for_test(deepcopy(model), config)
model_conv = get_all_modules_by_type(model, 'Conv2d')
quant_model_conv = get_all_modules_by_type(quant_model.get_nncf_wrapped_model(), 'NNCFConv2d')
assert len(model_conv) == len(quant_model_conv)
for module_scope, _ in model_conv.items():
quant_scope = deepcopy(module_scope) # type: Scope
quant_scope.pop()
quant_scope.push(ScopeElement('NNCFConv2d', 'conv'))
assert quant_scope in quant_model_conv.keys()
store = []
for op in quant_model_conv[quant_scope].pre_ops.values():
if isinstance(op, (UpdateInputs, UpdateWeight)) and isinstance(op.operand, SymmetricQuantizer):
assert op.__class__.__name__ not in store
store.append(op.__class__.__name__)
assert UpdateWeight.__name__ in store
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_can_create_magnitude_sparse_algo__with_defaults():
model = MagnitudeTestModel()
config = get_basic_magnitude_sparsity_config()
config['compression']['params'] = \
{'schedule': 'multistep'}
sparse_model, compression_ctrl = create_compressed_model_and_algo_for_test(
deepcopy(model), config)
assert isinstance(compression_ctrl, MagnitudeSparsityController)
assert compression_ctrl.sparsity_level == approx(0.1)
assert len(list(sparse_model.modules())) == 12
_, sparse_model_conv = check_correct_nncf_modules_replacement(
model, sparse_model)
i = 0
for sparse_module in sparse_model_conv.values():
store = []
ref_mask = torch.ones_like(
sparse_module.weight) if i == 0 else ref_mask_2
i += 1
for op in sparse_module.pre_ops.values():
if isinstance(op, UpdateWeight) and isinstance(
op.operand, BinaryMask):
assert compression_ctrl.threshold == approx(0.24, 0.1)
assert torch.allclose(op.operand.binary_mask, ref_mask)
assert isinstance(compression_ctrl.weight_importance,
type(normed_magnitude))
assert op.__class__.__name__ not in store
store.append(op.__class__.__name__)
def test_can_not_set_sparsity_more_than_one_for_magnitude_sparse_algo():
config = get_basic_magnitude_sparsity_config()
_, compression_ctrl = create_compressed_model_and_algo_for_test(
MagnitudeTestModel(), config)
with pytest.raises(AttributeError):
compression_ctrl.set_sparsity_level(1)
compression_ctrl.set_sparsity_level(1.2)
def test_can_choose_scheduler(algo, schedule_type, scheduler_class):
config = get_empty_config()
config['compression']['algorithm'] = algo
config['compression']["params"]["schedule"] = schedule_type
_, compression_ctrl = create_compressed_model_and_algo_for_test(
MockModel(), config)
assert isinstance(compression_ctrl.scheduler, scheduler_class)
def activation_quantizers_dumping_worker(current_gpu, config, tmp_path):
model = resnet50(pretrained=False)
quant_model, _ = create_compressed_model_and_algo_for_test(model, config)
path = get_path_to_keys(tmp_path, current_gpu)
print(path)
with open(path, 'w') as f:
f.writelines("%s\n" % key for key in quant_model.activation_quantizers.keys())
def test_can_create_rb_algo__with_adaptive_scheduler():
config = get_empty_config()
config['compression']['algorithm'] = 'rb_sparsity'
config['compression']["params"]["schedule"] = 'adaptive'
_, compression_ctrl = create_compressed_model_and_algo_for_test(
MockModel(), config)
assert isinstance(compression_ctrl.scheduler, AdaptiveSparsityScheduler)
def test_input_info_specification_from_config(mocker, input_info_test_struct):
stub_fn = mocker.stub()
mock_model = MockModel(stub_fn)
config = get_basic_quantization_config("symmetric", None)
input_info_config_entry = input_info_test_struct[0]
target_argument_info = input_info_test_struct[
1] # type: List[ModelInputInfo]
config["input_info"] = input_info_config_entry
_, _ = create_compressed_model_and_algo_for_test(mock_model, config)
forward_call_args = stub_fn.call_args[0]
forward_call_kwargs = stub_fn.call_args[1]
ref_args_info = list(
filter(lambda x: x.keyword is None, target_argument_info))
ref_kw_vs_arg_info = {
x.keyword: x
for x in target_argument_info if x.keyword is not None
}
def check_arg(arg: torch.Tensor, ref_arg_info: ModelInputInfo):
assert arg.shape == ref_arg_info.shape
assert arg.dtype == ref_arg_info.type
assert len(forward_call_args) == len(ref_args_info)
assert len(forward_call_kwargs) == len(ref_kw_vs_arg_info)
assert set(forward_call_kwargs.keys()) == set(ref_kw_vs_arg_info.keys())
for idx, arg in enumerate(forward_call_args):
check_arg(arg, ref_args_info[idx])
for keyword, arg in forward_call_kwargs.items():
check_arg(arg, ref_kw_vs_arg_info[keyword])
def test_can_not_create_magnitude_algo__without_steps():
config = get_basic_magnitude_sparsity_config()
config['compression']['params'] = {
'schedule': 'multistep',
'multistep_sparsity_levels': [0.1]
}
with pytest.raises(AttributeError):
_, _ = create_compressed_model_and_algo_for_test(MockModel(), config)
def test_sparse_algo_can_collect_sparse_layers():
model = TwoConvTestModel()
config = get_basic_sparsity_config()
_, compression_ctrl = create_compressed_model_and_algo_for_test(
model, config)
assert len(compression_ctrl.sparsified_module_info) == 2
def test_can_export_compressed_model(self, tmp_path, config_provider, model_provider):
test_path = str(tmp_path.joinpath('test.onnx'))
model = model_provider()
config = config_provider()
_, compression_ctrl = create_compressed_model_and_algo_for_test(model, config)
compression_ctrl.export_model(test_path)
assert os.path.exists(test_path)
def test_quantize_network(self, desc: ModelDesc, _case_config):
model = desc.model_builder()
config = get_basic_quantization_config(
_case_config.quant_type,
input_sample_sizes=desc.input_sample_sizes)
compressed_model, _ = \
create_compressed_model_and_algo_for_test(model, config, dummy_forward_fn=desc.dummy_forward_fn)
check_model_graph(compressed_model, desc.dot_filename,
_case_config.graph_dir)
def test_magnitude_sparse_algo_sets_threshold(weight_importance, sparsity_level, threshold):
model = MagnitudeTestModel()
config = get_basic_magnitude_sparsity_config()
config['compression']['weight_importance'] = weight_importance
config['compression']['params'] = {'schedule': 'multistep'}
_, 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)
def test_get_last_pruned_layers(model, ref_last_module_names):
config = get_basic_pruning_config(input_sample_size=(1, 1, 8, 8))
config['compression']['algorithm'] = 'filter_pruning'
pruned_model, _ = create_compressed_model_and_algo_for_test(model(), config)
first_pruned_modules = get_last_pruned_modules(pruned_model,
FilterPruningBuilder(config).get_types_of_pruned_modules())
ref_last_modules = [getattr(pruned_model, module_name) for module_name in ref_last_module_names]
assert set(first_pruned_modules) == set(ref_last_modules)
def test_resnet18__with_not_qinput(_case_config):
model = test_models.ResNet18()
input_shape = (1, 3, 32, 32)
config = get_basic_quantization_config(_case_config.quant_type, input_sample_size=input_shape)
config["compression"].update({"quantize_inputs": False})
compressed_model, _ = create_compressed_model_and_algo_for_test(model, config)
check_model_graph(compressed_model, 'resnet18_no_qinput.dot', _case_config.graph_dir)
def test_can_create_magnitude_algo__without_levels():
config = get_basic_magnitude_sparsity_config()
config['compression']['params'] = {
'schedule': 'multistep',
'multistep_steps': [1]
}
_, compression_ctrl = create_compressed_model_and_algo_for_test(
MockModel(), config)
assert compression_ctrl.sparsity_level == approx(0.1)
def test_output_quantization(_case_config):
model = test_models.UNet()
input_shape = (1, 3, 360, 480)
config = get_basic_quantization_config(_case_config.quant_type, input_sample_size=input_shape)
config["compression"].update({"quantize_outputs": True})
compressed_model, _ = create_compressed_model_and_algo_for_test(model, config)
check_model_graph(compressed_model, 'unet_qoutput.dot', _case_config.graph_dir)
def test_can_create_quant_loss_and_scheduler():
config = get_basic_quantization_config()
_, compression_ctrl = create_compressed_model_and_algo_for_test(MockModel(), config)
loss = compression_ctrl.loss
assert isinstance(loss, CompressionLoss)
scheduler = compression_ctrl.scheduler
assert isinstance(scheduler, CompressionScheduler)
