U-Net: 医用画像セグメンテーションネットワークの初期提案

1 U-Net の概要

U-Net [Ronneberger et al., 2015] は,セマンティックセグメンテーション向けのCNN構造である.(プロジェクトHP)

完全畳み込みネットワーク (Fully Convolutional Network) [Long et al., 2015] のネットワークをもとに,後半Decoder部分をプーリング層から(学習可能な)アップサンプリング層に改善した構造を提案し,Encoder-Decoder構成の対称型のネットワークになっている.

たとえば細胞核画像入力からの細胞領域の前景背景識別や,レントゲン画像からの臓器領域セグメンテーション,脳MRI画像からの病変疑い領域マスクの推定など,バイオ・医療系の画像での応用目的で用いられることが多いネットワーク構造でもある.

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

2. U-Net

2.1 U-Net の構造

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

U-Netの各層は,アルファベットの「U」の形をしたEncoder-Decoder形式の対称構造を成している (図1).U-Netの論文中では,前半のEncoderのことを縮小パス(contracting path),後半のDecoderブロックのことを拡張パス(expansive path)と呼んでいる.

入力画像はサイズ$w \times h$のグレー画像(1チャンネル)で,出力は各画素が2クラス識別確率である$w ¥times h$の2チャンネルの出力マップである.

図1のネットワーク構造は,元論文のネットワーク構造を模したものではある.しかし実際には,VGGNet的な「プーリング層間ではpaddingにより空間サイズを変更しない設計」にすることで,Encoder-Decoder全体を完全対称構造にする「実用的な」U-Netの例を示している.

U-Netのネットワーク構造は,以下の3ブロックから構成される対称構造である(図1):

  • Encoder:ダウンサンプリングを行うEncoder (VGGNet的な構造).
    • 4ステップの 3 x 3 畳み込み層.
    • 最大値プーリング後ごとに,畳み込みチャンネル数を64, 128, 256, 512と増加させていく.
    • 最大値プーリング後ごとの特徴マップを,Decoderの同一空間サイズの特徴マップに結合
  • Bottleneck:空間サイズは変化させないまま畳み込むボトルネック部.
  • Decoder:アップサンプリングを行うEncoder (逆VGGNet的な構造).
    • 4ステップの 3 x 3 畳み込み層.
    • スキップ接続でもらってきたEncoderからの特徴を結合し,チャンネル数を畳み込みで調整.
    • アップ畳み込み層(up-convolution layer)ごとに,畳み込みチャンネル数を512, 256, 128, 64と減少させていく

このU型の対称ネットワーク+特徴マップ構造に,スキップ接続が加わることで,Encoderブロック側の「局所情報」と「空間位置情報」をDecoder側へ受け渡して,Decode時も使用する.

2.2 スキップ接続によるEncoder-Decoder間での情報の受け渡し

前半で計算した特徴マップを後半ブロックにそのまま受け渡し結合することで,後半ブロックの特徴マップを2倍のチャンネル数の特徴に拡張している.

スケールごとのスキップ接続を用いた学習により,序盤のlow-level特徴やmid-level特徴が,特徴および位置情報が消されることなく後半ブロックへ受け渡すことができる.

おかげで,序盤や中盤の特徴・位置も生き残った「複数スケールの特徴」を最終的に用いて,画素レベルのクラス識別できるようになり,高精度のセマンティックセグメンテーションが可能となった.

2.3 SegNet との比較

同時期に並行研究されていた SegNet [Badrinarayanan et al., 2015] は U-Netと類似している.双方とも,EncoderのみであったFCN[Long et al., 2015]に,Decoderを追加することを提案してセマンティックセグメンテーションを解く,砂時計型のEncoder-Decoder CNNである.

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

  • U-Netでは,空間解像度を元画像サイズまで戻していくDecoder側において,畳み込み層の逆操作を担当する(学習可能な) Up-convolution層を導入し,プーリングと逆の「解像度を2倍にする」操作を担当させる (※ 他の論文ではup-convolutionのことを,deconvolution や transposed convoution とも呼ぶ).SegNetでは unpooling (poolingの逆操作)が代わりに行われる.
  • SegNetでは,スキップ接続後に特徴マップ同士を加算して合成するが,U-Netではスキップ結合して,加算せず両者を結合する設計である.
  • SegNetの「前半ブロックで最大プーリングした座標インデックスを,後半ブロックのUnpoolingでサイズ拡大時に使用する仕組み」は,U-Netには無い.その代わりにU-Netでは,単純に同じ空間スケール間でスキップ接続する(よって,実装的にも構造理解的にもシンプル).

以上のように,U-Netは当時の初期手法と比べて,シンプルで対称的なネットワーク構造を持ち理解しやすい.医用画像中心にセマンティックセグメンテーションでよく用いられるようになった.

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

3. まとめ

初期のセマンティックセグメンテーションネットワークの1つとして提案されたU-Netについて紹介した.U-Netは医用画像向けによく用いられるネットワークとなった.

U-Netはdeconvolutionup層をDecoder側に導入したうえで,同空間サイズ同士をスキップ接続したことで,low-level, mid-level,の各抽象化度の特徴量も取り込んだ特徴マップを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

参照外部リンク

関連記事