Save mask for the unstructured pruned model (for ft-attack evaluation). `neg_prune`: - if `args.neg_prune` is False (bottom pruning), save the mask as True for the weights not pruned. - if `args.neg_prune` is True (top pruning), save the mask as True for the pruned weights.
(model, neg_prune=False)
| 2039 | |
| 2040 | |
| 2041 | def get_mask(model, neg_prune=False): |
| 2042 | """ |
| 2043 | Save mask for the unstructured pruned model (for ft-attack evaluation). |
| 2044 | `neg_prune`: |
| 2045 | - if `args.neg_prune` is False (bottom pruning), save the mask as True for the weights not pruned. |
| 2046 | - if `args.neg_prune` is True (top pruning), save the mask as True for the pruned weights. |
| 2047 | """ |
| 2048 | use_cache = model.config.use_cache |
| 2049 | model.config.use_cache = False |
| 2050 | |
| 2051 | mask = {} |
| 2052 | |
| 2053 | mask_num = 0 |
| 2054 | total_num = 0 |
| 2055 | for name, module in model.named_modules(): |
| 2056 | if hasattr(module, "weight"): |
| 2057 | mask[name] = module.weight.data.abs().lt(1e-8).to("cpu").detach() |
| 2058 | if neg_prune is False: |
| 2059 | mask[name] = ~mask[name] |
| 2060 | |
| 2061 | mask_num += mask[name].eq(True).int().sum() |
| 2062 | total_num += mask[name].numel() |
| 2063 | |
| 2064 | print(f"{(100 * mask_num / total_num):.2f}% entries are True in mask.") |
| 2065 | return mask |
| 2066 | |
| 2067 | |
| 2068 | def prune_attention_head( |