修剪模型后，在Pytorch中删除weight_orig

Question

在Pytorch中修剪模型后，保存的模型同时包含修剪的权重和weight_orig。这将导致修剪的模型大小大于未修剪的模型。 有没有办法删除weight_orig并减少​​修剪的模型大小？

Answer 1

如offcial documentation中所述，您可以为此目的使用torch.nn.utils.prune.remove()。
remove()根据weight_orig和weight_mask删除重新参数化，并删除forward_pre_hook。 您将像这样使用它：

for module in model.modules():
    if isinstance(module, torch.nn.Conv2d):
        prune.remove(module,'weight')
    # etc...