PruningContainer

类 torch.nn.utils.prune.PruningContainer(*args)[source][source] ¶

容器包含一系列用于迭代剪枝的剪枝方法。

跟踪剪枝方法应用的顺序，并处理连续的剪枝调用。

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

add_pruning_method(method)

向容器中添加一个子剪枝 method 。

参数:

方法（BasePruningMethod 的子类）- 要添加到容器中的子剪枝方法。

@classmethod apply(module, name, *args, importance_scores=None, **kwargs)[source] ¶

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

添加了前向预钩子，该钩子使动态剪枝和张量的重新参数化（基于原始张量和剪枝掩码）成为可能。

参数:

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

• 名称（str）- 在 module 中参数名称，剪枝将在此参数上执行。

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

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

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

apply_mask(module)

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

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

参数:

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

返回值:

输入张量的修剪版本

返回类型:

pruned_tensor (torch.Tensor)

compute_mask(t, default_mask)

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

新的局部掩码应计算在未被 default_mask 清零的条目或通道上。新的掩码将从张量 t 的哪些部分计算取决于 PRUNING_TYPE （由类型处理器处理）：

• 对于‘非结构化’，掩码将从非掩码条目的展开列表中计算；

• 对于‘结构化’，掩码将从张量中的非掩码通道计算；

• 对于‘全局’，掩码将计算所有条目。

参数:

• t (torch.Tensor) – 表示要剪枝的参数的 tensor（与 default_mask 的维度相同）。

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

返回值:

新的掩码，结合了 default_mask 和当前剪枝迭代的新的掩码 method 的效果（与 default_mask 和 t 的维度相同）。

返回类型:

掩码（torch.Tensor）

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

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

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

参数:

• t (torch.Tensor) – 要剪枝的 tensor（与 default_mask 维度相同）。

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

• default_mask (torch.Tensor, 可选) – 前一次剪枝迭代的掩码，如果有的话。在确定剪枝应作用的 tensor 部分时需要考虑。如果为 None，则默认为全一的掩码。

返回值:

张量 t 的修剪版本。

remove(module)[源代码] ¶

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

被修剪的参数 name 将永久修剪，而参数 name+'_orig' 将从参数列表中移除。同样，缓冲区 name+'_mask' 也将从缓冲区中移除。

注意

自剪枝操作不可撤销或逆转！