9d468d2c742491af2d2f506c648ddc95ffea6a64,src/sdk/pynni/nni/compression/torch/builtin_pruners.py,SensitivityPruner,calc_mask,#SensitivityPruner#Any#Any#Any#,107

Before Change

        mask = self.mask_list.get(op_name, torch.ones(weight.shape).type_as(weight))
        // if we want to generate new mask, we should update weight first
        weight = weight * mask
        target_sparsity = config["sparsity"] * torch.std(weight).item()
        k = int(weight.numel() * target_sparsity)
        if k == 0:
            return mask

After Change

        self.if_init_list = {}

    def calc_mask(self, weight, config, op_name, **kwargs):
        if self.if_init_list.get(op_name, True):
            w_abs = weight.abs()
            k = int(weight.numel() * config["sparsity"])
            if k == 0:
                return torch.ones(weight.shape).type_as(weight)
            threshold = torch.topk(w_abs.view(-1), k, largest=False).values.max()
            mask = torch.gt(w_abs, threshold).type_as(weight)
            self.mask_list.update({op_name: mask})
            self.if_init_list.update({op_name: False})
        else:
            mask = self.mask_list[op_name]
        return mask


class AGP_Pruner(Pruner):

Instances