Support custom functions when computing pruned substructure of parameter

1. Support custom criterion of importance
2. Support custom importance sorting strategy

paddleslim/prune/__init__.py

@@ -27,7 +27,11 @@ from .prune_io import *
 from ..prune import prune_io
 from .group_param import *
 from ..prune import group_param
+from .criterion import *
+from ..prune import criterion
 
+from .importance_sort import *
+from ..prune import importance_sort
 __all__ = []
 
 __all__ += pruner.__all__
@@ -37,3 +41,5 @@ __all__ += sensitive.__all__
 __all__ += prune_walker.__all__
 __all__ += prune_io.__all__
 __all__ += group_param.__all__
+__all__ += criterion.__all__
+__all__ += importance_sort.__all__

paddleslim/prune/criterion.py

+"""Define some functions to compute the importance of structure to be pruned.
+"""
+# Copyright (c) 2020  PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import numpy as np
+from ..common import get_logger
+
+__all__ = ["l1_norm"]
+
+_logger = get_logger(__name__, level=logging.INFO)
+
+
+def l1_norm(group):
+    """Compute l1-norm scores of parameter on given axis.
+
+    This function return a list of parameters' l1-norm scores on given axis.
+    Each element of list is a tuple with format (name, axis, score) in which 'name' is parameter's name
+    and 'axis' is the axis reducing on and `score` is a np.array storing the l1-norm of strucure on `axis`.
+
+    Args:
+       group(list): A group of parameters. The first parameter of the group is convolution layer's weight
+                    while the others are parameters affected by pruning the first one. Each parameter in group
+                    is represented as tuple '(name, values, axis)' in which `name` is the parameter's name and
+                    and `values` is the values of parameter and `axis` is the axis reducing on pruning on.
+    Returns:
+       list: A list of tuple storing l1-norm on given axis.
+    """
+    scores = []
+    for name, value, axis in group:
+
+        reduce_dims = [i for i in range(len(value.shape)) if i != axis]
+        score = np.sum(np.abs(value), axis=tuple(reduce_dims))
+        scores.append((name, axis, score))
+
+    return scores

paddleslim/prune/importance_sort.py

+"""Define some functions to sort substructures of parameter by importance.
+"""
+# Copyright (c) 2020  PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from ..core import GraphWrapper
+from ..common import get_logger
+
+__all__ = ["channel_score_sort", "batch_norm_scale_sort"]
+
+
+def channel_score_sort(group, graph):
+    """Sort channels of convolution by importance.
+
+    This function return a list of parameters' sorted indexes on given axis.
+    Each element of list is a tuple with format (name, axis, indexes) in which 'name' is parameter's name
+    and 'axis' is the axis pruning on and `indexes` is sorted indexes.
+
+    The sorted indexes is computed by below steps:
+
+    step1: Find the first convolution layer in given group.
+    step2: Get the scores of first convolution's channels.
+    step3: Get sorted indexes by calling scores.argsort().
+    step4: All the parameters in group share the same sorted indexes computed in step3.
+
+    Args:
+       group(list): A group of parameters. The first parameter of the group is convolution layer's weight
+                    while the others are parameters affected by pruning the first one. Each parameter in group
+                    is represented as tuple '(name, axis, score)' in which `name` is the parameter's name and
+                    `axis` is the axis pruning on and `score` is a np.array storing the importance of strucure
+                    on `axis`. Show as below:
+
+                    .. code-block: text
+
+                       [("conv1_weights", 0, [0.7, 0.5, 0.6]), ("conv1_bn.scale", 0, [0.1, 0.2, 0.4])]
+
+                    The shape of "conv1_weights" is `[out_channel, in_channel, filter_size, filter_size]`, so
+                    `[0.7, 0.5, 0.6]` are the importance sores of each output channel in "conv1_weights"
+                    while axis is 0.
+     
+
+       graph(GraphWrapper): The graph is an auxiliary for sorting. It won't be used in this function.
+
+    Returns:
+
+       list: sorted indexes
+
+    """
+    assert (isinstance(graph, GraphWrapper))
+    name, axis, score = group[
+        0]  # sort channels by the first convolution's score
+    sorted_idx = score.argsort()
+    idxs = []
+    for name, axis, score in group:
+        idxs.append((name, axis, sorted_idx))
+    return idxs
+
+
+def batch_norm_scale_sort(group, graph):
+    """Sort channels of convolution by scales in batch norm layer.
+
+    This function return a list of parameters' sorted indexes on given axis.
+    Each element of list is a tuple with format (name, axis, indexes) in which 'name' is parameter's name
+    and 'axis' is the axis pruning on and `indexes` is sorted indexes.
+
+    The sorted indexes is computed by below steps:
+
+    step1: Find the batch norm layer after the first convolution in given group.
+    step2: Get the scales of the batch norm layer.
+    step3: Get sorted indexes by calling `scales.argsort()`.
+    step4: All the parameters in group share the same sorted indexes computed in step3.
+
+    Args:
+       group(list): A group of parameters. The first parameter of the group is convolution layer's weight
+                    while the others are parameters affected by pruning the first one. Each parameter in group
+                    is represented as tuple '(name, axis, score)' in which `name` is the parameter's name and
+                    `axis` is the axis pruning on and `score` is a np.array storing the importance of strucure
+                    on `axis`. Show as below:
+
+                    .. code-block: text
+
+                       [("conv1_weights", 0, [0.7, 0.5, 0.6]), ("conv1_bn.scale", 0, [0.1, 0.2, 0.4])]
+
+                    The shape of "conv1_weights" is `[out_channel, in_channel, filter_size, filter_size]`, so
+                    `[0.7, 0.5, 0.6]` are the importance sores of each output channel in "conv1_weights"
+                    while axis is 0.
+     
+
+       graph(GraphWrapper): The graph is an auxiliary for sorting. It is used to find
+                            the batch norm layer after given convolution layer.
+
+    Returns:
+       list: sorted indexes
+    """
+    assert (isinstance(graph, GraphWrapper))
+    # step1: Get first convolution
+    conv_weight, axis, score = group[0]
+    param_var = graph.var(conv_weight)
+    conv_op = param_var.outputs()[0]
+
+    # step2: Get bn layer after first convolution
+    conv_output = conv_op.outputs("Output")[0]
+    bn_op = conv_output.outputs()[0]
+    if bn_op is not None:
+        bn_scale_param = bn_op.inputs("Scale")[0].name()
+    else:
+        raise SystemExit("Can't find BatchNorm op after Conv op in Network.")
+
+    # steps3: Find score of bn and compute sorted indexes 
+    sorted_idx = None
+    for name, axis, score in group:
+        if name == bn_scale_param:
+            sorted_idx = score.argsort()
+            break
+
+    # step4: Share the sorted indexes with all the parameter in group
+    idxs = []
+    if sorted_idx is not None:
+        for name, axis, score in group:
+            idxs.append((name, axis, sorted_idx))
+    return idxs

paddleslim/prune/pruner.py

@@ -18,7 +18,9 @@ import numpy as np
 import paddle.fluid as fluid
 import copy
 from ..core import VarWrapper, OpWrapper, GraphWrapper
-from .prune_walker import conv2d as conv2d_walker
+from .group_param import collect_convs
+from .criterion import l1_norm
+from .importance_sort import channel_score_sort, batch_norm_scale
 from ..common import get_logger
 
 __all__ = ["Pruner"]
@@ -30,12 +32,21 @@ class Pruner():
     """The pruner used to prune channels of convolution.
 
     Args:
-        criterion(str): the criterion used to sort channels for pruning. It only supports 'l1_norm' currently.
+        criterion(str|function): the criterion used to sort channels for pruning.
+        channel_sortor(str|function): 
 
     """
 
-    def __init__(self, criterion="l1_norm"):
+    def __init__(self, criterion="l1_norm", channel_sortor="channel_score"):
         self.criterion = criterion
+        self.channel_sortor = channel_sortor
+        if criterion == "l1_norm":
+            self.criterion = l1_norm
+
+        if channel_sortor == "channel_score":
+            self.channel_sortor = channel_score_sort
+        elif channel_sortor == "batch_norm_scale":
+            self.channel_sortor = batch_norm_scale_sort
 
     def prune(self,
               program,
@@ -75,26 +86,40 @@ class Pruner():
         visited = {}
         pruned_params = []
         for param, ratio in zip(params, ratios):
+            group = collect_convs([param], graph)[0]  # [(name, axis)]
             if only_graph:
+
                 param_v = graph.var(param)
                 pruned_num = int(round(param_v.shape()[0] * ratio))
                 pruned_idx = [0] * pruned_num
+                for name, aixs in group:
+                    pruned_params.append((name, axis, pruned_idx))
+
             else:
-                pruned_idx = self._cal_pruned_idx(
-                    graph, scope, param, ratio, axis=0)
-            param = graph.var(param)
-            conv_op = param.outputs()[0]
-            walker = conv2d_walker(
-                conv_op, pruned_params=pruned_params, visited=visited)
-            walker.prune(param, pruned_axis=0, pruned_idx=pruned_idx)
+
+                group_values = []
+                for name, axis in group:
+                    values = np.array(scope.find_var(name).get_tensor())
+                    group_values.append((name, values, axis))
+
+                scores = self.criterion(
+                    group_with_values)  # [(name, axis, score)]
+
+                group_idx = self.channel_sortor(
+                    scores, graph=graph)  # [(name, axis, soted_idx)]
+                for param, pruned_axis, pruned_idx in group_idx:
+                    pruned_num = len(pruned_idx) * ratio
+                    pruned_params.append((
+                        param, pruned_axis,
+                        pruned_idx[:pruned_num]))  # [(name, axis, pruned_idx)]
 
         merge_pruned_params = {}
         for param, pruned_axis, pruned_idx in pruned_params:
-            if param.name() not in merge_pruned_params:
-                merge_pruned_params[param.name()] = {}
-            if pruned_axis not in merge_pruned_params[param.name()]:
-                merge_pruned_params[param.name()][pruned_axis] = []
-            merge_pruned_params[param.name()][pruned_axis].append(pruned_idx)
+            if param not in merge_pruned_params:
+                merge_pruned_params[param] = {}
+            if pruned_axis not in merge_pruned_params[param]:
+                merge_pruned_params[param][pruned_axis] = []
+            merge_pruned_params[param][pruned_axis].append(pruned_idx)
 
         for param_name in merge_pruned_params:
             for pruned_axis in merge_pruned_params[param_name]: