Shortcuts

PruningContainer

class torch.nn.utils.prune.PruningContainer(*args)[源代码]

用于迭代剪枝的剪枝方法序列容器。

跟踪修剪方法应用的顺序,并处理连续修剪调用的组合。

接受一个 BasePruningMethod 实例或它们的可迭代对象作为参数。

add_pruning_method(method)[源代码]

向容器中添加一个子节点剪枝的 方法

Parameters

方法 (子类 属于 BasePruningMethod) – 要添加到容器中的子剪枝方法。

classmethod apply(module, name, *args, importance_scores=None, **kwargs)

动态添加剪枝和重新参数化张量。

添加前向预钩子,以实现动态剪枝并 在原始张量和剪枝掩码的基础上重新参数化张量。

Parameters
  • 模块 (nn.Module) – 包含要剪枝的张量的模块

  • 名称 (字符串) – 在 模块 中进行剪枝操作的参数名称。

  • args – 传递给 BasePruningMethod 子类的参数

  • importance_scores (torch.Tensor) – 重要性分数的张量(与模块参数形状相同),用于计算剪枝的掩码。 该张量中的值表示被剪枝参数中相应元素的重要性。 如果未指定或为None,则将使用参数本身。

  • kwargs – 传递给 BasePruningMethod 子类的关键字参数

apply_mask(module)

简单处理待剪枝参数与生成的掩码之间的乘法。

从模块中获取掩码和原始张量,并返回张量的剪枝版本。

Parameters

模块 (nn.Module) – 包含要剪枝的张量的模块

Returns

输入张量的剪枝版本

Return type

修剪的张量(torch.Tensor

compute_mask(t, default_mask)[源代码]

应用最新的 method 方法,通过计算新的部分掩码并返回其与 default_mask 的组合。

新的部分掩码应该在那些未被default_mask置零的条目或通道上进行计算。 新的掩码将根据张量t的哪些部分进行计算,取决于PRUNING_TYPE(由类型处理器处理):

  • 对于‘unstructured’,掩码将从非掩码条目的展平列表中计算得出;

  • 对于‘structured’,掩码将从张量中的非掩码通道计算得出;

  • 对于‘global’,掩码将在所有条目中计算。

Parameters
  • t (torch.Tensor) – 表示要剪枝的参数的张量 (与default_mask具有相同的维度)。

  • default_mask (torch.Tensor) – 来自上一次剪枝迭代的掩码。

Returns

新的掩码,结合了default_mask的效果和当前剪枝method中的新掩码(与default_maskt具有相同的维度)。

Return type

掩码 (torch.Tensor)

prune(t, default_mask=None, importance_scores=None)

计算并返回输入张量 t 的修剪版本。

根据compute_mask()中指定的剪枝规则。

Parameters
  • t (torch.Tensor) – 要剪枝的张量(与 default_mask 具有相同的维度)。

  • importance_scores (torch.Tensor) – 重要性分数的张量(与t形状相同),用于计算剪枝的掩码t。 该张量中的值表示正在剪枝的t中相应元素的重要性。 如果未指定或为None,则将使用张量t

  • default_mask (torch.Tensor, 可选) – 来自先前剪枝迭代的掩码(如果有)。在确定剪枝应作用于张量的哪一部分时需要考虑。如果为 None,则默认为全为 1 的掩码。

Returns

修剪后的张量版本 t

remove(module)

从模块中移除修剪重参数化。

名为 name 的参数永久地被剪枝, 并且名为 name+'_orig' 的参数从参数列表中移除。 同样,名为 name+'_mask' 的缓冲区从缓冲区中移除。

注意

剪枝本身不会被撤销或逆转!

优云智算