如何通过pytorch修剪k%的最低权重？ [关闭]

Question

这里我从名为 Deep compression [Han et. al.] 使用 resnet18

我还使用以下代码，权重乘以掩码，以便它是被 k% 最低权重修剪为零的 after_weight。但是该代码对我不起作用。 任何有效的解决方案？

prune = float(0.1)
def prune_weights(torchweights):
    weights=np.abs(torchweights.cpu().numpy());
    weightshape=weights.shape
    rankedweights=weights.reshape(weights.size).argsort()#.reshape(weightshape)
    
    num = weights.size
    prune_num = int(np.round(num*prune))
    count=0
    masks = np.zeros_like(rankedweights)
    for n, rankedweight in enumerate(rankedweights):
        if rankedweight > prune_num:
            masks[n]=1
        else: count+=1
    print("total weights:", num)
    print("weights pruned:",count)
    
    masks=masks.reshape(weightshape)
    weights=masks*weights
    
    return torch.from_numpy(weights).cuda(), masks

# prune weights
# The pruned weight location is saved in the addressbook and maskbook.
# These will be used during training to keep the weights zero.
addressbook=[]
maskbook=[]
for k, v in net.state_dict().items():
    if "conv2" in k:
        addressbook.append(k)
        print("pruning layer:",k)
        weights=v
        weights, masks = prune_weights(weights)
        maskbook.append(masks)
        checkpoint['net'][k] = weights
        
checkpoint['address'] = addressbook
checkpoint['mask'] = maskbook
net.load_state_dict(checkpoint['net'])

Answer 1

您可以使用torch.nn.utils.prune。

您似乎想删除每个Conv2D 层的 10%。如果是这种情况，您可以这样做：

import torch
import torch.nn.utils.prune as prune

# load your model
net = ?

# in your example, you want to remove 10%
prune_perc = 0.1

for name, module in net.named_modules():
    if isinstance(module, torch.nn.Conv2d):
        prune.l1_unstructured(module, name='weight', amount=prune_perc)