diff --git a/docs/en_US/Compressor/Pruner.md b/docs/en_US/Compressor/Pruner.md index b9dcf3d9d1..939868c813 100644 --- a/docs/en_US/Compressor/Pruner.md +++ b/docs/en_US/Compressor/Pruner.md @@ -11,9 +11,9 @@ We provide several pruning algorithms that support fine-grained weight pruning a * [FPGM Pruner](#fpgm-pruner) * [L1Filter Pruner](#l1filter-pruner) * [L2Filter Pruner](#l2filter-pruner) -* [APoZ Rank Pruner](#activationapozrankfilterpruner) -* [Activation Mean Rank Pruner](#activationmeanrankfilterpruner) -* [Taylor FO On Weight Pruner](#taylorfoweightfilterpruner) +* [Activation APoZ Rank Filter Pruner](#activationAPoZRankFilter-pruner) +* [Activation Mean Rank Filter Pruner](#activationmeanrankfilter-pruner) +* [Taylor FO On Weight Pruner](#taylorfoweightfilter-pruner) **Pruning Schedule** * [AGP Pruner](#agp-pruner) @@ -51,16 +51,16 @@ pruner.compress() #### User configuration for Level Pruner -##### Tensorflow +##### PyTorch ```eval_rst -.. autoclass:: nni.compression.tensorflow.LevelPruner +.. autoclass:: nni.compression.torch.LevelPruner ``` -##### PyTorch +##### Tensorflow ```eval_rst -.. autoclass:: nni.compression.torch.LevelPruner +.. autoclass:: nni.compression.tensorflow.LevelPruner ``` @@ -140,17 +140,15 @@ pruner.compress() #### User configuration for FPGM Pruner -##### Tensorflow -```eval_rst -.. autoclass:: nni.compression.tensorflow.FPGMPruner -``` - ##### PyTorch ```eval_rst .. autoclass:: nni.compression.torch.FPGMPruner ``` -*** +##### Tensorflow +```eval_rst +.. autoclass:: nni.compression.tensorflow.FPGMPruner +``` ## L1Filter Pruner @@ -311,7 +309,6 @@ pruner = TaylorFOWeightFilterPruner(model, config_list, statistics_batch_num=1) pruner.compress() ``` -You can view [example](https://github.com/microsoft/nni/blob/master/examples/model_compress/model_prune_torch.py) for more information. #### User configuration for TaylorFOWeightFilterPruner @@ -323,13 +320,16 @@ You can view [example](https://github.com/microsoft/nni/blob/master/examples/mod ## AGP Pruner + This is an iterative pruner, In [To prune, or not to prune: exploring the efficacy of pruning for model compression](https://arxiv.org/abs/1710.01878), authors Michael Zhu and Suyog Gupta provide an algorithm to prune the weight gradually. >We introduce a new automated gradual pruning algorithm in which the sparsity is increased from an initial sparsity value si (usually 0) to a final sparsity value sf over a span of n pruning steps, starting at training step t0 and with pruning frequency ∆t: ![](../../img/agp_pruner.png) ->The binary weight masks are updated every ∆t steps as the network is trained to gradually increase the sparsity of the network while allowing the network training steps to recover from any pruning-induced loss in accuracy. In our experience, varying the pruning frequency ∆t between 100 and 1000 training steps had a negligible impact on the final model quality. Once the model achieves the target sparsity sf , the weight masks are no longer updated. The intuition behind this sparsity function in equation + +>The binary weight masks are updated every ∆t steps as the network is trained to gradually increase the sparsity of the network while allowing the network training steps to recover from any pruning-induced loss in accuracy. In our experience, varying the pruning frequency ∆t between 100 and 1000 training steps had a negligible impact on the final model quality. Once the model achieves the target sparsity sf , the weight masks are no longer updated. The intuition behind this sparsity function in equation (1). ### Usage + You can prune all weight from 0% to 80% sparsity in 10 epoch with the code below. PyTorch code @@ -372,7 +372,7 @@ PyTorch code ```python pruner.update_epoch(epoch) ``` -You can view example for more information. +You can view [example](https://github.com/microsoft/nni/blob/master/examples/model_compress/model_prune_torch.py) for more information. #### User configuration for AGP Pruner @@ -382,6 +382,12 @@ You can view example for more information. .. autoclass:: nni.compression.torch.AGP_Pruner ``` +##### Tensorflow + +```eval_rst +.. autoclass:: nni.compression.tensorflow.AGP_Pruner +``` + *** ## NetAdapt Pruner @@ -569,8 +575,11 @@ The above configuration means that there are 5 times of iterative pruning. As th #### User configuration for LotteryTicketPruner -* **prune_iterations:** The number of rounds for the iterative pruning, i.e., the number of iterative pruning. -* **sparsity:** The final sparsity when the compression is done. +##### PyTorch + +```eval_rst +.. autoclass:: nni.compression.torch.ADMMPruner +``` ### Reproduced Experiment diff --git a/docs/img/agp_pruner.png b/docs/img/agp_pruner.png index 889f42e764..98e991a0f9 100644 Binary files a/docs/img/agp_pruner.png and b/docs/img/agp_pruner.png differ diff --git a/src/sdk/pynni/nni/compression/tensorflow/builtin_pruners.py b/src/sdk/pynni/nni/compression/tensorflow/builtin_pruners.py index be70503de2..83f383bc40 100644 --- a/src/sdk/pynni/nni/compression/tensorflow/builtin_pruners.py +++ b/src/sdk/pynni/nni/compression/tensorflow/builtin_pruners.py @@ -41,23 +41,21 @@ def calc_mask(self, layer, config): class AGP_Pruner(Pruner): - """An automated gradual pruning algorithm that prunes the smallest magnitude - weights to achieve a preset level of network sparsity. - Michael Zhu and Suyog Gupta, "To prune, or not to prune: exploring the - efficacy of pruning for model compression", 2017 NIPS Workshop on Machine - Learning of Phones and other Consumer Devices, - https://arxiv.org/pdf/1710.01878.pdf + """ + Parameters + ---------- + model : torch.nn.module + Model to be pruned. + config_list : listlist + Supported keys: + - initial_sparsity: This is to specify the sparsity when compressor starts to compress. + - final_sparsity: This is to specify the sparsity when compressor finishes to compress. + - start_epoch: This is to specify the epoch number when compressor starts to compress, default start from epoch 0. + - end_epoch: This is to specify the epoch number when compressor finishes to compress. + - frequency: This is to specify every *frequency* number epochs compressor compress once, default frequency=1. """ def __init__(self, model, config_list): - """ - config_list: supported keys: - - initial_sparsity - - final_sparsity: you should make sure initial_sparsity <= final_sparsity - - start_epoch: start epoch numer begin update mask - - end_epoch: end epoch number stop update mask - - frequency: if you want update every 2 epoch, you can set it 2 - """ super().__init__(model, config_list) self.mask_list = {} self.if_init_list = {} diff --git a/src/sdk/pynni/nni/compression/torch/pruning/lottery_ticket.py b/src/sdk/pynni/nni/compression/torch/pruning/lottery_ticket.py index ace2b36154..c1805e7390 100644 --- a/src/sdk/pynni/nni/compression/torch/pruning/lottery_ticket.py +++ b/src/sdk/pynni/nni/compression/torch/pruning/lottery_ticket.py @@ -13,33 +13,22 @@ class LotteryTicketPruner(Pruner): """ - This is a Pytorch implementation of the paper "The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks", - following NNI model compression interface. - - 1. Randomly initialize a neural network f(x;theta_0) (where theta_0 follows D_{theta}). - 2. Train the network for j iterations, arriving at parameters theta_j. - 3. Prune p% of the parameters in theta_j, creating a mask m. - 4. Reset the remaining parameters to their values in theta_0, creating the winning ticket f(x;m*theta_0). - 5. Repeat step 2, 3, and 4. + Parameters + ---------- + model : pytorch model + The model to be pruned + config_list : list + Supported keys: + - prune_iterations : The number of rounds for the iterative pruning. + - sparsity : The final sparsity when the compression is done. + optimizer : pytorch optimizer + The optimizer for the model + lr_scheduler : pytorch lr scheduler + The lr scheduler for the model if used + reset_weights : bool + Whether reset weights and optimizer at the beginning of each round. """ - def __init__(self, model, config_list, optimizer=None, lr_scheduler=None, reset_weights=True): - """ - Parameters - ---------- - model : pytorch model - The model to be pruned - config_list : list - Supported keys: - - prune_iterations : The number of rounds for the iterative pruning. - - sparsity : The final sparsity when the compression is done. - optimizer : pytorch optimizer - The optimizer for the model - lr_scheduler : pytorch lr scheduler - The lr scheduler for the model if used - reset_weights : bool - Whether reset weights and optimizer at the beginning of each round. - """ # save init weights and optimizer self.reset_weights = reset_weights if self.reset_weights: diff --git a/src/sdk/pynni/nni/compression/torch/pruning/net_adapt_pruner.py b/src/sdk/pynni/nni/compression/torch/pruning/net_adapt_pruner.py index d5209cc50a..545057d25a 100644 --- a/src/sdk/pynni/nni/compression/torch/pruning/net_adapt_pruner.py +++ b/src/sdk/pynni/nni/compression/torch/pruning/net_adapt_pruner.py @@ -55,7 +55,7 @@ def short_term_fine_tuner(model, epoch=3): function to evaluate the masked model. This function should include `model` as the only parameter, and returns a scalar value. Example:: - + def evaluator(model): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") val_loader = ... diff --git a/src/sdk/pynni/nni/compression/torch/pruning/one_shot.py b/src/sdk/pynni/nni/compression/torch/pruning/one_shot.py index daa5b1d17d..4593b60b41 100644 --- a/src/sdk/pynni/nni/compression/torch/pruning/one_shot.py +++ b/src/sdk/pynni/nni/compression/torch/pruning/one_shot.py @@ -176,8 +176,8 @@ class FPGMPruner(_StructuredFilterPruner): - sparsity : This is to specify the sparsity operations to be compressed to. - op_types : Only Conv2d is supported in FPGM Pruner. """ - def __init__(self, model, config_list, optimizer=None): - super().__init__(model, config_list, pruning_algorithm='fpgm', optimizer=optimizer) + def __init__(self, model, config_list): + super().__init__(model, config_list, pruning_algorithm='fpgm') class TaylorFOWeightFilterPruner(_StructuredFilterPruner): """ @@ -204,8 +204,8 @@ class ActivationAPoZRankFilterPruner(_StructuredFilterPruner): - sparsity : How much percentage of convolutional filters are to be pruned. - op_types : Only Conv2d is supported in ActivationAPoZRankFilterPruner. """ - def __init__(self, model, config_list, optimizer=None, activation='relu', statistics_batch_num=1): - super().__init__(model, config_list, pruning_algorithm='apoz', optimizer=optimizer, \ + def __init__(self, model, config_list, activation='relu', statistics_batch_num=1): + super().__init__(model, config_list, pruning_algorithm='apoz', optimizer=None, \ activation=activation, statistics_batch_num=statistics_batch_num) class ActivationMeanRankFilterPruner(_StructuredFilterPruner): @@ -219,6 +219,6 @@ class ActivationMeanRankFilterPruner(_StructuredFilterPruner): - sparsity : How much percentage of convolutional filters are to be pruned. - op_types : Only Conv2d is supported in ActivationMeanRankFilterPruner. """ - def __init__(self, model, config_list, optimizer=None, activation='relu', statistics_batch_num=1): - super().__init__(model, config_list, pruning_algorithm='mean_activation', optimizer=optimizer, \ + def __init__(self, model, config_list, activation='relu', statistics_batch_num=1): + super().__init__(model, config_list, pruning_algorithm='mean_activation', optimizer=None, \ activation=activation, statistics_batch_num=statistics_batch_num) diff --git a/src/sdk/pynni/tests/test_compressor.py b/src/sdk/pynni/tests/test_compressor.py index 7641ae7d25..6a8727c9e4 100644 --- a/src/sdk/pynni/tests/test_compressor.py +++ b/src/sdk/pynni/tests/test_compressor.py @@ -88,9 +88,8 @@ def test_torch_quantizer_modules_detection(self): def test_torch_level_pruner(self): model = TorchModel() - optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5) configure_list = [{'sparsity': 0.8, 'op_types': ['default']}] - torch_compressor.LevelPruner(model, configure_list, optimizer).compress() + torch_compressor.LevelPruner(model, configure_list).compress() @tf2 def test_tf_level_pruner(self): @@ -129,7 +128,7 @@ def test_torch_fpgm_pruner(self): model = TorchModel() config_list = [{'sparsity': 0.6, 'op_types': ['Conv2d']}, {'sparsity': 0.2, 'op_types': ['Conv2d']}] - pruner = torch_compressor.FPGMPruner(model, config_list, torch.optim.SGD(model.parameters(), lr=0.01)) + pruner = torch_compressor.FPGMPruner(model, config_list) model.conv2.module.weight.data = torch.tensor(w).float() masks = pruner.calc_mask(model.conv2) @@ -315,7 +314,7 @@ def test_torch_QAT_quantizer(self): def test_torch_pruner_validation(self): # test bad configuraiton pruner_classes = [torch_compressor.__dict__[x] for x in \ - ['LevelPruner', 'SlimPruner', 'FPGMPruner', 'L1FilterPruner', 'L2FilterPruner', 'AGP_Pruner', \ + ['LevelPruner', 'SlimPruner', 'FPGMPruner', 'L1FilterPruner', 'L2FilterPruner', \ 'ActivationMeanRankFilterPruner', 'ActivationAPoZRankFilterPruner']] bad_configs = [ @@ -337,11 +336,10 @@ def test_torch_pruner_validation(self): ] ] model = TorchModel() - optimizer = torch.optim.SGD(model.parameters(), lr=0.01) for pruner_class in pruner_classes: for config_list in bad_configs: try: - pruner_class(model, config_list, optimizer) + pruner_class(model, config_list) print(config_list) assert False, 'Validation error should be raised for bad configuration' except schema.SchemaError: diff --git a/src/sdk/pynni/tests/test_pruners.py b/src/sdk/pynni/tests/test_pruners.py index e43f657534..7ac59733db 100644 --- a/src/sdk/pynni/tests/test_pruners.py +++ b/src/sdk/pynni/tests/test_pruners.py @@ -192,7 +192,9 @@ def pruners_test(pruner_names=['level', 'agp', 'slim', 'fpgm', 'l1', 'l2', 'tayl pruner = prune_config[pruner_name]['pruner_class'](model, config_list, trainer=prune_config[pruner_name]['trainer']) elif pruner_name == 'autocompress': pruner = prune_config[pruner_name]['pruner_class'](model, config_list, trainer=prune_config[pruner_name]['trainer'], evaluator=prune_config[pruner_name]['evaluator'], dummy_input=x) - else: + elif pruner_name in ['level', 'slim', 'fpgm', 'l1', 'l2', 'mean_activation', 'apoz']: + pruner = prune_config[pruner_name]['pruner_class'](model, config_list) + else: # 'agp', 'taylorfo' pruner = prune_config[pruner_name]['pruner_class'](model, config_list, optimizer) pruner.compress()