biendata U-RISC 神經元識別大賽簡單/複雜雙賽道第三方案

比賽頁面

鏈接: 比賽頁面.

github

鏈接: github.

任務簡介

本次比賽要求選手根據對小鼠視網膜的電鏡成像圖片，像素級地標註（即用不規則多邊形，而是精確描繪出神經元的形狀）每個神經元的外廓（只標註神經元的最外層細胞膜，神經元內部不進行填充標註）。

解決方案

總體思路爲使用了一個U型卷積神經網絡，對細胞圖片進行像素級的語義分割。
下面爲詳細的步驟介紹：

數據預處理

其中簡單賽道的訓練數據爲10241024分辨率的圖片，複雜賽道的訓練數據爲99589958分辨率的圖片。針對複雜賽道的訓練數據，將大圖切割爲10241024的小圖，滑動步長爲10241024，即小圖沒有重合。（考慮劃分驗證集時的數據泄露問題，也考慮到複雜賽道有較多的訓練數據，故小圖之間沒有重合部分，以便解決數據泄露的問題，更好的劃分本地的驗證集。）對簡單賽道和複雜賽道切分以後的數據做以下數據增強：

• 隨機從10241024的原圖中截取512512的圖片
• 以80%的概率調整對比度和亮度，調整的對比度範圍爲-20% ~ +20%，調整的對比度範圍爲-20% ~ +20%
• 以80%的概率調整伽馬值，調整的對比度範圍爲80~120
• 以50%的概率水平翻轉圖像
• 以50%的概率豎直翻轉圖像
• 以50%的概率將圖像隨機旋轉90°的整數倍
• 以50%的概率轉置圖像
• 進行歸一化處理
• 以上每個步驟的發生概率相互獨立

經過一次數據增強的效果如下所示

網絡架構

本次方案融合了多個模型，網絡結構分爲兩種類型：

• 狹義U-net網絡[1]（下文中簡稱U-net）
• 廣義U型網絡（下文中簡稱U型網絡）

U-net

利用改進的U-net網絡[1]（下文中簡稱U-net），網絡結構如下圖所示，爲了增強U-net網絡的性能，針對原始的U-net做了一下修改：

• 把卷積層替換成殘差結構[2]
• 在每個參差結構後面跟上一個instance normalization[3]
• 把下采樣由池化層替換爲跨步卷積
• 因爲原圖分割很精細，需要增加淺層權重，網絡通道全改成100
• 利用了一個迭代的思想[4]，把前一次的分割結果利用到第二次的訓練中，精細化分割結果，迭代三次。示意圖如下圖所示。詳細的一次迭代流程爲，網絡的輸入的通道數由原來的3（圖片的RGB通道）變爲4（RGB+label）。第一輪輸入的時候初始化的label爲全0，之後經過網絡預測得到第一次輸出的結果label’，將原圖RGB+label’再次輸入網絡，進行第二輪預測。
  
  

U型網絡

整體結構爲U型網絡（下文中簡稱U型網絡），將網絡分爲encode模塊和decode模塊，其中encode模塊採用efficientnet-b7[5]網絡，decode和U-net結構相似，結合encode的跳層鏈接上採樣，通道數分別爲1024,512,256,128,64，並且加入了空間和通道注意力機制scse模塊[6]，整體結構如下圖所示。

其中關於Backbone的選擇我們嘗試了densenet，dpn，inception，resnet，se_resnet，se_resnext等一系列網絡，本地驗證集結果如下圖所示最終efficientnet-b7網絡的性能在本地的驗證集上表現最佳。

損失函數

由於比賽的指標爲F1-score，因此引入Dice損失函數[7]直接優化F1-score，由於Dice損失函數的優化不夠平滑，和考慮到任務中正類和負類樣本極度不平衡，因此引入Focal Loss[8]。經過多組控制變量實驗，得出focalloss中伽馬的值爲2.5，Dice loss和focal loss的比值爲2：1時，可以在本地驗證集上得到最好的效果。下圖爲控制focal loss中伽馬變量的實驗結果。（Dice loss係數實驗結果由於在另外一臺電腦上，暫時無法展示，開學回學校以後可以補上）最後的損失函數爲2倍的(1-Dice)和Focal loss求和。
$loss=2*(1-Dice) + Focal$

優化器和訓練過程

優化器選擇方面也進行了一系列的實驗，包括基礎的SGD，Momentum-SGD，NAG(Nesterov accelerated gradient)，自適應優化器Adagrad，Adadelta，RMSprop，，Adam，以及Adam優化器的各種變體如NAdam，RAdam。更進一步，在優化效果較好的RAdam優化器上，嘗試了不同的優化策略，如針對大batchsize優化的LARS 技術和帶有嘗試思想的LookAHead技術。最終RAdam + LARS + LookAHead結合的優化器可以達到最優的本地驗證集指標。（優化器控制變量實驗結果由於在另外一臺電腦上，暫時無法展示，開學回學校以後可以補上）
具體的優化器和訓練策略選擇如下：

• U-net：使用RAdam[10] + LARS[11] + LookAHead[12]優化器，初始學習率3e-4，採用自適應學習策略（當損失連續沒有降低超過一定次數時減小學習率）
• U型網絡：使用RAdam[10] + LARS[11] + LookAHead[12]優化器，初始學習率爲1e-2，採用自適應學習策略

模型融合

獲取網絡輸出的每個像素點的概率進行求和平均，增加模型的魯棒性和精度，融合流程如下圖所示。

其他trick

膨脹預測

由於複雜賽道最後的圖的尺寸較大，因此需要切分爲小圖再進行拼接（和我19年暑假參加的一個比賽非常類似，因此拿那個比賽的示意圖來做示意），考慮到邊緣對於網絡的預測有較大的影響，因此拼接的時候採用了膨脹預測，即預測2048*2048的尺寸的小圖，會捨去周圍100個像素寬度的數據，僅僅拿中間的數據進行拼接，減少邊緣對網絡的影響。減少邊界拼接時的錯誤。示意圖如下圖所示。

梯度累計+同步BN

由於efficientnet-b7網絡較大，導致BatchSize較小（單卡1080ti，11G顯存只能跑BatchSize=2），導致BN的統計量不夠準確，收斂非常隨機等問題，大大制約了訓練結果的性能。因此引入多卡同步BN和梯度累計解決這個問題。

通過NVIDIA的APEX包實現的多卡之間的通信使得常規僅僅計算一張卡上的數據的BN可以一起計算多張卡上的數據從而得出更加穩定的BN統計量。

梯度累計即爲計算一個BatchSize的數據之後梯度不反傳，直到累積到設定的N個batchsize之後在進行反傳，效果等同於將BatchSize擴大了N倍。在本次efficientnet-b7網絡的訓練中，設置了N=20，兩張1080ti的環境下相當於BatchSize達到了2022=80，大大提高了模型訓練的穩定性。

測試時增強（TTA）

用了ttach模塊

       Input
         |           # input batch of images 
    / / /|\ \ \      # apply augmentations (flips, rotation, scale, etc.)
   | | | | | | |     # pass augmented batches through model
   | | | | | | |     # reverse transformations for each batch of masks/labels
    \ \ \ / / /      # merge predictions (mean, max, gmean, etc.)
         |           # output batch of masks/labels
       Output

用了ttach模塊由於簡單賽道時間比較充足，做了水平，垂直，旋轉[0, 90, 180, 270]的測試時數據增強，複雜賽道做了只做了水平，垂直，結果取平均值。

結論/討論

• 本次比賽中，考慮到細胞膜的寬度只佔幾個像素，非常細，因此pspnet，deeplab等這種在下采樣的尺寸上進行語義分割之後依靠雙線性插值算法直接對語義圖進行上採樣的算法效果非常差。這類算法由於decode沒有恢復到原始的分辨率，在下采樣4倍或者16倍的空間分辨率下進行語義分割，會導致細胞膜分割不準確。U型網絡在原始分辨率上進行語義分割，結合了高層的語義信息，低層紋理信息也使用跳層加以利用，可以獲得更加精細的細胞膜分割效果。
• 本次比賽中的數據，正負類別極度不平衡，如果不做數據平衡，網絡往往不願意做正類（細胞膜）的預測，因此交叉熵不能取得很好的效果，Dice loss直接優化F1-score，因此數據不平衡不會對Dice有影響，有較好的效果。
• 本次比賽經過我們的實驗，發現輸入圖像的尺寸是非常關鍵的一個因素，經過實驗得出，輸入圖像的尺寸越大可以去的越好的效果，我們認爲這個應該和比例（一個像素對應的實際面積）有關。尺寸越大，網絡可以學習到的高層語義信息就越多，但是考慮到輸入圖像尺寸越大，訓練數據的重複性就越高，考慮到泛化問題，最終確定尺寸爲512*512較爲合適。
• 暑假做的那個比賽後處理提分挺多的，但是這個比賽使用常規的開閉操作，或者移除較小的聯通分量都不能取得較好的結果，因此本次比賽後處理沒有用上，有個後處理的思路，我們發現非常多網絡預測的結果，都有如圖10所示的問題，即細胞膜不連續的問題，可以使用後處理將這個問題解決的話應該可以提升不少分數，嘗試了非常多的包括CRF，霍夫變換等傳統的手段，都沒有很好的解決這個問題。有點遺憾。
  

參考文獻

[1] Ronneberger O, Fischer P, Brox T. U-net: Convolutional networks for biomedical image segmentation[C]//International Conference on Medical image computing and computer-assisted intervention. Springer, Cham, 2015: 234-241.
[2] He K, Zhang X, Ren S, et al. Deep residual learning for image recognition[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2016: 770-778.
[3] Ulyanov D, Vedaldi A, Lempitsky V. Instance normalization: The missing ingredient for fast stylization[J]. arXiv preprint arXiv:1607.08022, 2016.
[4] Le T, Li Y, Duan Y. RED-NET: A Recursive Encoder-Decoder Network for Edge Detection[J]. arXiv preprint arXiv:1912.02914, 2019.
[5] Tan M, Le Q V. Efficientnet: Rethinking model scaling for convolutional neural networks[J]. arXiv preprint arXiv:1905.11946, 2019.
[6] Roy A G, Navab N, Wachinger C. Concurrent spatial and channel ‘squeeze & excitation’in fully convolutional networks[C]//International Conference on Medical Image Computing and Computer-Assisted Intervention. Springer, Cham, 2018: 421-429.
[7] https://zh.wikipedia.org/wiki/Dice%E7%B3%BB%E6%95%B0
[8] Lin T Y, Goyal P, Girshick R, et al. Focal loss for dense object detection[C]//Proceedings of the IEEE international conference on computer vision. 2017: 2980-2988.
[9] Kingma D P, Ba J. Adam: A method for stochastic optimization[J]. arXiv preprint arXiv:1412.6980, 2014.
[10] Liu L, Jiang H, He P, et al. On the variance of the adaptive learning rate and beyond[J]. arXiv preprint arXiv:1908.03265, 2019.
[11] Ginsburg B, Gitman I, You Y. Large batch training of convolutional networks with layer-wise adaptive rate scaling[J]. 2018.
[12] Zhang M, Lucas J, Ba J, et al. Lookahead Optimizer: k steps forward, 1 step back[C]//Advances in Neural Information Processing Systems. 2019: 9593-9604.