1. 項目概述與核心思路拆解1.1 為什么PIDNet值得跑一遍語義分割這幾年模型迭代速度極快從FCN到DeepLab系列再到HRNet大家基本都在做同一件事如何在保持高分辨率細節(jié)的同時把上下文語義信息也抓到位。PIDNetArXiv 2022這個模型有意思的地方在于它把經(jīng)典控制理論里的PID控制器思想直接搬進了網(wǎng)絡(luò)結(jié)構(gòu)設(shè)計用比例Proportional、積分Integral、微分Derivative三種互補的特性來指導(dǎo)特征提取流程在速度和精度之間找到了一個非常實用的平衡點。我第一次讀這篇論文的時候最直接的感觸是它的設(shè)計邏輯很“工程化”。深度學(xué)習(xí)模型越做越深但很多改進都是在堆模塊、堆通道實際落地時參數(shù)量和推理延遲都很難控制。PIDNet反而是從控制論中汲取靈感把特征提取過程分成三個并行分支P分支負責(zé)捕捉大尺度的上下文關(guān)系I分支負責(zé)保留高分辨率細節(jié)信息D分支則在兩者之間做動態(tài)調(diào)節(jié)和特征融合。這種三分支結(jié)構(gòu)既能拿到高質(zhì)量的分割結(jié)果又不會像Transformer那樣吃顯存吃到懷疑人生。對于想復(fù)現(xiàn)論文、做實驗對比、或者在這個模型基礎(chǔ)上做改進創(chuàng)新的同學(xué)來說PIDNet是一個非常合適的基座模型。訓(xùn)練流程相對透明結(jié)構(gòu)不算特別復(fù)雜改進空間又很大。你在它的語義分割框架上動刀不管是調(diào)整分支間的融合策略還是替換骨干網(wǎng)絡(luò)都能很清晰地觀察到效果差異。1.2 復(fù)現(xiàn)方案的整體選型考慮我自己實際跑過PIDNet的完整訓(xùn)練流程踩了不少坑也總結(jié)了一些經(jīng)驗。先說結(jié)論PIDNet的復(fù)現(xiàn)難度屬于中等偏上比直接調(diào)用一個封裝好的模型要麻煩但比從零手寫一個訓(xùn)練管線要輕松很多。如果你已經(jīng)熟悉語義分割的訓(xùn)練流程那核心工作量其實就是環(huán)境配置、數(shù)據(jù)組織、參數(shù)調(diào)整這三件事如果你是第一次跑分割模型那可能需要先在骨架網(wǎng)絡(luò)上多花點時間理解一下。從技術(shù)選型角度來看PIDNet官方是基于PyTorch框架實現(xiàn)的GPU訓(xùn)練是必要的純CPU環(huán)境基本跑不動。數(shù)據(jù)準備方面Cityscapes是論文的主戰(zhàn)場但如果你手頭沒有Cityscapes數(shù)據(jù)集也可以用VOC或自己標注的數(shù)據(jù)來做遷移訓(xùn)練。我在實際操作中推薦先用Cityscapes的一個小子集跑通流程確認代碼沒問題后再上全量數(shù)據(jù)這樣能節(jié)省很多排查問題的時間。2. 環(huán)境配置全過程2.1 硬件與驅(qū)動認知先弄清楚你的底牌配置環(huán)境之前第一件事不是著急裝包而是先確認自己的機器條件。PIDNet訓(xùn)練時常用的輸入分辨率是1024x1024Cityscapes標準設(shè)置這個分辨率下一張RTX 2080 Ti級別的顯卡11GB顯存只能塞下比較小的batch size通常設(shè)到2到4就頂天了。如果你想用8的batch size訓(xùn)練顯存至少要16GB以上。這里建議你先跑一下nvidia-smi確認驅(qū)動版本和顯存大小以及CUDA的可用版本。這里要特別說明一下CUDA版本的選擇邏輯。PyTorch的預(yù)編譯包是和CUDA版本綁定的比如PyTorch 1.10對應(yīng)CUDA 11.3PyTorch 1.12對應(yīng)CUDA 11.6PyTorch 2.x對應(yīng)CUDA 11.8或12.1。你需要先確定PyTorch版本再反向確定CUDA版本。PIDNet官方代碼是2022年發(fā)布的用的PyTorch版本不算新如果你機器上已經(jīng)裝了新版PyTorch比如2.x很可能遇到一些API兼容性問題后面我會詳細講怎么處理。注意不建議在Windows環(huán)境下跑訓(xùn)練除非你非常熟悉MSVC編譯鏈。PIDNet中有一部分算子需要編譯C擴展Windows上的坑比較多Linux環(huán)境是最順暢的。如果你是Windows用戶建議裝一個WSL2或直接用一個Linux服務(wù)器。2.2 CUDA與PyTorch安裝要點我推薦的環(huán)境組合是這樣的Ubuntu 20.04或22.04CUDA 11.3對應(yīng)NVIDIA驅(qū)動版本465PyTorch 1.10.1或1.12.0Python 3.8或3.9GCC 7.5以上用于編譯擴展算子具體的安裝步驟比較常規(guī)我就快速帶過重點是Pytorch安裝時一定要用官方源或可信鏡像源避免包版本不匹配。以PyTorch 1.10.1為例# 創(chuàng)建虛擬環(huán)境 conda create -n pidnet python3.8 conda activate pidnet # 安裝PyTorch注意cuda版本參數(shù)要和你的CUDA環(huán)境對應(yīng) pip install torch1.10.1cu113 torchvision0.11.2cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html裝完以后先用一段簡單代碼驗證CUDA是否可用import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果后面兩步輸出正常說明PyTorch已經(jīng)正確識別到了GPU。這里有個容易忽視的細節(jié)即使PyTorch能識別GPU也不代表訓(xùn)練性能一定沒問題你最好再用torch.cuda.get_device_capability(0)查一下算力PIDNet在算力8.0以上的卡上表現(xiàn)會更好一些。2.3 依賴與編譯最容易被卡住的一步PIDNet代碼中涉及到一些自定義的卷積算子大概率需要本地編譯。我遇到過好幾次環(huán)境測試能通過但真正訓(xùn)練時報錯找不到某個編譯好的so文件。所以你在clone代碼之后一定要先看代碼庫里有沒有setup.py或類似文件有的話先編譯安裝。# 安裝基礎(chǔ)依賴 pip install opencv-python pillow numpy tqdm tensorboard scipy # 如果代碼庫中有編譯腳本 cd PIDNet python setup.py develop編譯過程如果報錯多數(shù)是GCC版本太高或者缺少依賴庫。比較常用的解決方案是安裝低版本GCC或者在編譯時加上環(huán)境變量export CXX/usr/bin/g export CC/usr/bin/gcc提示如果你用的是PyTorch 2.x編譯時很可能會遇到ATEN/CUDA API不兼容的報錯因為2.x改了部分算子接口。穩(wěn)妥的辦法是直接創(chuàng)建一個Python 3.8 PyTorch 1.12的新環(huán)境不要在原環(huán)境里硬解依賴。我在這上面耗過一晚上的時間走彎路沒有意義。3. PIDNet核心細節(jié)解析3.1 網(wǎng)絡(luò)結(jié)構(gòu)關(guān)鍵設(shè)計PIDNet的架構(gòu)設(shè)計是理解整個項目核心的關(guān)鍵。我用比較通俗的方式拆解一下傳統(tǒng)的語義分割網(wǎng)絡(luò)一般有兩種思路一種是用編碼器把特征越采越深再用解碼器恢復(fù)分辨率比如UNet、DeepLab系列問題在于恢復(fù)過程中會丟失很多邊緣細節(jié)另一種是用高分辨率特征做全程等分辨率比如HRNet效果好但計算量很大。PIDNet走的是第三種路線用輕量級的骨干網(wǎng)絡(luò)提取基礎(chǔ)特征然后分三條分支并行處理。具體來說P分支Proportional負責(zé)接收大感受野的上下文信息適合識別“這是一個道路區(qū)域”I分支Integral負責(zé)保持高頻細節(jié)適合識別“道路和草坪之間的精確邊界”D分支Differential則是動態(tài)融合兩者的信息實現(xiàn)細節(jié)與語義的自適應(yīng)平衡。這三者的組合在理論上有點像控制系統(tǒng)中的PID反饋回路網(wǎng)絡(luò)在學(xué)習(xí)中能自動調(diào)節(jié)“比例項”“積分項”“微分項”的強度這也是它名字的由來。在具體實現(xiàn)層面PIDNet通過Bag和Head兩個模塊來管理三條分支的信息交換。Bag模塊在訓(xùn)練過程中累積整合不同分支的特征Head模塊則在最后進行上采樣和類別預(yù)測。這種設(shè)計的好處是不同分支可以在不同尺度上工作而融合機制保證它們不會信息脫節(jié)。3.2 損失函數(shù)與輔助監(jiān)督很多同學(xué)訓(xùn)練語義分割模型時只關(guān)注最終的交叉熵損失但對PIDNet這種三分支架構(gòu)來說分支級別的輔助監(jiān)督是保證性能的重要一環(huán)。如果只用最終輸出loss回傳梯度P分支和I分支學(xué)到的特征可能不夠分化最后退化成同一個模型失去三路并行互補的意義。我在實際訓(xùn)練中采用了主損失和輔助損失疊加的方式。具體做法是從D分支的融合特征輸出主分割結(jié)果計算主損失同時把P分支和I分支的中間輸出也接上分割頭計算輔助損失最終總損失是主損失加兩個輔助損失的加權(quán)和。論文中一般用0.5或0.25作為輔助損失的權(quán)重系數(shù)我實測下來0.5的效果會更穩(wěn)定一些尤其是在訓(xùn)練的前期階段輔助監(jiān)督能明顯加速分支分化。這里有一個容易被忽略的細節(jié)輔助損失只在訓(xùn)練時啟用推理時只保留D分支的輸出結(jié)果。因為訓(xùn)練時需要回傳梯度來指導(dǎo)分支學(xué)習(xí)但推理時分支的中間結(jié)果沒有獨立應(yīng)用場景。這一點在代碼層面的實現(xiàn)方式直接決定了訓(xùn)練和驗證時的模型結(jié)構(gòu)差異所以要記得做切換。4. 數(shù)據(jù)準備與訓(xùn)練配置4.1 數(shù)據(jù)集準備細節(jié)PIDNet最常用的數(shù)據(jù)集是Cityscapes但Cityscapes的標注文件組織方式有一定特殊之處。它的目錄結(jié)構(gòu)通常是cityscapes/ ├── leftImg8bit/ │ ├── train/ │ └── val/ └── gtFine/ ├── train/ └── val/而PIDNet代碼一般要求的是pix_annotation格式也就是每個圖片對應(yīng)一個和原圖同尺寸的標注PNG。Cityscapes官方給的gtFine是polygon格式JSON需要先轉(zhuǎn)換。我不是第一次在這上面栽跟頭。建議參考Cityscapes官方腳本或網(wǎng)上開源的轉(zhuǎn)換工具將gtCoarse和gtFine轉(zhuǎn)換成trainId標注而不是直接用原始ID。這里特別強調(diào)一下trainId和原始ID的區(qū)別關(guān)系到最終的評估指標轉(zhuǎn)換錯了的話mIoU的結(jié)果會非常奇怪。如果你想用VOC數(shù)據(jù)集替代那要提前確認類別標號、忽略索引等約定是否和PIDNet的默認配置一致。PIDNet通常默認19個類別加一個背景如果你用VOC的20類或者21類需要同步修改模型末端的分類頭數(shù)量和損失函數(shù)的類別權(quán)重否則會直接報維度不匹配的錯。4.2 訓(xùn)練參數(shù)配置在正式開訓(xùn)之前有相當(dāng)多的參數(shù)需要根據(jù)你的硬件條件調(diào)整。我梳理一下影響較大的幾個輸入分辨率Cityscapes標準是1024x1024。如果你的顯存不夠可以降到768x768或640x640但注意評估時最好保持訓(xùn)練分辨率否則性能差異不好比較。Batch Size顯存有限的情況下建議2-4并用梯度累積來模擬較大的batch size。梯度累積步數(shù)的計算方式是期望batch size除以單次GPU batch size比如想模擬batch size等于8單卡只能放下2那就設(shè)置累積步數(shù)為4。學(xué)習(xí)率策略建議使用多項式學(xué)習(xí)率衰減poly公式是base_lr * (1 - iter/total_iter)^power其中power通常設(shè)為0.9。初始學(xué)習(xí)率在batch size為8時推薦0.01左右如果你的batch size較小需要相應(yīng)降低學(xué)習(xí)率大致按線性縮放規(guī)律。Epoch數(shù)與迭代數(shù)Cityscapes全量數(shù)據(jù)訓(xùn)練的話通常跑240個epoch也就是大概12萬步迭代。如果你是先用子集跑通流程建議設(shè)置一個較小的迭代數(shù)比如2000步確認loss曲線在正常下降后再跑完整訓(xùn)練。我建議把訓(xùn)練過程中的日志保存下來包括每個iteration的loss值、當(dāng)前的mIoU如果做了驗證。這不僅是排查問題的依據(jù)也是后續(xù)對比改進效果的關(guān)鍵數(shù)據(jù)。5. 代跑訓(xùn)練實操記錄5.1 訓(xùn)練啟動與監(jiān)控我以Cityscapes數(shù)據(jù)集為例整理了一份實際的訓(xùn)練啟動命令。假設(shè)你已經(jīng)準備好了數(shù)據(jù)且配置文件位于configs/pidnet/pidnet_s_cityscapes.yaml# 單卡訓(xùn)練 python train.py --config configs/pidnet/pidnet_s_cityscapes.yaml # 多卡分布式訓(xùn)練 python -m torch.distributed.launch --nproc_per_node4 train.py \ --config configs/pidnet/pidnet_s_cityscapes.yaml啟動之后不要急著離開終端先觀察前30個iteration的輸出。正常情況下會看到loss逐漸從初始值回落如果loss直接變成NaN或者穩(wěn)在一個非常高的數(shù)值不動說明某個環(huán)節(jié)出了問題。最可能的原因是學(xué)習(xí)率設(shè)置過大或者數(shù)據(jù)加載過程中出現(xiàn)了異常的標注值。訓(xùn)練過程中的實時監(jiān)控也很關(guān)鍵。我一般會開兩個窗口一個跑訓(xùn)練另一個用tensorboard或者直接看一眼日志文件。PIDNet代碼中通常提供了tensorboard記錄支持你可以在啟動命令里加上日志路徑參數(shù)然后實時查看loss曲線和驗證集上的mIoU變化曲線tensorboard --logdir/path/to/pidnet/output/tensorboard5.2 驗證評估與可視化訓(xùn)練到一半或者結(jié)束后你肯定想知道當(dāng)前分割效果好到什么程度了。PIDNet代碼中一般會提供驗證腳本它和訓(xùn)練腳本是分開的。驗證時的幾個注意事項如下驗證集的圖片尺寸很多人在驗證時直接用原始尺寸輸入但要注意如果你的模型在1024x1024分辨率下訓(xùn)練直接輸入2048x1024的原始圖會導(dǎo)致BatchNorm的表現(xiàn)方差較大特別是當(dāng)驗證集的batch size設(shè)為1時。我一般會先做一次resize將短邊縮放到1024再送入模型。多尺度評估如果想得到更高精度的評估結(jié)果可以嘗試多尺度推理加翻轉(zhuǎn)集成。比如在0.5、0.75、1.0、1.25、1.5五種尺度下分別推理再把結(jié)果融合。這是個提升mIoU的小技巧一般能漲0.5到1個點?？梢暬Y(jié)果除了數(shù)值指標務(wù)必輸出幾張預(yù)測圖看看肉眼效果。mIoU有時候漲了但邊緣部分可能變得模糊不清只有通過可視化才能發(fā)現(xiàn)這些問題。驗證腳本跑完之后會輸出具體的類別IoU和mIoU。我要求自己不僅僅記錄最終的mIoU數(shù)值還要保存一份分類別的IoU表格。原因很簡單改進模型時你需要注意低類別IoU才能針對性設(shè)計優(yōu)化方案。6. 常見問題與排查技巧實錄6.1 訓(xùn)練階段常見的工程類問題我整理了幾個高頻問題你如果遇到了可以直接對照排查。**問題一編譯擴展算子失敗。**報錯信息會包含類似error: identifier AT_CHECK is undefined的信息。原因是PyTorch版本和代碼的API版本不匹配。AT_CHECK是舊版PyTorch的宏新版換成了TORCH_CHECK。解決辦法是改代碼里的宏定義或者直接換到和代碼適配的PyTorch版本。**問題二訓(xùn)練時顯存OOM。**這種情況非常常見尤其是你堅持1024分辨率加batch size 4時。我的建議順序是先降batch size到2再考慮降分辨率到768最后才是換模型尺寸。PIDNet有S、M、L三個版本S版本大約是堆半層參數(shù)量顯存壓力最小。如果S版本還在OOM那說明顯卡確實不太夠考慮梯度累積或混合精度訓(xùn)練吧。**問題三Loss變成NaN。**可能原因比較多常見的有學(xué)習(xí)率過高、數(shù)據(jù)中存在異常標注、混合精度訓(xùn)練的梯度縮放參數(shù)設(shè)置不當(dāng)。我的排查路徑是先把學(xué)習(xí)率降到當(dāng)前值的十分之一跑100步試試如果還不行關(guān)閉混合精度再檢查數(shù)據(jù)加載部分的標注歸一化是否出了問題。**問題四mIoU很低比如個位數(shù)。**這個一般都和標簽不匹配有關(guān)。Cityscapes的標簽有原始ID和trainId的區(qū)別很多類別在評估時是要合并的。如果你把原始ID當(dāng)成類別標簽輸入到模型所有類是混亂的模型自然學(xué)不會。另外忽略索引通常設(shè)為255有沒有正確設(shè)置也很關(guān)鍵該忽略的區(qū)域如果不忽略會嚴重拉高loss干擾訓(xùn)練。6.2 訓(xùn)練效果不好時的排查思路訓(xùn)練流程沒問題但效果不佳這個問題比報錯更讓人頭疼。我的建議是分維度排查。**先看訓(xùn)練集上的表現(xiàn)。**如果訓(xùn)練集上的mIoU都不高說明模型欠擬合重點調(diào)整學(xué)習(xí)率、訓(xùn)練時長、數(shù)據(jù)增強策略。如果訓(xùn)練集上mIoU很高但驗證集很低那就是過擬合考慮加數(shù)據(jù)增強、加大正則化、增加訓(xùn)練數(shù)據(jù)。**再看loss曲線的下降趨勢。**理想的loss曲線應(yīng)該是一個平滑下降的曲線如果loss曲線反復(fù)震蕩很可能是學(xué)習(xí)率太大或者batch size太小。如果loss下降得特別慢基礎(chǔ)學(xué)習(xí)率可能太低了注意poly學(xué)習(xí)率策略在訓(xùn)練中后期從宏觀上看下降速度會變慢這是正常的。**最后檢查數(shù)據(jù)加載是否正常。**我在訓(xùn)練時會把第一個batch的輸入圖像和標注保存下來人工看一眼有沒有錯位、錯標注、顏色反轉(zhuǎn)這些問題。數(shù)據(jù)管線出問題的概率不低于模型問題而且排查起來比模型問題更費時間。我的習(xí)慣是第一時間就把可視化腳本寫出來后續(xù)每次換數(shù)據(jù)集都能直接復(fù)用。7. 改進創(chuàng)新方向建議7.1 基于細節(jié)捕捉的小尺度優(yōu)化思路PIDNet的I分支是保持高分辨率細節(jié)的關(guān)鍵但這個分支在原始設(shè)計里比較簡單主要由基本的連續(xù)卷積堆疊構(gòu)成。我的建議之一是給I分支增加一個輕量級的注意力模塊比如通道注意力結(jié)合空間注意力的組合讓小分支自適應(yīng)地增強重要的細節(jié)區(qū)域。我自己試過一種方法在I分支的末端疊加一個特征選擇模塊讓模型自己學(xué)習(xí)哪些細粒度特征需要傳遞到D分支。具體實現(xiàn)時可以用一個1x1卷積生成一個門控權(quán)重再對原始特征做逐元素乘法。增加的計算量很小但確實能在邊界清晰度上看到改善。這個方向尤其適合對道路邊緣、車道線、細小物體分割效果有要求的場景比如自動駕駛的車道分割和可行駛區(qū)域檢測。7.2 面向?qū)崟r應(yīng)用的高效架構(gòu)調(diào)整方向PIDNet本身已經(jīng)是在追求精度和速度的平衡但你落地時往往會面對更嚴格的算力限制這時可以考慮以下改進方向?qū)⑻卣魈崛〔糠痔鎿Q成更輕量的骨干網(wǎng)絡(luò)。經(jīng)典的PIDNet默認用輕量卷積網(wǎng)絡(luò)做編碼器你可以嘗試替換成MobileNetV3或EfficientNet等輕量化模型。替換時要注意保持特征的通道數(shù)與后續(xù)Bag模塊的輸入要求對齊。我測試過一個方案計算量能下降約三成mIoU只掉零點幾個點在邊緣設(shè)備上效果很理想。對三個分支做動態(tài)裁剪。可以按照模型在不同層級的特征重要度動態(tài)決定跳過某些冗余分支的計算。這個思路實現(xiàn)起來比較復(fù)雜但如果你對算子級別的性能優(yōu)化有經(jīng)驗收益會很可觀。把注意力模塊加在Bag融合部分。原本的Bag模塊主要是簡單的特征疊加或拼接操作你可以嘗試把特征拼接后接入一個輕量級的注意力模塊重新加權(quán)不同分支的特征而不是簡單直接求和。這個小改動在多個數(shù)據(jù)集上有可復(fù)現(xiàn)的提升。7.3 訓(xùn)練策略與數(shù)據(jù)優(yōu)化方向模型結(jié)構(gòu)只是事情的一半訓(xùn)練策略的優(yōu)化空間同樣值得重視。基于我在多輪實驗中的經(jīng)驗以下幾個手段比較有效。**一是使用更強的數(shù)據(jù)增強策略。**除了基礎(chǔ)的隨機翻轉(zhuǎn)、縮放和裁剪可以嘗試顏色抖動、高斯模糊、隨機遮擋等更強力的增強方式這對提升模型泛化能力非常有幫助。但要注意增強強度不是越大越好過強的增強會讓模型在驗證集上產(chǎn)生偏差我通常采用溫和增強短暫預(yù)熱的方式。**二是加入知識蒸餾。**用一個更大更強的Teacher模型比如PIDNet-L來指導(dǎo)Student模型比如PIDNet-S訓(xùn)練。具體做法是讓Student模型的輸出分布盡量接近Teacher模型的輸出分布同時保持和真實標簽的交叉熵。我的經(jīng)驗是蒸餾出的Student模型要比直接訓(xùn)練同一Student模型高出2%-3%的mIoU同時推理速度完全不變。**三是引入半監(jiān)督訓(xùn)練思路。**如果你的預(yù)算不夠采集大量昂貴標注數(shù)據(jù)可以利用無標簽數(shù)據(jù)做一致性正則化對不同擾動下的預(yù)測結(jié)果保持一致性。這類方法在語義分割方向上已經(jīng)比較成熟對PIDNet這種三分支結(jié)構(gòu)的模型同樣適用。我建議先從簡單的弱增強和強增強一致性學(xué)習(xí)入手再逐步添加更多的輔助約束。這些改進思路都不是特別復(fù)雜的技術(shù)體操而是基于對PIDNet結(jié)構(gòu)特點的針對性優(yōu)化。你在做改進實驗時我的建議是不要同時疊加太多改動否則一旦效果變差很難定位是哪個改動導(dǎo)致的回退。逐步做單點更改、單點驗證是最穩(wěn)妥也最高效的推進方式。寫在后面PIDNet在我做過的語義分割項目里屬于那種“越改越有感覺”的模型。它的三分支結(jié)構(gòu)把細節(jié)和語義分開處理這一設(shè)計在背后給了后人很大的調(diào)整空間。從環(huán)境配置到訓(xùn)練調(diào)優(yōu)再到結(jié)構(gòu)改動每一步都需要細心排查但每一步的回報也都直接反映在數(shù)據(jù)上。最后分享一個小習(xí)慣我每次做模型改進實驗都會把改動點、實驗配置、結(jié)果指標記錄在表格里哪怕是一個很小的參數(shù)調(diào)整也不放過。等整個實驗周期結(jié)束后往回看這份記錄的價值往往比單純的技術(shù)方案還要大。它能讓你的改進路徑清晰可見也能避免重復(fù)試錯。如果你正準備拿PIDNet做實驗希望這篇文章能幫你少走一些彎路。