pruner.py 7.4 KB
Newer Older
W
wanghaoshuang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
# Copyright (c) 2019  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.

15
import logging
W
wanghaoshuang 已提交
16 17
import numpy as np
import paddle.fluid as fluid
18 19
import copy
from ..core import VarWrapper, OpWrapper, GraphWrapper
W
whs 已提交
20
from .prune_walker import conv2d as conv2d_walker
21
from ..common import get_logger
W
wanghaoshuang 已提交
22

23
__all__ = ["Pruner"]
W
wanghaoshuang 已提交
24

25 26
_logger = get_logger(__name__, level=logging.INFO)

W
wanghaoshuang 已提交
27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44

class Pruner():
    def __init__(self, criterion="l1_norm"):
        """
        Args:
            criterion(str): the criterion used to sort channels for pruning.
                            It only supports 'l1_norm' currently.
        """
        self.criterion = criterion

    def prune(self,
              program,
              scope,
              params,
              ratios,
              place=None,
              lazy=False,
              only_graph=False,
W
wanghaoshuang 已提交
45 46
              param_backup=False,
              param_shape_backup=False):
W
wanghaoshuang 已提交
47 48 49 50 51 52 53 54 55 56 57 58
        """
        Pruning the given parameters.
        Args:
            program(fluid.Program): The program to be pruned.
            scope(fluid.Scope): The scope storing paramaters to be pruned.
            params(list<str>): A list of parameter names to be pruned.
            ratios(list<float>): A list of ratios to be used to pruning parameters.
            place(fluid.Place): The device place of filter parameters. Defalut: None.
            lazy(bool): True means setting the pruned elements to zero.
                        False means cutting down the pruned elements. Default: False.
            only_graph(bool): True means only modifying the graph.
                              False means modifying graph and variables in scope. Default: False.
W
wanghaoshuang 已提交
59 60
            param_backup(bool): Whether to return a dict to backup the values of parameters. Default: False.
            param_shape_backup(bool): Whether to return a dict to backup the shapes of parameters. Default: False.
W
wanghaoshuang 已提交
61 62
        Returns:
            Program: The pruned program.
W
wanghaoshuang 已提交
63 64
            param_backup: A dict to backup the values of parameters.
            param_shape_backup: A dict to backup the shapes of parameters.
W
wanghaoshuang 已提交
65 66 67 68
        """

        self.pruned_list = []
        graph = GraphWrapper(program.clone())
W
wanghaoshuang 已提交
69 70
        param_backup = {} if param_backup else None
        param_shape_backup = {} if param_shape_backup else None
W
wanghaoshuang 已提交
71

W
wanghaoshuang 已提交
72
        visited = {}
W
whs 已提交
73
        pruned_params = []
W
wanghaoshuang 已提交
74
        for param, ratio in zip(params, ratios):
W
whs 已提交
75 76 77 78 79 80 81
            if only_graph:
                param_v = graph.var(param)
                pruned_num = int(round(param_v.shape()[0] * ratio))
                pruned_idx = [0] * pruned_num
            else:
                param_t = np.array(scope.find_var(param).get_tensor())
                pruned_idx = self._cal_pruned_idx(param_t, ratio, axis=0)
W
wanghaoshuang 已提交
82
            param = graph.var(param)
W
whs 已提交
83
            conv_op = param.outputs()[0]
W
whs 已提交
84 85
            walker = conv2d_walker(
                conv_op, pruned_params=pruned_params, visited=visited)
W
whs 已提交
86 87 88 89 90 91 92 93 94 95 96 97
            walker.prune(param, pruned_axis=0, pruned_idx=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)

        for param_name in merge_pruned_params:
            for pruned_axis in merge_pruned_params[param_name]:
W
whs 已提交
98 99
                pruned_idx = np.concatenate(merge_pruned_params[param_name][
                    pruned_axis])
W
whs 已提交
100
                param = graph.var(param_name)
W
whs 已提交
101 102 103 104 105 106 107 108 109
                if not lazy:
                    _logger.debug("{}\t{}\t{}".format(param.name(
                    ), pruned_axis, len(pruned_idx)))
                    if param_shape_backup is not None:
                        origin_shape = copy.deepcopy(param.shape())
                        param_shape_backup[param.name()] = origin_shape
                    new_shape = list(param.shape())
                    new_shape[pruned_axis] -= len(pruned_idx)
                    param.set_shape(new_shape)
W
whs 已提交
110 111
                if not only_graph:
                    param_t = scope.find_var(param.name()).get_tensor()
W
whs 已提交
112 113 114 115
                    if param_backup is not None and (
                            param.name() not in param_backup):
                        param_backup[param.name()] = copy.deepcopy(
                            np.array(param_t))
W
whs 已提交
116 117 118 119 120 121 122
                    try:
                        pruned_param = self._prune_tensor(
                            np.array(param_t),
                            pruned_idx,
                            pruned_axis=pruned_axis,
                            lazy=lazy)
                    except IndexError as e:
W
whs 已提交
123 124 125
                        _logger.error("Pruning {}, but get [{}]".format(
                            param.name(), e))

W
whs 已提交
126
                    param_t.set(pruned_param, place)
W
whs 已提交
127
        graph.update_groups_of_conv()
W
whs 已提交
128
        return graph.program, param_backup, param_shape_backup
W
wanghaoshuang 已提交
129

W
whs 已提交
130
    def _cal_pruned_idx(self, param, ratio, axis):
W
wanghaoshuang 已提交
131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176
        """
        Calculate the index to be pruned on axis by given pruning ratio.
        Args:
            name(str): The name of parameter to be pruned.
            param(np.array): The data of parameter to be pruned.
            ratio(float): The ratio to be pruned.
            axis(int): The axis to be used for pruning given parameter.
                       If it is None, the value in self.pruning_axis will be used.
                       default: None.
        Returns:
            list<int>: The indexes to be pruned on axis.
        """
        prune_num = int(round(param.shape[axis] * ratio))
        reduce_dims = [i for i in range(len(param.shape)) if i != axis]
        if self.criterion == 'l1_norm':
            criterions = np.sum(np.abs(param), axis=tuple(reduce_dims))
        pruned_idx = criterions.argsort()[:prune_num]
        return pruned_idx

    def _prune_tensor(self, tensor, pruned_idx, pruned_axis, lazy=False):
        """
        Pruning a array by indexes on given axis.
        Args:
            tensor(numpy.array): The target array to be pruned.
            pruned_idx(list<int>): The indexes to be pruned.
            pruned_axis(int): The axis of given array to be pruned on. 
            lazy(bool): True means setting the pruned elements to zero.
                        False means remove the pruned elements from memory.
                        default: False.
        Returns:
            numpy.array: The pruned array.
        """
        mask = np.zeros(tensor.shape[pruned_axis], dtype=bool)
        mask[pruned_idx] = True

        def func(data):
            return data[~mask]

        def lazy_func(data):
            data[mask] = 0
            return data

        if lazy:
            return np.apply_along_axis(lazy_func, pruned_axis, tensor)
        else:
            return np.apply_along_axis(func, pruned_axis, tensor)