修剪容器

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

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

记录剪枝方法的应用顺序，并处理连续剪枝调用的合并。

接受一个 BasePruningMethod 实例或其可迭代集合作为参数。

add_pruning_method(method)[源代码]

在容器中添加一个子元素修剪的方法。

参数

method (subclass of BasePruningMethod) – 要添加到容器中的剪枝方法。

classmethodapply(module, name, *args, importance_scores=None, **kwargs)

添加动态剪枝功能，并对张量进行重新参数化。

添加一个前向预处理钩子，在运行时启用剪枝功能，并根据原始张量和剪枝掩码来重新参数化张量。

参数

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

• name (str) – 进行剪枝操作的模块中的参数名称。

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

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

• kwargs - 传递给 BasePruningMethod 子类的关键词参数

apply_mask(模块)

简单地处理参数和生成的掩码之间的乘法运算。

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

参数

module (nn.Module) – 包含待修剪张量的模块

返回值

输入张量的简化版

返回类型

pruned_tensor (torch.Tensor)

compute_mask(t, default_mask)[源代码]

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

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

• 对于“unstructured”，掩码将根据未被遮盖条目的扁平化列表进行计算。

• 对于 'structured'，掩码将基于张量中未被掩码的通道进行计算。

• 对于“global”，掩码将跨所有条目进行计算。

参数

• t (torch.Tensor) – 表示要进行剪枝操作的参数张量（与default_mask具有相同的维度）。

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

返回值

一个新的掩码，结合了 default_mask 和当前剪枝方法中新的掩码效果（与 default_mask 和 t 具有相同的维度）。

返回类型

掩码 (torch.Tensor)

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

计算并返回输入张量t的简化版本。

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

参数

• t (torch.Tensor) – 需要进行剪枝操作的张量（与default_mask具有相同的维度）。

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

• default_mask (torch.Tensor, 可选) – 上一次剪枝迭代的掩码（如果有）。在确定剪枝应作用于张量的哪一部分时，需要考虑这个掩码。如果为 None，则默认使用全1掩码。

返回值

张量t的简化版本。

remove(module)

从模块中删除剪枝重参数化。

名为 name 的剪枝参数将永久保持被剪枝状态，而名为 name+'_orig' 的参数将从参数列表中删除。同样地，名为 name+'_mask' 的缓冲区也将从缓冲区列表中移除。

注意

剪枝操作本身无法撤销或逆转！