如何在 PyTorch 中修剪小于阈值的权重？

Question

如何修剪小于阈值的 CNN（卷积神经网络）模型的权重（让我们考虑修剪所有

对于在 pytorch 中保存为 .pth 格式的权重文件，我们如何实现这一点？

Answer 1

PyTorch 因为1.4.0 提供开箱即用的模型修剪，see official tutorial。

由于目前 PyTorch 中没有 threshold 修剪方法，因此您必须自己实现它，不过一旦您掌握了整体思路，这会很容易。

阈值剪枝方法

下面是执行剪枝的代码：

from torch.nn.utils import prune


class ThresholdPruning(prune.BasePruningMethod):
    PRUNING_TYPE = "unstructured"

    def __init__(self, threshold):
        self.threshold = threshold

    def compute_mask(self, tensor, default_mask):
        return torch.abs(tensor) > self.threshold

解释：

• PRUNING_TYPE 可以是 global、structured、unstructured 之一。 global 作用于整个模块（例如，删除权重最小的20%），structured 作用于整个通道/模块。我们需要unstructured，因为我们想修改特定参数张量中的每个连接（比如weight 或bias）
• __init__ - 在这里传递任何你想要或需要让它工作的东西，普通的东西
• compute_mask - 用于修剪特定张量的掩码。在我们的例子中，所有低于阈值的参数都应该为零。我是用绝对值来做的，因为它更有意义。 default_mask 在这里不需要，但保留为命名参数，因为这是 API 需要 atm 的。

此外，从prune.BasePruningMethod 继承定义了将掩码应用于每个参数、使修剪永久化等方法。有关详细信息，请参阅base class docs。

示例模块

没什么花哨的，你可以放任何你想要的东西：

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.first = torch.nn.Linear(50, 30)
        self.second = torch.nn.Linear(30, 10)

    def forward(self, inputs):
        return self.second(torch.relu(self.first(inputs)))


module = MyModule()

您也可以通过module = torch.load('checkpoint.pth') 加载您的模块 如果你需要，没关系。

修剪模块的参数

我们应该定义应该修剪我们模块的哪个参数（以及它是weight 还是bias），如下所示：

parameters_to_prune = ((module.first, "weight"), (module.second, "weight"))

现在，我们可以将globally 我们的unstructured 修剪应用到所有定义的parameters（threshold 作为kwarg 传递给__init__ 的ThresholdPruning）：

prune.global_unstructured(
    parameters_to_prune, pruning_method=ThresholdPruning, threshold=0.1
)

结果

weight属性

要查看效果，只需检查first子模块的权重：

print(module.first.weight)

这是一个应用了我们的剪枝技术的权重，但是请注意它不再是torch.nn.Parameter了！现在它只是我们模型的一个属性，因此它不会占用目前参与培训或评估。

weight_mask

我们可以通过module.first.weight_mask 检查创建的掩码以查看一切是否正确（在这种情况下它将是二进制的）。

weight_orig

应用剪枝会创建一个新的torch.nn.Parameter，其原始权重名为name + _orig，在本例中为weight_orig，让我们看看：

print(module.first.weight_orig)

目前将在训练和评估期间使用此参数！。通过上述方法应用pruning 后，添加了forward_pre_hooks，将原来的weight“切换”到weight_orig。

由于这种方法，您可以在training 或inference 的任何部分定义和应用您的修剪，而不会“破坏”原始权重。

永久应用修剪

如果您想永久应用修剪，只需发出：

prune.remove(module.first, "weight")

现在我们的module.first.weight 再次成为参数，并适当地修剪了条目，module.first.weight_mask 被删除，module.first.weight_orig 也被删除。 这可能是您所追求的。

您可以迭代 children 以使其永久化：

for child in module.children():
    prune.remove(child, "weight")

您可以使用相同的逻辑定义parameters_to_prune：

parameters_to_prune = [(child, "weight") for child in module.children()]

或者，如果您只想修剪 convolution 层（或其他任何东西）：

parameters_to_prune = [
    (child, "weight")
    for child in module.children()
    if isinstance(child, torch.nn.Conv2d)
]

优势

• 使用“PyTorch 修剪方式”，因此可以更轻松地将您的意图传达给其他程序员
• 在每个张量的基础上定义修剪，单一职责而不是遍历所有内容
• 仅限于预定义的方式
• 修剪不是永久性的，因此您可以在需要时从中恢复。可以使用修剪掩码和原始权重保存模块，以便为您留出一些空间来恢复最终的错误（例如，threshold 太高，现在您所有的权重都为零，渲染结果毫无意义）
• 在forward 调用期间使用原始权重，除非您想最终更改为修剪后的版本（简单调用remove）

缺点

• IMO 修剪 API 可能更清晰
• 你可以做得更短（由Shai提供）
• 可能会让那些不知道 PyTorch “定义”这种东西的人感到困惑（仍然有教程和文档，所以我认为这不是一个大问题）

Answer 2

您可以直接处理保存在state_dict 中的值：

sd = torch.load('saved_weights.pth')  # load the state dicd
for k in sd.keys():
  if not 'weight' in k:
    continue  # skip biases and other saved parameters
  w = sd[k]
  sd[k] = w * (w > thr)  # set to zero weights smaller than thr 
torch.save(sd, 'pruned_weights.pth')