Compression doc structure refactor

docs/en_US/Compressor/AutoCompression.md

@@ -84,7 +84,7 @@ config_list_agp = [{'initial_sparsity': 0, 'final_sparsity': conv0_sparsity,
                    {'initial_sparsity': 0, 'final_sparsity': conv1_sparsity,
                     'start_epoch': 0, 'end_epoch': 3,
                     'frequency': 1,'op_name': 'conv1' },]
-PRUNERS = {'level':LevelPruner(model, config_list_level), 'agp':AGP_Pruner(model, config_list_agp)}
+PRUNERS = {'level':LevelPruner(model, config_list_level), 'agp':AGPPruner(model, config_list_agp)}
 pruner = PRUNERS(params['prune_method']['_name'])
 pruner.compress()
 ... # fine tuning

examples/model_compress/README.md

@@ -22,7 +22,7 @@ configure_list = [{
     'frequency': 1,
     'op_types': ['default']
 }]
-pruner = AGP_Pruner(configure_list)
+pruner = AGPPruner(configure_list)
 ```
 
 When ```pruner(model)``` is called, your model is injected with masks as embedded operations. For example, a layer takes a weight as input, we will insert an operation between the weight and the layer, this operation takes the weight as input and outputs a new weight applied by the mask. Thus, the masks are applied at any time the computation goes through the operations. You can fine-tune your model **without** any modifications.

examples/model_compress/model_prune_torch.py

@@ -10,7 +10,7 @@
 from models.cifar10.vgg import VGG
 import nni
 from nni.compression.torch import LevelPruner, SlimPruner, FPGMPruner, L1FilterPruner, \
-    L2FilterPruner, AGP_Pruner, ActivationMeanRankFilterPruner, ActivationAPoZRankFilterPruner
+    L2FilterPruner, AGPPruner, ActivationMeanRankFilterPruner, ActivationAPoZRankFilterPruner
 
 prune_config = {
     'level': {
@@ -25,7 +25,7 @@
     'agp': {
         'dataset_name': 'mnist',
         'model_name': 'naive',
-        'pruner_class': AGP_Pruner,
+        'pruner_class': AGPPruner,
         'config_list': [{
             'initial_sparsity': 0.,
             'final_sparsity': 0.8,

src/sdk/pynni/nni/compression/tensorflow/builtin_pruners.py

@@ -6,17 +6,23 @@
 import tensorflow as tf
 from .compressor import Pruner
 
-__all__ = ['LevelPruner', 'AGP_Pruner', 'FPGMPruner']
+__all__ = ['LevelPruner', 'AGPPruner', 'FPGMPruner']
 
 _logger = logging.getLogger(__name__)
 
 
 class LevelPruner(Pruner):
+    """
+    Parameters
+    ----------
+    model : tensorflow model
+        Model to be pruned
+    config_list : list
+        Supported keys:
+            - sparsity : This is to specify the sparsity operations to be compressed to.
+            - op_types : Operation types to prune.
+    """
     def __init__(self, model, config_list):
-        """
-        config_list: supported keys:
-            - sparsity
-        """
         super().__init__(model, config_list)
         self.mask_list = {}
         self.if_init_list = {}
@@ -34,24 +40,22 @@ def calc_mask(self, layer, config):
         return mask
 
 
-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
+class AGPPruner(Pruner):
+    """
+    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 = {}
@@ -102,23 +106,19 @@ def update_epoch(self, epoch, sess):
         for k in self.if_init_list:
             self.if_init_list[k] = True
 
+
 class FPGMPruner(Pruner):
     """
-    A filter pruner via geometric median.
-    "Filter Pruning via Geometric Median for Deep Convolutional Neural Networks Acceleration",
-    https://arxiv.org/pdf/1811.00250.pdf
+    Parameters
+    ----------
+    model : tensorflow model
+        Model to be pruned
+    config_list : list
+        Supported keys:
+            - sparsity : percentage of convolutional filters to be pruned.
+            - op_types : Only Conv2d is supported in FPGM Pruner.
     """
-
     def __init__(self, model, config_list):
-        """
-        Parameters
-        ----------
-        model : pytorch model
-            the model user wants to compress
-        config_list: list
-            support key for each list item:
-                - sparsity: percentage of convolutional filters to be pruned.
-        """
         super().__init__(model, config_list)
         self.mask_dict = {}
         self.assign_handler = []

src/sdk/pynni/nni/compression/torch/pruning/admm_pruner.py

@@ -15,58 +15,50 @@
 
 class ADMMPruner(OneshotPruner):
     """
-    This is a Pytorch implementation of ADMM Pruner algorithm.
+    A Pytorch implementation of ADMM Pruner algorithm.
+
+    Parameters
+    ----------
+    model : torch.nn.Module
+        Model to be pruned.
+    config_list : list
+        List on pruning configs.
+    trainer : function
+        Function used for the first subproblem.
+        Users should write this function as a normal function to train the Pytorch model
+        and include `model, optimizer, criterion, epoch, callback` as function arguments.
+        Here `callback` acts as an L2 regulizer as presented in the formula (7) of the original paper.
+        The logic of `callback` is implemented inside the Pruner,
+        users are just required to insert `callback()` between `loss.backward()` and `optimizer.step()`.
+        Example::
+
+            def trainer(model, criterion, optimizer, epoch, callback):
+                device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+                train_loader = ...
+                model.train()
+                for batch_idx, (data, target) in enumerate(train_loader):
+                    data, target = data.to(device), target.to(device)
+                    optimizer.zero_grad()
+                    output = model(data)
+                    loss = criterion(output, target)
+                    loss.backward()
+                    # callback should be inserted between loss.backward() and optimizer.step()
+                    if callback:
+                        callback()
+                    optimizer.step()
+    num_iterations : int
+        Total number of iterations.
+    training_epochs : int
+        Training epochs of the first subproblem.
+    row : float
+        Penalty parameters for ADMM training.
+    base_algo : str
+        Base pruning algorithm. `level`, `l1` or `l2`, by default `l1`. Given the sparsity distribution among the ops,
+        the assigned `base_algo` is used to decide which filters/channels/weights to prune.
 
-    Alternating Direction Method of Multipliers (ADMM) is a mathematical optimization technique,
-    by decomposing the original nonconvex problem into two subproblems that can be solved iteratively.
-    In weight pruning problem, these two subproblems are solved via 1) gradient descent algorithm and 2) Euclidean projection respectively.
-    This solution framework applies both to non-structured and different variations of structured pruning schemes.
-
-    For more details, please refer to the paper: https://arxiv.org/abs/1804.03294.
     """
 
     def __init__(self, model, config_list, trainer, num_iterations=30, training_epochs=5, row=1e-4, base_algo='l1'):
-        """
-        Parameters
-        ----------
-        model : torch.nn.module
-            Model to be pruned
-        config_list : list
-            List on pruning configs
-        trainer : function
-            Function used for the first subproblem.
-            Users should write this function as a normal function to train the Pytorch model
-            and include `model, optimizer, criterion, epoch, callback` as function arguments.
-            Here `callback` acts as an L2 regulizer as presented in the formula (7) of the original paper.
-            The logic of `callback` is implemented inside the Pruner,
-            users are just required to insert `callback()` between `loss.backward()` and `optimizer.step()`.
-            Example::
-            ```
-            >>> def trainer(model, criterion, optimizer, epoch, callback):
-            >>>     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            >>>     train_loader = ...
-            >>>     model.train()
-            >>>     for batch_idx, (data, target) in enumerate(train_loader):
-            >>>         data, target = data.to(device), target.to(device)
-            >>>         optimizer.zero_grad()
-            >>>         output = model(data)
-            >>>         loss = criterion(output, target)
-            >>>         loss.backward()
-            >>>         # callback should be inserted between loss.backward() and optimizer.step()
-            >>>         if callback:
-            >>>             callback()
-            >>>         optimizer.step()
-            ```
-        num_iterations : int
-            Total number of iterations.
-        training_epochs : int
-            Training epochs of the first subproblem.
-        row : float
-            Penalty parameters for ADMM training.
-        base_algo : str
-            Base pruning algorithm. `level`, `l1` or `l2`, by default `l1`. Given the sparsity distribution among the ops,
-            the assigned `base_algo` is used to decide which filters/channels/weights to prune.
-        """
         self._base_algo = base_algo
 
         super().__init__(model, config_list)
@@ -83,7 +75,7 @@ def validate_config(self, model, config_list):
         """
         Parameters
         ----------
-        model : torch.nn.module
+        model : torch.nn.Module
             Model to be pruned
         config_list : list
             List on pruning configs

src/sdk/pynni/nni/compression/torch/pruning/agp.py

@@ -1,41 +1,46 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+"""
+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.
+"""
+
 import logging
 import torch
 from schema import And, Optional
 from .constants import MASKER_DICT
 from ..utils.config_validation import CompressorSchema
 from ..compressor import Pruner
 
-__all__ = ['AGP_Pruner']
+__all__ = ['AGPPruner']
 
 logger = logging.getLogger('torch pruner')
 
-class AGP_Pruner(Pruner):
+class AGPPruner(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.
+    optimizer: torch.optim.Optimizer
+        Optimizer used to train model.
+    pruning_algorithm: str
+        Algorithms being used to prune model,
+        choose from `['level', 'slim', 'l1', 'l2', 'fpgm', 'taylorfo', 'apoz', 'mean_activation']`, by default `level`
     """
 
     def __init__(self, model, config_list, optimizer, pruning_algorithm='level'):
-        """
-        Parameters
-        ----------
-        model : torch.nn.module
-            Model to be pruned
-        config_list : list
-            List on pruning configs
-        optimizer: torch.optim.Optimizer
-            Optimizer used to train model
-        pruning_algorithm: str
-            algorithms being used to prune model
-        """
-
         super().__init__(model, config_list, optimizer)
         assert isinstance(optimizer, torch.optim.Optimizer), "AGP pruner is an iterative pruner, please pass optimizer of the model to it"
         self.masker = MASKER_DICT[pruning_algorithm](model, self)
@@ -47,7 +52,7 @@ def validate_config(self, model, config_list):
         """
         Parameters
         ----------
-        model : torch.nn.module
+        model : torch.nn.Module
             Model to be pruned
         config_list : list
             List on pruning configs

src/sdk/pynni/nni/compression/torch/pruning/apply_compression.py

@@ -14,7 +14,7 @@ def apply_compression_results(model, masks_file, map_location=None):
 
     Parameters
     ----------
-    model : torch.nn.module
+    model : torch.nn.Module
         The model to be compressed
     masks_file : str
         The path of the mask file