構(gòu)化剪枝實(shí)戰(zhàn)：從稀疏化訓(xùn)練到部署加速)
先講一個(gè)我前陣子真實(shí)遇到的部署場(chǎng)景。模型用YOLOv8s在GPU上訓(xùn)得漂漂亮亮精度也在線可一旦挪到邊緣盒子上單幀推理直接干到四十多毫秒真到了要拼實(shí)時(shí)性的時(shí)候幀率死活穩(wěn)不住。調(diào)NMS、降分辨率、換解碼方式折騰了一圈效果都有限最后真正把時(shí)延壓下來(lái)的是剪枝。剪枝這個(gè)詞大家都不陌生可一旦落到Y(jié)OLOv8s上情況就跟網(wǎng)上一堆舊教程里的YOLOv5源碼完全不一樣了。yolov8s在結(jié)構(gòu)上把C3換成了C2fdetect head也改成了decoupled結(jié)構(gòu)直接把v5時(shí)代的剪枝腳本搬過(guò)來(lái)十有八九跑不通。我花了大概兩周時(shí)間把整條鏈路完整跑通從稀疏化訓(xùn)練、通道裁剪、結(jié)構(gòu)重建到微調(diào)恢復(fù)精度全部基于ultralytics YOLOv8s源碼改造。這篇博文就是把一套能落地的yolov8s剪枝源碼思路拆開(kāi)來(lái)講每一步為什么要這么做核心代碼應(yīng)該寫(xiě)在哪有哪些坑我替你先踩過(guò)了。內(nèi)容適合已經(jīng)在用YOLOv8做項(xiàng)目、對(duì)PyTorch不陌生、但還沒(méi)系統(tǒng)接觸過(guò)模型壓縮的開(kāi)發(fā)者。1. YOLOv8s結(jié)構(gòu)解析先弄清楚算力燒在了哪里1.1 YOLOv8s各模塊的算力分布剪枝之前我建議你先別急著翻代碼而是花半天時(shí)間把你手里這個(gè)模型的算力分布看明白。YOLOv8s在640×640輸入分辨率下總計(jì)算量約在28GFLOPs左右參數(shù)量約11.2M含檢測(cè)頭。這個(gè)體量在YOLOv8家族里屬于small檔看起來(lái)不大但對(duì)邊緣設(shè)備來(lái)說(shuō)依然偏重。從結(jié)構(gòu)上看YOLOv8s由三大部分組成Backbone負(fù)責(zé)從輸入圖像提取多尺度特征Neck是PAN-FPN結(jié)構(gòu)負(fù)責(zé)把不同層的特征融合起來(lái)最后是Decoupled檢測(cè)頭分別輸出分類(lèi)和回歸分支。我自己的實(shí)測(cè)統(tǒng)計(jì)里Backbone占據(jù)了接近一半的計(jì)算量Neck占了剩下三分之一左右檢測(cè)頭因?yàn)榇罅渴褂?×1卷積FLOPs占比其實(shí)沒(méi)有想象中高。這里要重點(diǎn)說(shuō)的是Backbone和Neck里大量的3×3卷積。對(duì)邊緣設(shè)備而言3×3卷積在CPU和多數(shù)NPU上都是計(jì)算熱點(diǎn)剪掉一層輸出通道省下的時(shí)間非?？捎^。所以剪枝的核心目標(biāo)很明確把你的裁剪預(yù)算優(yōu)先花在Backbone和Neck的C2f模塊上檢測(cè)頭做一個(gè)克制的輸入通道裁剪即可最后的輸出層盡量保持不動(dòng)。1.2 C2f的split機(jī)制剪枝時(shí)為什么比v5麻煩YOLOv8s和YOLOv5s在結(jié)構(gòu)上最大的變化就是C3換成了C2f。C2f模塊里有一個(gè)很關(guān)鍵的split操作def forward(self, x): y list(self.cv1(x).chunk(2, 1)) # 沿通道維度對(duì)半拆開(kāi) y.extend(m(y[-1]) for m in self.m) # 后半路依次經(jīng)過(guò)多個(gè)Bottleneck return self.cv2(torch.cat(y, 1))cv1先把輸入通道映射到2c然后沿通道方向?qū)Π氩鸪蓛煞?。前半份直接送去concat后半份經(jīng)過(guò)一組Bottleneck之后再送去concat最終由cv2把(n1)c壓縮回目標(biāo)輸出通道。這個(gè)split機(jī)制直接導(dǎo)致了一個(gè)問(wèn)題C2f內(nèi)部所有分支的通道組必須保持對(duì)齊才能保證最終concat之后通道位置嚴(yán)格對(duì)應(yīng)。如果照搬YOLOv5時(shí)代的“逐個(gè)BN層剪γ”方案幾乎必然在C2f模塊的torch.cat處報(bào)shape mismatch或者更隱蔽的——通道錯(cuò)位。我在寫(xiě)剪枝源碼時(shí)把C2f當(dāng)成一個(gè)整體去生成統(tǒng)一的通道m(xù)ask而不是單獨(dú)處理每個(gè)Bottleneck這一點(diǎn)是YOLOv8s剪枝區(qū)別于v5的核心難點(diǎn)后面會(huì)專(zhuān)門(mén)講。1.3 結(jié)構(gòu)化剪枝與非結(jié)構(gòu)化剪枝想清楚再動(dòng)手模型剪枝分為兩大類(lèi)非結(jié)構(gòu)化剪枝unstructured pruning剪的是單個(gè)權(quán)重稀疏度分布不均勻模型文件里全是零值但通用硬件上幾乎得不到加速收益除非配套稀疏矩陣專(zhuān)用庫(kù)。結(jié)構(gòu)化剪枝structured pruning以通道或?yàn)V波器為基本單位整體裁剪直接改變卷積層的輸入輸出維度任何推理框架都能拿到實(shí)實(shí)在在的加速。我見(jiàn)過(guò)不少朋友把SparseGPT這類(lèi)非結(jié)構(gòu)化方案用在YOLOv8s上精度保住了但導(dǎo)出到TensorRT或ONNX Runtime之后推理速度紋絲不動(dòng)。原因很簡(jiǎn)單民用級(jí)硬件基本不為稀疏權(quán)重做優(yōu)化。所以你如果目的是部署加速直接選結(jié)構(gòu)化通道剪枝不要去碰非結(jié)構(gòu)化路線。2. 剪枝方案選型基于BN層γ系數(shù)做通道重要性判斷2.1 Network Slimmingγ系數(shù)就是天然的通道重要度指標(biāo)通道剪枝要回答一個(gè)核心問(wèn)題怎么判定哪些通道不重要最經(jīng)典也最省事的方案是ICCV 2017年的Network Slimming方法。它利用的是卷積層后面那個(gè)幾乎被人忽略的BatchNorm層。BN層的計(jì)算公式是y γ * (x - μ) / √(σ2 ε) β其中γ和β都是可學(xué)習(xí)的仿射參數(shù)。γ是縮放系數(shù)如果訓(xùn)練后某個(gè)通道的γ變得非常接近0那么該通道的輸出就幾乎恒定在一個(gè)由β決定的固定值上信息量趨近于零這個(gè)通道就可以被安全地剪掉。關(guān)鍵就在于普通訓(xùn)練得到的γ數(shù)值分布不夠“極端”很多通道的γ值都落在0.5到1.5之間你沒(méi)法說(shuō)誰(shuí)更重要。所以需要人為給γ加上L1正則約束讓它稀疏化。訓(xùn)練時(shí)的總損失從L變成L λ * Σ|γ|其中λ是稀疏化懲罰系數(shù)。這個(gè)正則項(xiàng)會(huì)鼓勵(lì)γ往0靠攏訓(xùn)練結(jié)束后γ的分布會(huì)形成明顯的雙峰一堆極度接近0剩下的保持相對(duì)大的數(shù)值。普通訓(xùn)練后直接按γ剪和稀疏化訓(xùn)練之后按γ剪效果差別非常大。前者誤傷率高剪完精度掉得一塌糊涂微調(diào)也救不回來(lái)后者才能精準(zhǔn)識(shí)別出真正冗余的通道。這是整個(gè)剪枝方案成立的前提千萬(wàn)別跳過(guò)。2.2 為什么不考慮低秩分解和基于梯度的通道評(píng)估你可能還會(huì)看到一些方案用通道注意力權(quán)重、梯度敏感度或者低秩分解來(lái)選擇剪枝通道。低秩分解能把卷積核拆成小矩陣序列但從工程角度講拆完之后的模型結(jié)構(gòu)要重寫(xiě)底層算子才能發(fā)揮加速效果和普通推理框架的兼容性很一般?；谔荻鹊姆椒ㄊ歉?xì)的敏感度分析但要在每個(gè)剪枝候選上做一次反向傳播算力開(kāi)銷(xiāo)不低性?xún)r(jià)比反而不如γ稀疏化。對(duì)YOLOv8s這種層數(shù)不算極端、結(jié)構(gòu)相對(duì)規(guī)整的模型來(lái)說(shuō)γ系數(shù)方案是平衡效果與工程復(fù)雜度最好的選擇。它不需要額外訓(xùn)練一個(gè)評(píng)判網(wǎng)絡(luò)不需要逐層搜索敏感度只需要在訓(xùn)練loss上加一個(gè)正則項(xiàng)剪枝時(shí)拿一把全局閾值來(lái)篩通道整個(gè)流程完全可控。2.3 開(kāi)工前的基礎(chǔ)配置我建議你基于ultralytics官方的YOLOv8s源碼來(lái)做版本選8.0.x或8.1.x都行但1.x分支之后的代碼結(jié)構(gòu)有些調(diào)整注意保持一致。環(huán)境方面PyTorch 1.13以上、CUDA 11.7以上就行不需要額外裝稀疏計(jì)算庫(kù)剪枝前的稀疏化訓(xùn)練就是常規(guī)訓(xùn)練。另外準(zhǔn)備一個(gè)你自己的、已經(jīng)收斂到可用狀態(tài)的YOLOv8s權(quán)重文件。最佳實(shí)踐是先正常訓(xùn)練到float32的精度峰值再?gòu)倪@個(gè)權(quán)重開(kāi)始做稀疏化訓(xùn)練不要從頭就在loss里加稀疏正則否則模型收斂速度和精度上限都會(huì)受影響。3. 稀疏化訓(xùn)練源碼改造給YOLOv8s訓(xùn)練循環(huán)加上γ正則3.1 在Ultralytics訓(xùn)練循環(huán)中插入稀疏化損失Ultralytics YOLOv8的訓(xùn)練邏輯集中在trainer.py的_do_train方法里核心循環(huán)對(duì)每個(gè)batch做前向和反向。為了加稀疏化損失我的做法是在trainer里覆寫(xiě)一個(gè)loss計(jì)算hook。不同版本源碼結(jié)構(gòu)有差異但思路是通用的在模型前向算出的主loss之后、loss.backward()之前加入BN層γ的L1正則項(xiàng)。# sparse_loss.py import torch import torch.nn as nn SPARSITY_LAMBDA 1e-4 def sparsity_loss(model): reg_loss 0.0 for m in model.modules(): if isinstance(m, torch.nn.BatchNorm2d): reg_loss torch.abs(m.weight).sum() return SPARSITY_LAMBDA * reg_loss然后在訓(xùn)練循環(huán)的loss計(jì)算處調(diào)用# 在Ultralytics trainer的_do_train循環(huán)內(nèi) loss, loss_items self.model(batch) # 原有的任務(wù)損失 loss loss sparsity_loss(self.model) # 加上稀疏化正則 self.loss loss self.optimizer.zero_grad() loss.backward() self.optimizer.step()實(shí)際接入時(shí)要注意一個(gè)細(xì)節(jié)Ultralytics在不同版本里對(duì)AMP混合精度和梯度縮放的處理方式不同你需要確保sparsity_loss返回的是一個(gè)普通tensorfloat16和float32的類(lèi)型轉(zhuǎn)換正確即可。如果加了正則之后出現(xiàn)梯度異常先檢查一下是不是這里的類(lèi)型沒(méi)對(duì)齊。3.2 怎么判斷稀疏化訓(xùn)練到位了稀疏化訓(xùn)練不能瞎跑幾個(gè)epoch就算數(shù)。我訓(xùn)練過(guò)程中會(huì)周期性記錄所有BN層γ值的分布。具體做法是寫(xiě)一個(gè)小腳本加載當(dāng)前模型把每一個(gè)BatchNorm2d的weight.data收集起來(lái)然后打印直方圖。你希望看到的分布是明顯的雙峰左側(cè)在0附近聚集大量通道右側(cè)還有一個(gè)相對(duì)較高的峰。如果訓(xùn)練了幾十個(gè)epoch之后γ分布基本沒(méi)變化說(shuō)明SPARSITY_LAMBDA太小了懲罰力度不夠反過(guò)來(lái)如果γ值幾乎全部壓到了0附近說(shuō)明λ太大模型表達(dá)能力已經(jīng)被破壞這種狀態(tài)下剪完模型基本廢掉。我自己的經(jīng)驗(yàn)閾值YOLOv8s在COCO風(fēng)格數(shù)據(jù)集上SPARSITY_LAMBDA從1e-4起步每跑30個(gè)epoch觀察一次γ分布不夠稀疏就加到5e-4上限不要超過(guò)1e-3。要注意這個(gè)參數(shù)跟數(shù)據(jù)集、batch size都有關(guān)系換數(shù)據(jù)集后必須重新觀察分布不能一個(gè)數(shù)值通吃。3.3 稀疏化訓(xùn)練的完整參數(shù)參考我實(shí)際跑通的一套配置大致如下供參考參數(shù)建議值備注SPARSITY_LAMBDA1e-4 ~ 1e-3以γ分布雙峰為標(biāo)準(zhǔn)稀疏化訓(xùn)練epoch100 ~ 300在預(yù)訓(xùn)練權(quán)重基礎(chǔ)上繼續(xù)初始學(xué)習(xí)率0.01配合余弦衰減批量大小64或32一致即可權(quán)重衰減5e-4保持常規(guī)訓(xùn)練習(xí)慣數(shù)據(jù)增強(qiáng)與正式訓(xùn)練一致盡量保持一致稀疏化訓(xùn)練結(jié)束之后不要急著剪。先保存一份稀疏化權(quán)重然后花十分鐘驗(yàn)證一下加載模型遍歷所有BN層統(tǒng)計(jì)γ絕對(duì)值小于1e-3的通道占比。這個(gè)比例基本決定了你后續(xù)剪枝的性?xún)r(jià)比比如占比超過(guò)60%說(shuō)明模型有充足的水分可以擠。4. 剪枝源碼實(shí)現(xiàn)定位BN層、算mask、重建C2f與SPPF4.1 剪枝腳本的完整主流程剪枝腳本是整個(gè)方案的核心工程。我的思路是不直接在原模型對(duì)象上做in-place修改而是新建一個(gè)同構(gòu)的模型通過(guò)計(jì)算好的保留通道索引逐層重建最后把新state_dict拷貝進(jìn)去。這樣最安全出錯(cuò)時(shí)可以隨時(shí)回到原始權(quán)重。主流程可以拆成五步加載稀疏化訓(xùn)練好的模型權(quán)重。遍歷模型所有BatchNorm2d層收集γ絕對(duì)值。按剪枝比例算出全局γ閾值。結(jié)合C2f、SPPF、Concat、shortcut的結(jié)構(gòu)約束對(duì)每層生成輸入通道和輸出通道的保留索引。重建所有Conv和BN層把原始權(quán)重按照保留索引拷貝到新層里保存新模型結(jié)構(gòu)和權(quán)重。4.2 單層Conv-BN通道裁剪的基準(zhǔn)代碼先寫(xiě)最基礎(chǔ)的通道級(jí)裁剪函數(shù)。假設(shè)某個(gè)卷積層后面跟著一個(gè)BN層我們已知輸入通道保留索引keep_idx_in和輸出通道保留索引keep_idx_out需要生成新的Conv和BN。import torch import torch.nn as nn def prune_conv_bn(conv, bn, keep_idx_in, keep_idx_out): # 裁剪卷積層 new_conv nn.Conv2d( in_channelskeep_idx_in.size(0), out_channelskeep_idx_out.size(0), kernel_sizeconv.kernel_size, strideconv.stride, paddingconv.padding, dilationconv.dilation, groupsconv.groups, biasconv.bias is not None ) # 權(quán)重形狀 [out_channels, in_channels/groups, kH, kW] new_conv.weight.data conv.weight.data[keep_idx_out][:, keep_idx_in, :, :] if conv.bias is not None: new_conv.bias.data conv.bias.data[keep_idx_out] # 裁剪BN層 new_bn nn.BatchNorm2d(keep_idx_out.size(0)) new_bn.weight.data bn.weight.data[keep_idx_out] new_bn.bias.data bn.bias.data[keep_idx_out] new_bn.running_mean.data bn.running_mean.data[keep_idx_out] new_bn.running_var.data bn.running_var.data[keep_idx_out] new_bn.num_batches_tracked bn.num_batches_tracked return new_conv, new_bn這個(gè)函數(shù)是整個(gè)剪枝源碼的地基。YOLOv8s里幾乎所有的Conv后面都跟著B(niǎo)N所以只要把每一對(duì)Conv-BN的keep_idx算出來(lái)就能用這個(gè)函數(shù)逐層重建。真正難的是keep_idx怎么算尤其是C2f模塊。4.3 C2f的mask對(duì)齊shortcut和concat是兩個(gè)繞不開(kāi)的約束C2f模塊的剪枝不能只考慮單個(gè)BN的γ值必須把結(jié)構(gòu)約束考慮進(jìn)去?；氐?.2的forward代碼C2f產(chǎn)生的分支輸出最終會(huì)在通道維度上concat因此所有參與concat的分支必須共享同一套保留通道索引否則拼接后通道位置就對(duì)不上了。更棘手的是Bottleneck內(nèi)部的shortcut。Bottleneck的forward邏輯是這樣的x經(jīng)過(guò)cv1再經(jīng)cv2如果shortcut開(kāi)啟且輸入輸出維度一致就把輸入x直接加到輸出上。這意味著cv2輸出通道的保留mask必須和該Bottleneck輸入通道的保留mask一致否則逐元素相加必然維度沖突。所以實(shí)際處理C2f時(shí)我是把整個(gè)模塊作為一個(gè)通道組來(lái)對(duì)待def compute_c2f_mask(module, global_threshold): 收集C2f模塊內(nèi)所有相關(guān)BN的γ絕對(duì)值合并出一組通道保留mask。 返回的是一個(gè)與C2f內(nèi)部通道維度等長(zhǎng)的布爾mask。 gamma_list [] for m in module.modules(): if isinstance(m, nn.BatchNorm2d): gamma_list.append(m.weight.data.abs()) all_gamma torch.cat(gamma_list) # 也可以用與全局thr的比較來(lái)做但為了保險(xiǎn)推薦按模塊內(nèi)γ中位數(shù) mask all_gamma global_threshold # 需要把所有對(duì)齊約束體現(xiàn)在返回的mask上 return mask更嚴(yán)格的做法是對(duì)一個(gè)具體的C2f把cv1輸出通道對(duì)應(yīng)的γ前c和后c兩部分、每個(gè)Bottleneck里cv1和cv2的γ全部收集起來(lái)Reshape成對(duì)齊形狀逐位置取最大值。如果同一個(gè)位置在任何一個(gè)BN里的γ高于閾值就保留這個(gè)位置。這樣可以保證concat和shortcut兩條約束同時(shí)滿足。這段邏輯寫(xiě)起來(lái)確實(shí)繁瑣但它是yolov8s剪枝質(zhì)量的關(guān)鍵。如果你偷懶跳過(guò)只按γ排序裁通道裁出來(lái)的模型根本沒(méi)法前向推理或者loss直接爆掉。4.4 SPPF和檢測(cè)頭的處理策略SPPF模塊在剪枝時(shí)也是個(gè)特殊位置。它內(nèi)部結(jié)構(gòu)是cv1 - 三個(gè)5×5 MaxPool - concat - cv2池化層不改變通道數(shù)concat把4份特征拼在一起。由于SPPF本身沒(méi)有BN層不能直接用γ來(lái)判斷通道重要性。我采取的辦法是用它前后相鄰的BN層mask來(lái)約束配合Backbone主干和Neck的關(guān)聯(lián)mask同時(shí)裁剪cv1的輸出通道和cv2的輸入通道。檢測(cè)頭的處理需要克制。YOLOv8s的檢測(cè)頭是Decoupled結(jié)構(gòu)包含cls分支和box分支box分支還帶DFL輸出維度與reg_max和類(lèi)別數(shù)綁定。cls和box的最終輸出通道都不能隨便剪否則損失函數(shù)和部署后處理都會(huì)出問(wèn)題。但檢測(cè)頭內(nèi)部的1×1卷積輸入通道來(lái)自Neck這些輸入通道可以跟著Neck的mask一起裁剪。實(shí)際操作中我只對(duì)檢測(cè)頭做輸入通道的同步裁剪輸出通道一律不動(dòng)。5. 微調(diào)與精度恢復(fù)剪完不是結(jié)束重訓(xùn)才是開(kāi)始5.1 為什么剪完必須微調(diào)剪枝不是在模型文件上做減法那么簡(jiǎn)單。通道被剪掉之后模型權(quán)重從一個(gè)預(yù)訓(xùn)練分布中被硬生生截取了一部分剩余權(quán)重之間的協(xié)同關(guān)系已經(jīng)被破壞。更關(guān)鍵的是BN層的running_mean和running_var統(tǒng)計(jì)量也失效了——它們統(tǒng)計(jì)的是剪枝前通道的分布通道被刪掉后這些統(tǒng)計(jì)值完全沒(méi)有意義。如果不微調(diào)剪完的模型幾乎必然在驗(yàn)證集上崩盤(pán)。微調(diào)的本質(zhì)是讓剩余通道重新適應(yīng)新的模型結(jié)構(gòu)把損失掉的精度找回來(lái)。這一步不能省也急不得。5.2 微調(diào)階段的參數(shù)設(shè)置微調(diào)最忌諱的就是按正常訓(xùn)練的學(xué)習(xí)率來(lái)跑。我的一般做法參數(shù)建議值說(shuō)明初始學(xué)習(xí)率1e-4約為正常訓(xùn)練的1/10學(xué)習(xí)率調(diào)度余弦或階梯下降后期手動(dòng)降一次學(xué)習(xí)率訓(xùn)練的epoch30 ~ 50不需要長(zhǎng)訓(xùn)重點(diǎn)是恢復(fù)而非繼續(xù)提升稀疏化正則關(guān)閉微調(diào)階段不加γ正則BN層保持開(kāi)啟讓統(tǒng)計(jì)量自適應(yīng)新結(jié)構(gòu)微調(diào)階段不要凍結(jié)BN層。剪枝后BN的統(tǒng)計(jì)量本來(lái)就需要重新估計(jì)凍結(jié)它會(huì)讓模型在驗(yàn)證集上表現(xiàn)很差。另外微調(diào)數(shù)據(jù)集建議用訓(xùn)練集的全量數(shù)據(jù)不要只取子集否則采樣偏差會(huì)影響統(tǒng)計(jì)量重估。5.3 剪枝比例與精度之間的取舍參考以一個(gè)在自有數(shù)據(jù)集上mAP約90%的YOLOv8s模型為例參考我的實(shí)操經(jīng)驗(yàn)剪枝比例和精度變化大致如下剪枝比例剩余通道比例參數(shù)量估算mAP變化部署加速效果30%70%約5.5M幾乎無(wú)下降約20%~30%50%50%約4.2M下降0.5~1.5個(gè)百分點(diǎn)約40%~50%70%30%約2.6M下降2~4個(gè)百分點(diǎn)約60%~70%85%15%約1.8M大幅下降恢復(fù)困難約75%以上注意這個(gè)表是基于我自己的工程估算不是標(biāo)準(zhǔn)答案。不同數(shù)據(jù)集、不同檢測(cè)難度、不同稀疏化訓(xùn)練質(zhì)量都會(huì)讓數(shù)字偏移。但有一個(gè)規(guī)律是通用的剪枝比例超過(guò)50%之后精度衰減會(huì)加速變快微調(diào)挽回的空間越來(lái)越有限。部署目標(biāo)很緊的話建議把50%當(dāng)優(yōu)選檔先跑通再繼續(xù)往下壓。5.4 導(dǎo)出部署剪枝與量化疊加的效果剪枝是在結(jié)構(gòu)上減通道量化是把float32換成int8或float16這兩者是正交的可以疊加。實(shí)際上絕大多數(shù)邊緣設(shè)備上真正的大頭收益來(lái)自?xún)烧呓Y(jié)合剪枝之后通道數(shù)變少int8量化后的有效計(jì)算量大幅度縮小帶寬壓力也顯著降低。我常用的路徑是微調(diào)結(jié)束之后把剪枝模型導(dǎo)出為ONNX再用TensorRT或OpenVINO做推理優(yōu)化。由于剪枝后的模型已經(jīng)在結(jié)構(gòu)上變窄TensorRT在構(gòu)建engine時(shí)會(huì)自動(dòng)裁剪不必要的計(jì)算實(shí)測(cè)在CPU上使用OpenVINO能比原始模型快2到3倍在NVIDIA邊緣設(shè)備上配合int8量化再疊加一層加速。6. 踩坑實(shí)錄從稀疏化到部署的六個(gè)經(jīng)典問(wèn)題6.1 SPARSITY_LAMBDA設(shè)太大γ全部歸零我第一輪稀疏化訓(xùn)練就踩了這個(gè)坑。跑了20個(gè)epoch后發(fā)現(xiàn)γ分布幾乎全部塌到0模型輸出的loss還在往下降但驗(yàn)證集mAP已經(jīng)慘不忍睹。排查鏈路其實(shí)很短先打印γ的均值和標(biāo)準(zhǔn)差發(fā)現(xiàn)均值只有1e-5量級(jí)基本確定懲罰過(guò)強(qiáng)。解決方法是把SPARSITY_LAMBDA從5e-4降到1e-4重新從預(yù)訓(xùn)練權(quán)重起步訓(xùn)練。記住一個(gè)原則γ稀疏化是讓部分通道趨近0而不是把所有通道都打死分布圖應(yīng)該始終保留一個(gè)非零的峰。6.2 剪完模型loss變成NaN剪枝之后前向推理正常但微調(diào)時(shí)loss直接跳NaN。我排查的思路是先檢查模型輸出的shape是否正常再逐層檢查卷積權(quán)重里有沒(méi)有NaN值。最后發(fā)現(xiàn)原因出在剪枝后的某個(gè)C2f模塊它的concat輸入通道被剪成了0導(dǎo)致某個(gè)Bottleneck的輸出是空張量乘到loss上就變成NaN。解決方法是給每個(gè)mask加一個(gè)強(qiáng)制約束任何卷積層的輸出通道數(shù)不得低于預(yù)設(shè)下限比如至少保留16通道。這一步在生成mask的代碼里直接做兜底判斷即可。6.3 shortcut維度對(duì)不上報(bào)size mismatch這是C2f剪枝里最容易出現(xiàn)的問(wèn)題。報(bào)錯(cuò)信息會(huì)指向某個(gè)Bottleneck的加法操作說(shuō)兩個(gè)張量shape不一致。根因幾乎都是同一個(gè)沒(méi)有對(duì)shortcut兩側(cè)的mask做對(duì)齊。修復(fù)方式是我在4.3里描述過(guò)的方案生成mask時(shí)把同一個(gè)C2f模塊內(nèi)所有需要對(duì)齊的位置統(tǒng)一計(jì)算。排查的時(shí)候不要只看報(bào)錯(cuò)的那一層要從這個(gè)C2f的cv1開(kāi)始順藤摸瓜把整個(gè)模塊的channel流向畫(huà)出來(lái)再對(duì)照mask。6.4 剪完保存的模型文件反而變大了這是一個(gè)非常迷惑的現(xiàn)象。模型通道數(shù)減小了按理說(shuō)文件應(yīng)該變小但如果你用torch.save(model)直接保存整個(gè)模型對(duì)象會(huì)有很多冗余——包括模型結(jié)構(gòu)定義、優(yōu)化器狀態(tài)、各類(lèi)緩存字段。解決方法是剪枝后只保存模型的state_dict或者更好一點(diǎn)把剪枝后的結(jié)構(gòu)寫(xiě)成一個(gè)新的YAML配置文件再單獨(dú)保存對(duì)應(yīng)的state_dict。這樣剪枝后的模型文件才會(huì)真正反映通道數(shù)的減少導(dǎo)出ONNX后大小也才可信。6.5 PyTorch里測(cè)速?zèng)]有提升覺(jué)得白剪了剪枝之后用PyTorch的CPU直接推理發(fā)現(xiàn)速度幾乎沒(méi)變第一反應(yīng)往往是方案沒(méi)效果。但我排查后發(fā)現(xiàn)這是PyTorch在CPU上的卷積算子對(duì)小通道數(shù)沒(méi)有做充分優(yōu)化通道減少了但算子調(diào)度開(kāi)銷(xiāo)還在。真正的提速必須依靠部署側(cè)的優(yōu)化。把剪枝模型導(dǎo)出為ONNX再用OpenVINO或TensorRT跑一遍速度差異立刻出來(lái)了。所以測(cè)速一定要在目標(biāo)部署框架里測(cè)不要在PyTorch里過(guò)早下結(jié)論。6.6 state_dict鍵名錯(cuò)亂加載模型報(bào)錯(cuò)剪枝重建模型過(guò)程中最容易翻車(chē)的環(huán)節(jié)是拼接新的state_dict。因?yàn)榕f模型的鍵名是model.0.cv1.conv.weight這種形式你重建的新C2f層如果命名規(guī)則稍有不同加載時(shí)就會(huì)報(bào)unexpected key或者missing key。我的習(xí)慣是在剪枝腳本末尾加一個(gè)驗(yàn)證函數(shù)把新舊state_dict逐層做形狀比對(duì)發(fā)現(xiàn)不匹配立即打印出錯(cuò)層的鍵名。寧可多花幾分鐘跑校驗(yàn)也別在部署階段才發(fā)現(xiàn)權(quán)重對(duì)不上。最后再分享一個(gè)我在實(shí)操中一直在用的小原則剪枝前一定保存好原始權(quán)重和稀疏化權(quán)重兩份備份每剪完一個(gè)比例立刻做一次完整的前向驗(yàn)證和mAP評(píng)估。剪枝是一個(gè)遞增式的過(guò)程不要一次性壓到目標(biāo)比例先剪30%驗(yàn)證流程沒(méi)問(wèn)題再逐步往上加。這樣即便出問(wèn)題也能清楚地知道是哪個(gè)環(huán)節(jié)引入的回退和重來(lái)的成本都低得多。 YOLOv8s結(jié)構(gòu)化剪枝實(shí)戰(zhàn)：從稀疏化訓(xùn)練到部署加速
先講一個(gè)我前陣子真實(shí)遇到的部署場(chǎng)景。模型用YOLOv8s在GPU上訓(xùn)得漂漂亮亮精度也在線可一旦挪到邊緣盒子上單幀推理直接干到四十多毫秒真到了要拼實(shí)時(shí)性的時(shí)候幀率死活穩(wěn)不住。調(diào)NMS、降分辨率、換解碼方式折騰了一圈效果都有限最后真正把時(shí)延壓下來(lái)的是剪枝。剪枝這個(gè)詞大家都不陌生可一旦落到Y(jié)OLOv8s上情況就跟網(wǎng)上一堆舊教程里的YOLOv5源碼完全不一樣了。yolov8s在結(jié)構(gòu)上把C3換成了C2fdetect head也改成了decoupled結(jié)構(gòu)直接把v5時(shí)代的剪枝腳本搬過(guò)來(lái)十有八九跑不通。我花了大概兩周時(shí)間把整條鏈路完整跑通從稀疏化訓(xùn)練、通道裁剪、結(jié)構(gòu)重建到微調(diào)恢復(fù)精度全部基于ultralytics YOLOv8s源碼改造。這篇博文就是把一套能落地的yolov8s剪枝源碼思路拆開(kāi)來(lái)講每一步為什么要這么做核心代碼應(yīng)該寫(xiě)在哪有哪些坑我替你先踩過(guò)了。內(nèi)容適合已經(jīng)在用YOLOv8做項(xiàng)目、對(duì)PyTorch不陌生、但還沒(méi)系統(tǒng)接觸過(guò)模型壓縮的開(kāi)發(fā)者。1. YOLOv8s結(jié)構(gòu)解析先弄清楚算力燒在了哪里1.1 YOLOv8s各模塊的算力分布剪枝之前我建議你先別急著翻代碼而是花半天時(shí)間把你手里這個(gè)模型的算力分布看明白。YOLOv8s在640×640輸入分辨率下總計(jì)算量約在28GFLOPs左右參數(shù)量約11.2M含檢測(cè)頭。這個(gè)體量在YOLOv8家族里屬于small檔看起來(lái)不大但對(duì)邊緣設(shè)備來(lái)說(shuō)依然偏重。從結(jié)構(gòu)上看YOLOv8s由三大部分組成Backbone負(fù)責(zé)從輸入圖像提取多尺度特征Neck是PAN-FPN結(jié)構(gòu)負(fù)責(zé)把不同層的特征融合起來(lái)最后是Decoupled檢測(cè)頭分別輸出分類(lèi)和回歸分支。我自己的實(shí)測(cè)統(tǒng)計(jì)里Backbone占據(jù)了接近一半的計(jì)算量Neck占了剩下三分之一左右檢測(cè)頭因?yàn)榇罅渴褂?×1卷積FLOPs占比其實(shí)沒(méi)有想象中高。這里要重點(diǎn)說(shuō)的是Backbone和Neck里大量的3×3卷積。對(duì)邊緣設(shè)備而言3×3卷積在CPU和多數(shù)NPU上都是計(jì)算熱點(diǎn)剪掉一層輸出通道省下的時(shí)間非?？捎^。所以剪枝的核心目標(biāo)很明確把你的裁剪預(yù)算優(yōu)先花在Backbone和Neck的C2f模塊上檢測(cè)頭做一個(gè)克制的輸入通道裁剪即可最后的輸出層盡量保持不動(dòng)。1.2 C2f的split機(jī)制剪枝時(shí)為什么比v5麻煩YOLOv8s和YOLOv5s在結(jié)構(gòu)上最大的變化就是C3換成了C2f。C2f模塊里有一個(gè)很關(guān)鍵的split操作def forward(self, x): y list(self.cv1(x).chunk(2, 1)) # 沿通道維度對(duì)半拆開(kāi) y.extend(m(y[-1]) for m in self.m) # 后半路依次經(jīng)過(guò)多個(gè)Bottleneck return self.cv2(torch.cat(y, 1))cv1先把輸入通道映射到2c然后沿通道方向?qū)Π氩鸪蓛煞?。前半份直接送去concat后半份經(jīng)過(guò)一組Bottleneck之后再送去concat最終由cv2把(n1)c壓縮回目標(biāo)輸出通道。這個(gè)split機(jī)制直接導(dǎo)致了一個(gè)問(wèn)題C2f內(nèi)部所有分支的通道組必須保持對(duì)齊才能保證最終concat之后通道位置嚴(yán)格對(duì)應(yīng)。如果照搬YOLOv5時(shí)代的“逐個(gè)BN層剪γ”方案幾乎必然在C2f模塊的torch.cat處報(bào)shape mismatch或者更隱蔽的——通道錯(cuò)位。我在寫(xiě)剪枝源碼時(shí)把C2f當(dāng)成一個(gè)整體去生成統(tǒng)一的通道m(xù)ask而不是單獨(dú)處理每個(gè)Bottleneck這一點(diǎn)是YOLOv8s剪枝區(qū)別于v5的核心難點(diǎn)后面會(huì)專(zhuān)門(mén)講。1.3 結(jié)構(gòu)化剪枝與非結(jié)構(gòu)化剪枝想清楚再動(dòng)手模型剪枝分為兩大類(lèi)非結(jié)構(gòu)化剪枝unstructured pruning剪的是單個(gè)權(quán)重稀疏度分布不均勻模型文件里全是零值但通用硬件上幾乎得不到加速收益除非配套稀疏矩陣專(zhuān)用庫(kù)。結(jié)構(gòu)化剪枝structured pruning以通道或?yàn)V波器為基本單位整體裁剪直接改變卷積層的輸入輸出維度任何推理框架都能拿到實(shí)實(shí)在在的加速。我見(jiàn)過(guò)不少朋友把SparseGPT這類(lèi)非結(jié)構(gòu)化方案用在YOLOv8s上精度保住了但導(dǎo)出到TensorRT或ONNX Runtime之后推理速度紋絲不動(dòng)。原因很簡(jiǎn)單民用級(jí)硬件基本不為稀疏權(quán)重做優(yōu)化。所以你如果目的是部署加速直接選結(jié)構(gòu)化通道剪枝不要去碰非結(jié)構(gòu)化路線。2. 剪枝方案選型基于BN層γ系數(shù)做通道重要性判斷2.1 Network Slimmingγ系數(shù)就是天然的通道重要度指標(biāo)通道剪枝要回答一個(gè)核心問(wèn)題怎么判定哪些通道不重要最經(jīng)典也最省事的方案是ICCV 2017年的Network Slimming方法。它利用的是卷積層后面那個(gè)幾乎被人忽略的BatchNorm層。BN層的計(jì)算公式是y γ * (x - μ) / √(σ2 ε) β其中γ和β都是可學(xué)習(xí)的仿射參數(shù)。γ是縮放系數(shù)如果訓(xùn)練后某個(gè)通道的γ變得非常接近0那么該通道的輸出就幾乎恒定在一個(gè)由β決定的固定值上信息量趨近于零這個(gè)通道就可以被安全地剪掉。關(guān)鍵就在于普通訓(xùn)練得到的γ數(shù)值分布不夠“極端”很多通道的γ值都落在0.5到1.5之間你沒(méi)法說(shuō)誰(shuí)更重要。所以需要人為給γ加上L1正則約束讓它稀疏化。訓(xùn)練時(shí)的總損失從L變成L λ * Σ|γ|其中λ是稀疏化懲罰系數(shù)。這個(gè)正則項(xiàng)會(huì)鼓勵(lì)γ往0靠攏訓(xùn)練結(jié)束后γ的分布會(huì)形成明顯的雙峰一堆極度接近0剩下的保持相對(duì)大的數(shù)值。普通訓(xùn)練后直接按γ剪和稀疏化訓(xùn)練之后按γ剪效果差別非常大。前者誤傷率高剪完精度掉得一塌糊涂微調(diào)也救不回來(lái)后者才能精準(zhǔn)識(shí)別出真正冗余的通道。這是整個(gè)剪枝方案成立的前提千萬(wàn)別跳過(guò)。2.2 為什么不考慮低秩分解和基于梯度的通道評(píng)估你可能還會(huì)看到一些方案用通道注意力權(quán)重、梯度敏感度或者低秩分解來(lái)選擇剪枝通道。低秩分解能把卷積核拆成小矩陣序列但從工程角度講拆完之后的模型結(jié)構(gòu)要重寫(xiě)底層算子才能發(fā)揮加速效果和普通推理框架的兼容性很一般?；谔荻鹊姆椒ㄊ歉?xì)的敏感度分析但要在每個(gè)剪枝候選上做一次反向傳播算力開(kāi)銷(xiāo)不低性?xún)r(jià)比反而不如γ稀疏化。對(duì)YOLOv8s這種層數(shù)不算極端、結(jié)構(gòu)相對(duì)規(guī)整的模型來(lái)說(shuō)γ系數(shù)方案是平衡效果與工程復(fù)雜度最好的選擇。它不需要額外訓(xùn)練一個(gè)評(píng)判網(wǎng)絡(luò)不需要逐層搜索敏感度只需要在訓(xùn)練loss上加一個(gè)正則項(xiàng)剪枝時(shí)拿一把全局閾值來(lái)篩通道整個(gè)流程完全可控。2.3 開(kāi)工前的基礎(chǔ)配置我建議你基于ultralytics官方的YOLOv8s源碼來(lái)做版本選8.0.x或8.1.x都行但1.x分支之后的代碼結(jié)構(gòu)有些調(diào)整注意保持一致。環(huán)境方面PyTorch 1.13以上、CUDA 11.7以上就行不需要額外裝稀疏計(jì)算庫(kù)剪枝前的稀疏化訓(xùn)練就是常規(guī)訓(xùn)練。另外準(zhǔn)備一個(gè)你自己的、已經(jīng)收斂到可用狀態(tài)的YOLOv8s權(quán)重文件。最佳實(shí)踐是先正常訓(xùn)練到float32的精度峰值再?gòu)倪@個(gè)權(quán)重開(kāi)始做稀疏化訓(xùn)練不要從頭就在loss里加稀疏正則否則模型收斂速度和精度上限都會(huì)受影響。3. 稀疏化訓(xùn)練源碼改造給YOLOv8s訓(xùn)練循環(huán)加上γ正則3.1 在Ultralytics訓(xùn)練循環(huán)中插入稀疏化損失Ultralytics YOLOv8的訓(xùn)練邏輯集中在trainer.py的_do_train方法里核心循環(huán)對(duì)每個(gè)batch做前向和反向。為了加稀疏化損失我的做法是在trainer里覆寫(xiě)一個(gè)loss計(jì)算hook。不同版本源碼結(jié)構(gòu)有差異但思路是通用的在模型前向算出的主loss之后、loss.backward()之前加入BN層γ的L1正則項(xiàng)。# sparse_loss.py import torch import torch.nn as nn SPARSITY_LAMBDA 1e-4 def sparsity_loss(model): reg_loss 0.0 for m in model.modules(): if isinstance(m, torch.nn.BatchNorm2d): reg_loss torch.abs(m.weight).sum() return SPARSITY_LAMBDA * reg_loss然后在訓(xùn)練循環(huán)的loss計(jì)算處調(diào)用# 在Ultralytics trainer的_do_train循環(huán)內(nèi) loss, loss_items self.model(batch) # 原有的任務(wù)損失 loss loss sparsity_loss(self.model) # 加上稀疏化正則 self.loss loss self.optimizer.zero_grad() loss.backward() self.optimizer.step()實(shí)際接入時(shí)要注意一個(gè)細(xì)節(jié)Ultralytics在不同版本里對(duì)AMP混合精度和梯度縮放的處理方式不同你需要確保sparsity_loss返回的是一個(gè)普通tensorfloat16和float32的類(lèi)型轉(zhuǎn)換正確即可。如果加了正則之后出現(xiàn)梯度異常先檢查一下是不是這里的類(lèi)型沒(méi)對(duì)齊。3.2 怎么判斷稀疏化訓(xùn)練到位了稀疏化訓(xùn)練不能瞎跑幾個(gè)epoch就算數(shù)。我訓(xùn)練過(guò)程中會(huì)周期性記錄所有BN層γ值的分布。具體做法是寫(xiě)一個(gè)小腳本加載當(dāng)前模型把每一個(gè)BatchNorm2d的weight.data收集起來(lái)然后打印直方圖。你希望看到的分布是明顯的雙峰左側(cè)在0附近聚集大量通道右側(cè)還有一個(gè)相對(duì)較高的峰。如果訓(xùn)練了幾十個(gè)epoch之后γ分布基本沒(méi)變化說(shuō)明SPARSITY_LAMBDA太小了懲罰力度不夠反過(guò)來(lái)如果γ值幾乎全部壓到了0附近說(shuō)明λ太大模型表達(dá)能力已經(jīng)被破壞這種狀態(tài)下剪完模型基本廢掉。我自己的經(jīng)驗(yàn)閾值YOLOv8s在COCO風(fēng)格數(shù)據(jù)集上SPARSITY_LAMBDA從1e-4起步每跑30個(gè)epoch觀察一次γ分布不夠稀疏就加到5e-4上限不要超過(guò)1e-3。要注意這個(gè)參數(shù)跟數(shù)據(jù)集、batch size都有關(guān)系換數(shù)據(jù)集后必須重新觀察分布不能一個(gè)數(shù)值通吃。3.3 稀疏化訓(xùn)練的完整參數(shù)參考我實(shí)際跑通的一套配置大致如下供參考參數(shù)建議值備注SPARSITY_LAMBDA1e-4 ~ 1e-3以γ分布雙峰為標(biāo)準(zhǔn)稀疏化訓(xùn)練epoch100 ~ 300在預(yù)訓(xùn)練權(quán)重基礎(chǔ)上繼續(xù)初始學(xué)習(xí)率0.01配合余弦衰減批量大小64或32一致即可權(quán)重衰減5e-4保持常規(guī)訓(xùn)練習(xí)慣數(shù)據(jù)增強(qiáng)與正式訓(xùn)練一致盡量保持一致稀疏化訓(xùn)練結(jié)束之后不要急著剪。先保存一份稀疏化權(quán)重然后花十分鐘驗(yàn)證一下加載模型遍歷所有BN層統(tǒng)計(jì)γ絕對(duì)值小于1e-3的通道占比。這個(gè)比例基本決定了你后續(xù)剪枝的性?xún)r(jià)比比如占比超過(guò)60%說(shuō)明模型有充足的水分可以擠。4. 剪枝源碼實(shí)現(xiàn)定位BN層、算mask、重建C2f與SPPF4.1 剪枝腳本的完整主流程剪枝腳本是整個(gè)方案的核心工程。我的思路是不直接在原模型對(duì)象上做in-place修改而是新建一個(gè)同構(gòu)的模型通過(guò)計(jì)算好的保留通道索引逐層重建最后把新state_dict拷貝進(jìn)去。這樣最安全出錯(cuò)時(shí)可以隨時(shí)回到原始權(quán)重。主流程可以拆成五步加載稀疏化訓(xùn)練好的模型權(quán)重。遍歷模型所有BatchNorm2d層收集γ絕對(duì)值。按剪枝比例算出全局γ閾值。結(jié)合C2f、SPPF、Concat、shortcut的結(jié)構(gòu)約束對(duì)每層生成輸入通道和輸出通道的保留索引。重建所有Conv和BN層把原始權(quán)重按照保留索引拷貝到新層里保存新模型結(jié)構(gòu)和權(quán)重。4.2 單層Conv-BN通道裁剪的基準(zhǔn)代碼先寫(xiě)最基礎(chǔ)的通道級(jí)裁剪函數(shù)。假設(shè)某個(gè)卷積層后面跟著一個(gè)BN層我們已知輸入通道保留索引keep_idx_in和輸出通道保留索引keep_idx_out需要生成新的Conv和BN。import torch import torch.nn as nn def prune_conv_bn(conv, bn, keep_idx_in, keep_idx_out): # 裁剪卷積層 new_conv nn.Conv2d( in_channelskeep_idx_in.size(0), out_channelskeep_idx_out.size(0), kernel_sizeconv.kernel_size, strideconv.stride, paddingconv.padding, dilationconv.dilation, groupsconv.groups, biasconv.bias is not None ) # 權(quán)重形狀 [out_channels, in_channels/groups, kH, kW] new_conv.weight.data conv.weight.data[keep_idx_out][:, keep_idx_in, :, :] if conv.bias is not None: new_conv.bias.data conv.bias.data[keep_idx_out] # 裁剪BN層 new_bn nn.BatchNorm2d(keep_idx_out.size(0)) new_bn.weight.data bn.weight.data[keep_idx_out] new_bn.bias.data bn.bias.data[keep_idx_out] new_bn.running_mean.data bn.running_mean.data[keep_idx_out] new_bn.running_var.data bn.running_var.data[keep_idx_out] new_bn.num_batches_tracked bn.num_batches_tracked return new_conv, new_bn這個(gè)函數(shù)是整個(gè)剪枝源碼的地基。YOLOv8s里幾乎所有的Conv后面都跟著B(niǎo)N所以只要把每一對(duì)Conv-BN的keep_idx算出來(lái)就能用這個(gè)函數(shù)逐層重建。真正難的是keep_idx怎么算尤其是C2f模塊。4.3 C2f的mask對(duì)齊shortcut和concat是兩個(gè)繞不開(kāi)的約束C2f模塊的剪枝不能只考慮單個(gè)BN的γ值必須把結(jié)構(gòu)約束考慮進(jìn)去?；氐?.2的forward代碼C2f產(chǎn)生的分支輸出最終會(huì)在通道維度上concat因此所有參與concat的分支必須共享同一套保留通道索引否則拼接后通道位置就對(duì)不上了。更棘手的是Bottleneck內(nèi)部的shortcut。Bottleneck的forward邏輯是這樣的x經(jīng)過(guò)cv1再經(jīng)cv2如果shortcut開(kāi)啟且輸入輸出維度一致就把輸入x直接加到輸出上。這意味著cv2輸出通道的保留mask必須和該Bottleneck輸入通道的保留mask一致否則逐元素相加必然維度沖突。所以實(shí)際處理C2f時(shí)我是把整個(gè)模塊作為一個(gè)通道組來(lái)對(duì)待def compute_c2f_mask(module, global_threshold): 收集C2f模塊內(nèi)所有相關(guān)BN的γ絕對(duì)值合并出一組通道保留mask。 返回的是一個(gè)與C2f內(nèi)部通道維度等長(zhǎng)的布爾mask。 gamma_list [] for m in module.modules(): if isinstance(m, nn.BatchNorm2d): gamma_list.append(m.weight.data.abs()) all_gamma torch.cat(gamma_list) # 也可以用與全局thr的比較來(lái)做但為了保險(xiǎn)推薦按模塊內(nèi)γ中位數(shù) mask all_gamma global_threshold # 需要把所有對(duì)齊約束體現(xiàn)在返回的mask上 return mask更嚴(yán)格的做法是對(duì)一個(gè)具體的C2f把cv1輸出通道對(duì)應(yīng)的γ前c和后c兩部分、每個(gè)Bottleneck里cv1和cv2的γ全部收集起來(lái)Reshape成對(duì)齊形狀逐位置取最大值。如果同一個(gè)位置在任何一個(gè)BN里的γ高于閾值就保留這個(gè)位置。這樣可以保證concat和shortcut兩條約束同時(shí)滿足。這段邏輯寫(xiě)起來(lái)確實(shí)繁瑣但它是yolov8s剪枝質(zhì)量的關(guān)鍵。如果你偷懶跳過(guò)只按γ排序裁通道裁出來(lái)的模型根本沒(méi)法前向推理或者loss直接爆掉。4.4 SPPF和檢測(cè)頭的處理策略SPPF模塊在剪枝時(shí)也是個(gè)特殊位置。它內(nèi)部結(jié)構(gòu)是cv1 - 三個(gè)5×5 MaxPool - concat - cv2池化層不改變通道數(shù)concat把4份特征拼在一起。由于SPPF本身沒(méi)有BN層不能直接用γ來(lái)判斷通道重要性。我采取的辦法是用它前后相鄰的BN層mask來(lái)約束配合Backbone主干和Neck的關(guān)聯(lián)mask同時(shí)裁剪cv1的輸出通道和cv2的輸入通道。檢測(cè)頭的處理需要克制。YOLOv8s的檢測(cè)頭是Decoupled結(jié)構(gòu)包含cls分支和box分支box分支還帶DFL輸出維度與reg_max和類(lèi)別數(shù)綁定。cls和box的最終輸出通道都不能隨便剪否則損失函數(shù)和部署后處理都會(huì)出問(wèn)題。但檢測(cè)頭內(nèi)部的1×1卷積輸入通道來(lái)自Neck這些輸入通道可以跟著Neck的mask一起裁剪。實(shí)際操作中我只對(duì)檢測(cè)頭做輸入通道的同步裁剪輸出通道一律不動(dòng)。5. 微調(diào)與精度恢復(fù)剪完不是結(jié)束重訓(xùn)才是開(kāi)始5.1 為什么剪完必須微調(diào)剪枝不是在模型文件上做減法那么簡(jiǎn)單。通道被剪掉之后模型權(quán)重從一個(gè)預(yù)訓(xùn)練分布中被硬生生截取了一部分剩余權(quán)重之間的協(xié)同關(guān)系已經(jīng)被破壞。更關(guān)鍵的是BN層的running_mean和running_var統(tǒng)計(jì)量也失效了——它們統(tǒng)計(jì)的是剪枝前通道的分布通道被刪掉后這些統(tǒng)計(jì)值完全沒(méi)有意義。如果不微調(diào)剪完的模型幾乎必然在驗(yàn)證集上崩盤(pán)。微調(diào)的本質(zhì)是讓剩余通道重新適應(yīng)新的模型結(jié)構(gòu)把損失掉的精度找回來(lái)。這一步不能省也急不得。5.2 微調(diào)階段的參數(shù)設(shè)置微調(diào)最忌諱的就是按正常訓(xùn)練的學(xué)習(xí)率來(lái)跑。我的一般做法參數(shù)建議值說(shuō)明初始學(xué)習(xí)率1e-4約為正常訓(xùn)練的1/10學(xué)習(xí)率調(diào)度余弦或階梯下降后期手動(dòng)降一次學(xué)習(xí)率訓(xùn)練的epoch30 ~ 50不需要長(zhǎng)訓(xùn)重點(diǎn)是恢復(fù)而非繼續(xù)提升稀疏化正則關(guān)閉微調(diào)階段不加γ正則BN層保持開(kāi)啟讓統(tǒng)計(jì)量自適應(yīng)新結(jié)構(gòu)微調(diào)階段不要凍結(jié)BN層。剪枝后BN的統(tǒng)計(jì)量本來(lái)就需要重新估計(jì)凍結(jié)它會(huì)讓模型在驗(yàn)證集上表現(xiàn)很差。另外微調(diào)數(shù)據(jù)集建議用訓(xùn)練集的全量數(shù)據(jù)不要只取子集否則采樣偏差會(huì)影響統(tǒng)計(jì)量重估。5.3 剪枝比例與精度之間的取舍參考以一個(gè)在自有數(shù)據(jù)集上mAP約90%的YOLOv8s模型為例參考我的實(shí)操經(jīng)驗(yàn)剪枝比例和精度變化大致如下剪枝比例剩余通道比例參數(shù)量估算mAP變化部署加速效果30%70%約5.5M幾乎無(wú)下降約20%~30%50%50%約4.2M下降0.5~1.5個(gè)百分點(diǎn)約40%~50%70%30%約2.6M下降2~4個(gè)百分點(diǎn)約60%~70%85%15%約1.8M大幅下降恢復(fù)困難約75%以上注意這個(gè)表是基于我自己的工程估算不是標(biāo)準(zhǔn)答案。不同數(shù)據(jù)集、不同檢測(cè)難度、不同稀疏化訓(xùn)練質(zhì)量都會(huì)讓數(shù)字偏移。但有一個(gè)規(guī)律是通用的剪枝比例超過(guò)50%之后精度衰減會(huì)加速變快微調(diào)挽回的空間越來(lái)越有限。部署目標(biāo)很緊的話建議把50%當(dāng)優(yōu)選檔先跑通再繼續(xù)往下壓。5.4 導(dǎo)出部署剪枝與量化疊加的效果剪枝是在結(jié)構(gòu)上減通道量化是把float32換成int8或float16這兩者是正交的可以疊加。實(shí)際上絕大多數(shù)邊緣設(shè)備上真正的大頭收益來(lái)自?xún)烧呓Y(jié)合剪枝之后通道數(shù)變少int8量化后的有效計(jì)算量大幅度縮小帶寬壓力也顯著降低。我常用的路徑是微調(diào)結(jié)束之后把剪枝模型導(dǎo)出為ONNX再用TensorRT或OpenVINO做推理優(yōu)化。由于剪枝后的模型已經(jīng)在結(jié)構(gòu)上變窄TensorRT在構(gòu)建engine時(shí)會(huì)自動(dòng)裁剪不必要的計(jì)算實(shí)測(cè)在CPU上使用OpenVINO能比原始模型快2到3倍在NVIDIA邊緣設(shè)備上配合int8量化再疊加一層加速。6. 踩坑實(shí)錄從稀疏化到部署的六個(gè)經(jīng)典問(wèn)題6.1 SPARSITY_LAMBDA設(shè)太大γ全部歸零我第一輪稀疏化訓(xùn)練就踩了這個(gè)坑。跑了20個(gè)epoch后發(fā)現(xiàn)γ分布幾乎全部塌到0模型輸出的loss還在往下降但驗(yàn)證集mAP已經(jīng)慘不忍睹。排查鏈路其實(shí)很短先打印γ的均值和標(biāo)準(zhǔn)差發(fā)現(xiàn)均值只有1e-5量級(jí)基本確定懲罰過(guò)強(qiáng)。解決方法是把SPARSITY_LAMBDA從5e-4降到1e-4重新從預(yù)訓(xùn)練權(quán)重起步訓(xùn)練。記住一個(gè)原則γ稀疏化是讓部分通道趨近0而不是把所有通道都打死分布圖應(yīng)該始終保留一個(gè)非零的峰。6.2 剪完模型loss變成NaN剪枝之后前向推理正常但微調(diào)時(shí)loss直接跳NaN。我排查的思路是先檢查模型輸出的shape是否正常再逐層檢查卷積權(quán)重里有沒(méi)有NaN值。最后發(fā)現(xiàn)原因出在剪枝后的某個(gè)C2f模塊它的concat輸入通道被剪成了0導(dǎo)致某個(gè)Bottleneck的輸出是空張量乘到loss上就變成NaN。解決方法是給每個(gè)mask加一個(gè)強(qiáng)制約束任何卷積層的輸出通道數(shù)不得低于預(yù)設(shè)下限比如至少保留16通道。這一步在生成mask的代碼里直接做兜底判斷即可。6.3 shortcut維度對(duì)不上報(bào)size mismatch這是C2f剪枝里最容易出現(xiàn)的問(wèn)題。報(bào)錯(cuò)信息會(huì)指向某個(gè)Bottleneck的加法操作說(shuō)兩個(gè)張量shape不一致。根因幾乎都是同一個(gè)沒(méi)有對(duì)shortcut兩側(cè)的mask做對(duì)齊。修復(fù)方式是我在4.3里描述過(guò)的方案生成mask時(shí)把同一個(gè)C2f模塊內(nèi)所有需要對(duì)齊的位置統(tǒng)一計(jì)算。排查的時(shí)候不要只看報(bào)錯(cuò)的那一層要從這個(gè)C2f的cv1開(kāi)始順藤摸瓜把整個(gè)模塊的channel流向畫(huà)出來(lái)再對(duì)照mask。6.4 剪完保存的模型文件反而變大了這是一個(gè)非常迷惑的現(xiàn)象。模型通道數(shù)減小了按理說(shuō)文件應(yīng)該變小但如果你用torch.save(model)直接保存整個(gè)模型對(duì)象會(huì)有很多冗余——包括模型結(jié)構(gòu)定義、優(yōu)化器狀態(tài)、各類(lèi)緩存字段。解決方法是剪枝后只保存模型的state_dict或者更好一點(diǎn)把剪枝后的結(jié)構(gòu)寫(xiě)成一個(gè)新的YAML配置文件再單獨(dú)保存對(duì)應(yīng)的state_dict。這樣剪枝后的模型文件才會(huì)真正反映通道數(shù)的減少導(dǎo)出ONNX后大小也才可信。6.5 PyTorch里測(cè)速?zèng)]有提升覺(jué)得白剪了剪枝之后用PyTorch的CPU直接推理發(fā)現(xiàn)速度幾乎沒(méi)變第一反應(yīng)往往是方案沒(méi)效果。但我排查后發(fā)現(xiàn)這是PyTorch在CPU上的卷積算子對(duì)小通道數(shù)沒(méi)有做充分優(yōu)化通道減少了但算子調(diào)度開(kāi)銷(xiāo)還在。真正的提速必須依靠部署側(cè)的優(yōu)化。把剪枝模型導(dǎo)出為ONNX再用OpenVINO或TensorRT跑一遍速度差異立刻出來(lái)了。所以測(cè)速一定要在目標(biāo)部署框架里測(cè)不要在PyTorch里過(guò)早下結(jié)論。6.6 state_dict鍵名錯(cuò)亂加載模型報(bào)錯(cuò)剪枝重建模型過(guò)程中最容易翻車(chē)的環(huán)節(jié)是拼接新的state_dict。因?yàn)榕f模型的鍵名是model.0.cv1.conv.weight這種形式你重建的新C2f層如果命名規(guī)則稍有不同加載時(shí)就會(huì)報(bào)unexpected key或者missing key。我的習(xí)慣是在剪枝腳本末尾加一個(gè)驗(yàn)證函數(shù)把新舊state_dict逐層做形狀比對(duì)發(fā)現(xiàn)不匹配立即打印出錯(cuò)層的鍵名。寧可多花幾分鐘跑校驗(yàn)也別在部署階段才發(fā)現(xiàn)權(quán)重對(duì)不上。最后再分享一個(gè)我在實(shí)操中一直在用的小原則剪枝前一定保存好原始權(quán)重和稀疏化權(quán)重兩份備份每剪完一個(gè)比例立刻做一次完整的前向驗(yàn)證和mAP評(píng)估。剪枝是一個(gè)遞增式的過(guò)程不要一次性壓到目標(biāo)比例先剪30%驗(yàn)證流程沒(méi)問(wèn)題再逐步往上加。這樣即便出問(wèn)題也能清楚地知道是哪個(gè)環(huán)節(jié)引入的回退和重來(lái)的成本都低得多。