def test_groups(test_input_info_struct_: GroupPruningModulesTestStruct): model = test_input_info_struct_.model not_pruned_modules = test_input_info_struct_.not_pruned_modules pruned_groups = test_input_info_struct_.pruned_groups prune_first, prune_last, prune_downsample = test_input_info_struct_.prune_params model = model() nncf_config = get_basic_pruning_config(input_sample_size=[1, 1, 8, 8]) nncf_config['compression']['algorithm'] = 'filter_pruning' nncf_config['compression']['params']['prune_first_conv'] = prune_first nncf_config['compression']['params']['prune_last_conv'] = prune_last nncf_config['compression']['params']['prune_downsample_convs'] = prune_downsample compressed_model, compression_ctrl = create_compressed_model_and_algo_for_test(model, nncf_config) # 1. Check all not pruned modules clusters = compression_ctrl.pruned_module_groups_info all_pruned_modules_info = clusters.get_all_nodes() all_pruned_modules = [info.module for info in all_pruned_modules_info] print([minfo.module_scope for minfo in all_pruned_modules_info]) for module_name in not_pruned_modules: module = compressed_model.get_module_by_scope(Scope.from_str(module_name)) assert module not in all_pruned_modules # 2. Check that all pruned groups are valid for group in pruned_groups: first_node_scope = Scope.from_str(group[0]) cluster = clusters.get_cluster_by_node_id(first_node_scope) cluster_modules = [n.module for n in cluster.nodes] group_modules = [compressed_model.get_module_by_scope(Scope.from_str(module_scope)) for module_scope in group] assert cluster_modules == group_modules
def test_can_choose_scheduler(algo, scheduler, scheduler_class): config = get_basic_pruning_config() config['compression']['algorithm'] = algo config['compression']['params']['schedule'] = scheduler model = PruningTestModel() _, compression_ctrl = create_compressed_model_and_algo_for_test( model, config) scheduler = compression_ctrl.scheduler assert isinstance(scheduler, scheduler_class)
def test_both_targets_assert(): config = get_basic_pruning_config() config['compression']['algorithm'] = 'filter_pruning' config['compression']['params']['pruning_target'] = 0.3 config['compression']['params']['pruning_flops_target'] = 0.5 model = PruningTestModel() with pytest.raises(ValueError): create_compressed_model_and_algo_for_test(model, config)
def create_nncf_model_and_builder(model, config_params): nncf_config = get_basic_pruning_config(input_sample_size=[1, 1, 8, 8]) nncf_config['compression']['algorithm'] = 'filter_pruning' for key, value in config_params.items(): nncf_config['compression']['params'][key] = value nncf_model, algo_builders_list = create_nncf_model_and_algo_builder(model, nncf_config) assert len(algo_builders_list) == 1 algo_builder = algo_builders_list[0] return nncf_model, algo_builder
def test_init_params_for_flops_calculation(model, ref_params): config = get_basic_pruning_config() config['compression']['algorithm'] = 'filter_pruning' config['compression']['params']['pruning_flops_target'] = 0.3 config['compression']['params']['prune_first_conv'] = True model = model() _, compression_ctrl = create_compressed_model_and_algo_for_test(model, config) for key, value in ref_params.items(): assert getattr(compression_ctrl, key) == value
def test_check_default_scheduler_params(algo, ref_scheduler, ref_scheduler_params): config = get_basic_pruning_config() config['compression']['algorithm'] = algo model = PruningTestModel() _, compression_ctrl = create_compressed_model_and_algo_for_test( model, config) scheduler = compression_ctrl.scheduler assert isinstance(scheduler, ref_scheduler) for key, value in ref_scheduler_params.items(): assert getattr(scheduler, key) == value
def test_pruning_export_simple_model(tmp_path): model = BigPruningTestModel() nncf_config = get_basic_pruning_config(input_sample_size=[1, 1, 8, 8]) nncf_config['compression']['algorithm'] = 'filter_pruning' onnx_model_proto = load_exported_onnx_version(nncf_config, model, path_to_storage_dir=tmp_path) # Check that conv2 + BN were pruned by output filters check_bias_and_weight_shape('nncf_module.conv2', onnx_model_proto, [16, 16, 3, 3], [16]) check_bias_and_weight_shape('nncf_module.bn', onnx_model_proto, [16], [16]) # Check that conv3 was pruned by input filters check_bias_and_weight_shape('nncf_module.conv3', onnx_model_proto, [1, 16, 5, 5], [1])
def test_pruning_export_eltwise_model(tmp_path, prune_first, prune_last, ref_shapes): model = PruningTestModelEltwise() nncf_config = get_basic_pruning_config(input_sample_size=[1, 1, 8, 8]) nncf_config['compression']['algorithm'] = 'filter_pruning' nncf_config['compression']['params']['prune_first_conv'] = prune_first nncf_config['compression']['params']['prune_last_conv'] = prune_last onnx_model_proto = load_exported_onnx_version(nncf_config, model, path_to_storage_dir=tmp_path) for i in range(1, 5): conv_name = "nncf_module.conv{}".format(i) check_bias_and_weight_shape(conv_name, onnx_model_proto, *ref_shapes[i - 1])
def test_prune_flops_param(pruning_target, pruning_flops_target, prune_flops_ref, pruning_target_ref): config = get_basic_pruning_config() config['compression']['algorithm'] = 'filter_pruning' if pruning_target: config['compression']['params']['pruning_target'] = pruning_target if pruning_flops_target: config['compression']['params']['pruning_flops_target'] = pruning_flops_target config['compression']['params']['prune_first_conv'] = True model = PruningTestModel() _, compression_ctrl = create_compressed_model_and_algo_for_test(model, config) assert compression_ctrl.prune_flops is prune_flops_ref assert compression_ctrl.scheduler.pruning_target == pruning_target_ref
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_op_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_get_bn_for_module_scope(): 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( BigPruningTestModel(), config) conv1_scope = Scope.from_str('BigPruningTestModel/NNCFConv2d[conv1]') bn = get_bn_for_module_scope(pruned_model, conv1_scope) assert bn is None conv2_scope = Scope.from_str('BigPruningTestModel/NNCFConv2d[conv2]') bn = get_bn_for_module_scope(pruned_model, conv2_scope) assert bn == pruned_model.bn conv3_scope = Scope.from_str('BigPruningTestModel/NNCFConv2d[conv3]') bn = get_bn_for_module_scope(pruned_model, conv3_scope) assert bn is None
def test_pruning_export_simple_model(tmp_path): model = BigPruningTestModel() nncf_config = get_basic_pruning_config(input_sample_size=[1, 1, 8, 8]) nncf_config['compression']['pruning_init'] = 0.5 nncf_config['compression']['algorithm'] = 'filter_pruning' onnx_model_proto = load_exported_onnx_version(nncf_config, model, path_to_storage_dir=tmp_path) # Check that conv2 + BN were pruned by output filters # WARNING: starting from at least torch 1.7.0, torch.onnx.export will fuses BN into previous # convs if torch.onnx.export is done with `training=False`, so this test might fail. check_bias_and_weight_shape('nncf_module.conv2', onnx_model_proto, [16, 16, 3, 3], [16]) check_bias_and_weight_shape('nncf_module.bn', onnx_model_proto, [16], [16]) # Check that up was pruned by input filters check_bias_and_weight_shape('nncf_module.up', onnx_model_proto, [16, 32, 3, 3], [32]) # Check that conv3 was pruned by input filters check_bias_and_weight_shape('nncf_module.conv3', onnx_model_proto, [1, 32, 5, 5], [1])