Compression doc structure refactor

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

@@ -12,11 +12,17 @@
 
 
 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 = {}
@@ -35,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 = {}
@@ -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)

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

@@ -1,6 +1,14 @@
 # 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
@@ -14,28 +22,24 @@
 
 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.
+    optimizer: torch.optim.Optimizer
+        Optimizer used to train model.
+    pruning_algorithm: str
+        Algorithms being used to prune model.
     """
 
     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)