論文地址：[2402.02491] VM-UNet: Vision Mamba UNet for Medical Image Segmentation (arxiv.org)
代碼地址：JCruan519/VM-UNet: (ARXIV24) This is the official code repository for "VM-UNet: Vision Mamba UNet for Medical Image Segmentation". (github.com)
摘要：
在醫(yī)學(xué)圖像分割領(lǐng)域，基于細胞神經(jīng)網(wǎng)絡(luò)和基于變換器的模型都得到了廣泛的探索。然而，細胞神經(jīng)網(wǎng)絡(luò)在長期建模能力方面表現(xiàn)出局限性，而變壓器則受到其二次計算復(fù)雜性的阻礙。最近，以曼巴為例的狀態(tài)空間模型（SSM）已成為一種很有前途的方法。它們不僅擅長對長程相互作用進行建模，而且保持線性計算復(fù)雜性。在本文中，利用狀態(tài)空間模型，我們提出了一種用于醫(yī)學(xué)圖像分割的Ushape架構(gòu)模型，稱為Vision Mamba UNet（VM UNet）。具體地，引入視覺狀態(tài)空間（VSS）塊作為基礎(chǔ)塊來捕獲廣泛的上下文信息，并構(gòu)建了一種不對稱的編碼器-解碼器結(jié)構(gòu)。我們在ISIC17、ISIC18和Synapse數(shù)據(jù)集上進行了全面的實驗，結(jié)果表明，VM UNet在醫(yī)學(xué)圖像分割任務(wù)中具有競爭力。據(jù)我們所知，這是第一個基于純SSM模型構(gòu)建的醫(yī)學(xué)圖像分割模型。我們的目標(biāo)是建立一個基線，并為未來開發(fā)更高效、更有效的基于SSM的分割系統(tǒng)提供有價值的見解。
技術(shù)現(xiàn)狀
近年來，狀態(tài)空間模型(State Space Models)引起了研究者的極大興趣。在經(jīng)典SSM研究的基礎(chǔ)上，現(xiàn)代SSM（例如Mamba）不僅建立了長距離依賴關(guān)系，而且在輸入大小方面表現(xiàn)出線性復(fù)雜性。此外，基于SSM的模型在許多領(lǐng)域得到了大量研究，包括語言理解、一般視覺等。特別是，U-Mamba最近推出了一種新的SSM-CNN混合模型，標(biāo)志著它首次應(yīng)用于醫(yī)學(xué)圖像分割任務(wù)。SegMamba在編碼器部分結(jié)合了SSM，同時在解碼器部分仍然使用CNN，這提出了用于3D腦腫瘤分割任務(wù)的SSM-CNN混合模型。盡管上述工作已經(jīng)將SSM用于醫(yī)學(xué)圖像分割任務(wù)，但純基于SSM的模型的性能仍有待探索。
文章貢獻
1）提出了VM-UNet，這標(biāo)志著首次探索了純基于SSM的模型在醫(yī)學(xué)圖像分割中的潛在應(yīng)用。
2） 在三個數(shù)據(jù)集上進行了綜合實驗，結(jié)果表明VM UNet表現(xiàn)出相當(dāng)大的競爭力。
3） 在醫(yī)學(xué)圖像分割任務(wù)中為純基于SSM的模型建立了基線，提供了有價值的見解，為開發(fā)更高效、更有效的基于SSM分割方法鋪平了道路。
整體結(jié)構(gòu)
① Vision Mamba UNet (VM-UNet)
VM-UNet包括Patch Embedding layer, an encoder, a decoder, a Final Projection layer, and skip connections。采用了不對稱結(jié)構(gòu)設(shè)計。Patch Embedding layer將輸入圖像劃分為4*4的不重疊補丁，隨后將圖像的維度映射到C（默認96），該過程得到嵌入圖像H4×W/4×C，然后編碼器的4個stage進行特征提取，前3個stage結(jié)束時應(yīng)用補丁合并操作以減少輸入特征的高度和寬度，同時增加通信數(shù)量。每個stage使用2個VSS塊。解碼器部分與編碼器類似。
② VSS block
輸入數(shù)據(jù)通過層歸一化后，被拆分為兩個分支。在第一個分支中，輸入通過線性層和激活函數(shù)；第二個分支通過線性層、深度可分離卷積和激活函數(shù)、2D選擇性掃描模塊2D-Selective-Scan（SS2D）。隨后進行層歸一化，與第一個分支的輸出執(zhí)行逐元素生成。最后進行一個殘差連接操作得到VSS block的輸出。結(jié)構(gòu)中的損失函數(shù)選擇的SiLU
SS2D由3部分組成，a scan expanding operation, an S6 block, and a scan merging operation
scan expanding operation將輸入圖像沿著四個不同的方向（從左上到右下、從右下到左上、從右上到左下、從左下到右上）展開為序列。然后S6塊對這些序列進行處理，以進行特征提取，確保對來自各個方向的信息進行徹底的掃描，從而獲取不同的特征。隨后，合并序列，將輸入圖像恢復(fù)到與輸入相同的大小。
S6塊源自Mamba，通過根據(jù)輸入調(diào)整SSM的參數(shù)，在S4之上引入了一種選擇性機制。這使得模型能夠區(qū)分和保留相關(guān)信息，同時過濾掉不相關(guān)的信息。算法1中給出了S6塊的偽代碼
結(jié)論
在本文中，我們首次引入了一種純的基于SSM的醫(yī)學(xué)圖像分割模型，將VM-UNet作為基線。為了利用基于SSM的模型的能力，我們使用VSS塊構(gòu)建VM UNet，并使用預(yù)訓(xùn)練的VMamba-S初始化其權(quán)重。在皮膚損傷和多器官分割數(shù)據(jù)集上進行的綜合實驗表明，純基于SSM模型在醫(yī)學(xué)圖像分割任務(wù)中具有很強的競爭力，值得在未來進行深入探索。
代碼復(fù)現(xiàn)
注意目前只支持linux系統(tǒng)復(fù)現(xiàn)文章來源：http://www.zghlxwxcb.cn/news/detail-842959.html
conda create -n vmunet python=3.8
conda activate vmunet
pip install torch==1.13.0 torchvision==0.14.0 torchaudio==0.13.0 --extra-index-url https://download.pytorch.org/whl/cu117
pip install packaging
pip install timm==0.4.12
pip install pytest chardet yacs termcolor
pip install submitit tensorboardX
pip install triton==2.0.0
pip install causal_conv1d==1.0.0 # causal_conv1d-1.0.0+cu118torch1.13cxx11abiFALSE-cp38-cp38-linux_x86_64.whl
pip install mamba_ssm==1.0.1
# 如果安不上按下面步驟
# 點擊報錯中鏈接
# pip install mamba_ssm-1.0.1+cu118torch1.13cxx11abiFALSE-cp38-cp38-linux_x86_64.whl
pip install scikit-learn matplotlib thop h5py SimpleITK scikit-image medpy yacs
#預(yù)訓(xùn)練權(quán)重下載地址：https://pan.baidu.com/s/144nc45k28NBB_8k9hf_OeQ?pwd=xuns
文章來源地址http://www.zghlxwxcb.cn/news/detail-842959.html
到了這里，關(guān)于《VM-UNet: Vision Mamba UNet for Medical Image Segmentation》論文閱讀及代碼復(fù)現(xiàn)的文章就介紹完了。如果您還想了解更多內(nèi)容，請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！
