點擊上方“CVer”，選擇加"星標"或“置頂”

重磅干貨，第一時間送達

作者：?街道口扛把子

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

本文已由原作者授權，不得擅自二次轉載

Non-local U-Nets for Biomedical Image Segmentation

論文地址：https://arxiv.org/abs/1812.04103

開源代碼地址(tensorflow代碼)：

https://github.com/divelab/Non-local-U-Nets

以上代碼是tensorflow的，用于3D圖像，根據這個，我做了一個pytorch的2D圖像的代碼。

https://github.com/Whu-wxy/Non-local-U-Nets-2D-block

(這篇很用心寫的，點個贊哦親)

簡介

自2015年以來，在生物醫學圖像分割領域，U-Net得到了廣泛的應用，至今，U-Net已經有了很多變體。U-Net如下圖所示，是一個encoder-decoder結構，左邊一半的encoder包括若干卷積，池化，把圖像進行下采樣，右邊的decoder進行上采樣，恢復到原圖的形狀，給出每個像素的預測。

這篇文章的作者提出了一個新的U-Net模型:?Non-local U-Nets，推理速度更快，精度更高，參數量更少。提出了一個新的上/下采樣方法：Global Aggregation Block，把self-attention和上/下采樣相結合，在上/下采樣的同時考慮全圖(Non-local)信息。

更多的U-Net可以在以下代碼庫了解：

https://github.com/ShawnBIT/UNet-family

存在的不足

作者首先分析了U-Net存在的不足。

U-Net的encoder部分由若干卷積層和池化層組成，由于他們都是local的運算，只能看到局部的信息，因此需要通過堆疊多層來提取長距離信息，這種方式較為低效，參數量大，計算量也大。過多的下采樣導致更多空間信息的損失(U-Net下采樣16倍)，圖像分割要求對每個像素進行準確地預測，空間信息的損失會導致分割圖不準確。

decoder的形式與encoder部分正好相反，包括若干個上采樣運算，使用反卷積或插值方法，他們也都是local的方法。

創新點

為了解決以上問題，作者基于self-attention提出了一個Non-local的結構，global aggregation block用于在上/下采樣時可以看到全圖的信息，這樣會使得到更精準的分割圖。

簡化U-Net，減少參數量，提高推理速度，上采樣和下采樣使用global aggregation block，使分割更準確。

global aggregation block

global aggregation block如上圖所示，圖看起來很復雜，實際上運算過程很簡單。

1.該結構與Attention Is All You Need這篇文章的形式很相似。輸入圖為X(B*H*W*C)，經過QueryTransform和1*1卷積，轉換為Q(B*Hq*Wq*Ck)，K(B*H*W*Ck)，V(B*H*W*Cv)。QueryTransform可以為卷積，反卷積，插值等你喜歡的方法，最后的輸出結果的H與W將與這個值一致。

2.代碼里在Unfold之前有Multi-Head(Attention Is All You Need)的操作，不過在論文中沒有說明，實際上是把通道分為N等份。Unfold是把Batch,height,width,N通道合并,Q(B*Hq*Wq*N*ck)，K(B*H*W*N*ck)，V(B*H*W*N*cv)。

3.接下來是經典的點積attention操作，得到一個權值矩陣A((B*Hq*Wq*N)*(B*H*W*N))，用于self-attention的信息加權，分母Ck是通道數，作用是調節矩陣的數值不要過大，使訓練更穩定(這個也是Attention Is All You Need提出的)。最后權值矩陣A和V點乘，得到最終的結果((B*Hq*Wq*N)*cv)，可見輸出的height和width由Q決定，通道數由V決定。

4.我最近看了兩篇上采樣的論文，DUpsample和CARAFE，現在很多上采樣相關的論文關注于在上采樣時用CNN擴大感受野，增加圖像局部信息。這篇文章提出的global aggregation block是一個將注意力機制和上/下采樣相結合的方法，關注全圖信息，感受野更大，可以在其他任務上試用一下，效果如何還是要看實踐的結果。

Non-local U-Nets

文章提出的Non-local U-Nets如上圖所示。

相比U-Net，卷積層數減少，圖像下采樣倍率從16倍變成4倍，保留了更多空間信息。encoder和decoder之間的skip connections用相加的方式，不是拼接的方式，讓模型推理速度更快。

上下采樣使用global aggregation block，使分割圖更準確。

實驗結果

實際上這篇文章提出的東西沒有很高深，但是它數據很多，很漂亮。

醫學圖像分割的評價指標不太了解哈，不過看數據效果挺好。

參數量減少了很多，推理速度更快。

結果圖如下：(訓練數據直譯是3D多模態等強度嬰兒腦MR圖像)

參考文獻

Non-local U-Nets for Biomedical Image Segmentation

U-Net: Convolutional Networks for Biomedical Image Segmentation

Attention Is All You Need

Non-local Neural Networks

Decoders Matter for Semantic Segmentation: Data-Dependent Decoding Enables Flexible Feature Aggregation

CARAFE: Content-Aware ReAssembly of FEatures

推薦閱讀

2020年AI算法崗求職群來了(含準備攻略、面試經驗、內推和學習資料等)

重磅！CVer-醫療影像和圖像分割交流群已成立

掃碼添加CVer助手，可申請加入CVer-醫療影像或者圖像分割微信交流群，目前已滿500+人，旨在交流SLAM的學習、科研、工程項目等內容。

一定要備注：研究方向+地點+學校/公司+昵稱(如醫療影像/圖像分割+上海+上交+卡卡)，根據格式備注，可更快被通過且邀請進群

▲長按加群

▲長按關注我們

麻煩給我一個在看！

總結

以上是生活随笔為你收集整理的tensorflow图像分割unet_AAAI2020 | 医学图像分割的Nonlocal UNets的全部內容，希望文章能夠幫你解決所遇到的問題。

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