BasePruningMethod

class torch.nn.utils.prune.BasePruningMethod

创建新剪枝技术的抽象基类。

提供了一个用于定制的骨架，需要重写诸如 compute_mask() 和 apply() 等方法。

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)

abstract compute_mask(t, default_mask)

计算并返回输入张量 t 的掩码。

从基础 default_mask 开始（如果张量尚未剪枝，则应为全一掩码），根据具体的剪枝方法生成一个随机掩码应用于 default_mask 之上。

参数

• t (torch.Tensor) – 表示要

• 剪枝 (参数的) –

• default_mask (torch.Tensor) – 来自先前剪枝的基础掩码

• 迭代 –

• 是 (新掩码应用后需要保留的部分。) –

• t. (与 维度相同) –

返回

要应用于 t 的掩码，维度与 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' 的缓冲区从缓冲区中移除。

注意

剪枝本身无法撤销或反转！