c7d58033db0e25736d33406ed262cb5232d366e8,src/sdk/pynni/nni/compression/torch/pruners.py,AGP_Pruner,calc_mask,#AGP_Pruner#Any#Any#,88

Before Change

        

        weight = layer.module.weight.data
        op_name = layer.name
        start_epoch = config.get("start_epoch", 0)
        freq = config.get("frequency", 1)
        if self.now_epoch >= start_epoch and self.if_init_list.get(op_name, True) \
                and (self.now_epoch - start_epoch) % freq == 0:
            mask = self.mask_dict.get(op_name, {"weight": 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)[0].max()
            new_mask = {"weight": torch.gt(w_abs, threshold).type_as(weight)}
            self.mask_dict.update({op_name: new_mask})
            self.if_init_list.update({op_name: False})
        else:
            new_mask = self.mask_dict.get(op_name, {"weight": torch.ones(weight.shape).type_as(weight)})
        return new_mask

    def compute_target_sparsity(self, config):
        

After Change

        freq = config.get("frequency", 1)

        if_calculated = kwargs["if_calculated"]
        if if_calculated:
            return None
        if not (self.now_epoch >= start_epoch and (self.now_epoch - start_epoch) % freq == 0):
            return None

        mask = {"weight": 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()
        threshold = torch.topk(w_abs.view(-1), k, largest=False)[0].max()
        new_mask = {"weight": torch.gt(w_abs, threshold).type_as(weight)}
        if_calculated.copy_(torch.tensor(1)) // pylint: disable=not-callable

        return new_mask

Instances