BasePruningMethod

class torch.nn.utils.prune.BasePruningMethod

用於建立新剪枝技術的抽象基底類別。

提供自訂的骨架，需要覆寫 compute_mask() 和 apply() 等方法。

classmethod apply(module, name, *args, importance_scores=None, **kwargs)

動態地新增剪枝（pruning）和張量的重新參數化（reparametrization）。

新增 forward pre-hook，以啟用動態的剪枝，並根據原始張量和剪枝遮罩（pruning mask）重新參數化張量。

參數

• 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 開始（如果張量尚未被剪枝，則應為全 1 遮罩），根據特定的剪枝方法規則生成隨機遮罩，以應用於 default_mask 之上。

參數

• t (torch.Tensor) – 表示重要性分數的張量

• prune. (parameter to) –

• default_mask (torch.Tensor) – 來自先前剪枝的基本遮罩

• iterations –

• is (that need to be respected after the new mask) –

• t. (applied. Same dims as) –

返回

要應用於 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 中相應元素的重要性，該 t 正在被剪枝。如果未指定或為 None，則將使用張量 t 本身。

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

返回

張量 t 的剪枝版本。

remove(module)

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

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

注意

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