[論文メモ] Fully Transformer Networks for Semantic Image Segmentation

arxiv.org

図・表は論文から引用しています。

概要

NLPで使われているtransformerがsemantic segmentationでどれだけ有効なのかを調査した論文。
encoder-decoder basedなFully Transformer Networks(FTN)を提案。

Transformer basedな手法としては3種類ほどある(図1を参照)。
f:id:Ninhydrin:20210906093814p:plain

1)Transformer-CNN architecture(図1(a))
transformerのencoderとCNNのdecoderを使う。画像をパッチに切り出し埋め込みベクトルにして処理していく。ViTに近い。

2) Hybrid-CNN architecture(図1(b))
decoderはCNN。encoderもほとんどがResNetなどのCNNだが、出力に近い低解像度の部分を埋め込みベクトルとしてtransformerのモジュールに入れる。

3) Hybrid-Transformer architecture(図1(c))
Hybrid-CNNのdecoderをtransformerにした。

今回提案するのはすべてをtransformerにしたアーキテクチャ(図1(d))

手法

Pyramid Group Transformer encoderとFeature Pyramid Transformer decoderで構成される(図2)。
f:id:Ninhydrin:20210906094906p:plain

Pyramid Group Transformer(PGT)

PGTはCNNアーキテクチャの様に階層的な特徴を学習する。
PGTの各ステージでは似たような構造をしている。
Patch Transforme Layerでは画像のパッチをN次元に埋め込む。
各ステージではN_i個のPGTブロックがあり、最後のPGTブロックで2x2のオーバーラップなしの結合が行われる(conv? pooling?)。これにより解像度を下げる。
各ステージの出力をdecoder(FPT)に渡してもとの解像度のセグメンテーションを得る。

PGTは徐々に解像度を下げることでCNNで言う受容野を徐々に広げていく。
これはViTが提案した方法だが、これだと画像全体でself-attention(SA)を行うのでコストが高い。なのでグループに分け、その中でSAを行う。
図3を参照。赤い枠がSAを行う範囲で、ViTが画像全体に行うのに対してPGTではグリッドレベルでSAが行われる。
ステージ3, 4はストライド16, 32で、十分にリッチな特徴が得られるので、globalレベルでのSAを適用する。
f:id:Ninhydrin:20210906100335p:plain



l番目のPGTブロックの計算は式1。
f:id:Ninhydrin:20210906100908p:plain
Pyramid Group Multi-Self Attention(PG-MSA)は式2。Hはヘッドの数、iはヘッドのインデックス、Gがグループ数、jがグループのインデックス。
f:id:Ninhydrin:20210906101027p:plain

PGTに関しては4つデザインした。表1を参照。
iはステージ、Pはトークンを減らすファクター、Cはチャンネル、NはPGTブロック数、Gはグループ数、Hはヘッド数、EはMLPでの次元の拡大率。
f:id:Ninhydrin:20210906101512p:plain

Feature Pyramid Transformer(FPT)

PGTの各レベル(解像度)の特徴を集約してsemantic segmentationのマップを出力する(図2を参照)。
CNNベースのFPNを参考に意味情報と空間情報を高解像度なレイヤーに集約していく。
ここで重要な点は2つ
1) トップダウン(低解像の特徴から高解像の特徴)接続は意味情報を直接拡大して高解像に加えているので意味情報が正確にローカライズされている?
2) 各レベルの特徴がspatial reduction transformer blockとbilinearによって徐々にアップサンプリングされてsumされる
これらの2つにより各レベルの各解像度の特徴を扱うことができる。

実験

ImageNet-1Kのクラス分類

同じレベルのパラメータ・GFLOPsのモデルと比較して(RegNetY-4G vs ResNet-50 vs PGT-S)優位。
小さいモデルは特に顕著(PVT-Tiny vs PGT-T)。
大きいモデルではSOTAなSwin-Bにもわずかに勝利(PGT-L)
f:id:Ninhydrin:20210907085336p:plain

Semantic Image Segmentation

PASCAL Context、ADE20K、COCO-Stuffの3つ

PASCAL Context

同じパラメータレベルでの比較は勝利。
Swin-BのUperNetとFTN(PGT-L )を比較して大幅に改善。
ImageNetで事前学習したSwin-LをバックボーンにしたらSOTA。

ADE20K

CNNと比べ少ないパラメータで優位(FTN(PGT-S))。pretrained Swin-Lをバックボーンにすると強い。

COCO-Stuff

Swin transformersより大体改善。pretrained Swin-LバックボーンはSOTA(48.48%)。

f:id:Ninhydrin:20210907085859p:plain

Ablation Strudy

Effectiveness of PGT

PGTの効果を確認するため、segmentationタスクにて ViT、PVT、Swin TransformerにPGTを使ってPASCAL Contextを学習。
大体改善。特にViT-L/16よりパラメータがかなり少ないPGT-Lの方が高精度。
f:id:Ninhydrin:20210907091352p:plain

Effectiveness of FPT

表4の右端の列(Decoder.Transformer.FPT)を参照。僅かな改善量だけど、全てで改善。

Group Choices of PGT

PGTの受容野の大きさはGroupの大きさで調整される。大きいと計算量が多く、小さいと適切にコンテクストを読み取れない。
比較結果が図4
16-4-1-1がベストだけど、メモリコストがかなり大きいので64-16-1-1を採用。
f:id:Ninhydrin:20210907091858p:plain

所感

最近ViT系を読んでいて読んだけど思ったよりシンプルだった。
SAをConvと同じ用にするための制約(ここではGroup)を入れるタイプが多く、クラス分類と違ってImage to ImageのSemantic Segmentation等のタスクでは画像という巨大な時系列を扱う場合メモリや計算量が問題になるようだ。Pixel Shufflerとかを使った事例はあるのかな(Swinのdownsamplingは近いものを感じる)。
ただ、Image to ImageもTransformerでなんとかなりそうにはなってきたが、事前学習や学習速度の問題が実際のところどうなのかが気になる。実用的なのだろうか。