From 39c0a12274ee86c39cb3700b76386e2acece7b02 Mon Sep 17 00:00:00 2001 From: J-shang Date: Mon, 16 Aug 2021 17:56:29 +0800 Subject: [PATCH 01/13] update config list key --- .../v2/pytorch/pruning/basic_pruner.py | 28 +++----- .../v2/pytorch/pruning/tools/base.py | 2 +- .../pruning/tools/sparsity_allocator.py | 4 +- .../v2/pytorch/utils/config_validation.py | 4 +- .../compression/v2/pytorch/utils/pruning.py | 72 +++++++++++++++++++ 5 files changed, 85 insertions(+), 25 deletions(-) create mode 100644 nni/algorithms/compression/v2/pytorch/utils/pruning.py diff --git a/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py b/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py index ef4c2247f7..e5d7bd3cd8 100644 --- a/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py +++ b/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py @@ -13,6 +13,7 @@ from nni.algorithms.compression.v2.pytorch.base.pruner import Pruner from nni.algorithms.compression.v2.pytorch.utils.config_validation import PrunerSchema +from nni.algorithms.compression.v2.pytorch.utils.pruning import config_list_canonical from .tools import ( DataCollector, @@ -46,24 +47,13 @@ class OneShotPruner(Pruner): def __init__(self, model: Module, config_list: List[Dict]): + config_list = config_list_canonical(model, config_list) self.data_collector: DataCollector = None self.metrics_calculator: MetricsCalculator = None self.sparsity_allocator: SparsityAllocator = None - self._convert_config_list(config_list) super().__init__(model, config_list) - def _convert_config_list(self, config_list: List[Dict]): - """ - Convert `sparsity` in config to `sparsity_per_layer`. - """ - 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') - def reset(self, model: Optional[Module], config_list: Optional[List[Dict]]): super().reset(model=model, config_list=config_list) self.reset_tools() @@ -117,7 +107,7 @@ def __init__(self, model: Module, config_list: List[Dict]): def validate_config(self, model: Module, config_list: List[Dict]): schema = PrunerSchema([{ - SchemaOptional('sparsity_per_layer'): And(float, lambda n: 0 < n < 1), + SchemaOptional('total_sparsity'): And(float, lambda n: 0 <= n < 1), SchemaOptional('op_types'): [str], SchemaOptional('op_names'): [str], SchemaOptional('exclude'): bool @@ -173,7 +163,7 @@ def __init__(self, model: Module, config_list: List[Dict], p: int, def validate_config(self, model: Module, config_list: List[Dict]): schema = PrunerSchema([{ - SchemaOptional('sparsity_per_layer'): And(float, lambda n: 0 < n < 1), + SchemaOptional('total_sparsity'): And(float, lambda n: 0 <= n < 1), SchemaOptional('op_types'): ['Conv2d', 'Linear'], SchemaOptional('op_names'): [str], SchemaOptional('exclude'): bool @@ -293,7 +283,7 @@ def __init__(self, model: Module, config_list: List[Dict], def validate_config(self, model: Module, config_list: List[Dict]): schema = PrunerSchema([{ - SchemaOptional('sparsity_per_layer'): And(float, lambda n: 0 < n < 1), + SchemaOptional('total_sparsity'): And(float, lambda n: 0 <= n < 1), SchemaOptional('op_types'): ['Conv2d', 'Linear'], SchemaOptional('op_names'): [str], SchemaOptional('exclude'): bool @@ -378,8 +368,7 @@ def trainer(model: Module, optimizer: Optimizer, criterion: Callable[[Tensor, Te def validate_config(self, model: Module, config_list: List[Dict]): schema = PrunerSchema([{ - SchemaOptional('sparsity_per_layer'): And(float, lambda n: 0 < n < 1), - SchemaOptional('total_sparsity'): And(float, lambda n: 0 < n < 1), + SchemaOptional('total_sparsity'): And(float, lambda n: 0 <= n < 1), SchemaOptional('max_sparsity_per_layer'): And(float, lambda n: 0 < n < 1), SchemaOptional('op_types'): ['BatchNorm2d'], SchemaOptional('op_names'): [str], @@ -479,7 +468,7 @@ def trainer(model: Module, optimizer: Optimizer, criterion: Callable[[Tensor, Te def validate_config(self, model: Module, config_list: List[Dict]): schema = PrunerSchema([{ - SchemaOptional('sparsity_per_layer'): And(float, lambda n: 0 < n < 1), + SchemaOptional('total_sparsity'): And(float, lambda n: 0 <= n < 1), SchemaOptional('op_types'): ['Conv2d', 'Linear'], SchemaOptional('op_names'): [str], SchemaOptional('exclude'): bool @@ -605,8 +594,7 @@ def trainer(model: Module, optimizer: Optimizer, criterion: Callable[[Tensor, Te def validate_config(self, model: Module, config_list: List[Dict]): schema = PrunerSchema([{ - SchemaOptional('sparsity_per_layer'): And(float, lambda n: 0 < n < 1), - SchemaOptional('total_sparsity'): And(float, lambda n: 0 < n < 1), + SchemaOptional('total_sparsity'): And(float, lambda n: 0 <= n < 1), SchemaOptional('max_sparsity_per_layer'): And(float, lambda n: 0 < n < 1), SchemaOptional('op_types'): ['Conv2d', 'Linear'], SchemaOptional('op_names'): [str], diff --git a/nni/algorithms/compression/v2/pytorch/pruning/tools/base.py b/nni/algorithms/compression/v2/pytorch/pruning/tools/base.py index 59a5f8190f..8a0be9a0b9 100644 --- a/nni/algorithms/compression/v2/pytorch/pruning/tools/base.py +++ b/nni/algorithms/compression/v2/pytorch/pruning/tools/base.py @@ -437,6 +437,6 @@ def _compress_mask(self, mask: Tensor) -> Tensor: mask = mask.unfold(i, step, step) ein_expression += lower_case_letters[i] ein_expression = '...{},{}'.format(ein_expression, ein_expression) - mask = torch.einsum(ein_expression, mask, torch.ones(self.block_sparse_size)) + mask = torch.einsum(ein_expression, mask, torch.ones(self.block_sparse_size).to(mask.device)) return (mask != 0).type_as(mask) diff --git a/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py b/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py index b959d3a5eb..8ae6c76328 100644 --- a/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py +++ b/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py @@ -20,7 +20,7 @@ class NormalSparsityAllocator(SparsityAllocator): def generate_sparsity(self, metrics: Dict[str, Tensor]) -> Dict[str, Dict[str, Tensor]]: masks = {} for name, wrapper in self.pruner.get_modules_wrapper().items(): - sparsity_rate = wrapper.config['sparsity_per_layer'] + sparsity_rate = wrapper.config['total_sparsity'] assert name in metrics, 'Metric of %s is not calculated.' metric = metrics[name] * self._compress_mask(wrapper.weight_mask) @@ -108,7 +108,7 @@ def generate_sparsity(self, metrics: Dict) -> Dict[str, Dict[str, Tensor]]: for _, group_metric_dict in grouped_metrics.items(): group_metric = self._group_metric_calculate(group_metric_dict) - sparsities = {name: self.pruner.get_modules_wrapper()[name].config['sparsity_per_layer'] for name in group_metric_dict.keys()} + sparsities = {name: self.pruner.get_modules_wrapper()[name].config['total_sparsity'] for name in group_metric_dict.keys()} min_sparsity = min(sparsities.values()) conv2d_groups = [self.group_depen[name] for name in group_metric_dict.keys()] diff --git a/nni/algorithms/compression/v2/pytorch/utils/config_validation.py b/nni/algorithms/compression/v2/pytorch/utils/config_validation.py index c673cd55af..4df68a6529 100644 --- a/nni/algorithms/compression/v2/pytorch/utils/config_validation.py +++ b/nni/algorithms/compression/v2/pytorch/utils/config_validation.py @@ -53,8 +53,8 @@ def validate(self, data): self.compressor_schema.validate(data) def validate_exclude_sparsity(data): - if not ('exclude' in data or 'sparsity_per_layer' in data or 'total_sparsity' in data): - raise SchemaError('One of [sparsity_per_layer, total_sparsity, exclude] should be specified.') + if not ('exclude' in data or 'total_sparsity' in data): + raise SchemaError('One of [total_sparsity, exclude] should be specified.') return True def validate_exclude_quant_types_quant_bits(data): diff --git a/nni/algorithms/compression/v2/pytorch/utils/pruning.py b/nni/algorithms/compression/v2/pytorch/utils/pruning.py new file mode 100644 index 0000000000..d28d274045 --- /dev/null +++ b/nni/algorithms/compression/v2/pytorch/utils/pruning.py @@ -0,0 +1,72 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from copy import deepcopy +from typing import Dict, List + +from torch.nn import Module + + +def config_list_canonical(model: Module, config_list: List[Dict]) -> List[Dict]: + 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') + + 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: + new_config = deepcopy(config) + new_config['op_names'] = [op_name] + new_config['total_sparsity'] = sparsity_per_layer + new_config_list.append(new_config) + else: + new_config_list.append(config) + + return new_config_list + + +def unfold_config_list(model: Module, config_list: List[Dict]) -> List[Dict]: + ''' + unfold config_list to op_names level + ''' + unfolded_config_list = [] + for config in config_list: + op_names = [] + for module_name, module in model.named_modules(): + module_type = type(module).__name__ + if 'op_types' in config and module_type not in config['op_types']: + continue + if 'op_names' in config and module_name not in config['op_names']: + continue + op_names.append(module_name) + unfolded_config = deepcopy(config) + unfolded_config['op_names'] = op_names + unfolded_config_list.append(unfolded_config) + return unfolded_config_list + + +def dedupe_config_list(config_list: List[Dict]) -> List[Dict]: + ''' + dedupe the op_names in unfolded config_list + ''' + exclude = set() + exclude_idxes = [] + config_list = deepcopy(config_list) + for idx, config in reversed(list(enumerate(config_list))): + if 'exclude' in config: + exclude.update(config['op_names']) + exclude_idxes.append(idx) + continue + config['op_names'] = sorted(list(set(config['op_names']).difference(exclude))) + exclude.update(config['op_names']) + for idx in sorted(exclude_idxes, reverse=True): + config_list.pop(idx) + return config_list From 038f2ad551f3e5716aecb5cfa12b0006aafcbd90 Mon Sep 17 00:00:00 2001 From: J-shang Date: Mon, 16 Aug 2021 18:04:57 +0800 Subject: [PATCH 02/13] update docstr --- .../compression/v2/pytorch/utils/pruning.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/nni/algorithms/compression/v2/pytorch/utils/pruning.py b/nni/algorithms/compression/v2/pytorch/utils/pruning.py index d28d274045..0b81cc327f 100644 --- a/nni/algorithms/compression/v2/pytorch/utils/pruning.py +++ b/nni/algorithms/compression/v2/pytorch/utils/pruning.py @@ -8,6 +8,10 @@ def config_list_canonical(model: Module, config_list: List[Dict]) -> List[Dict]: + ''' + 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']. + ''' for config in config_list: if 'sparsity' in config: if 'sparsity_per_layer' in config: @@ -23,10 +27,10 @@ def config_list_canonical(model: Module, config_list: List[Dict]) -> List[Dict]: sparsity_per_layer = config.pop('sparsity_per_layer') op_names = config.pop('op_names') for op_name in op_names: - new_config = deepcopy(config) - new_config['op_names'] = [op_name] - new_config['total_sparsity'] = sparsity_per_layer - new_config_list.append(new_config) + sub_config = deepcopy(config) + sub_config['op_names'] = [op_name] + sub_config['total_sparsity'] = sparsity_per_layer + new_config_list.append(sub_config) else: new_config_list.append(config) @@ -35,7 +39,7 @@ def config_list_canonical(model: Module, config_list: List[Dict]) -> List[Dict]: def unfold_config_list(model: Module, config_list: List[Dict]) -> List[Dict]: ''' - unfold config_list to op_names level + Unfold config_list to op_names level. ''' unfolded_config_list = [] for config in config_list: @@ -55,7 +59,7 @@ def unfold_config_list(model: Module, config_list: List[Dict]) -> List[Dict]: def dedupe_config_list(config_list: List[Dict]) -> List[Dict]: ''' - dedupe the op_names in unfolded config_list + Dedupe the op_names in unfolded config_list. ''' exclude = set() exclude_idxes = [] From cdc2036cbd784d909253d3855d6961a01635849a Mon Sep 17 00:00:00 2001 From: J-shang Date: Tue, 17 Aug 2021 10:08:48 +0800 Subject: [PATCH 03/13] fix total_prune_num == 0 --- .../v2/pytorch/pruning/tools/sparsity_allocator.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py b/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py index 8ae6c76328..8b9ec9c890 100644 --- a/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py +++ b/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py @@ -77,8 +77,10 @@ def _calculate_threshold(self, group_metric_dict: Dict[str, Tensor]) -> Tuple[fl assert total_sparsity <= max_sparsity_per_layer, 'total_sparsity should less than max_sparsity_per_layer.' total_prune_num = int(total_sparsity * total_weight_num) - - threshold = torch.topk(torch.cat(metric_list).view(-1), total_prune_num, largest=False)[0].max().item() + if total_prune_num == 0: + threshold = torch.cat(metric_list).min().item() - 1 + else: + threshold = torch.topk(torch.cat(metric_list).view(-1), total_prune_num, largest=False)[0].max().item() return threshold, sub_thresholds From c28f6cfc60ffc141f998b9989fa5ff3d212e2d95 Mon Sep 17 00:00:00 2001 From: J-shang Date: Tue, 17 Aug 2021 11:04:43 +0800 Subject: [PATCH 04/13] add convert max_sparsity_per_layer to min_retention_numel --- .../v2/pytorch/pruning/tools/sparsity_allocator.py | 14 ++++++++------ .../compression/v2/pytorch/utils/pruning.py | 12 ++++++++++++ 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py b/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py index 8b9ec9c890..397a9d4007 100644 --- a/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py +++ b/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import math from typing import Any, Dict, List, Tuple, Union import numpy as np @@ -58,24 +59,25 @@ def _calculate_threshold(self, group_metric_dict: Dict[str, Tensor]) -> Tuple[fl temp_wrapper_config = self.pruner.get_modules_wrapper()[list(group_metric_dict.keys())[0]].config total_sparsity = temp_wrapper_config['total_sparsity'] - max_sparsity_per_layer = temp_wrapper_config.get('max_sparsity_per_layer', 1.0) + min_retention_numel = temp_wrapper_config.get('min_retention_numel', {}) for name, metric in group_metric_dict.items(): wrapper = self.pruner.get_modules_wrapper()[name] metric = metric * self._compress_mask(wrapper.weight_mask) - print(metric) layer_weight_num = wrapper.module.weight.data.numel() - stay_num = int(metric.numel() * max_sparsity_per_layer) + + retention_numel = 0 if name not in min_retention_numel else min_retention_numel[name] + removed_metric_num = math.floor(retention_numel / (wrapper.weight_mask.numel() / metric.numel())) + stay_metric_num = metric.numel() - removed_metric_num # Remove the weight parts that must be left - stay_metric = torch.topk(metric.view(-1), stay_num, largest=False)[0] + stay_metric = torch.topk(metric.view(-1), stay_metric_num, largest=False)[0] sub_thresholds[name] = stay_metric.max() expend_times = int(layer_weight_num / metric.numel()) if expend_times > 1: - stay_metric = stay_metric.expand(stay_num, int(layer_weight_num / metric.numel())).view(-1) + stay_metric = stay_metric.expand(stay_metric_num, int(layer_weight_num / metric.numel())).view(-1) metric_list.append(stay_metric) total_weight_num += layer_weight_num - assert total_sparsity <= max_sparsity_per_layer, 'total_sparsity should less than max_sparsity_per_layer.' total_prune_num = int(total_sparsity * total_weight_num) if total_prune_num == 0: threshold = torch.cat(metric_list).min().item() - 1 diff --git a/nni/algorithms/compression/v2/pytorch/utils/pruning.py b/nni/algorithms/compression/v2/pytorch/utils/pruning.py index 0b81cc327f..950d2e84a3 100644 --- a/nni/algorithms/compression/v2/pytorch/utils/pruning.py +++ b/nni/algorithms/compression/v2/pytorch/utils/pruning.py @@ -2,10 +2,13 @@ # Licensed under the MIT license. from copy import deepcopy +import math from typing import Dict, List from torch.nn import Module +from nni.compression.pytorch.utils import get_module_by_name + def config_list_canonical(model: Module, config_list: List[Dict]) -> List[Dict]: ''' @@ -31,6 +34,15 @@ def config_list_canonical(model: Module, config_list: List[Dict]) -> List[Dict]: sub_config['op_names'] = [op_name] sub_config['total_sparsity'] = sparsity_per_layer new_config_list.append(sub_config) + elif 'max_sparsity_per_layer' in config: + min_retention_per_layer = (1 - config.pop('max_sparsity_per_layer')) + op_names = config.get('op_names', []) + min_retention_numel = {} + for op_name in op_names: + total_element_num = get_module_by_name(model, op_name)[0].weight.numel() + min_retention_numel[op_name] = math.floor(total_element_num * min_retention_per_layer) + config['min_retention_numel'] = min_retention_numel + new_config_list.append(config) else: new_config_list.append(config) From dc272a6e42ca9771cbbc204119f969286c82b1fc Mon Sep 17 00:00:00 2001 From: J-shang Date: Tue, 17 Aug 2021 17:38:19 +0800 Subject: [PATCH 05/13] fix typing --- nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py b/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py index e5d7bd3cd8..f17537b9f6 100644 --- a/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py +++ b/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py @@ -603,7 +603,7 @@ def validate_config(self, model: Module, config_list: List[Dict]): schema.validate(config_list) - def _collector(self, buffer: List, weight_tensor: Tensor) -> Callable[[Module, Tensor, Tensor], None]: + def _collector(self, buffer: List, weight_tensor: Tensor) -> Callable[[Tensor], None]: def collect_taylor(grad: Tensor): if len(buffer) < self.training_batches: buffer.append(self._calculate_taylor_expansion(weight_tensor, grad)) From 816b5ecad1af8a1dc0477eaa422fdd251300240a Mon Sep 17 00:00:00 2001 From: J-shang Date: Wed, 18 Aug 2021 10:56:01 +0800 Subject: [PATCH 06/13] update validation --- .../v2/pytorch/pruning/basic_pruner.py | 115 ++++++++++-------- .../v2/pytorch/utils/config_validation.py | 82 ++++++------- 2 files changed, 101 insertions(+), 96 deletions(-) diff --git a/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py b/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py index f17537b9f6..93a110b88f 100644 --- a/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py +++ b/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py @@ -1,10 +1,11 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from copy import deepcopy import logging from typing import List, Dict, Tuple, Callable, Optional -from schema import And, Optional as SchemaOptional +from schema import And, Or, Optional as SchemaOptional import torch from torch import Tensor import torch.nn as nn @@ -12,7 +13,7 @@ from torch.optim import Optimizer from nni.algorithms.compression.v2.pytorch.base.pruner import Pruner -from nni.algorithms.compression.v2.pytorch.utils.config_validation import PrunerSchema +from nni.algorithms.compression.v2.pytorch.utils.config_validation import CompressorSchema from nni.algorithms.compression.v2.pytorch.utils.pruning import config_list_canonical from .tools import ( @@ -44,16 +45,41 @@ __all__ = ['LevelPruner', 'L1NormPruner', 'L2NormPruner', 'FPGMPruner', 'SlimPruner', 'ActivationPruner', 'ActivationAPoZRankPruner', 'ActivationMeanRankPruner', 'TaylorFOWeightPruner'] +NORMAL_SCHEMA = { + Or('sparsity', 'sparsity_per_layer'): And(float, lambda n: 0 <= n < 1), + SchemaOptional('op_types'): [str], + SchemaOptional('op_names'): [str] +} + +GLOBAL_SCHEMA = { + 'total_sparsity': And(float, lambda n: 0 <= n < 1), + 'max_sparsity_per_layer': And(float, lambda n: 0 < n <= 1), + SchemaOptional('op_types'): [str], + SchemaOptional('op_names'): [str] +} + +EXCLUDE_SCHEMA = { + 'exclude': bool, + SchemaOptional('op_types'): [str], + SchemaOptional('op_names'): [str] +} + class OneShotPruner(Pruner): def __init__(self, model: Module, config_list: List[Dict]): - config_list = config_list_canonical(model, config_list) self.data_collector: DataCollector = None self.metrics_calculator: MetricsCalculator = None self.sparsity_allocator: SparsityAllocator = None super().__init__(model, config_list) + def validate_config(self, model: Module, config_list: List[Dict]): + self._validate_config_before_canonical(model, config_list) + self.config_list = config_list_canonical(model, config_list) + + def _validate_config_before_canonical(self, model: Module, config_list: List[Dict]): + pass + def reset(self, model: Optional[Module], config_list: Optional[List[Dict]]): super().reset(model=model, config_list=config_list) self.reset_tools() @@ -105,14 +131,8 @@ def __init__(self, model: Module, config_list: List[Dict]): self.mode = 'normal' super().__init__(model, config_list) - def validate_config(self, model: Module, config_list: List[Dict]): - schema = PrunerSchema([{ - SchemaOptional('total_sparsity'): And(float, lambda n: 0 <= n < 1), - SchemaOptional('op_types'): [str], - SchemaOptional('op_names'): [str], - SchemaOptional('exclude'): bool - }], model, _logger) - + def _validate_config_before_canonical(self, model: Module, config_list: List[Dict]): + schema = CompressorSchema([deepcopy(NORMAL_SCHEMA), deepcopy(EXCLUDE_SCHEMA)], model, _logger) schema.validate(config_list) def reset_tools(self): @@ -161,13 +181,12 @@ def __init__(self, model: Module, config_list: List[Dict], p: int, self.dummy_input = dummy_input super().__init__(model, config_list) - def validate_config(self, model: Module, config_list: List[Dict]): - schema = PrunerSchema([{ - SchemaOptional('total_sparsity'): And(float, lambda n: 0 <= n < 1), - SchemaOptional('op_types'): ['Conv2d', 'Linear'], - SchemaOptional('op_names'): [str], - SchemaOptional('exclude'): bool - }], model, _logger) + def _validate_config_before_canonical(self, model: Module, config_list: List[Dict]): + normal_schema = deepcopy(NORMAL_SCHEMA) + normal_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] + exclude_schema = deepcopy(EXCLUDE_SCHEMA) + exclude_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] + schema = CompressorSchema([normal_schema, exclude_schema], model, _logger) schema.validate(config_list) @@ -281,13 +300,12 @@ def __init__(self, model: Module, config_list: List[Dict], self.dummy_input = dummy_input super().__init__(model, config_list) - def validate_config(self, model: Module, config_list: List[Dict]): - schema = PrunerSchema([{ - SchemaOptional('total_sparsity'): And(float, lambda n: 0 <= n < 1), - SchemaOptional('op_types'): ['Conv2d', 'Linear'], - SchemaOptional('op_names'): [str], - SchemaOptional('exclude'): bool - }], model, _logger) + def _validate_config_before_canonical(self, model: Module, config_list: List[Dict]): + normal_schema = deepcopy(NORMAL_SCHEMA) + normal_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] + exclude_schema = deepcopy(EXCLUDE_SCHEMA) + exclude_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] + schema = CompressorSchema([normal_schema, exclude_schema], model, _logger) schema.validate(config_list) @@ -366,14 +384,15 @@ def trainer(model: Module, optimizer: Optimizer, criterion: Callable[[Tensor, Te self._scale = scale super().__init__(model, config_list) - def validate_config(self, model: Module, config_list: List[Dict]): - schema = PrunerSchema([{ - SchemaOptional('total_sparsity'): And(float, lambda n: 0 <= n < 1), - SchemaOptional('max_sparsity_per_layer'): And(float, lambda n: 0 < n < 1), - SchemaOptional('op_types'): ['BatchNorm2d'], - SchemaOptional('op_names'): [str], - SchemaOptional('exclude'): bool - }], model, _logger) + def _validate_config_before_canonical(self, model: Module, config_list: List[Dict]): + if self.mode == 'global': + normal_schema = deepcopy(GLOBAL_SCHEMA) + else: + normal_schema = deepcopy(NORMAL_SCHEMA) + normal_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] + exclude_schema = deepcopy(EXCLUDE_SCHEMA) + exclude_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] + schema = CompressorSchema([normal_schema, exclude_schema], model, _logger) schema.validate(config_list) @@ -466,13 +485,12 @@ def trainer(model: Module, optimizer: Optimizer, criterion: Callable[[Tensor, Te self._activation = self._choose_activation(activation) super().__init__(model, config_list) - def validate_config(self, model: Module, config_list: List[Dict]): - schema = PrunerSchema([{ - SchemaOptional('total_sparsity'): And(float, lambda n: 0 <= n < 1), - SchemaOptional('op_types'): ['Conv2d', 'Linear'], - SchemaOptional('op_names'): [str], - SchemaOptional('exclude'): bool - }], model, _logger) + def _validate_config_before_canonical(self, model: Module, config_list: List[Dict]): + normal_schema = deepcopy(NORMAL_SCHEMA) + normal_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] + exclude_schema = deepcopy(EXCLUDE_SCHEMA) + exclude_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] + schema = CompressorSchema([normal_schema, exclude_schema], model, _logger) schema.validate(config_list) @@ -592,14 +610,15 @@ def trainer(model: Module, optimizer: Optimizer, criterion: Callable[[Tensor, Te self.training_batches = training_batches super().__init__(model, config_list) - def validate_config(self, model: Module, config_list: List[Dict]): - schema = PrunerSchema([{ - SchemaOptional('total_sparsity'): And(float, lambda n: 0 <= n < 1), - SchemaOptional('max_sparsity_per_layer'): And(float, lambda n: 0 < n < 1), - SchemaOptional('op_types'): ['Conv2d', 'Linear'], - SchemaOptional('op_names'): [str], - SchemaOptional('exclude'): bool - }], model, _logger) + def _validate_config_before_canonical(self, model: Module, config_list: List[Dict]): + if self.mode == 'global': + normal_schema = deepcopy(GLOBAL_SCHEMA) + else: + normal_schema = deepcopy(NORMAL_SCHEMA) + normal_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] + exclude_schema = deepcopy(EXCLUDE_SCHEMA) + exclude_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] + schema = CompressorSchema([normal_schema, exclude_schema], model, _logger) schema.validate(config_list) diff --git a/nni/algorithms/compression/v2/pytorch/utils/config_validation.py b/nni/algorithms/compression/v2/pytorch/utils/config_validation.py index 4df68a6529..3b9200cc4b 100644 --- a/nni/algorithms/compression/v2/pytorch/utils/config_validation.py +++ b/nni/algorithms/compression/v2/pytorch/utils/config_validation.py @@ -1,75 +1,61 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from logging import Logger +from typing import Dict, List from schema import Schema, And, SchemaError -def validate_op_names(model, op_names, logger): - found_names = set(map(lambda x: x[0], model.named_modules())) - - not_found_op_names = list(set(op_names) - found_names) - if not_found_op_names: - logger.warning('op_names %s not found in model', not_found_op_names) - - return True - -def validate_op_types(model, op_types, logger): - found_types = set(['default']) | set(map(lambda x: type(x[1]).__name__, model.named_modules())) +from torch.nn import Module - not_found_op_types = list(set(op_types) - found_types) - if not_found_op_types: - logger.warning('op_types %s not found in model', not_found_op_types) - - return True - -def validate_op_types_op_names(data): - if not ('op_types' in data or 'op_names' in data): - raise SchemaError('Either op_types or op_names must be specified.') - return True class CompressorSchema: - def __init__(self, data_schema, model, logger): + def __init__(self, data_schema: List[Dict], model: Module, logger: Logger): assert isinstance(data_schema, list) and len(data_schema) <= 1 self.data_schema = data_schema self.compressor_schema = Schema(self._modify_schema(data_schema, model, logger)) - def _modify_schema(self, data_schema, model, logger): + def _modify_schema(self, data_schema: List[Dict], model: Module, logger: Logger) -> List[Dict]: if not data_schema: return data_schema - for k in data_schema[0]: - old_schema = data_schema[0][k] - if k == 'op_types' or (isinstance(k, Schema) and k._schema == 'op_types'): - new_schema = And(old_schema, lambda n: validate_op_types(model, n, logger)) - data_schema[0][k] = new_schema - if k == 'op_names' or (isinstance(k, Schema) and k._schema == 'op_names'): - new_schema = And(old_schema, lambda n: validate_op_names(model, n, logger)) - data_schema[0][k] = new_schema + for sub_schema in data_schema: + for k, old_schema in sub_schema.items(): + if k == 'op_types' or (isinstance(k, Schema) and k._schema == 'op_types'): + new_schema = And(old_schema, lambda n: validate_op_types(model, n, logger)) + sub_schema[k] = new_schema + if k == 'op_names' or (isinstance(k, Schema) and k._schema == 'op_names'): + new_schema = And(old_schema, lambda n: validate_op_names(model, n, logger)) + sub_schema[k] = new_schema - data_schema[0] = And(data_schema[0], lambda d: validate_op_types_op_names(d)) + sub_schema = And(data_schema[0], lambda d: validate_op_types_op_names(d)) return data_schema def validate(self, data): self.compressor_schema.validate(data) -def validate_exclude_sparsity(data): - if not ('exclude' in data or 'total_sparsity' in data): - raise SchemaError('One of [total_sparsity, exclude] should be specified.') + +def validate_op_names(model, op_names, logger): + found_names = set(map(lambda x: x[0], model.named_modules())) + + not_found_op_names = list(set(op_names) - found_names) + if not_found_op_names: + logger.warning('op_names %s not found in model', not_found_op_names) + return True -def validate_exclude_quant_types_quant_bits(data): - if not ('exclude' in data or ('quant_types' in data and 'quant_bits' in data)): - raise SchemaError('Either (quant_types and quant_bits) or exclude must be specified.') + +def validate_op_types(model, op_types, logger): + found_types = set(['default']) | set(map(lambda x: type(x[1]).__name__, model.named_modules())) + + not_found_op_types = list(set(op_types) - found_types) + if not_found_op_types: + logger.warning('op_types %s not found in model', not_found_op_types) + return True -class PrunerSchema(CompressorSchema): - def _modify_schema(self, data_schema, model, logger): - data_schema = super()._modify_schema(data_schema, model, logger) - data_schema[0] = And(data_schema[0], lambda d: validate_exclude_sparsity(d)) - return data_schema -class QuantizerSchema(CompressorSchema): - def _modify_schema(self, data_schema, model, logger): - data_schema = super()._modify_schema(data_schema, model, logger) - data_schema[0] = And(data_schema[0], lambda d: validate_exclude_quant_types_quant_bits(d)) - return data_schema +def validate_op_types_op_names(data): + if not ('op_types' in data or 'op_names' in data): + raise SchemaError('Either op_types or op_names must be specified.') + return True From 7a26ebe258d18de7c34c09d9d56c8fec2c486c54 Mon Sep 17 00:00:00 2001 From: J-shang Date: Wed, 18 Aug 2021 13:21:50 +0800 Subject: [PATCH 07/13] fix warning --- .../compression/v2/pytorch/pruning/tools/metrics_calculator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nni/algorithms/compression/v2/pytorch/pruning/tools/metrics_calculator.py b/nni/algorithms/compression/v2/pytorch/pruning/tools/metrics_calculator.py index afe12de3b9..311fef5b71 100644 --- a/nni/algorithms/compression/v2/pytorch/pruning/tools/metrics_calculator.py +++ b/nni/algorithms/compression/v2/pytorch/pruning/tools/metrics_calculator.py @@ -120,7 +120,7 @@ def calculate_metrics(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]: metric = torch.ones(*reorder_tensor.size()[:len(keeped_dim)], device=reorder_tensor.device) across_dim = list(range(len(keeped_dim), len(reorder_dim))) - idxs = metric.nonzero() + idxs = metric.nonzero(as_tuple=False) for idx in idxs: other = reorder_tensor for i in idx: From e0761826a3a381547fc7795eec1d3027089ff2a2 Mon Sep 17 00:00:00 2001 From: J-shang Date: Wed, 18 Aug 2021 13:57:42 +0800 Subject: [PATCH 08/13] fix type --- .../compression/v2/pytorch/pruning/tools/metrics_calculator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nni/algorithms/compression/v2/pytorch/pruning/tools/metrics_calculator.py b/nni/algorithms/compression/v2/pytorch/pruning/tools/metrics_calculator.py index 311fef5b71..309a73eba3 100644 --- a/nni/algorithms/compression/v2/pytorch/pruning/tools/metrics_calculator.py +++ b/nni/algorithms/compression/v2/pytorch/pruning/tools/metrics_calculator.py @@ -161,7 +161,7 @@ def calculate_metrics(self, data: Dict[str, List[Tensor]]) -> Dict[str, Tensor]: for dim, dim_size in enumerate(_eq_zero.size()): if dim not in keeped_dim: total_size *= dim_size - _apoz = torch.sum(_eq_zero, dim=across_dim, dtype=torch.float64) / total_size + _apoz = torch.sum(_eq_zero, dim=across_dim).type_as(activations) / total_size # NOTE: the metric is (1 - apoz) because we assume the smaller metric value is more needed to be pruned. metrics[name] = torch.ones_like(_apoz) - _apoz return metrics From 8ce0d88c2a16a03a9da597c41e33f644b46e2b0d Mon Sep 17 00:00:00 2001 From: J-shang Date: Wed, 18 Aug 2021 15:58:03 +0800 Subject: [PATCH 09/13] add internal_config --- .../v2/pytorch/pruning/basic_pruner.py | 63 ++++++++++--------- .../pruning/tools/sparsity_allocator.py | 6 +- .../compression/v2/pytorch/utils/pruning.py | 22 ++++--- 3 files changed, 49 insertions(+), 42 deletions(-) diff --git a/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py b/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py index 93a110b88f..9ab6b7f617 100644 --- a/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py +++ b/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py @@ -53,7 +53,7 @@ GLOBAL_SCHEMA = { 'total_sparsity': And(float, lambda n: 0 <= n < 1), - 'max_sparsity_per_layer': And(float, lambda n: 0 < n <= 1), + SchemaOptional('max_sparsity_per_layer'): And(float, lambda n: 0 < n <= 1), SchemaOptional('op_types'): [str], SchemaOptional('op_names'): [str] } @@ -64,6 +64,13 @@ SchemaOptional('op_names'): [str] } +INTERNAL_SCHEMA = { + '_sparsity': And(float, lambda n: 0 <= n < 1), + SchemaOptional('_min_retention_numel'): {str: int}, + SchemaOptional('op_types'): [str], + SchemaOptional('op_names'): [str] +} + class OneShotPruner(Pruner): def __init__(self, model: Module, config_list: List[Dict]): @@ -132,7 +139,8 @@ def __init__(self, model: Module, config_list: List[Dict]): super().__init__(model, config_list) def _validate_config_before_canonical(self, model: Module, config_list: List[Dict]): - schema = CompressorSchema([deepcopy(NORMAL_SCHEMA), deepcopy(EXCLUDE_SCHEMA)], model, _logger) + schema_list = [deepcopy(NORMAL_SCHEMA), deepcopy(EXCLUDE_SCHEMA), deepcopy(INTERNAL_SCHEMA)] + schema = CompressorSchema(schema_list, model, _logger) schema.validate(config_list) def reset_tools(self): @@ -182,11 +190,10 @@ def __init__(self, model: Module, config_list: List[Dict], p: int, super().__init__(model, config_list) def _validate_config_before_canonical(self, model: Module, config_list: List[Dict]): - normal_schema = deepcopy(NORMAL_SCHEMA) - normal_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] - exclude_schema = deepcopy(EXCLUDE_SCHEMA) - exclude_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] - schema = CompressorSchema([normal_schema, exclude_schema], model, _logger) + schema_list = [deepcopy(NORMAL_SCHEMA), deepcopy(EXCLUDE_SCHEMA), deepcopy(INTERNAL_SCHEMA)] + for sub_shcema in schema_list: + sub_shcema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] + schema = CompressorSchema(schema_list, model, _logger) schema.validate(config_list) @@ -301,11 +308,10 @@ def __init__(self, model: Module, config_list: List[Dict], super().__init__(model, config_list) def _validate_config_before_canonical(self, model: Module, config_list: List[Dict]): - normal_schema = deepcopy(NORMAL_SCHEMA) - normal_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] - exclude_schema = deepcopy(EXCLUDE_SCHEMA) - exclude_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] - schema = CompressorSchema([normal_schema, exclude_schema], model, _logger) + schema_list = [deepcopy(NORMAL_SCHEMA), deepcopy(EXCLUDE_SCHEMA), deepcopy(INTERNAL_SCHEMA)] + for sub_shcema in schema_list: + sub_shcema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] + schema = CompressorSchema(schema_list, model, _logger) schema.validate(config_list) @@ -385,14 +391,14 @@ def trainer(model: Module, optimizer: Optimizer, criterion: Callable[[Tensor, Te super().__init__(model, config_list) def _validate_config_before_canonical(self, model: Module, config_list: List[Dict]): + schema_list = [deepcopy(EXCLUDE_SCHEMA), deepcopy(INTERNAL_SCHEMA)] if self.mode == 'global': - normal_schema = deepcopy(GLOBAL_SCHEMA) + schema_list.append(deepcopy(GLOBAL_SCHEMA)) else: - normal_schema = deepcopy(NORMAL_SCHEMA) - normal_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] - exclude_schema = deepcopy(EXCLUDE_SCHEMA) - exclude_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] - schema = CompressorSchema([normal_schema, exclude_schema], model, _logger) + schema_list.append(deepcopy(NORMAL_SCHEMA)) + for sub_shcema in schema_list: + sub_shcema[SchemaOptional('op_types')] = ['BatchNorm2d'] + schema = CompressorSchema(schema_list, model, _logger) schema.validate(config_list) @@ -486,11 +492,10 @@ def trainer(model: Module, optimizer: Optimizer, criterion: Callable[[Tensor, Te super().__init__(model, config_list) def _validate_config_before_canonical(self, model: Module, config_list: List[Dict]): - normal_schema = deepcopy(NORMAL_SCHEMA) - normal_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] - exclude_schema = deepcopy(EXCLUDE_SCHEMA) - exclude_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] - schema = CompressorSchema([normal_schema, exclude_schema], model, _logger) + schema_list = [deepcopy(NORMAL_SCHEMA), deepcopy(EXCLUDE_SCHEMA), deepcopy(INTERNAL_SCHEMA)] + for sub_shcema in schema_list: + sub_shcema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] + schema = CompressorSchema(schema_list, model, _logger) schema.validate(config_list) @@ -611,14 +616,14 @@ def trainer(model: Module, optimizer: Optimizer, criterion: Callable[[Tensor, Te super().__init__(model, config_list) def _validate_config_before_canonical(self, model: Module, config_list: List[Dict]): + schema_list = [deepcopy(EXCLUDE_SCHEMA), deepcopy(INTERNAL_SCHEMA)] if self.mode == 'global': - normal_schema = deepcopy(GLOBAL_SCHEMA) + schema_list.append(deepcopy(GLOBAL_SCHEMA)) else: - normal_schema = deepcopy(NORMAL_SCHEMA) - normal_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] - exclude_schema = deepcopy(EXCLUDE_SCHEMA) - exclude_schema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] - schema = CompressorSchema([normal_schema, exclude_schema], model, _logger) + schema_list.append(deepcopy(NORMAL_SCHEMA)) + for sub_shcema in schema_list: + sub_shcema[SchemaOptional('op_types')] = ['Conv2d', 'Linear'] + schema = CompressorSchema(schema_list, model, _logger) schema.validate(config_list) diff --git a/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py b/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py index 397a9d4007..7145888594 100644 --- a/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py +++ b/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py @@ -21,7 +21,7 @@ class NormalSparsityAllocator(SparsityAllocator): def generate_sparsity(self, metrics: Dict[str, Tensor]) -> Dict[str, Dict[str, Tensor]]: masks = {} for name, wrapper in self.pruner.get_modules_wrapper().items(): - sparsity_rate = wrapper.config['total_sparsity'] + sparsity_rate = wrapper.config['_sparsity'] assert name in metrics, 'Metric of %s is not calculated.' metric = metrics[name] * self._compress_mask(wrapper.weight_mask) @@ -58,7 +58,7 @@ def _calculate_threshold(self, group_metric_dict: Dict[str, Tensor]) -> Tuple[fl total_weight_num = 0 temp_wrapper_config = self.pruner.get_modules_wrapper()[list(group_metric_dict.keys())[0]].config - total_sparsity = temp_wrapper_config['total_sparsity'] + total_sparsity = temp_wrapper_config['_sparsity'] min_retention_numel = temp_wrapper_config.get('min_retention_numel', {}) for name, metric in group_metric_dict.items(): @@ -112,7 +112,7 @@ def generate_sparsity(self, metrics: Dict) -> Dict[str, Dict[str, Tensor]]: for _, group_metric_dict in grouped_metrics.items(): group_metric = self._group_metric_calculate(group_metric_dict) - sparsities = {name: self.pruner.get_modules_wrapper()[name].config['total_sparsity'] for name in group_metric_dict.keys()} + sparsities = {name: self.pruner.get_modules_wrapper()[name].config['_sparsity'] for name in group_metric_dict.keys()} min_sparsity = min(sparsities.values()) conv2d_groups = [self.group_depen[name] for name in group_metric_dict.keys()] diff --git a/nni/algorithms/compression/v2/pytorch/utils/pruning.py b/nni/algorithms/compression/v2/pytorch/utils/pruning.py index 950d2e84a3..382d47d217 100644 --- a/nni/algorithms/compression/v2/pytorch/utils/pruning.py +++ b/nni/algorithms/compression/v2/pytorch/utils/pruning.py @@ -13,7 +13,7 @@ def config_list_canonical(model: Module, config_list: List[Dict]) -> List[Dict]: ''' 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 set the sub_config['_sparsity'] = config['sparsity_per_layer']. ''' for config in config_list: if 'sparsity' in config: @@ -32,16 +32,18 @@ def config_list_canonical(model: Module, config_list: List[Dict]) -> List[Dict]: for op_name in op_names: sub_config = deepcopy(config) sub_config['op_names'] = [op_name] - sub_config['total_sparsity'] = sparsity_per_layer + sub_config['_sparsity'] = sparsity_per_layer new_config_list.append(sub_config) - elif 'max_sparsity_per_layer' in config: - min_retention_per_layer = (1 - config.pop('max_sparsity_per_layer')) - op_names = config.get('op_names', []) - min_retention_numel = {} - for op_name in op_names: - total_element_num = get_module_by_name(model, op_name)[0].weight.numel() - min_retention_numel[op_name] = math.floor(total_element_num * min_retention_per_layer) - config['min_retention_numel'] = min_retention_numel + elif 'total_sparsity' in config: + if 'max_sparsity_per_layer' in config: + min_retention_per_layer = (1 - config.pop('max_sparsity_per_layer')) + op_names = config.get('op_names', []) + min_retention_numel = {} + for op_name in op_names: + total_element_num = get_module_by_name(model, op_name)[0].weight.numel() + min_retention_numel[op_name] = math.floor(total_element_num * min_retention_per_layer) + config['_min_retention_numel'] = min_retention_numel + config['_sparsity'] = config.pop('total_sparsity') new_config_list.append(config) else: new_config_list.append(config) From ad9fa8c6f1a617f1a42bd907f381cc0c7b73c4b2 Mon Sep 17 00:00:00 2001 From: J-shang Date: Wed, 18 Aug 2021 16:31:43 +0800 Subject: [PATCH 10/13] fix bug --- .../compression/v2/pytorch/utils/config_validation.py | 2 +- nni/algorithms/compression/v2/pytorch/utils/pruning.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/nni/algorithms/compression/v2/pytorch/utils/config_validation.py b/nni/algorithms/compression/v2/pytorch/utils/config_validation.py index 3b9200cc4b..100a81dfe9 100644 --- a/nni/algorithms/compression/v2/pytorch/utils/config_validation.py +++ b/nni/algorithms/compression/v2/pytorch/utils/config_validation.py @@ -10,7 +10,7 @@ class CompressorSchema: def __init__(self, data_schema: List[Dict], model: Module, logger: Logger): - assert isinstance(data_schema, list) and len(data_schema) <= 1 + assert isinstance(data_schema, list) self.data_schema = data_schema self.compressor_schema = Schema(self._modify_schema(data_schema, model, logger)) diff --git a/nni/algorithms/compression/v2/pytorch/utils/pruning.py b/nni/algorithms/compression/v2/pytorch/utils/pruning.py index 382d47d217..9f04e8d608 100644 --- a/nni/algorithms/compression/v2/pytorch/utils/pruning.py +++ b/nni/algorithms/compression/v2/pytorch/utils/pruning.py @@ -40,7 +40,7 @@ def config_list_canonical(model: Module, config_list: List[Dict]) -> List[Dict]: op_names = config.get('op_names', []) min_retention_numel = {} for op_name in op_names: - total_element_num = get_module_by_name(model, op_name)[0].weight.numel() + total_element_num = get_module_by_name(model, op_name)[1].weight.numel() min_retention_numel[op_name] = math.floor(total_element_num * min_retention_per_layer) config['_min_retention_numel'] = min_retention_numel config['_sparsity'] = config.pop('total_sparsity') From 2d3ab7ee1b690d28b961b273671dd673277ba5d0 Mon Sep 17 00:00:00 2001 From: J-shang Date: Thu, 19 Aug 2021 21:05:38 +0800 Subject: [PATCH 11/13] change _sparsity to toral_sparsity --- .../v2/pytorch/pruning/basic_pruner.py | 4 ++-- .../pruning/tools/sparsity_allocator.py | 13 +++++----- .../compression/v2/pytorch/utils/pruning.py | 24 +++++++------------ 3 files changed, 18 insertions(+), 23 deletions(-) diff --git a/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py b/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py index 9ab6b7f617..e373816a2b 100644 --- a/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py +++ b/nni/algorithms/compression/v2/pytorch/pruning/basic_pruner.py @@ -65,8 +65,8 @@ } INTERNAL_SCHEMA = { - '_sparsity': And(float, lambda n: 0 <= n < 1), - SchemaOptional('_min_retention_numel'): {str: int}, + 'total_sparsity': And(float, lambda n: 0 <= n < 1), + SchemaOptional('max_sparsity_per_layer'): {str: float}, SchemaOptional('op_types'): [str], SchemaOptional('op_names'): [str] } diff --git a/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py b/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py index 7145888594..fe9ca02037 100644 --- a/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py +++ b/nni/algorithms/compression/v2/pytorch/pruning/tools/sparsity_allocator.py @@ -21,7 +21,7 @@ class NormalSparsityAllocator(SparsityAllocator): def generate_sparsity(self, metrics: Dict[str, Tensor]) -> Dict[str, Dict[str, Tensor]]: masks = {} for name, wrapper in self.pruner.get_modules_wrapper().items(): - sparsity_rate = wrapper.config['_sparsity'] + sparsity_rate = wrapper.config['total_sparsity'] assert name in metrics, 'Metric of %s is not calculated.' metric = metrics[name] * self._compress_mask(wrapper.weight_mask) @@ -58,16 +58,17 @@ def _calculate_threshold(self, group_metric_dict: Dict[str, Tensor]) -> Tuple[fl total_weight_num = 0 temp_wrapper_config = self.pruner.get_modules_wrapper()[list(group_metric_dict.keys())[0]].config - total_sparsity = temp_wrapper_config['_sparsity'] - min_retention_numel = temp_wrapper_config.get('min_retention_numel', {}) + total_sparsity = temp_wrapper_config['total_sparsity'] + max_sparsity_per_layer = temp_wrapper_config.get('max_sparsity_per_layer', {}) for name, metric in group_metric_dict.items(): wrapper = self.pruner.get_modules_wrapper()[name] metric = metric * self._compress_mask(wrapper.weight_mask) layer_weight_num = wrapper.module.weight.data.numel() - retention_numel = 0 if name not in min_retention_numel else min_retention_numel[name] - removed_metric_num = math.floor(retention_numel / (wrapper.weight_mask.numel() / metric.numel())) + retention_ratio = 1 - max_sparsity_per_layer.get(name, 1) + retention_numel = math.ceil(retention_ratio * layer_weight_num) + removed_metric_num = math.ceil(retention_numel / (wrapper.weight_mask.numel() / metric.numel())) stay_metric_num = metric.numel() - removed_metric_num # Remove the weight parts that must be left stay_metric = torch.topk(metric.view(-1), stay_metric_num, largest=False)[0] @@ -112,7 +113,7 @@ def generate_sparsity(self, metrics: Dict) -> Dict[str, Dict[str, Tensor]]: for _, group_metric_dict in grouped_metrics.items(): group_metric = self._group_metric_calculate(group_metric_dict) - sparsities = {name: self.pruner.get_modules_wrapper()[name].config['_sparsity'] for name in group_metric_dict.keys()} + sparsities = {name: self.pruner.get_modules_wrapper()[name].config['total_sparsity'] for name in group_metric_dict.keys()} min_sparsity = min(sparsities.values()) conv2d_groups = [self.group_depen[name] for name in group_metric_dict.keys()] diff --git a/nni/algorithms/compression/v2/pytorch/utils/pruning.py b/nni/algorithms/compression/v2/pytorch/utils/pruning.py index 9f04e8d608..d104d38dbf 100644 --- a/nni/algorithms/compression/v2/pytorch/utils/pruning.py +++ b/nni/algorithms/compression/v2/pytorch/utils/pruning.py @@ -2,18 +2,15 @@ # Licensed under the MIT license. from copy import deepcopy -import math from typing import Dict, List from torch.nn import Module -from nni.compression.pytorch.utils import get_module_by_name - def config_list_canonical(model: Module, config_list: List[Dict]) -> List[Dict]: ''' Split the config by op_names if 'sparsity' or 'sparsity_per_layer' in config, - and set the sub_config['_sparsity'] = config['sparsity_per_layer']. + and set the sub_config['total_sparsity'] = config['sparsity_per_layer']. ''' for config in config_list: if 'sparsity' in config: @@ -32,18 +29,15 @@ def config_list_canonical(model: Module, config_list: List[Dict]) -> List[Dict]: for op_name in op_names: sub_config = deepcopy(config) sub_config['op_names'] = [op_name] - sub_config['_sparsity'] = sparsity_per_layer + sub_config['total_sparsity'] = sparsity_per_layer new_config_list.append(sub_config) - elif 'total_sparsity' in config: - if 'max_sparsity_per_layer' in config: - min_retention_per_layer = (1 - config.pop('max_sparsity_per_layer')) - op_names = config.get('op_names', []) - min_retention_numel = {} - for op_name in op_names: - total_element_num = get_module_by_name(model, op_name)[1].weight.numel() - min_retention_numel[op_name] = math.floor(total_element_num * min_retention_per_layer) - config['_min_retention_numel'] = min_retention_numel - config['_sparsity'] = config.pop('total_sparsity') + elif 'max_sparsity_per_layer' in config and isinstance(config['max_sparsity_per_layer'], float): + op_names = config.get('op_names', []) + max_sparsity_per_layer = {} + max_sparsity = config['max_sparsity_per_layer'] + for op_name in op_names: + max_sparsity_per_layer[op_name] = max_sparsity + config['max_sparsity_per_layer'] = max_sparsity_per_layer new_config_list.append(config) else: new_config_list.append(config) From 0e92ce43a8bd66f9fe8243daac54ad7c69a5af66 Mon Sep 17 00:00:00 2001 From: J-shang Date: Fri, 20 Aug 2021 21:03:42 +0800 Subject: [PATCH 12/13] fix bug --- .../compression/v2/pytorch/utils/config_validation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nni/algorithms/compression/v2/pytorch/utils/config_validation.py b/nni/algorithms/compression/v2/pytorch/utils/config_validation.py index 100a81dfe9..e51e598fa0 100644 --- a/nni/algorithms/compression/v2/pytorch/utils/config_validation.py +++ b/nni/algorithms/compression/v2/pytorch/utils/config_validation.py @@ -27,7 +27,7 @@ def _modify_schema(self, data_schema: List[Dict], model: Module, logger: Logger) new_schema = And(old_schema, lambda n: validate_op_names(model, n, logger)) sub_schema[k] = new_schema - sub_schema = And(data_schema[0], lambda d: validate_op_types_op_names(d)) + sub_schema = And(sub_schema, lambda d: validate_op_types_op_names(d)) return data_schema From 68b537d8e89daee1f27594a09c672098bbf7b9d2 Mon Sep 17 00:00:00 2001 From: J-shang Date: Mon, 23 Aug 2021 11:43:06 +0800 Subject: [PATCH 13/13] fix bug --- .../compression/v2/pytorch/utils/config_validation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nni/algorithms/compression/v2/pytorch/utils/config_validation.py b/nni/algorithms/compression/v2/pytorch/utils/config_validation.py index e51e598fa0..d1b2a96f75 100644 --- a/nni/algorithms/compression/v2/pytorch/utils/config_validation.py +++ b/nni/algorithms/compression/v2/pytorch/utils/config_validation.py @@ -18,7 +18,7 @@ def _modify_schema(self, data_schema: List[Dict], model: Module, logger: Logger) if not data_schema: return data_schema - for sub_schema in data_schema: + for i, sub_schema in enumerate(data_schema): for k, old_schema in sub_schema.items(): if k == 'op_types' or (isinstance(k, Schema) and k._schema == 'op_types'): new_schema = And(old_schema, lambda n: validate_op_types(model, n, logger)) @@ -27,7 +27,7 @@ def _modify_schema(self, data_schema: List[Dict], model: Module, logger: Logger) new_schema = And(old_schema, lambda n: validate_op_names(model, n, logger)) sub_schema[k] = new_schema - sub_schema = And(sub_schema, lambda d: validate_op_types_op_names(d)) + data_schema[i] = And(sub_schema, lambda d: validate_op_types_op_names(d)) return data_schema