語(yǔ)義分割模型 SegNeXt方法概述

語(yǔ)義分割是對(duì)圖像中的每個(gè)像素進(jìn)行識(shí)別的一種算法，可以對(duì)圖像進(jìn)行像素級(jí)別的理解。作為計(jì)算機(jī)視覺(jué)中的基礎(chǔ)任務(wù)之一，其不僅僅在學(xué)術(shù)界廣受關(guān)注，也在無(wú)人駕駛、工業(yè)檢測(cè)、輔助診斷等領(lǐng)域有著廣泛的應(yīng)用。

近期，計(jì)圖團(tuán)隊(duì)與南開(kāi)大學(xué)程明明教授團(tuán)隊(duì)、非十科技劉政寧博士等合作，提出了一種全新的語(yǔ)義分割模型 SegNeXt，該方法大幅提高了當(dāng)前語(yǔ)義分割方法的性能，并在Pascal VOC 分割排行榜上名列第一。該論文已被 NeurIPS 2022 接收。

Part1

語(yǔ)義分割模型SegNeXt

研究背景 自2015年FCN[2] 被提出以來(lái)，語(yǔ)義分割開(kāi)始逐漸走向深度學(xué)習(xí)算法，其常用架構(gòu)為編碼-解碼器結(jié)構(gòu)(Encoder-Decoder)。在 vision transformer 被提出之前，人們通常采用卷積神經(jīng)網(wǎng)絡(luò)(如 ResNet、VGGNet、GoogleNet 等) 作為其編碼器部分；最近，由于vision transformer 在視覺(jué)領(lǐng)域的成功，語(yǔ)義分割編碼器部分開(kāi)始逐漸被換成基于vision transformer的模型(如 ViT、SegFormer、HRFormer等)。但是，基于 vision transformer編碼器的方法真的比基于卷積神經(jīng)網(wǎng)絡(luò)的方法更好么？為了回答這個(gè)問(wèn)題，Jittor團(tuán)隊(duì)重新思考了語(yǔ)義分割任務(wù)對(duì)神經(jīng)網(wǎng)絡(luò)的要求，并針對(duì)語(yǔ)義分割的任務(wù)專(zhuān)門(mén)設(shè)計(jì)了一個(gè)基于卷積神經(jīng)網(wǎng)絡(luò)的編碼器MSCAN 和一個(gè)語(yǔ)義分割模型 SegNeXt。

圖1. SegNeXt 和其他語(yǔ)義分割方法的性能對(duì)比，其中紅色為SegNeXt

方法概述

論文首先分析了語(yǔ)義分割任務(wù)本身以及之前的相關(guān)工作，總結(jié)出四點(diǎn)語(yǔ)義分割任務(wù)所需的關(guān)鍵因素。1）強(qiáng)大的骨干網(wǎng)絡(luò)作為編碼器。與之前基于 CNN 的模型相比，基于Transformer 的模型的性能提升主要來(lái)自更強(qiáng)大的骨干網(wǎng)絡(luò)。2）多尺度信息交互。與主要識(shí)別單個(gè)對(duì)象的圖像分類(lèi)任務(wù)不同，語(yǔ)義分割是一項(xiàng)密集的預(yù)測(cè)任務(wù)，因此需要在單個(gè)圖像中處理不同大小的對(duì)象，這就使得針對(duì)語(yǔ)義分割任務(wù)的網(wǎng)絡(luò)需要多尺度信息的交互。3）注意力機(jī)制：注意力可以使得模型關(guān)注到重點(diǎn)的部分，并且可以使得網(wǎng)絡(luò)獲得自適應(yīng)性。4）低計(jì)算復(fù)雜度：這對(duì)于常常處理高分辨率圖像的語(yǔ)義分割任務(wù)來(lái)說(shuō)至關(guān)重要。

表 1 不同方法所具有的的屬性對(duì)比

為了滿(mǎn)足上述四點(diǎn)要求，作者設(shè)計(jì)了一種簡(jiǎn)單的多尺度卷積注意力機(jī)制 （MSCA）。如圖 2 所示，MSCA 主要是采用大卷積核分解、多分支并行架構(gòu)以及類(lèi)似VAN[3]的注意力機(jī)制。這使得 MSCA 可以獲得大感受野、多尺度信息以及自適應(yīng)性等有益屬性。基于 MSCA，該論文搭建了一種層次化神經(jīng)網(wǎng)絡(luò) MSCAN 作為SegNeXt 的編碼器部分。除此之外，作者采用了 UNet 架構(gòu)，并選擇了HamNet[4] 作為 SegNeXt 的解碼器部分。分析和實(shí)驗(yàn)證明，MSCAN和 Ham 優(yōu)勢(shì)互補(bǔ)，兩者相互配合，使得 SegNeXt 實(shí)現(xiàn)了優(yōu)異的性能。

圖 2：多尺度卷積注意力（MSCA） 示意圖 實(shí)驗(yàn)結(jié)果 本文在五個(gè)常見(jiàn)分割數(shù)據(jù)集上 ADE20K, Cityscapes,COCO-Stuff, Pascal VOC, Pascal Context 和一個(gè)遙感分割數(shù)據(jù)集 iSAID做了測(cè)評(píng)，SegNeXt均超過(guò)了之前的方法。限于篇幅，我們僅展示部分結(jié)果。

表2：在 ADE20K、Cityscapes, COCO-Stuff 上的實(shí)驗(yàn)結(jié)果

表 3 SegNeXt 在遙感數(shù)據(jù)集上的實(shí)驗(yàn)結(jié)果

Part2 計(jì)圖語(yǔ)義分割算法庫(kù)JSeg Jittor團(tuán)隊(duì)基于自主深度學(xué)習(xí)框架Jittor[5]，并借鑒MMSegmentation語(yǔ)義分割算法庫(kù)的特點(diǎn)，開(kāi)發(fā)了語(yǔ)義分割算法庫(kù)JSeg。MMSegmentation是廣泛使用的功能強(qiáng)大的語(yǔ)義分割算法庫(kù)，新推出的JSeg可以直接加載MMSegmentation的模型，同時(shí)借助Jittor深度學(xué)習(xí)平臺(tái)的優(yōu)勢(shì)，使其更高效、穩(wěn)定運(yùn)行，可以實(shí)現(xiàn)訓(xùn)練和推理快速的從PyTorch向Jittor遷移。 目前JSeg已經(jīng)支持4個(gè)模型、4個(gè)數(shù)據(jù)集，其中模型包括在Pascal VOC test dataset斬獲第一的SegNeXt模型，數(shù)據(jù)集包括經(jīng)典的ADE20K Dataset、CityScapes Dataset以及遙感分割中的iSAID Dataset等，后續(xù)JSeg也將支持更多的模型和數(shù)據(jù)集！ 性能提升 我們使用SegNeX-Tiny模型，與Pytorch實(shí)現(xiàn)的版本在NVIDIA TITAN RTX上進(jìn)行了對(duì)比，可以顯著縮短模型訓(xùn)練所需要的時(shí)間。

表1JSeg和mmseg（PyTorch）的訓(xùn)練時(shí)間對(duì)比

易用性提升

由于Jittor動(dòng)態(tài)編譯的特性及code算子對(duì)python內(nèi)聯(lián)C++及CUDA的支持，JSeg在不同環(huán)境下無(wú)需對(duì)任何算子進(jìn)行手動(dòng)編譯，即可輕松運(yùn)行不同模型，免去了用戶(hù)對(duì)不同模型分別配置環(huán)境的負(fù)擔(dān)，同時(shí)方便用戶(hù)對(duì)不同方法進(jìn)行更公平的比較。此外，JSeg的設(shè)計(jì)易于拓展，用戶(hù)可以基于JSeg已有的模型和功能方便地開(kāi)展進(jìn)一步的研究和開(kāi)發(fā)。

實(shí)踐案例

下面，我們將簡(jiǎn)要介紹如何使用JSeg訓(xùn)練一個(gè)基礎(chǔ)模型。

首先，下載數(shù)據(jù)集到原始數(shù)據(jù)集目錄。

通過(guò)tools/convert_datasets下的數(shù)據(jù)處理腳本對(duì)原始數(shù)據(jù)進(jìn)行預(yù)處理，得到處理后的數(shù)據(jù)集。然后即可對(duì)模型進(jìn)行單卡或者多卡訓(xùn)練、評(píng)估和測(cè)試，同時(shí)提供了推理接口，用戶(hù)可以使用10行代碼完成一張圖片的語(yǔ)義分割，盡可能地降低了用戶(hù)的使用成本。