Repository: microsoft/nni
File: nni/algorithms/compression/v2/pytorch/utils/pruning.py (Python)

Code before:

    Example::
        model = models.resnet18()
        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8, 'op_partial_names': ['conv1']}]
        pruner = L1NormPruner(model, config_list)
        pruner.compress()
        pruner.show_pruned_weights()

    In this process, the config_list will implicitly convert to the following:

    [{'op_types': ['Conv2d'], 'sparsity_per_layer': 0.8,
        'op_names': ['conv1', 'layer1.0.conv1', 'layer1.1.conv1',
        'layer2.0.conv1', 'layer2.1.conv1', 'layer3.0.conv1', 'layer3.1.conv1',
        'layer4.0.conv1', 'layer4.1.conv1']}]
    '''
    for config in config_list:
        if 'sparsity' in config:
            if 'sparsity_per_layer' in config:
                raise ValueError("'sparsity' and 'sparsity_per_layer' have the same semantics, can not set both in one config.")
            else:
                config['sparsity_per_layer'] = config.pop('sparsity')

    for config in config_list:
        if 'op_types' in config:
            if 'default' in config['op_types']:
                config['op_types'] = weighted_modules

    for config in config_list:
        if 'op_partial_names' in config:
            op_names = []
            for partial_name in config['op_partial_names']:
                for name, _ in model.named_modules():
                    if partial_name in name:
                        op_names.append(name)
            if 'op_names' in config:
                config['op_names'].extend(op_names)
                config['op_names'] = list(set(config['op_names']))
            else:
                config['op_names'] = op_names
            config.pop('op_partial_names')

    config_list = dedupe_config_list(unfold_config_list(model, config_list))
    new_config_list = []

    for config in config_list:
        if 'sparsity_per_layer' in config:
            sparsity_per_layer = config.pop('sparsity_per_layer')
            op_names = config.pop('op_names')
            for op_name in op_names:
                sub_config = deepcopy(config)
                sub_config['op_names'] = [op_name]

Review comment (suggestion):
please using `config['op_types'].remove('default')` and `config['op_types'].extend(weighted_modules)` to avoid:
`config = {'op_types': ['default', 'customized_layer']}`

Code after:
    '''
    Split the config by op_names if 'sparsity' or 'sparsity_per_layer' in config,
    and set the sub_config['total_sparsity'] = config['sparsity_per_layer'].
    And every item in 'op_partial_names' will match corresponding 'op_names' in model,
    then convert 'op_partial_names' to 'op_names' in config.

    Example::
        model = models.resnet18()
        config_list = [{'op_types': ['Conv2d'], 'sparsity': 0.8, 'op_partial_names': ['conv1']}]
        pruner = L1NormPruner(model, config_list)
        pruner.compress()
        pruner.show_pruned_weights()

    In this process, the config_list will implicitly convert to the following:

    [{'op_types': ['Conv2d'], 'sparsity_per_layer': 0.8,
        'op_names': ['conv1', 'layer1.0.conv1', 'layer1.1.conv1',
        'layer2.0.conv1', 'layer2.1.conv1', 'layer3.0.conv1', 'layer3.1.conv1',
        'layer4.0.conv1', 'layer4.1.conv1']}]
    '''
    for config in config_list:
        if 'sparsity' in config:
            if 'sparsity_per_layer' in config:
                raise ValueError("'sparsity' and 'sparsity_per_layer' have the same semantics, can not set both in one config.")
            else:
                config['sparsity_per_layer'] = config.pop('sparsity')

    for config in config_list:
        if 'op_types' in config:
            if 'default' in config['op_types']:
                config['op_types'].remove('default')
                config['op_types'].extend(weighted_modules)

    for config in config_list:
        if 'op_partial_names' in config:
            op_names = []
            for partial_name in config['op_partial_names']:
                for name, _ in model.named_modules():
                    if partial_name in name:
                        op_names.append(name)
            if 'op_names' in config:
                config['op_names'].extend(op_names)
                config['op_names'] = list(set(config['op_names']))
            else:
                config['op_names'] = op_names
            config.pop('op_partial_names')

    config_list = dedupe_config_list(unfold_config_list(model, config_list))
    new_config_list = []

    for config in config_list: