CustomFromMask

class torch.nn.utils.prune.CustomFromMask(mask)

classmethod apply(module, name, mask)

動態地新增剪枝並重新參數化張量。

新增前向預鉤 (forward pre-hook)，以啟用動態剪枝，並根據原始張量和剪枝遮罩 (mask) 重新參數化張量。

參數

• module (nn.Module) – 包含要剪枝的張量的模組

• name (str) – module 內將執行剪枝的參數名稱。

apply_mask(module)

簡單地處理被剪枝的參數和產生的遮罩之間的乘法。

從模組中提取遮罩和原始張量，並返回張量的剪枝版本。

參數

module (nn.Module) – 包含要剪枝的張量的模組

返回

輸入張量的剪枝版本

返回類型

pruned_tensor (torch.Tensor)

prune(t, default_mask=None, importance_scores=None)

計算並返回輸入張量 t 的剪枝版本。

根據 compute_mask() 中指定的剪枝規則。

參數

• t (torch.Tensor) – 要剪枝的張量 (與 default_mask 具有相同維度)。

• importance_scores (torch.Tensor) – 重要性分數 (importance scores) 的張量 (與 t 具有相同形狀)，用於計算剪枝 t 的遮罩。此張量中的值表示正在剪枝的 t 中對應元素的重要性。如果未指定或為 None，將使用張量 t 代替。

• default_mask (torch.Tensor, optional) – 先前剪枝迭代的遮罩 (如果有的話)。在確定剪枝應作用於張量的哪一部分時要考慮。如果為 None，則預設為全 1 的遮罩。

返回

張量 t 的剪枝版本。

remove(module)

從模組中移除剪枝重新參數化。

名為 name 的剪枝參數將永久保持剪枝狀態，並且名為 name+'_orig' 的參數將從參數列表中移除。同樣，名為 name+'_mask' 的緩衝區 (buffer) 將從緩衝區中移除。

注意

剪枝本身不會被撤銷或還原！