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

投稿日投稿者Masaki Hayashi

1. U-Net の概要

U-Net は，医用画像のセマンティックセグメンテーション向けに提案された，CNN構造である [Ronneberger et al., 2015] (プロジェクトHP)

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

たとえば細胞核画像入力からの細胞領域の前景背景識別や，レントゲン画像からの臓器領域セグメンテーション，脳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の論文中では，前半のEncoderのことを縮小パス(contracting path)，後半のDecoderブロックのことを拡張パス(expansive path)と呼んでいる．

ちなみに，U-Netと似たようなネットワーク構造だが，物体検出向けに提案された Feature Pyramid Network の論文では，前半のEncoderをボトムアップパス(bottom-up path)と呼び，後半のDecoderをトップダウンパス(top-down path)と呼んでいる．このFPNの呼びかたの方が，伝統的な画像認識の呼びかたそのままであるので，砂時計全体でそれぞれが何を担当しているかをそのまま表していて分かりやすいかもしれない．

U-Netは，画像の縦横サイズが同じ，以下の入出力間のEncoder-Decoder ネットワークである：

  • 入力：サイズ$w \times h$のグレー画像（1チャンネル）．
  • 出力：各画素が2クラス識別確率である$w \times h$の2チャンネルの出力マップ．

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

U-Netのネットワーク構造は，以下の３ブロックから構成される対称構造である(図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 同じスケール同士で特徴をDecoder側へつなぐスキップ接続

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

この「同一スケール同士のスキップ接続」の導入により，序盤や中盤のlow-levelとmid-levelの特徴・位置も生き残った「複数スケールの特徴」を最終的に用いて，画素レベルのクラス識別でき，高精度なセマンティックセグメンテーションが可能となった．

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

以上のように，U-Netは当時の初期手法と比べて，シンプルで対称的なネットワーク構造を持ち理解しやすい．それもあって，U-Netのスキップ接続を利用したEncoder-Decoder構造は，医用画像目的だけに限らず，セマンティックセグメンテーション全般や，画像対画像変換全般において，その後広く用いられるようになった．

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

3. U-Net のまとめ

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

U-Netは，deconvolution層を，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

参照外部リンク

SNSでシェアやブックマークを行う