摘要：所謂模型剪枝，其實是一種從神經網路中移除"不必要"權重或偏差的模型壓縮技術。
本文分享自華為雲社群《模型壓縮-pytorch 中的模型剪枝方法實踐》，作者：嵌入式視覺。
一，剪枝分類
所謂模型剪枝，其實是一種從神經網路中移除"不必要"權重或偏差（weigths/bias）的模型壓縮技術。關於什麼引數才是“不必要的”，這是一個目前依然在研究的領域。
1.1，非結構化剪枝
非結構化剪枝（Unstructured Puning）是指修剪引數的單個元素，比如全連線層中的單個權重、卷積層中的單個卷積核引數元素或者自定義層中的浮點數（scaling floats）。其重點在於，剪枝權重物件是隨機的，沒有特定結構，因此被稱為非結構化剪枝。
1.2，結構化剪枝
與非結構化剪枝相反，結構化剪枝會剪枝整個引數結構。比如，丟棄整行或整列的權重，或者在卷積層中丟棄整個過濾器（Filter）。
1.3，本地與全域性修剪
剪枝可以在每層（區域性）或多層/所有層（全域性）上進行。
二，PyTorch 的剪枝
目前 PyTorch 框架支援的權重剪枝方法有:
- Random: 簡單地修剪隨機引數。
- Magnitude: 修剪權重最小的引數（例如它們的 L2 範數）
以上兩種方法實現簡單、計算容易，且可以在沒有任何資料的情況下應用。
2.1，pytorch 剪枝工作原理
剪枝功能在 torch.nn.utils.prune 類中實現，程式碼在檔案 torch/nn/utils/prune.py 中，主要剪枝類如下圖所示。
剪枝原理是基於張量（Tensor）的掩碼（Mask）實現。掩碼是一個與張量形狀相同的布林型別的張量，掩碼的值為 True 表示相應位置的權重需要保留，掩碼的值為 False 表示相應位置的權重可以被刪除。
Pytorch 將原始引數 <param> 複製到名為 <param>_original 的引數中，並建立一個緩衝區來儲存剪枝掩碼 <param>_mask。同時，其也會建立一個模組級的 forward_pre_hook 回撥函式（在模型前向傳播之前會被呼叫的回撥函式），將剪枝掩碼應用於原始權重。
pytorch 剪枝的 api 和教程比較混亂，我個人將做了如下表格，希望能將 api 和剪枝方法及分類總結好。
pytorch 中進行模型剪枝的工作流程如下：
- 選擇剪枝方法（或者子類化 BasePruningMethod 實現自己的剪枝方法）。
- 指定剪枝模組和引數名稱。
- 設定剪枝方法的引數，比如剪枝比例等。
2.2，區域性剪枝
Pytorch 框架中的區域性剪枝有非結構化和結構化剪枝兩種型別，值得注意的是結構化剪枝只支援區域性不支援全域性。
2.2.1，區域性非結構化剪枝
1，區域性非結構化剪枝（Locall Unstructured Pruning）對應函式原型如下：
def random_unstructured(module, name, amount)
1）函式功能：
用於對權重引數張量進行非結構化剪枝。該方法會在張量中隨機選擇一些權重或連線進行剪枝，剪枝率由使用者指定。
2）函式引數定義：
- module (nn.Module): 需要剪枝的網路層/模組，例如 nn.Conv2d() 和 nn.Linear()。
- name(str): 要剪枝的引數名稱，比如 “weight” 或 “bias”。
- amount (int or float): 指定要剪枝的數量，如果是 0~1 之間的小數，則表示剪枝比例；如果是證書，則直接剪去引數的絕對數量。比如amount=0.2 ，表示將隨機選擇 20% 的元素進行剪枝。
3）下面是 random_unstructured 函式的使用示例。
import torch import torch.nn.utils.prune as prune conv = torch.nn.Conv2d(1, 1, 4) prune.random_unstructured(conv, name="weight", amount=0.5) conv.weight """ tensor([[[[-0.1703, 0.0000, -0.0000, 0.0690], [ 0.1411, 0.0000, -0.0000, -0.1031], [-0.0527, 0.0000, 0.0640, 0.1666], [ 0.0000, -0.0000, -0.0000, 0.2281]]]], grad_fn=<MulBackward0>) """
可以看書輸出的 conv 層中權重值有一半比例為 0。
2.2.2，區域性結構化剪枝
區域性結構化剪枝（Locall Structured Pruning）有兩種函式，對應函式原型如下：
def random_structured(module, name, amount, dim)
def ln_structured(module, name, amount, n, dim, importance_scores=None)
1）函式功能
與非結構化移除的是連線權重不同，結構化剪枝移除的是整個通道權重。
2）引數定義
與區域性非結構化函式非常相似，唯一的區別是您必須定義 dim 引數(ln_structured 函式多了 n 引數)。
n 表示剪枝的範數，dim 表示剪枝的維度。
對於 torch.nn.Linear：
- dim = 0： 移除一個神經元。
- dim = 1：移除與一個輸入的所有連線。
對於 torch.nn.Conv2d：
- dim = 0(Channels) : 通道 channels 剪枝/過濾器 filters 剪枝
- dim = 1（Neurons）: 二維卷積核 kernel 剪枝，即與輸入通道相連線的 kernel
2.2.3，區域性結構化剪枝示例程式碼
在寫示例程式碼之前，我們先需要理解 Conv2d 函式引數、卷積核 shape、軸以及張量的關係。
首先，Conv2d 函式原型如下;
class torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True)
而 pytorch 中常規卷積的卷積核權重 shape 都為（C_out, C_in, kernel_height, kernel_width），所以在程式碼中卷積層權重 shape 為 [3, 2, 3, 3]，dim = 0 對應的是 shape [3, 2, 3, 3] 中的 3。這裡我們 dim 設定了哪個軸，那自然剪枝之後權重張量對應的軸機會發生變換。
理解了前面的關鍵概念，下面就可以實際使用了，dim=0 的示例如下所示。
conv = torch.nn.Conv2d(2, 3, 3) norm1 = torch.norm(conv.weight, p=1, dim=[1,2,3]) print(norm1) """ tensor([1.9384, 2.3780, 1.8638], grad_fn=<NormBackward1>) """ prune.ln_structured(conv, name="weight", amount=1, n=2, dim=0) print(conv.weight) """ tensor([[[[-0.0005, 0.1039, 0.0306], [ 0.1233, 0.1517, 0.0628], [ 0.1075, -0.0606, 0.1140]], [[ 0.2263, -0.0199, 0.1275], [-0.0455, -0.0639, -0.2153], [ 0.1587, -0.1928, 0.1338]]], [[[-0.2023, 0.0012, 0.1617], [-0.1089, 0.2102, -0.2222], [ 0.0645, -0.2333, -0.1211]], [[ 0.2138, -0.0325, 0.0246], [-0.0507, 0.1812, -0.2268], [-0.1902, 0.0798, 0.0531]]], [[[ 0.0000, -0.0000, -0.0000], [ 0.0000, -0.0000, -0.0000], [ 0.0000, -0.0000, 0.0000]], [[ 0.0000, 0.0000, 0.0000], [-0.0000, 0.0000, 0.0000], [-0.0000, -0.0000, -0.0000]]]], grad_fn=<MulBackward0>) """
從執行結果可以明顯看出，卷積層引數的最後一個通道引數張量被移除了（為 0 張量），其解釋參見下圖。
dim = 1 的情況：
conv = torch.nn.Conv2d(2, 3, 3) norm1 = torch.norm(conv.weight, p=1, dim=[0, 2,3]) print(norm1) """ tensor([3.1487, 3.9088], grad_fn=<NormBackward1>) """ prune.ln_structured(conv, name="weight", amount=1, n=2, dim=1) print(conv.weight) """ tensor([[[[ 0.0000, -0.0000, -0.0000], [-0.0000, 0.0000, 0.0000], [-0.0000, 0.0000, -0.0000]], [[-0.2140, 0.1038, 0.1660], [ 0.1265, -0.1650, -0.2183], [-0.0680, 0.2280, 0.2128]]], [[[-0.0000, 0.0000, 0.0000], [ 0.0000, 0.0000, -0.0000], [-0.0000, -0.0000, -0.0000]], [[-0.2087, 0.1275, 0.0228], [-0.1888, -0.1345, 0.1826], [-0.2312, -0.1456, -0.1085]]], [[[-0.0000, 0.0000, 0.0000], [ 0.0000, -0.0000, 0.0000], [ 0.0000, -0.0000, 0.0000]], [[-0.0891, 0.0946, -0.1724], [-0.2068, 0.0823, 0.0272], [-0.2256, -0.1260, -0.0323]]]], grad_fn=<MulBackward0>) """
很明顯，對於 dim=1的維度，其第一個張量的 L2 範數更小，所以shape 為 [2, 3, 3] 的張量中，第一個 [3, 3] 張量引數會被移除（即張量為 0 矩陣） 。
2.3，全域性非結構化剪枝
前文的 local 剪枝的物件是特定網路層，而 global 剪枝是將模型看作一個整體去移除指定比例（數量）的引數，同時 global 剪枝結果會導致模型中每層的稀疏比例是不一樣的。
全域性非結構化剪枝函式原型如下：
# v1.4.0 版本 def global_unstructured(parameters, pruning_method, **kwargs) # v2.0.0-rc2版本 def global_unstructured(parameters, pruning_method, importance_scores=None, **kwargs):
1）函式功能：
隨機選擇全域性所有引數（包括權重和偏置）的一部分進行剪枝，而不管它們屬於哪個層。
2）引數定義：
- parameters（(Iterable of (module, name) tuples)）: 修剪模型的引數列表，列表中的元素是 (module, name)。
- pruning_method（function）: 目前好像官方只支援 pruning_method=prune.L1Unstuctured，另外也可以是自己實現的非結構化剪枝方法函式。
- importance_scores: 表示每個引數的重要性得分，如果為 None，則使用預設得分。
- **kwargs: 表示傳遞給特定剪枝方法的額外引數。比如 amount 指定要剪枝的數量。
3）global_unstructured 函式的示例程式碼如下所示。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") class LeNet(nn.Module): def __init__(self): super(LeNet, self).__init__() # 1 input image channel, 6 output channels, 3x3 square conv kernel self.conv1 = nn.Conv2d(1, 6, 3) self.conv2 = nn.Conv2d(6, 16, 3) self.fc1 = nn.Linear(16 * 5 * 5, 120) # 5x5 image dimension self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) def forward(self, x): x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2)) x = F.max_pool2d(F.relu(self.conv2(x)), 2) x = x.view(-1, int(x.nelement() / x.shape[0])) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x model = LeNet().to(device=device) model = LeNet() parameters_to_prune = ( (model.conv1, 'weight'), (model.conv2, 'weight'), (model.fc1, 'weight'), (model.fc2, 'weight'), (model.fc3, 'weight'), ) prune.global_unstructured( parameters_to_prune, pruning_method=prune.L1Unstructured, amount=0.2, ) # 計算卷積層和整個模型的稀疏度 # 其實呼叫的是 Tensor.numel 內內函式，返回輸入張量中元素的總數 print( "Sparsity in conv1.weight: {:.2f}%".format( 100. * float(torch.sum(model.conv1.weight == 0)) / float(model.conv1.weight.nelement()) ) ) print( "Global sparsity: {:.2f}%".format( 100. * float( torch.sum(model.conv1.weight == 0) + torch.sum(model.conv2.weight == 0) + torch.sum(model.fc1.weight == 0) + torch.sum(model.fc2.weight == 0) + torch.sum(model.fc3.weight == 0) ) / float( model.conv1.weight.nelement() + model.conv2.weight.nelement() + model.fc1.weight.nelement() + model.fc2.weight.nelement() + model.fc3.weight.nelement() ) ) ) # 程式執行結果 """ Sparsity in conv1.weight: 3.70% Global sparsity: 20.00% """
執行結果表明，雖然模型整體（全域性）的稀疏度是 20%，但每個網路層的稀疏度不一定是 20%。
三，總結
另外，pytorch 框架還提供了一些幫助函式:
- torch.nn.utils.prune.is_pruned(module): 判斷模組 是否被剪枝。
- torch.nn.utils.prune.remove(module, name)： 用於將指定模組中指定引數上的剪枝操作移除，從而恢復該引數的原始形狀和數值。
雖然 PyTorch 提供了內建剪枝 API ，也支援了一些非結構化和結構化剪枝方法，但是 API 比較混亂，對應文件描述也不清晰，所以後面我還會結合微軟的開源 nni 工具來實現模型剪枝功能。
參考資料
- How to Prune Neural Networks with PyTorch
- PRUNING TUTORIAL
- PyTorch Pruning