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

Before Change

            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

After Change

        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