PruningContainer#
- class torch.nn.utils.prune.PruningContainer(*args)[source]#
用于迭代剪枝的剪枝方法容器。
跟踪剪枝方法的应用顺序，并处理连续剪枝调用的合并。
接受 BasePruningMethod 实例或其可迭代对象作为参数。
- add_pruning_method(method)[source]#
向容器中添加一个子剪枝方法
method。- 参数
method (subclass of BasePruningMethod) – 要添加到容器中的子剪枝方法。
- classmethod apply(module, name, *args, importance_scores=None, **kwargs)[source]#
即时添加剪枝和张量的重参数化。
添加前向预钩子以启用即时剪枝，并进行张量的重参数化，使其表示为原始张量和剪枝掩码。
- 参数
module (nn.Module) – 包含要剪枝张量的模块
name (str) –
module中将应用剪枝的参数名称。args – 传递给 BasePruningMethod 子类的参数
importance_scores (torch.Tensor) – 重要性得分张量（与模块参数形状相同），用于计算剪枝掩码。此张量中的值指示被剪枝参数相应元素的importance。如果未指定或为 None，则将使用参数本身。
kwargs – 传递给 BasePruningMethod 子类的关键字参数
- apply_mask(module)[source]#
简单处理要剪枝的参数与生成的掩码之间的乘法。
获取模块中的掩码和原始张量，并返回张量的剪枝版本。
- 参数
module (nn.Module) – 包含要剪枝张量的模块
- 返回
剪枝后的张量版本
- 返回类型
pruned_tensor (torch.Tensor)
- compute_mask(t, default_mask)[source]#
通过计算新的部分掩码来应用最新的
method，并将其与default_mask组合返回。新的部分掩码应根据
default_mask未归零的条目或通道进行计算。新的掩码将从张量t的哪些部分计算，取决于PRUNING_TYPE（由类型处理程序处理）。对于“非结构化”，掩码将从未掩码条目的展平列表中计算；
对于“结构化”，掩码将从张量中的非掩码通道计算；
对于“全局”，掩码将跨所有条目进行计算。
- 参数
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)[source]#
计算并返回输入张量
t的剪枝版本。根据
compute_mask()中指定的剪枝规则。- 参数
t（torch.Tensor）– 要剪枝的张量（维度与
default_mask相同）。importance_scores（torch.Tensor）– 重要性分数张量（形状与
t相同），用于计算剪枝t的掩码。此张量中的值指示正在剪枝的t中相应元素的重要性。如果未指定或为None，则将使用t张量代替。default_mask（torch.Tensor, 可选）– 前一个剪枝迭代的掩码（如果有）。在确定剪枝应作用于张量的哪个部分时需要考虑。如果为None，则默认为一个全为1的掩码。
- 返回
张量
t的修剪版本。