PruningContainer

class torch.nn.utils.prune.PruningContainer(*args)

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

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

接受 BasePruningMethod 的实例或其可迭代对象作为参数。

add_pruning_method(method)

向容器添加子剪枝 method。

参数

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

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

添加即时剪枝和张量的重新参数化。

添加启用即时剪枝的前向预钩子，以及根据原始张量和剪枝掩码对张量进行重新参数化。

参数

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

• name (str) – module 中剪枝将作用的参数名称。

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

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

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

apply_mask(module)

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

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

参数

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

返回值

输入张量的剪枝版本

返回值类型

pruned_tensor (torch.Tensor)

compute_mask(t, default_mask)

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

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

• 对于 ‘unstructured’，掩码将从非掩码条目的展平列表计算；

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

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

参数

• t (torch.Tensor) – 表示要剪枝参数的张量（与 default_mask 尺寸相同）。

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

返回值

新的掩码，结合了 default_mask 和当前剪枝 method 生成的新掩码的效果（与 default_mask 和 t 尺寸相同）。

返回值类型

mask (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，则默认为全一掩码。

返回值

张量 t 的剪枝版本。

remove(module)

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

名为 name 的剪枝参数将永久保持剪枝状态，名为 name+'_orig' 的参数将从参数列表中移除。类似地，名为 name+'_mask' 的缓冲区将从缓冲区中移除。

注意

剪枝本身不会被撤销或反转！