U-Net: 医用画像セマンティックセグメンテーションのCNN構造

投稿日投稿者Masaki Hayashi

1. U-Net の概要

U-Net は，医用画像のセマンティックセグメンテーション向けに提案された，CNN構造である [Ronneberger et al., 2015] ．FCN(Fully Convolutional Network) [Long et al., 2015] のネットワーク構造をもとに， 転置畳み込みやスキップ接続を用いて，特徴ピラミッドを形成する砂時計型のEncoder-Decoderネットワークを提案した．

プロジェクトHP

低解像度画像サイズのマップしか出力できなかったFCNと異なり，入力画像と全く同じ画像サイズで，密な画素ごとのクラスシ識別を学習できるネットワーク構造をU-Netは提案できた．よってその後はU-Netの構造が，密な推定タスク全般でも使用されることとなった (2.4節)．

バイオ・医療系の画像での応用だと，たとえば，細胞核画像入力からの細胞領域の前景背景識別や，レントゲン画像からの臓器領域セグメンテーション，脳MRI画像からの病変疑い領域マスクの推定などに用いられる．

U-Netでは，入力の[512 × 512] 画像から，同じ空間サイズ[512 × 512]の，セグメンテーション結果マップ画像を出力するように，学習をおこなう．ただし，Fully Convolutionalな構造なので，必ずしも入力画像サイズを固定する必要はない．

親記事：セマンティックセグメンテーション(Semantic Segmentation)

2. U-Net

2.1 U-Net 構造の概要

U-Netの構造
図1 U-Netの構造．

U-Netの各層は，アルファベットの「U」の形をした(砂時計型)Encoder-Decoder形式の，対称的な構造を成している(図1)．U-Netと似たようなネットワーク構造であるが，別目的(物体検出向け)に提案された Feature Pyramid Network(FPN) の論文では，前半のEncoderをボトムアップパス(bottom-up path)と呼び，後半のDecoderをトップダウンパス(top-down path)と呼んでいる．このFPNにおける呼びかたの方が，「伝統的な画像認識の呼びかた」そのままであるゆえ，砂時計全体でそれぞれが何を担当しているかを表しており分かりやすいと(管理人は)思う．また，U-NetのようなEncoder-Decoderでは，中央部で次元数が一番低くなった中間部分を，ボトルネック構造と呼ぶ (図1 の中央部，[16 x 16] の 3層部分)．

※ U-Netの元の論文では，前半のEncoder経路全体を縮小パス(contracting path)，後半のDecoder経路全体を拡張パス(expansive path)とそれぞれ呼んではいるが，それぞれは「ボトムアップ/トップダウン」で呼ぶほうが良い．

この図1 のU-Netは，(よく実装される形として)画像入出力サイズが同じとなるようにアレンジした，以下のEncoder-Decoder構成である：

  • 入力：画像サイズ$w \times h$のグレー画像（1チャンネル）．
  • ボトルネック:$16 \times 16$サイズの特徴マップ(× 3層)
  • 出力：各画素が2クラス識別確率である，画像サイズ$(w \times h)$の，2チャンネル出力マップ．

図1のネットワーク構造は，元論文のネットワーク構造を模したものではある．しかし実際には，VGGNetにならって「プーリング層間ではpaddingにより空間サイズを変更しない設計」にすることで，Encoder-Decoder全体を完全な対称構造にする「実用的な」U-Netの例を示している．元論文[Ronneberger et al., 2015] のスキップ接続においては，特徴マップの部分的なクロップをおこなったあとにコピーをするが，そのクロップ処理も簡単のために省いた構成であることに注意．

2.2 U-Net の詳細

より詳しくU-Netのネットワーク構造を示すと，以下の3ブロックから構成される対称構造である(図1)：

  • Encoder：ダウンサンプリングを行うEncoder (VGGNet的な構造)．
  • ボトルネック：空間サイズは変化させないまま畳み込む．
  • Decoder：アップサンプリングを行うEncoder (逆VGGNet的な構造)．
    • [3 x 3] 転置畳み込み層を4段実施する．
    • スキップ接続でEncoder側からコピーされた特徴マップを，後方チャンネルに結合．
    • チャンネル数は，その後の畳み込み層で調整していく．
    • 転置畳み込みごとに，畳み込みチャンネル数を512, 256, 128, 64と減少させていく

このように，「U型の対称ネットワーク + 順伝搬で得られていく各特徴マップ構造」に，スキップ接続が加わることで，Encoderブロック側の「局所情報」と「空間位置情報」を，Decoder側へショートカットして受け渡し，Decode時も使用する．

あとは，U-Net(の元論文)の場合，前景後景のマスクを教師に，画素ごとのバイナリークロスエントロピー和を損失学習で学習を行えば良い．教師マスクが，Nクラスのセマンティックセグメンテーション設定になっても，出力をNクラス特徴マップおよび損失関数にかえるだけで，あとは同じネットワーク構造のままで良い．このように類似の多問題への応用しやすい画像ゆえ，画像補間問題や，pix2pixなどの「画像対画像変換」問題の各種にも，このU-Net系構造がよく用いられることになる(2.4節で，再度詳しく)．

2.3 同じスケール同士で特徴をDecoder側へつなぐスキップ接続

関連記事：スキップ接続(skip connection)

前半で計算した特徴マップを後半ブロックにそのまま受け渡して結合することで，後半ブロックの特徴マップを，2倍のチャンネル数の特徴に拡張したマルチスケールでスキップ接続がされているネットワーク構造である(図1)．これにより，序盤層の例レベル特徴や，中盤層の中レベル特徴を，特徴・位置情報が消されることなく，後半ブロックへとスキップさせてそのままコピーし受け渡すことができる．

このU-Netの「同一スケール特徴マップ間でのスキップ接続」の導入により，低レベル/中レベルの特徴およびその位置情報も受け渡せて生き残った「複数スケールの特徴」を終盤層に用いることができ，画素レベルのクラス識別でき，高精度なセマンティックセグメンテーションが可能となった．

2.4 SegNet との比較．および密なタスク全般への展開

同時期に並行研究されていた SegNet [Badrinarayanan et al., 2015] は U-Netと類似している．双方とも，EncoderのみであったFCN (Fully Convolutional Networks) [Long et al., 2015] に，Decoderを追加することを提案して，セマンティックセグメンテーションを解くようにした，砂時計型Encoder-Decoder 構造のCNN(多くの場合，全結合層を使わないのでFCN)である．

ただし，U-NetとSegNetの両者は同時研究同士であるゆえ，U-Net論文中にはSegNetへの言及は一切ない．そこで，両者の異なる点を私の方で以下に整理しておく：

  • U-Netでは，Decoder側で空間解像度を元画像サイズまで戻していく際に，畳み込み層の逆操作を担当する(学習可能な) 転置畳み込み を導入し，プーリングと逆の「解像度を2倍にする」操作を担当させる．それに対して，SegNet のDecoderでは，unpooling （poolingの逆操作）が行われる．
  • SegNetでは，スキップ接続後に特徴マップ同士を加算して合成する．U-Netではスキップ結合して，加算はせず両者を結合する．
  • SegNetの「前半ブロックで最大プーリングした座標インデックスを，後半ブロックのUnpoolingでサイズ拡大時に使用する仕組み」はU-Netには無い．U-Netでは，代わりに，単純に同じ空間スケール同士をスキップ接続でつなぐ（従って，U-Netのほうが，実装的にも構造理解的にもシンプルに済む）．

以上のように，U-Netは当時の初期手法と比べて，シンプルで対称的なネットワーク構造を持つので，他の類似タスクにも応用が容易である．それもあって，U-Netのスキップ接続を利用した画像間Encoder-Decoder構造は，医用画像目的だけに限らず，セマンティックセグメンテーション全般や，画像対画像変換全般において，その後広く用いられるようになった．また，スキップ接続の記事でも，類似系としてまとめているように，FPN(Feature Pyramid Networks)で得るネットワーク構造も，U-Netの構造と同じ「U型 + 各解像度同士でスキップ接続」である．

ちなみに，画像対画像変換の pix2pix [Isola et al., 2017] の論文に「U-Netのスキップ接続を用いたEncoder-Decoder構造を解説したわかりやすい図」が描かれている．よって，このあたりのテーマを初めて勉強する方は，pix2pixの論文と図も読むと，とても参考になる（pix2pixのプロジェクトHP）．

3. U-Net のまとめ

初期のセマンティックセグメンテーションの１つとして提案されたU-Netについて紹介した．その「シンプルで対称的な(美しい)構造」により(2.1,2.2節)，元論文で提案された医用画像向けに限らず，「密な推定タスク」全般や，画像生成タスク・画像変換タスクによく用いられる定番Encoder-Decoder構造となった(2.4節)．

U-Netは，転置畳み込み層を，Decoder側に導入したうえで，同空間サイズ同士をスキップ接続した4レベルピラミッドEncoder-Decoderである．スキップ接続のおかげで，低レベル・中レベルのボトムアップ画像特徴量も，そのまま後半へコピーして受け渡し，Decodeすることができる．これにより，細かな変化や，中くらい領域の変化があっても，それらもきちんとコンテキストとして捉えた特徴マップをDecode時に使用できるゆえ，高精度な(医用画像の)セマンティックセグメンテーションが可能となった．

参考書籍

References

  • [Badrinarayanan et al., 2015] V. Badrinarayanan, A. Kendall, and R. Cipolla. Segnet: A deep convolutional encoder-decoder architecture for image segmentation. arXiv:1511.00561, 2015.
  • [Isola et al., 2017] Isola, P., Zhu, J.Y., Zhou, T., Efros, A.A.: Image-to-image translation with conditional adversarial networks. In CVPR, 2017.
  • [Long et al., 2015] J. Long, E. Shelhamer, and T. Darrell. Fully convolutional networks for semantic segmentation. In CVPR, 2015.
  • [Ronneberger et al., 2015] O. Ronneberger, P. Fischer, and T. Brox. U-Net: Convolutional Networks for Biomedical Image Segmentation, pages 234–241. Springer International Publishing, Cham, 2015

参照外部リンク