torch.nn.utils.prune.custom_from_mask¶
- torch.nn.utils.prune.custom_from_mask(module, name, mask)[源代码]¶
通过应用
mask
中的预先计算的掩码,修剪对应于module
中名为name
的参数的张量。通过以下方式修改模块就地(并返回修改后的模块):
添加一个名为
name+'_mask'
的命名缓冲区,对应于修剪方法应用于参数name
的二进制掩码。用其修剪后的版本替换参数
name
,而原始(未修剪)参数存储在名为name+'_orig'
的新参数中。
- 参数
- 返回值
输入模块的修改(即修剪)版本
- 返回类型
module (nn.Module)
示例
>>> from torch.nn.utils import prune >>> m = prune.custom_from_mask( ... nn.Linear(5, 3), name='bias', mask=torch.tensor([0, 1, 0]) ... ) >>> print(m.bias_mask) tensor([0., 1., 0.])