2de52a8976971da4836727ba9242fedcc7474878,src/sdk/pynni/nni/compression/torch/pruners.py,SlimPruner,calc_mask,#SlimPruner#Any#Any#,202

Before Change

        

        weight = layer.module.weight.data
        op_name = layer.name
        op_type = layer.type
        assert op_type == "BatchNorm2d", "SlimPruner only supports 2d batch normalization layer pruning"
        if op_name in self.mask_calculated_ops:
            assert op_name in self.mask_dict
            return self.mask_dict.get(op_name)
        base_mask = torch.ones(weight.size()).type_as(weight).detach()
        mask = {"weight": base_mask.detach(), "bias": base_mask.clone().detach()}
        try:
            filters = weight.size(0)

After Change

        op_type = layer.type
        if_calculated = kwargs["if_calculated"]
        assert op_type == "BatchNorm2d", "SlimPruner only supports 2d batch normalization layer pruning"
        if if_calculated:
            return None
        base_mask = torch.ones(weight.size()).type_as(weight).detach()
        mask = {"weight": base_mask.detach(), "bias": base_mask.clone().detach()}
        filters = weight.size(0)
        num_prune = int(filters * config.get("sparsity"))
        if filters >= 2 and num_prune >= 1:
            w_abs = weight.abs()
            mask_weight = torch.gt(w_abs, self.global_threshold).type_as(weight)
            mask_bias = mask_weight.clone()
            mask = {"weight": mask_weight.detach(), "bias": mask_bias.detach()}
        if_calculated.copy_(torch.tensor(True)) // pylint: disable=not-callable
        return mask

class LotteryTicketPruner(Pruner):

Instances