torch.nn.utils.prune.global_unstructured

torch.nn.utils.prune.global_unstructured(parameters, pruning_method, importance_scores=None, **kwargs)[源代码]

根据指定的 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) – 一个将 (模块, 名称) 元组映射到相应参数的重要性分数张量的字典。该张量应与参数具有相同的形状,并用于计算剪枝掩码。如果未指定或为 None,则使用参数本身作为其重要性分数。

  • kwargs – 其他关键字参数,例如:amount (int 或 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)
本页目录