[Unstructured prune]Support block sparse training (PaddlePaddle#1000)

Author: minghaoBD

demo/dygraph/unstructured_pruning/README.md

@@ -33,7 +33,7 @@ paddleslim>=2.2.0
 默认根据参数的绝对值大小进行稀疏化，且不稀疏归一化层参数。如果开发者想更改相应的逻辑，可按照下述操作：
 
 - 开发者可以通过重写`paddleslim.dygraph.prune.unstructured_pruner.py`中的`UnstructuredPruner.mask_parameters()`和`UnstructuredPruner.update_threshold()`来定义自己的非结构化稀疏策略（目前为剪裁掉绝对值小的parameters）。
-- 开发可以在初始化`UnstructuredPruner`时，传入自定义的`skip_params_func`，来定义哪些参数不参与剪裁。`skip_params_func`示例代码如下(路径：`paddleslim.dygraph.prune.unstructured_pruner._get_skip_params())`。默认为所有的归一化层的参数不参与剪裁。
+- 开发可以在初始化`UnstructuredPruner`时，传入自定义的`skip_params_func`，来定义哪些参数不参与剪裁。`skip_params_func`示例代码如下(路径：`paddleslim.dygraph.prune.unstructured_pruner._get_skip_params())`。默认为所有的归一化层的参数和 `bias` 不参与剪裁。
 
 ```python
 NORMS_ALL = [ 'BatchNorm', 'GroupNorm', 'LayerNorm', 'SpectralNorm', 'BatchNorm1D',
@@ -55,6 +55,8 @@ def _get_skip_params(model):
     for _, sub_layer in model.named_sublayers():
         if type(sub_layer).__name__.split('.')[-1] in NORMS_ALL:
             skip_params.add(sub_layer.full_name())
+        for param in sub_layer.parameters(include_sublayers=False):
+            if len(param.shape) == 1: skip_params.add(param.name)
     return skip_params
 ```
 

demo/dygraph/unstructured_pruning/train.py

@@ -35,6 +35,7 @@
 add_arg('threshold',            float,  0.01,               "The threshold to set zeros. Default: 0.01")
 add_arg('num_epochs',       int,  120,               "The number of total epochs. Default: 120")
 parser.add_argument('--step_epochs', nargs='+', type=int, default=[30, 60, 90], help="piecewise decay step")
+parser.add_argument('--sparse_block', nargs='+', type=int, default=[1, 1], help="There must be two integers inside this array. The array defines the shape of the block, the values within which are either sparsified to all zeros or kept original. [1, 1] means unstructured pruning. Default: [1, 1]")
 add_arg('data',             str, "imagenet",                 "Which data to use. 'cifar10' or 'imagenet'. Default: imagenet")
 add_arg('log_period',       int, 100,                 "Log period in batches. Default: 100")
 add_arg('test_period',      int, 5,                 "Test period in epoches. Default: 5")
@@ -100,13 +101,15 @@ def create_unstructured_pruner(model, args, configs=None):
             ratio=args.ratio,
             threshold=args.threshold,
             prune_params_type=args.prune_params_type,
-            local_sparsity=args.local_sparsity)
+            local_sparsity=args.local_sparsity,
+            sparse_block=args.sparse_block)
     else:
         return GMPUnstructuredPruner(
             model,
             ratio=args.ratio,
             prune_params_type=args.prune_params_type,
             local_sparsity=args.local_sparsity,
+            sparse_block=args.sparse_block,
             configs=configs)
 
 

demo/unstructured_prune/README.md

@@ -42,7 +42,7 @@ tar -xf MobileNetV1_pretrained.tar
 默认根据参数的绝对值大小进行稀疏化，且不稀疏归一化层参数。如果开发者想更改相应的逻辑，可按照下述操作：
 
 - 可以通过重写`paddleslim.prune.unstructured_pruner.py`中的`UnstructuredPruner.update_threshold()`来定义自己的非结构化稀疏策略（目前为剪裁掉绝对值小的parameters）。
-- 可以在初始化`UnstructuredPruner`时，传入自定义的`skip_params_func`，来定义哪些参数不参与剪裁。`skip_params_func`示例代码如下(路径：`paddleslim.prune.unstructured_pruner._get_skip_params()`)。默认为所有的归一化层的参数不参与剪裁。
+- 可以在初始化`UnstructuredPruner`时，传入自定义的`skip_params_func`，来定义哪些参数不参与剪裁。`skip_params_func`示例代码如下(路径：`paddleslim.prune.unstructured_pruner._get_skip_params()`)。默认为所有的归一化层的参数和 `bias` 不参与剪裁。
 
 ```python
 def _get_skip_params(program):
@@ -61,6 +61,9 @@ def _get_skip_params(program):
         if 'norm' in op.type() and 'grad' not in op.type():
             for input in op.all_inputs():
                 skip_params.add(input.name())
+    for param in program.all_parameters():
+        if len(param.shape) == 1:
+            skip_params.add(param.name)  
     return skip_params
 ```
 

demo/unstructured_prune/train.py

@@ -39,6 +39,7 @@
 add_arg('ratio',            float,  0.55,               "The ratio to set zeros, the smaller portion will be zeros. Default: 0.55")
 add_arg('num_epochs',       int,  120,               "The number of total epochs. Default: 120")
 parser.add_argument('--step_epochs', nargs='+', type=int, default=[30, 60, 90], help="piecewise decay step")
+parser.add_argument('--sparse_block', nargs='+', type=int, default=[1, 1], help="There must be two integers inside this array. The array defines the shape of the block, the values within which are either sparsified to all zeros or kept original. [1, 1] means unstructured pruning. Default: [1, 1]")
 add_arg('data',             str, "imagenet",                 "Which data to use. 'mnist', 'cifar10' or 'imagenet'. Default: imagenet")
 add_arg('log_period',       int, 100,                 "Log period in batches. Default: 100")
 add_arg('test_period',      int, 5,                 "Test period in epoches. Default: 5")
@@ -102,14 +103,16 @@ def create_unstructured_pruner(train_program, args, place, configs):
             threshold=args.threshold,
             prune_params_type=args.prune_params_type,
             place=place,
-            local_sparsity=args.local_sparsity)
+            local_sparsity=args.local_sparsity,
+            sparse_block=args.sparse_block)
     else:
         return GMPUnstructuredPruner(
             train_program,
             ratio=args.ratio,
             prune_params_type=args.prune_params_type,
             place=place,
             local_sparsity=args.local_sparsity,
+            sparse_block=args.sparse_block,
             configs=configs)
 
 
@@ -312,7 +315,6 @@ def train(epoch, program):
                 fetch_list=[avg_cost.name, acc_top1.name, acc_top5.name])
             # GMP pruner step 2: step() to update ratios and other internal states of the pruner.
             pruner.step()
-
             train_run_cost += time.time() - train_start
             total_samples += args.batch_size
             loss_n = np.mean(loss_n)

docs/zh_cn/api_cn/dygraph/pruners/unstructured_pruner.rst

@@ -4,7 +4,7 @@
 UnstructuredPruner
 ----------
 
-.. py:class:: paddleslim.UnstructuredPruner(model, mode, threshold=0.01, ratio=0.55, prune_params_type=None, skip_params_func=None, local_sparsity=False)
+.. py:class:: paddleslim.UnstructuredPruner(model, mode, threshold=0.01, ratio=0.55, prune_params_type=None, skip_params_func=None, local_sparsity=False, sparse_block=[1,1])
 
 `源代码 <https://github.com/PaddlePaddle/PaddleSlim/blob/develop/paddleslim/dygraph/prune/unstructured_pruner.py>`_
 
@@ -17,7 +17,7 @@ UnstructuredPruner
 - **ratio(float)** - 稀疏化比例期望，只有在 mode=='ratio' 时才会生效。
 - **threshold(float)** - 稀疏化阈值期望，只有在 mode=='threshold' 时才会生效。
 - **prune_params_type(String)** - 用以指定哪些类型的参数参与稀疏。目前只支持None和"conv1x1_only"两个选项，后者表示只稀疏化1x1卷积。而前者表示稀疏化除了归一化层的参数。
-- **skip_params_func(function)** - 一个指向function的指针，该function定义了哪些参数不应该被剪裁，默认（None）时代表所有归一化层参数不参与剪裁。示例代码如下：
+- **skip_params_func(function)** - 一个指向function的指针，该function定义了哪些参数不应该被剪裁，默认（None）时代表所有归一化层参数和 bias 不参与剪裁。示例代码如下：
 
 .. code-block:: python
 
@@ -28,7 +28,7 @@ UnstructuredPruner
   def _get_skip_params(model):
       """
       This function is used to check whether the given model's layers are valid to be pruned.
-      Usually, the convolutions are to be pruned while we skip the normalization-related parameters.
+      Usually, the convolutions are to be pruned while we skip the normalization-related parameters and bias.
       Deverlopers could replace this function by passing their own when initializing the UnstructuredPuner instance.
 
       Args:
@@ -40,11 +40,14 @@ UnstructuredPruner
       for _, sub_layer in model.named_sublayers():
           if type(sub_layer).__name__.split('.')[-1] in NORMS_ALL:
               skip_params.add(sub_layer.full_name())
+          for param in sub_layer.parameters(include_sublayers=False):
+              if len(param.shape) == 1: skip_params.add(param.name)
       return skip_params
 
 ..
 
 - **local_sparsity(bool)** - 剪裁比例（ratio）应用的范围：local_sparsity 开启时意味着每个参与剪裁的参数矩阵稀疏度均为 'ratio'， 关闭时表示只保证模型整体稀疏度达到'ratio'，但是每个参数矩阵的稀疏度可能存在差异。
+- **sparse_block(Array<Integer>)** - 一个含有两个正整数的数组，定义了稀疏化时候block的大小。即 sparse_block[0] x sparse_block[1]中的参数作为一个整体，要么被置0，要么保持不变。默认为 [1,1]，代表非结构化稀疏。
 
 **返回：** 一个UnstructuredPruner类的实例。
 
@@ -245,6 +248,7 @@ GMPUnstructuredPruner
 - **prune_params_type(str)** - 用以指定哪些类型的参数参与稀疏。目前只支持None和"conv1x1_only"两个选项，后者表示只稀疏化1x1卷积。而前者表示稀疏化除了归一化层的参数。
 - **skip_params_func(function)** - 一个指向function的指针，该function定义了哪些参数不应该被剪裁，默认（None）时代表所有归一化层参数不参与剪裁。
 - **local_sparsity(bool)** - 剪裁比例（ratio）应用的范围：local_sparsity 开启时意味着每个参与剪裁的参数矩阵稀疏度均为 'ratio'， 关闭时表示只保证模型整体稀疏度达到'ratio'，但是每个参数矩阵的稀疏度可能存在差异。
+- **sparse_block(Array<Integer>)** - 一个含有两个正整数的数组，定义了稀疏化时候block的大小。即 sparse_block[0] x sparse_block[1]中的参数作为一个整体，要么被置0，要么保持不变。默认为 [1,1]，代表非结构化稀疏。
 - **configs(Dict)** - 传入额外的训练超参用以指导GMP训练过程。各参数介绍如下：
 
 .. code-block:: python

docs/zh_cn/api_cn/static/prune/unstructured_prune_api.rst

@@ -4,7 +4,7 @@
 UnstrucuturedPruner
 ----------
 
-.. py:class:: paddleslim.prune.UnstructuredPruner(program, mode, ratio=0.55, threshold=1e-2, scope=None, place=None, prune_params_type, skip_params_func=None, local_sparsity=False)
+.. py:class:: paddleslim.prune.UnstructuredPruner(program, mode, ratio=0.55, threshold=1e-2, scope=None, place=None, prune_params_type, skip_params_func=None, local_sparsity=False, sparse_block=[1,1])
 
 `源代码 <https://github.com/PaddlePaddle/PaddleSlim/blob/develop/paddleslim/prune/unstructured_pruner.py>`_
 
@@ -19,14 +19,14 @@ UnstrucuturedPruner
 - **scope(paddle.static.Scope)** - 一个paddle.static.Scope对象，存储了所有变量的数值，默认（None）时表示paddle.static.global_scope。
 - **place(CPUPlace|CUDAPlace)** - 模型执行的设备，类型为CPUPlace或者CUDAPlace，默认（None）时代表CPUPlace。
 - **prune_params_type(String)** - 用以指定哪些类型的参数参与稀疏。目前只支持None和"conv1x1_only"两个选项，后者表示只稀疏化1x1卷积。而前者表示稀疏化除了归一化的参数。
-- **skip_params_func(function)** - 一个指向function的指针，该function定义了哪些参数不应该被剪裁，默认（None）时代表所有归一化层参数不参与剪裁。
+- **skip_params_func(function)** - 一个指向function的指针，该function定义了哪些参数不应该被剪裁，默认（None）时代表所有归一化层参数和 bias 不参与剪裁。
 
 .. code-block:: python
 
   def _get_skip_params(program):
       """
       The function is used to get a set of all the skipped parameters when performing pruning.
-      By default, the normalization-related ones will not be pruned.
+      By default, the normalization-related ones and bias will not be pruned.
       Developers could replace it by passing their own function when initializing the UnstructuredPruner instance.
       Args:
         - program(paddle.static.Program): the current model.
@@ -39,11 +39,15 @@ UnstrucuturedPruner
           if 'norm' in op.type() and 'grad' not in op.type():
               for input in op.all_inputs():
                   skip_params.add(input.name())
+      for param in program.all_parameters():
+          if len(param.shape) == 1:
+              skip_params.add(param.name)
       return skip_params
 
 ..
 
 - **local_sparsity(bool)** - 剪裁比例（ratio）应用的范围：local_sparsity 开启时意味着每个参与剪裁的参数矩阵稀疏度均为 'ratio'， 关闭时表示只保证模型整体稀疏度达到'ratio'，但是每个参数矩阵的稀疏度可能存在差异。
+- **sparse_block(Array<Integer>)** - 一个含有两个正整数的数组，定义了稀疏化时候block的大小。即 sparse_block[0] x sparse_block[1]中的参数作为一个整体，要么被置0，要么保持不变。默认为 [1,1]，代表非结构化稀疏。
 
 **返回：** 一个UnstructuredPruner类的实例
 
@@ -339,6 +343,8 @@ GMPUnstrucuturedPruner
 - **prune_params_type(String)** - 用以指定哪些类型的参数参与稀疏。目前只支持None和"conv1x1_only"两个选项，后者表示只稀疏化1x1卷积。而前者表示稀疏化除了归一化的参数。
 - **skip_params_func(function)** - 一个指向function的指针，该function定义了哪些参数不应该被剪裁，默认（None）时代表所有归一化层参数不参与剪裁。
 - **local_sparsity(bool)** - 剪裁比例（ratio）应用的范围：local_sparsity 开启时意味着每个参与剪裁的参数矩阵稀疏度均为 'ratio'， 关闭时表示只保证模型整体稀疏度达到'ratio'，但是每个参数矩阵的稀疏度可能存在差异。
+- **sparse_block(Array<Integer>)** - 一个含有两个正整数的数组，定义了稀疏化时候block的大小。即 sparse_block[0] x sparse_block[1]中的参数作为一个整体，要么被置0，要么保持不变。默认为 [1,1]，代表非结构化稀疏。
+
 - **configs(Dict)** - 传入额外的训练超参用以指导GMP训练过程。具体描述如下：
 
 .. code-block:: python