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)