torch.nn.utils.prune.global_unstructured
- torch.nn.utils.prune.global_unstructured(parameters, pruning_method, importance_scores=None, **kwargs)[源代码]
-
根据指定的
pruning_method
对所有参数对应的张量进行全局剪枝。通过以下方式就地修改模块:
-
添加一个名为
name+'_mask'
的缓冲区,该缓冲区对应于剪枝方法应用于参数name
的二进制掩码。 -
将参数
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)
-