基剪枝方法

torch.nn.utils.prune.BasePruningMethod[源代码]

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

提供了一个自定义的骨架,要求重写一些方法,如compute_mask()apply()

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)[源代码]

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

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

参数

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

返回值

输入张量的简化版

返回类型

pruned_tensor (torch.Tensor)

abstractcompute_mask(t, default_mask)[源代码]

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

从基础的 default_mask(如果张量尚未被剪枝,则应为全1的掩码)开始,根据特定的剪枝方法生成一个随机掩码,并将其应用到 default_mask 上。

参数
  • t (torch.Tensor) – 代表重要性分数的张量

  • 修剪。参数)–

  • default_mask (torch.Tensor) – 前一次剪枝所用的基础掩码

  • 迭代次数

  • (在应用新遮罩后需遵守的)–

  • t. (应用。相同维度) –

返回值

t具有相同维度的掩码,用于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' 的缓冲区也将从缓冲区列表中移除。

注意

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

本页目录