Pytorch網絡壓縮系列教程一:Prune你的模型

本文由林大佬原創，轉載請注明出處，來自騰訊、阿里等一線AI算法工程師組成的QQ交流群歡迎你的加入: 1037662480

深度學習模型取得了前所未有的巨大成功, 并且引領著新的科技潮流. 但學術界流傳的深度學習模型, 通常都是超參數化的, 在更低功耗, 便攜式的設備上部署時, 會遇到很多問題. 為了將我們的模型進一步壓縮, 同時保證模型具有和以前一樣的精確度, 一個新的領域開始備受矚目: 網絡壓縮與加速. 總的來說, 這里面分為一個不同的領域, 但是實際上做的是同一件事:

• 剪枝(Prune);
• 稀疏訓練;
• 蒸餾

對于Prune來說, 這是有很好的范例和理論依據來做這件事的, 現在越來越多的公司也在這個領域發力, 不斷的將深度學習運算的極限push到更高的值, 讓低功耗設備下的深度學習模型運算成為可能.

正式由于現有的深度學習模型都存在超參數化的現象, 對于一些對于整體網絡貢獻度不高的層, 或者參數, 我們可以將其裁剪掉, 進而使得網絡輕量化. 如果你正在研究, 學習這個方向, 這篇教程可以讓你足夠入門.

我們將結合pytorch和實際的例子來教大家如何來prune你的模型, 這篇文章不僅僅教你找到你要剪掉的層, 同時還會告訴你如何保存你減掉的模型, 并且這個模型你可以拿來推理, 可以導出簡化模型的onnx并部署.

01. Prune Lenet

我們先從一個十分簡單的例子開始, 比如LeNet. 這是一個只有兩個卷積層和三個全連接層的極簡網絡. 它的結構如下:

LeNet((conv1): Conv2d(1, 6, kernel_size=(3, 3), stride=(1, 1))(conv2): Conv2d(6, 16, kernel_size=(3, 3), stride=(1, 1))(fc1): Linear(in_features=400, out_features=120, bias=True)(fc2): Linear(in_features=120, out_features=84, bias=True)(fc3): Linear(in_features=84, out_features=10, bias=True) )

即便是這么一個網絡, 我們也是可以對它進行壓縮的. 事實上, 我們會先從pytorch自帶的Prune模塊入手, 教大家如何去尋找網絡的稀疏度, 但如果你想要更高層次的東西, 比如Distiller里面那種通過修改每一層再做evaluation, 測試每一層的影響, 并生成剪枝表的更復雜操作, 我們會在后面的教程跟大家講.

先來看看如何用pytorch來做.

實際上在pytorch里面, prune通過對權重進行掩碼來完成. 這個如何理解?

首先, 我們打印一下原始的conv1的權重看看:

module = model.conv1 print(list(module.named_modules())) print(list(module.named_buffers())) print(list(module.named_parameters()))

這里我列舉了后面可能會用到三個方法, 這個可以查看當前的module到底是一個啥情況.

然后我們著重觀察 named_parameters 因為參數都保存在這里, 打印完了之后你可以看到:

[('weight', Parameter containing: tensor([[[[-0.2312, 0.2133, -0.1313],[-0.2980, -0.1838, -0.2902],[-0.3006, 0.1338, -0.0980]]],[[[-0.1239, 0.1060, 0.3271],[-0.0301, -0.0245, 0.0493],[-0.0160, 0.0397, -0.1242]]]], device='cuda:0', requires_grad=True)), ('bias', Parameter containing: tensor([-0.2593, -0.0520, 0.0303, 0.0382, -0.0468, -0.1053], device='cuda:0',requires_grad=True))] ?

它有一個weight和一個bias, 這沒錯, 合乎常理. 我們甚至可以看看weights的尺寸是多少.

for a in module.named_parameters():print(a[1].shape)

輸出:

torch.Size([6, 1, 3, 3]) torch.Size([6])

這個其實就是卷積的維度了, 6指的是channel, 1值得還是stride, 3指的是kernel size.

然后重點來了, 我們要開始做prune了, 在pytorch里面操作也很簡單, 只需要一行代碼:

import torch.nn.utils.prune as prune ? prune.random_unstructured(module, name='weight', amount=0.3) ?

這個可以從眾多的剪枝方法中, 選擇一個很好的手段來完成同樣的目的. 然后我們再打印一下named_parameters:

[('bias', Parameter containing: tensor([-0.2281, 0.3085, 0.0937, -0.0540, 0.3295, 0.1107], device='cuda:0',requires_grad=True)), ('weight_orig', Parameter containing: tensor([[[[ 0.1934, -0.0172, -0.1957],[ 0.1655, 0.1669, -0.2448],[-0.2250, -0.0963, -0.0195]]], ?[[[-0.3154, 0.1868, 0.0103],[-0.2245, 0.1548, 0.2567],[ 0.0713, 0.1262, 0.1547]]]], device='cuda:0', requires_grad=True))] ?

唯一的變化就是 weights 變成了 weights_orig, prune之后通過掩碼的方式存放在了 named_buffers里面:

print(list(module.named_buffers()))

可以看到:

[('weight_mask', tensor([[[[1., 1., 1.],[0., 1., 0.],[0., 1., 1.]]],[[[0., 1., 0.],[0., 0., 1.],[0., 0., 1.]]]], device='cuda:0'))]

那么問題來了, 你只是把權重進行了掩碼, 那么我要知道你剪掉了哪幾個channel怎么辦? 而且你這個是剪的權重, 結構呢? 我怎么把這個結構找出來??

所以說這只是第一步, 接下來我們來看看結構化修剪. 結構化修剪講道理你可以知道你修剪了哪些結構.

一個比較好的結構化修建的例子是通過沿著Tensor的某個維度進行裁剪, 這樣你可以直接看到維度的變化. 關于結構化剪枝和非結構的區別, 大家可以看看我們的知乎這個回答:

https://www.zhihu.com/question/391195715/answer/1421770287

我們現在開始修剪模塊, 比如上面的LeNet的conv1層, 首先我們可以從prune層里面拿一個我們喜歡的技術, 比如基于 ln范數 的評判標準來進行結構話的裁剪.

prune.ln_structured(module, name='weight', amount=0.5, n=2, dim=0)

這個操作之后, 我們得到的將是一個新的權重, 和上面的非結構化的不同的地方在于, 這里是整個矩陣的一行為零, 上面我們用的dim=0, 那么就是channel這一個維度, 會有50%為零.

講到這里, 你們很容易復現, 也很容易理解, 我在編寫這篇教程的時候與參考了很多博客, 文章, 我發現即便是你看完了你還是無法這些讀者遇到的問題:

簡單來說, 大家其實關注的本身不在于你怎么去做的prune, 而是prune完了之后我要怎么用.

很可惜網上大部分教程到此戛然而止, 沒有人告訴你接下來要怎么做, 網絡如何保存? 如何訓練? 如何部署? 如何把新的結構提取出來?等等.

這就是你要關注本專欄或者本公眾號的根本原因: 我們做后20%的事情.

02. 保存裁剪后的模型

接下來是我們的重點. 如何將這些模型保存下來. 換句話說, 你剪枝完了之后, 你得到了一個mask的掩碼, 不管是結構化還是非結構化, 這個新的模型你要怎么樣才能導出呢? 試想一下, 如果只是得到這么一個掩碼, 你在推理的時候, 只是這一部分權重為零而已, 你推理的長寬高啥的都么有改變, so, 速度會變得更快嗎?

事實上善于谷歌的我們, 可以在stackoverflow輕松找到一位老哥試圖復現的代碼:

stackoverflow

這位老哥發現, 他嘗試不管是結構化剪枝還是非結構化剪枝, 速度和模型的大小都沒有變化. 然后另外一個老哥就告訴他真像了:

重要的是要了解非結構化修剪和結構化修剪之間的區別。
結構化修剪：通過刪除張量的整行/列來減小重量張量的尺寸。這轉化為去除神經元及其所有傳入和傳出連接（在密集層中）或整個卷積過濾器（在卷積層中）。
非結構化修剪：單個權重可以“去除”（歸零），而不受最終張量形狀的限制。這意味著刪除神經元之間的單個連接（在密集層中）或刪除卷積過濾器的單個權重（在卷積層中）。
請注意，生成的重量張量可以稀疏，但保持其原始形狀。目前，torch.nn.utils.prune僅支持非結構化修剪，這幾乎無法幫助降低推理成本，因為GPU并未針對稀疏矩陣乘法進行優化。雖然您可能希望減小權重張量的尺寸以減少浮點運算的數量，但非結構化修剪會產生帶有許多零的權重張量，但不會自動減小此類張量的大小。僅當去除許多負擔時，非結構化修剪才能幫助提高性能。在這種情況下，您可以依靠PyTorch稀疏操作，也可以嘗試查找包含全零的行/列，因此可以將其刪除。相反，如果要研究結構化修剪，可以看一下TorchPruner，這是我為研究目的而自行開發的一個庫，它提供了實用程序來查找最不重要的神經元并相應地對重量張量進行切片。

所以, 這個部分我們也應該知道, 事實上pytorch的prune里面做的事情, 只是找到哪些層是可以置零的而已, 請注意, 這里的層指的不是某個layer, 而是一個卷積層矩陣里面的某一整行或者某一整列.

那么問題來了, 聽起來我們要保存模型, 要得到速度加快, 模型變小的模型, 還需要再每一個prune的層里面找到某一被全置零的鏈接, 然后再拿來做推理.

可是這個怎么聽起來這么復雜呢? 事實上在實際應用中應該怎么做呢?

這就引入了我們的下一篇教程: 使用Distiller的工業應用模型剪枝教程.

更多

http://weixin.qq.com/r/TjgIEHTEzymwrTBz921c (二維碼自動識別)

如果你想學習人工智能，對前沿的AI技術比較感興趣，可以加入我們的知識星球，獲取第一時間資訊，前沿學術動態，業界新聞等等！你的支持將會鼓勵我們更頻繁的創作，我們也會幫助你開啟更深入的深度學習之旅！

往期文章

https://zhuanlan.zhihu.com/p/165009477

https://zhuanlan.zhihu.com/p/149398749

https://zhuanlan.zhihu.com/p/147622974

https://zhuanlan.zhihu.com/p/144727162

總結

以上是生活随笔為你收集整理的pytorch打印模型参数_Pytorch网络压缩系列教程一:Prune你的模型的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。