BasePruningMethod¶
- class torch.nn.utils.prune.BasePruningMethod[source]¶
创建新的剪枝技术的抽象基类。
提供了一个用于自定义的框架,需要重写
compute_mask()
和apply()
等方法。- classmethod apply(module, name, *args, importance_scores=None, **kwargs)[source]¶
动态添加剪枝并对张量进行重新参数化。
添加前向预钩子,该钩子可以动态启用剪枝,并根据原始张量和剪枝掩码对张量进行重新参数化。
- 参数
module (nn.Module) – 包含要剪枝的张量的模块
name (str) –
module
中将要进行剪枝操作的参数名称。args – 传递给
BasePruningMethod
子类的参数importance_scores (torch.Tensor) – 用于计算剪枝掩码的重要性分数张量(与模块参数形状相同)。此张量中的值表示要剪枝的参数中对应元素的重要性。如果未指定或为 None,则将使用参数代替。
kwargs – 传递给
BasePruningMethod
子类的关键字参数
- apply_mask(module)[source]¶
仅处理要剪枝的参数与生成的掩码之间的乘法。
从模块中获取掩码和原始张量,并返回张量的剪枝版本。
- 参数
module (nn.Module) – 包含要剪枝的张量的模块
- 返回值
输入张量的剪枝版本
- 返回类型
pruned_tensor (torch.Tensor)
- abstract compute_mask(t, default_mask)[source]¶
计算并返回输入张量
t
的掩码。从基本
default_mask
(如果张量尚未被剪枝,则应为全1掩码)开始,根据特定的剪枝方法方案生成一个随机掩码应用于default_mask
之上。- 参数
t (torch.Tensor) – 表示要
剪枝的。 (参数的) –
default_mask (torch.Tensor) – 来自先前剪枝的基本掩码
迭代次数 –
是 (需要在应用新掩码后保持的) –
t. (与) –
- 返回值
要应用于
t
的掩码,与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,则默认为全为一的掩码。
- 返回值
张量
t
的修剪版本。