fd551c8611033e2c9e9f350f13fc3bce8c7fd32e,src/sdk/pynni/nni/compression/torch/builtin_pruners.py,AGP_Pruner,calc_mask,#AGP_Pruner#,51

Before Change

        self.now_epoch = 1

    def calc_mask(self, weight, config, op_name, **kwargs):
        mask = self.mask_list.get(op_name, torch.ones(weight.shape))
        target_sparsity = self.compute_target_sparsity(config)
        k = int(weight.numel() * target_sparsity)
        if k == 0 or target_sparsity >= 1 or target_sparsity <= 0:
            return mask
        // if we want to generate new mask, we should update weigth first 
        w_abs = weight.abs()*mask
        threshold = torch.topk(w_abs.view(-1), k, largest = False).values.max()
        new_mask = torch.gt(w_abs, threshold).type(weight.type())
        self.mask_list[op_name] = new_mask
        return new_mask

After Change

        self.now_epoch = 1

    def calc_mask(self, weight, config, op_name, **kwargs):
        mask = self.mask_list.get(op_name, torch.ones(weight.shape).type_as(weight))
        target_sparsity = self.compute_target_sparsity(config)
        k = int(weight.numel() * target_sparsity)
        if k == 0 or target_sparsity >= 1 or target_sparsity <= 0:
            return mask
        // if we want to generate new mask, we should update weigth first 
        w_abs = weight.abs() * mask
        threshold = torch.topk(w_abs.view(-1), k, largest=False).values.max()
        new_mask = torch.gt(w_abs, threshold).type_as(weight)
        self.mask_list[op_name] = new_mask
        return new_mask

Instances