[論文紹介] U^2-Net

6 min read読了の目安(約5500字

[論文紹介] U^2-Net: Going Deeper with Nested U-Structure for Salient Object Detection

タイトル U^2-Net: Going Deeper with Nested U-Structure for Salient Object Detection
著者 Xuebin Qin, Zichen Zhang, Chenyang Huang, Masood Dehghan, Osmar R. Zaiane, Martin Jagersand
URL https://doi.org/10.1016/j.patcog.2020.107404.

概要

U-Net をネストした構造の U^2-Net を提唱し、画像の中から「興味の対象」を抽出する SOD(Salient Object Detection) タスクにおいて、SOTA を達成。また、パラメータ数を抑えた U^2-Net^\dagger では、パラメタサイズがたかだか 4.7M にもかかわらず、既存の手法と遜色のない精度を達成した。既存の手法と違い、 ImageNet による事前学習を行わず、フルスクラッチで訓練しているため、ImageNet と分布の異なる画像であっても利用可能。

検証画像 (c) および (d) を見ると、他の手法よりも精度良く物体を抽出できていることがわかる。

モデルサイズと精度 (max-F measure) の関係は以下の通り。

Introduction

Salitent Object Detection (SOD) と既存の手法

Salient Objet Detection (SOD) とは、画像の中から「興味のある対象」を抽出するタスクのことで、既存の手法の多くは ImageNet の分類タスクなどで訓練したネットワークを使って特徴量を抽出し、それをもとに物体の抽出を行う。

訓練済モデルを使うことで、質の良い特徴量を利用できるメリットがある反面、

  1. ネットワークが複雑になりがち
  2. ImageNet と分布が異なる場合、十分な性能が出ない
  3. 分類問題を解くように訓練されているため、必ずしも良い特徴量とは限らない。特に分類モデルは位置情報がそこまで重要ではないため、浅い層で pooling を行って解像度を落とすことが多いが、SODやSegmantationでは、位置情報は重要である

関連手法

既存の手法は、訓練済モデルからどうやって特徴量を抽出するかという観点で、以下の2種類に分けられる。

  1. Multi-level deep feature integration
    • 複数のレイヤーから特徴量を集める。各レイヤはそれぞれ違った特徴を学習しているため、複数のレイヤーの特徴を使うことで精度が期待される
  2. Multi-scale feature extraction
    • SODでは局所的な情報と大域的な情報の両方が重要であることが知られているため、訓練済モデルに新たな層を追加し、1つの層から局所的な情報を大局的な情報の両方を抽出する。以下のような手法が提案されている。
      • Pyramid pooling module
      • spatial and channel-wise attention module
      • inception-like contextual weighting module

提案手法

提案手法のネットワーク構成は、下図を見るとわかりやすい。 U-Net をネストさせることで、局所的な特徴と大域的な特徴の効率的な抽出を狙う。

Residual U-blocks

実際には、単純な U-Net のネストではなく、Residual U-block と呼ばれるものを使っている。下図は色々なモデルで使われている代表的な ブロック の構成

定性的な議論をすると、

  1. (a)、(b)、(c) は受容野が小さいため、ブロック単体では局所的な特徴しか抽出できない。
  2. (d) では dilated convolition を使うことで受容野が広くなったが、計算量が多い

そこで、U-Net に Residual block の概念を取り入れた、Residual U-Block (RSU) を提案する。 RSU-L(C_{in}, M, C_{out}) と書いた時、それぞれ

  • L: ブロック内のレイヤ数
  • C_{in}: 入力チャンネル数
  • M: 中間層のチャンネル数

を表すものとする。図(e)にある通り、RSUは以下の3つから構成されている

  1. input convolution layer
  2. U-Net like symmetric encoder-decoder
  3. residual connection which fuses local features and the multi-scale feature

U-Net 構造のおかげで、中間層のチャンネル数の与える計算量へのインパクトは、他のブロックと比べてとても小さい。

RSUブロックの結合

次に、RSUブロックをつなげて全体のネットワークを構築することを考える。ブロックのつなげ方には色々な方法が考えられる。とくに既存の手法(DocUNet や CU-Net)では、U-Netブロックを Sequential に結合しているが、その場合は計算量がブロック数 n に比例してしまう。

そこで、U-Netをネストした形の U^n-Net を提案する。

n は任意の正の整数を選べるが、本論文では n=2 とする。U^n-Net では、入力画像サイズの小さい En_5En_6De_5 では RSU ではなく dilated convolution を使うことで不要な計算を抑えている。

デコーダーの各層からは、上の層への出力に加えて、3x3 convoution と sigumoid を追加し、saliency map を出力するようにしている。この saliency map は、訓練時に用いる。

また、U-Net構造としたことで、深いネットワークにも関わらず、他の手法と比較して少ない計算量とメモリを実現できている。本論文では、構造は同じでパラメーター数の違う2つのモデル U^2-Net(176.3MB)と U^2-Net^\dagger(4.7MB) で精度検証を行っている。

実験に使用したパラメータ数は以下の通り

訓練

各レイヤの saliency map の損失と、最終的な出力の損失の合計をモデルの損失として訓練する。

\mathcal{L}=\sum_{m=1}^{M} w_{\mathrm{side}}^{(m)}l_{\mathrm{side}}^{(m)} + w_{\mathrm{fuse}}l_{\mathrm{fuse}}

ここで l_{\mathrm{side}}^{(m)} は各 saliency map の損失関数であり、それぞれの損失関数は binary crossentropy を利用している

l = - \sum_{r, c}^{H,W}[P_{G(r,c)}\log P_{S(r,c)} + (1 - P_{G(r,c)})\log (1 - P_{S(r,c)})]

P_{S(r,c)} は各ピクセルの出力値、 P_{G(r,c)} は各ピクセルの正解値を表す。

実験

データセット

訓練には DUTS データセットのサブセット、DUTS-TR (10553枚) を利用した。評価は複数のデータセットを使って行った。

評価データセット
DUTOMRON
DUTS-TE
HKU-IS
ECSSD
PASCAL-S
SOD

評価指標

以下の6つの指標を用いて総合的に評価した

  1. PR Curve
  2. F-measure
  3. MAE
  4. weighted F-measure
  5. S-measure
  6. relax boundary F-measure

Ablation Study

本論文では、

  1. RSU
  2. U-Netのネスト
  3. 事前学習不要

の3つの提案をおこなったので、それぞれについて検証を行った

RSUブロックの評価

外側の Encoder-Decoder アーキテクチャを固定して、RSUブロックの部分を他のものに置き換えた結果が以下のとおりであった(下図の Baseline U-Net 〜 PPM U-Net)。U^2-Net と U^2-Net^\dagger がともに他のブロックよりも良い精度であった

U-Netのネスト構造の評価

U-Netブロックをつなげる既存のネットワークがいくつかあるが、代表的なものは以下の2つである。

  1. HourglassNet
  2. CU-Net

正当な評価をするため、U^2-NetのRSUを (Residualでない) U-Net に置き換えたもの(NIV U^2-Net)と、それらを比較した結果、 NIV U^2-Net の精度が良かった。つまり、U-Netのネスト構造は有効であると考えられる。

事前学習の評価

エンコーダーの部分を学習済モデルに付け替えたものと比較した結果、事前学習を用いない U^2-Net のほうが精度が高かった。よって、U^2-Net においては必ずしも事前学習は必要ではないといえる。

State of the Arts との比較

各 SOTA 手法との比較結果は以下の通り。いくつかのタスクにおいては PoolNet に及ばない結果となったものの、多くのタスクでSOTA手法を上回った

まとめと感想

  1. 事前学習なしで、既存の事前学習済のSOTAメソッドと同等以上の精度を達成しているのはすごい
  2. 特に U^2-Net^\dagger は非常にモデルサイズが小さいにもかかわらず、他のSOTA手法と拮抗する精度を達成しているのもすごい
  3. モデルがシンプルで、実装しやすそうに感じた