剪枝模型大小与未剪枝模型相同[PyTorch]

Question

我正在尝试使用 ResNet18 在 PyTorch 上实现模型修剪。给定 ResNet18 的实例，我运行以下代码来加载预训练的模型，对其进行修剪并保存修剪后的模型：

def random_unstructured_pruning(pruning_rate: float, device, log_file):
    trained_model=retrieve_file(folder="./models", file_name='trained_model.pth')
    model=ResNet18(num_classes=10, input_channels=1).to(device)
    model.load_state_dict(torch.load(trained_model))
    
    modules_list=filter(lambda x: isinstance(x[1], (nn.Conv2d, nn.Linear, nn.BatchNorm2d)), model.named_modules())
    modules_list = map(lambda x: (x[1], 'weight'), modules_list)
    modules_list=tuple(modules_list)
 
  
    prune.global_unstructured(modules_list, pruning_method=prune.L1Unstructured, amount=0.8)
    for module in modules_list:
        prune.remove(module[0], module[1])
        
        
        
    pruning_rate_str= "{:02d}".format(int(pruning_rate * 10))
    path=f"{model_saving_path}pruned_{pruning_rate_str}.pth"
    # 
    torch.save(model.state_dict(), f"{path}")

在上述函数的末尾，.pth 文件与我在开始时加载的文件具有相同的尺寸，但我希望它更小，因为我修剪了 80% 的权重。

有人可以解释一下为什么会发生这种情况吗？我有什么错吗？ 谢谢！！

我认为问题出在函数的保存部分，似乎我总是保存在开始时重新加载的相同模型，并且修剪无效。

Answer 1

Pytorch 中的修剪实用程序充当接收修剪的层上的掩蔽包装器。这意味着您仍然可以访问原始模型权重，并且网络大小保持不变（如果不是更大的话），因为每个剪枝张量的掩码都已初始化。

如果您查看

prune.global_unstructured

的文档页面：

通过以下方式修改模块：

• 添加一个名为

  name+'_mask'

  的命名缓冲区，对应于二进制文件 通过修剪方法应用于参数名称的掩码。

• 用其修剪后的版本替换参数名称，而原始的 （未修剪的）参数存储在名为

  name+'_orig'

  的新参数中。

这是一个最小的例子，表明未剪枝的权重仍然可以访问：

net = nn.Sequential(OrderedDict(
    f1=nn.Linear(10, 5),
    f2=nn.Linear(5, 1)))

pruned = ((net.f1, 'weight'),)
prune.global_unstructured(pruned, prune.L1Unstructured, amount=0.8)

然后您可以访问修剪后的权重：

>>> net.f1.weight
tensor([[0.0000, 0.0000, 0.0000, -0.0000, 0.0000],
        [0.0000, 0.3599, 0.0000, -0.0000, 0.4034]])

原始未剪枝权重：

>>> net.f1.weight_orig
Parameter containing:
tensor([[ 0.1312,  0.1105,  0.0910, -0.2650,  0.3439],
        [ 0.0412,  0.3599,  0.2040, -0.2672,  0.4034]])

还有修剪面膜：

>>> net.f1.weight_mask
tensor([[0., 0., 0., 0., 0.],
        [0., 1., 0., 0., 1.]])