评价此页

torch.nn.utils.prune.global_unstructured#

torch.nn.utils.prune.global_unstructured(parameters, pruning_method, importance_scores=None, **kwargs)[source]#

全局地剪枝 parameters 中所有参数对应的张量,通过应用指定的 pruning_method 来实现。

通过以下方式就地修改模块:

  1. 添加一个名为 name+'_mask' 的命名缓冲区,对应于剪枝方法应用于参数 name 的二进制掩码。

  2. 用其剪枝后的版本替换参数 name,同时将原始(未剪枝)参数存储在新命名为 name+'_orig' 的参数中。

参数
  • parameters (Iterable of (module, name) tuples) – 要全局剪枝的模型参数,即在决定剪枝哪些权重之前聚合所有权重。module 必须是 nn.Module 类型,name 必须是字符串。

  • pruning_method (function) – 来自此模块的有效剪枝函数,或用户实现的自定义函数,该函数满足实现指南并具有 PRUNING_TYPE='unstructured'

  • importance_scores (dict) – 一个字典,将 (module, name) 元组映射到相应的参数的重要性分数张量。张量应与参数的形状相同,并用于计算剪枝掩码。如果未指定或为 None,则参数将与其重要性分数一起使用。

  • kwargs – 其他关键字参数,例如:amount (int or float): 要在指定参数中剪枝的参数数量。如果为 float,应在 0.0 和 1.0 之间,表示要剪枝的参数分数。如果为 int,则表示要剪枝的参数绝对数量。

引发

TypeError – 如果 PRUNING_TYPE != 'unstructured'

注意

由于全局结构化剪枝除非对范数按参数大小进行归一化,否则没有太大意义,因此我们现在将全局剪枝的范围限制在非结构化方法。

示例

>>> from torch.nn.utils import prune
>>> from collections import OrderedDict
>>> net = nn.Sequential(
...     OrderedDict(
...         [
...             ("first", nn.Linear(10, 4)),
...             ("second", nn.Linear(4, 1)),
...         ]
...     )
... )
>>> parameters_to_prune = (
...     (net.first, "weight"),
...     (net.second, "weight"),
... )
>>> prune.global_unstructured(
...     parameters_to_prune,
...     pruning_method=prune.L1Unstructured,
...     amount=10,
... )
>>> print(sum(torch.nn.utils.parameters_to_vector(net.buffers()) == 0))
tensor(10)