Skip to content
This repository has been archived by the owner on Sep 18, 2024. It is now read-only.

[#3507 follow up] update doc #3688

Merged
merged 8 commits into from
May 27, 2021
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions docs/en_US/Compression/DependencyAware.rst
Original file line number Diff line number Diff line change
Expand Up @@ -54,11 +54,11 @@ To enable the dependency-aware mode for ``L1FilterPruner``\ :
# for FPGMPruner
# pruner = FPGMPruner(model, config_list, dependency_aware=True, dummy_input=dummy_input)
# for ActivationAPoZRankFilterPruner
# pruner = ActivationAPoZRankFilterPruner(model, config_list, statistics_batch_num=1, , dependency_aware=True, dummy_input=dummy_input)
# pruner = ActivationAPoZRankFilterPruner(model, config_list, optimizer, trainer, criterion, statistics_batch_num=1, , dependency_aware=True, dummy_input=dummy_input)
# for ActivationMeanRankFilterPruner
# pruner = ActivationMeanRankFilterPruner(model, config_list, statistics_batch_num=1, dependency_aware=True, dummy_input=dummy_input)
# pruner = ActivationMeanRankFilterPruner(model, config_list, optimizer, trainer, criterion, statistics_batch_num=1, dependency_aware=True, dummy_input=dummy_input)
# for TaylorFOWeightFilterPruner
# pruner = TaylorFOWeightFilterPruner(model, config_list, statistics_batch_num=1, dependency_aware=True, dummy_input=dummy_input)
# pruner = TaylorFOWeightFilterPruner(model, config_list, optimizer, trainer, criterion, statistics_batch_num=1, dependency_aware=True, dummy_input=dummy_input)

pruner.compress()

Expand Down
6 changes: 3 additions & 3 deletions docs/en_US/Compression/Framework.rst
Original file line number Diff line number Diff line change
Expand Up @@ -29,8 +29,7 @@ Compressor is the base class for pruner and quntizer, it provides a unified inte
'op_types': ['Conv2d', 'Linear'],
}]

optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=1e-4)
pruner = LevelPruner(model, configure_list, optimizer)
pruner = LevelPruner(model, configure_list)
model = pruner.compress()

# model is ready for pruning, now start finetune the model,
Expand Down Expand Up @@ -103,7 +102,8 @@ Users can also remove this collector like this:
Pruner
------

A pruner receives ``model``\ , ``config_list`` and ``optimizer`` as arguments. It prunes the model per the ``config_list`` during training loop by adding a hook on ``optimizer.step()``.
A pruner receives ``model``\ , ``config_list`` as arguments.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

remove \

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed

Some pruners like ``TaylorFOWeightFilter Pruner`` prune the model per the ``config_list`` during training loop by adding a hook on ``optimizer.step()``.

Pruner class is a subclass of Compressor, so it contains everything in the Compressor class and some additional components only for pruning, it contains:

Expand Down
20 changes: 6 additions & 14 deletions docs/en_US/Compression/Pruner.rst
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ PyTorch code

from nni.algorithms.compression.pytorch.pruning import SlimPruner
config_list = [{ 'sparsity': 0.8, 'op_types': ['BatchNorm2d'] }]
pruner = SlimPruner(model, config_list)
pruner = SlimPruner(model, config_list, optimizer, trainer, criterion)
pruner.compress()

User configuration for Slim Pruner
Expand Down Expand Up @@ -269,7 +269,7 @@ PyTorch code
'sparsity': 0.5,
'op_types': ['Conv2d']
}]
pruner = ActivationAPoZRankFilterPruner(model, config_list, statistics_batch_num=1)
pruner = ActivationAPoZRankFilterPruner(model, config_list, optimizer, trainer, criterion, statistics_batch_num=1)
pruner.compress()

Note: ActivationAPoZRankFilterPruner is used to prune convolutional layers within deep neural networks, therefore the ``op_types`` field supports only convolutional layers.
Expand Down Expand Up @@ -304,7 +304,7 @@ PyTorch code
'sparsity': 0.5,
'op_types': ['Conv2d']
}]
pruner = ActivationMeanRankFilterPruner(model, config_list, statistics_batch_num=1)
pruner = ActivationMeanRankFilterPruner(model, config_list, optimizer, trainer, criterion, statistics_batch_num=1)
pruner.compress()

Note: ActivationMeanRankFilterPruner is used to prune convolutional layers within deep neural networks, therefore the ``op_types`` field supports only convolutional layers.
Expand Down Expand Up @@ -344,7 +344,7 @@ PyTorch code
'sparsity': 0.5,
'op_types': ['Conv2d']
}]
pruner = TaylorFOWeightFilterPruner(model, config_list, statistics_batch_num=1)
pruner = TaylorFOWeightFilterPruner(model, config_list, optimizer, trainer, criterion, statistics_batch_num=1)
pruner.compress()

User configuration for TaylorFOWeightFilter Pruner
Expand Down Expand Up @@ -389,7 +389,7 @@ PyTorch code
# optimizer.step(), so an optimizer is required to prune the model.
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=1e-4)

pruner = AGPPruner(model, config_list, optimizer, pruning_algorithm='level')
pruner = AGPPruner(model, config_list, optimizer, trainer, criterion, pruning_algorithm='level')
pruner.compress()

AGP pruner uses ``LevelPruner`` algorithms to prune the weight by default, however you can set ``pruning_algorithm`` parameter to other values to use other pruning algorithms:
Expand All @@ -404,14 +404,6 @@ AGP pruner uses ``LevelPruner`` algorithms to prune the weight by default, howev
* ``apoz``\ : ActivationAPoZRankFilterPruner
* ``mean_activation``\ : ActivationMeanRankFilterPruner

You should add code below to update epoch number when you finish one epoch in your training code.

PyTorch code

.. code-block:: python

pruner.update_epoch(epoch)


User configuration for AGP Pruner
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Expand Down Expand Up @@ -620,7 +612,7 @@ PyTorch code
'op_types': ['Conv2d'],
'op_names': ['conv2']
}]
pruner = ADMMPruner(model, config_list, trainer=trainer, num_iterations=30, epochs=5)
pruner = ADMMPruner(model, config_list, trainer, num_iterations=30, epochs_per_iteration=5)
pruner.compress()

You can view :githublink:`example <examples/model_compress/pruning/auto_pruners_torch.py>` for more information.
Expand Down
7 changes: 3 additions & 4 deletions docs/en_US/Compression/QuickStart.rst
Original file line number Diff line number Diff line change
Expand Up @@ -31,17 +31,16 @@ The specification of configuration can be found `here <./Tutorial.rst#specify-th
Step2. Choose a pruner and compress the model
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

First instantiate the chosen pruner with your model and configuration as arguments, then invoke ``compress()`` to compress your model. Note that, some algorithms may check gradients for compressing, so we also define an optimizer and pass it to the pruner.
First instantiate the chosen pruner with your model and configuration as arguments, then invoke ``compress()`` to compress your model. Note that, some algorithms may check gradients for compressing, so we may also define an optimizer and pass it to the pruner.

.. code-block:: python

from nni.algorithms.compression.pytorch.pruning import LevelPruner

optimizer_finetune = torch.optim.SGD(model.parameters(), lr=0.01)
pruner = LevelPruner(model, config_list, optimizer_finetune)
pruner = LevelPruner(model, config_list)
model = pruner.compress()

Then, you can train your model using traditional training approach (e.g., SGD), pruning is applied transparently during the training. Some pruners (e.g., L1FilterPruner, FPGMPruner) prune once at the beginning, the following training can be seen as fine-tune. Some pruners (e.g., AGPPruner) prune your model iteratively, the masks are adjusted epoch by epoch during training.
Some pruners (e.g., L1FilterPruner, FPGMPruner) prune once, some pruners (e.g., AGPPruner) prune your model iteratively, the masks are adjusted epoch by epoch during training.

Note that, ``pruner.compress`` simply adds masks on model weights, it does not include fine-tuning logic. If users want to fine tune the compressed model, they need to write the fine tune logic by themselves after ``pruner.compress``.

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@ class AutoCompressPruner(Pruner):
Function used for the first subproblem of ADMM Pruner.
Users should write this function as a normal function to train the Pytorch model
and include `model, optimizer, criterion, epoch` as function arguments.
criterion: function
Function used to calculate the loss between the target and the output. By default, we use CrossEntropyLoss.
QuanluZhang marked this conversation as resolved.
Show resolved Hide resolved
evaluator : function
function to evaluate the pruned model.
This function should include `model` as the only parameter, and returns a scalar value.
Expand Down Expand Up @@ -80,7 +82,7 @@ def evaluator(model):
PATH to store temporary experiment data.
"""

def __init__(self, model, config_list, trainer, criterion, evaluator, dummy_input,
def __init__(self, model, config_list, trainer, evaluator, dummy_input, criterion=torch.nn.CrossEntropyLoss(),
num_iterations=3, optimize_mode='maximize', base_algo='l1',
# SimulatedAnnealing related
start_temperature=100, stop_temperature=20, cool_down_rate=0.9, perturbation_magnitude=0.35,
Expand Down
26 changes: 17 additions & 9 deletions nni/algorithms/compression/pytorch/pruning/iterative_pruner.py
Original file line number Diff line number Diff line change
Expand Up @@ -461,6 +461,8 @@ class TaylorFOWeightFilterPruner(IterativePruner):
Function used to calculate the loss between the target and the output.
sparsity_training_epochs: int
The number of epochs to collect the contributions.
statistics_batch_num: int
The number of batches to statistic the activation.
dependency_aware: bool
If prune the model in a dependency-aware way. If it is `True`, this pruner will
prune the model according to the l2-norm of weights and the channel-dependency or
Expand All @@ -472,14 +474,14 @@ class TaylorFOWeightFilterPruner(IterativePruner):
dummy_input : torch.Tensor
The dummy input to analyze the topology constraints. Note that, the dummy_input
should on the same device with the model.

"""

def __init__(self, model, config_list, optimizer, trainer, criterion, sparsity_training_epochs=1, dependency_aware=False,
dummy_input=None):
def __init__(self, model, config_list, optimizer, trainer, criterion, sparsity_training_epochs=1,
statistics_batch_num=1, dependency_aware=False, dummy_input=None):
super().__init__(model, config_list, optimizer=optimizer, pruning_algorithm='taylorfo', trainer=trainer,
criterion=criterion, num_iterations=1, epochs_per_iteration=sparsity_training_epochs,
dependency_aware=dependency_aware, dummy_input=dummy_input)
criterion=criterion, statistics_batch_num=statistics_batch_num, num_iterations=1,
epochs_per_iteration=sparsity_training_epochs, dependency_aware=dependency_aware,
dummy_input=dummy_input)

def _supported_dependency_aware(self):
return True
Expand Down Expand Up @@ -507,6 +509,8 @@ class ActivationAPoZRankFilterPruner(IterativePruner):
The activation type.
sparsity_training_epochs: int
The number of epochs to statistic the activation.
statistics_batch_num: int
The number of batches to statistic the activation.
dependency_aware: bool
If prune the model in a dependency-aware way. If it is `True`, this pruner will
prune the model according to the l2-norm of weights and the channel-dependency or
Expand All @@ -522,10 +526,11 @@ class ActivationAPoZRankFilterPruner(IterativePruner):
"""

def __init__(self, model, config_list, optimizer, trainer, criterion, activation='relu',
sparsity_training_epochs=1, dependency_aware=False, dummy_input=None):
sparsity_training_epochs=1, statistics_batch_num=1, dependency_aware=False, dummy_input=None):
super().__init__(model, config_list, pruning_algorithm='apoz', optimizer=optimizer, trainer=trainer,
criterion=criterion, dependency_aware=dependency_aware, dummy_input=dummy_input,
activation=activation, num_iterations=1, epochs_per_iteration=sparsity_training_epochs)
activation=activation, statistics_batch_num=statistics_batch_num, num_iterations=1,
epochs_per_iteration=sparsity_training_epochs)
self.patch_optimizer(self.update_mask)

def _supported_dependency_aware(self):
Expand Down Expand Up @@ -554,6 +559,8 @@ class ActivationMeanRankFilterPruner(IterativePruner):
The activation type.
sparsity_training_epochs: int
The number of batches to statistic the activation.
statistics_batch_num: int
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why there is no such parameter before?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

workaround by using sparsifying_training_batches

The number of batches to statistic the activation.
dependency_aware: bool
If prune the model in a dependency-aware way. If it is `True`, this pruner will
prune the model according to the l2-norm of weights and the channel-dependency or
Expand All @@ -568,10 +575,11 @@ class ActivationMeanRankFilterPruner(IterativePruner):
"""

def __init__(self, model, config_list, optimizer, trainer, criterion, activation='relu',
sparsity_training_epochs=1, dependency_aware=False, dummy_input=None):
sparsity_training_epochs=1, statistics_batch_num=1, dependency_aware=False, dummy_input=None):
super().__init__(model, config_list, pruning_algorithm='mean_activation', optimizer=optimizer, trainer=trainer,
criterion=criterion, dependency_aware=dependency_aware, dummy_input=dummy_input,
activation=activation, num_iterations=1, epochs_per_iteration=sparsity_training_epochs)
activation=activation, statistics_batch_num=statistics_batch_num, num_iterations=1,
epochs_per_iteration=sparsity_training_epochs)
self.patch_optimizer(self.update_mask)

def _supported_dependency_aware(self):
Expand Down