U-Net 構造と特徴

U-Netとは

U-Netは2015年にOlaf Ronneberger氏らによって発表された論文""U-Net: Convolutional Networks for Biomedical Image Segmentation""に登場した畳み込みニューラルネットワーク(CNN) を使用したアーキテクチャです。
U-Netは主に医療用の画像処理モデルとして開発され、現在は様々な画像セグメンテーションタスクに応用されています。画像セグメンテーションとは、画像の１ピクセルごとに、対象かそうでないかなどのラベルを定義する操作をさします。医療分野では、"腫瘍部位"かそうでないかのセグメンテーションといった技術で応用されています。本記事ではU-Netのアーキテクチャを詳しく解説し、他のCNNアーキテクチャとの違いを見ていきます。

前提知識：畳み込み演算とCNNの基礎知識

引用:https://axa.biopapyrus.jp/deep-learning/cnn/convolution.html

畳み込み演算とは、画像などの2次元データに対し、フィルタやカーネルと呼ばれる小さな行列を用いて各位置で積和演算を実行する演算方法です。通常、フィルタを畳み込みたい行列の左上から順にスライドしながら計算していきます。これにより、以下のような利点を得ることが出来ます。

パラメータ数を効率的に削減することが出来る。
局所的な特徴を抽出できる
物体が画像内のどこにあっても特徴を捉えやすい（移動不変性）

移動不変性とは、物体がどの位置にあっても同じ物体として認識できる性質を表します。

※ここからの解説は単純化のためにpadding=0,stride=1としますn×nの二次元行列に対して、t×tの行列を用いて畳み込み演算を行うと、演算結果は(n-t+1)×(n-t+1)の行列になります。図では6×6の二次元行列に対して、3×3の行列を用いて畳み込み演算を行っているので、結果は4×4になります。畳み込み演算について詳しく知りたい方は引用元サイトをご参照ください。

次に、チャンネル数が入出力で異なり、複数あるときのことを考えます。ここでは、チャンネル数を3次元目の要素数と定義します。図はチャンネル数が３つのRGB画像を、チャンネル数が８つの特徴マップに変換することを表現しています。
図のように、"3層のフィルタ"を8つ用意します。これら合計24個の行列はすべて異なる行列です。 "3層のフィルタ"はそれぞれR,G,Bの画像データと畳み込み演算を行い、3つの行列データが作成します。これらの行列の和をとり、スカラー値であるバイアスを全ての行列要素に加算します。（※バイアスは加算しないモデルも多く存在します。）
これらをA~Hまで行うと、8つのチャンネルを持つ特徴マップを作成することが出来ます。つまり、A~Hまでの24個のフィルタが全結合層(Linear層) でいう重み行列と同じ機能を持ち、これらは訓練データによって最適化されていきます。

これらを一般化すると、チャンネル数(i)から、チャンネル数(j)に畳み込みニューラルネットワークを繋ぐとき、チャンネル数(i)のフィルタがj個必要になります。事前知識は以上です。

U-Netのアーキテクチャ

引用:https://arxiv.org/abs/1505.04597

UNetという名前は、アーキテクチャが「U」の形をしていることから名付けられました。図では、 inputのimage（画像）がoutputのsegmentation（分類）までどのように変換されているかが分かります。一般に図の左部分、下っていく操作部をDecoder, 上っていく操作部をEncoderといいます。図の深い場所ほど多くの畳み込みを行っているので、位置情報が薄れています。これはCNNと特徴ともいえます。以下に図の参照方法と各矢印で表される変換がどのような操作なのかを説明します。通常のCNNと大きく異なるのは全結合層を排除し、スキップ接続の概念を取り込んだことです。 U-Netは画像セグメンテーションに特化しており、全結合層を必要としません。その代わり、コピークロップとアップサンプリングを追加し、スキップ接続を可能にしました。

各ニューラルネットワークを開設する前に☝図の参照方式を説明します。水色で表される長方形を本記事では"Box"と表現します。これはデータのサイズを表しています。 Boxを繋ぐ矢印はニューラルネットワークを表しています。例えば、図の最も左上の572×572×1のBoxは矢印で表されるニューラルネットワークを介して570×570×64のBoxになります。

conv

convとは畳み込みニューラルネットワークのことで、conv 3×3ならサイズ3×3のフィルタを使って畳み込み演算を行います。

前述のとおり、畳み込みニューラルネットワークでは入出力チャンネル数に応じてフィルタを用意する必要があります。

ReLU

畳み込み演算を行ったあとはReLU関数を用いて非線形性を導入します。入力xに対し、ReLU関数は以下の式で表されます。

ReLU(x)=\max(0, x)

ReLU関数では負の値は0とし、正の値はそのまま返します。ReLU関数は他の層と比べて正の値の勾配が 1に固定されるので、勾配消失が起きにくいという性質があります。この辺りの非線形関数については別記事でまとめようと思います。

max pool 2×2

引用:https://axa.biopapyrus.jp/deep-learning/cnn/pooling.html

プーリング層では画像サイズを削減しています。図のように2×2のウィンドウの中で最大値を値として出力します。図のようなstride(ウィンドウの移動幅)が2の場合、画像サイズは半分になります。図では6×6のウィンドウが3×3になっています。このようにデータサイズを小さくするようなpooling操作を論文ではダウンサンプリングといいます。移動不変性を得るうえで非常に重要になります。

up-conv 2×2

引用:https://datascience.stackexchange.com/questions/6107/what-are-deconvolutional-layers

up-convはTransposed ConvolutionまたはDeconvolution、日本語では転置畳み込みまたは逆畳み込みと呼ばれる操作で、畳み込みの逆で行列サイズを増やす操作になります。図はTransposed Convolutionの説明に使用される有名な図です。上の緑色のマスが5×5の出力の行列、青いマスが3×3の入力の行列であり、白いマスはすべて0です。影になっている部分は畳み込み演算を行うウィンドウであり影になっている部分に3×3のフィルタを畳み込んで緑色の出力を作成しています。

論文では2×2のカーネルを使用しています。2×2のカーネルを使用して転置畳み込みを行うと、3×3の入力行列は 6×6の出力行列になります。同様の方法で、4倍の画像サイズの行列を論文では作成しています。

また、前提知識にあった畳み込みによるチャンネル数の操作と同様に逆畳み込みでチャンネル数を半分にします。 256チャンネル数の特徴マップを128チャンネル数にする場合は、256チャンネル数のフィルタを128個用意すれば良いです。転置畳み込みの操作をアップサンプリングといいます。

copy and crop

U-Netの特徴ともいえる層です。

コピーとは特徴マップをコピーする操作を指します。コピーした特徴マップをup-convの出力と "結合"したいのですが画像サイズが合わないと結合できません。そのため、画像の外側の部分をクロップして省きます。
例として、図上部の灰色矢印で表されるコピークロップを考えます。136×136×256のデータをコピーして保持します。図右部のup-convニューラルネットワークにより、104×104×256の出力が得られ、これに"結合"したいので画像サイズを合わせるために 136×136×256の特徴マップ形式の画像データの外側をクロップして104×104×256の画像データにし、チャンネル方向にたいして前に"結合"します。これを"スキップ接続"といいます。図の白い部分が結合された特徴マップです。

U-Netの特徴とまとめ

アーキテクチャの解説がかなり長くなってしまったので特徴を簡単にまとめます。

スキップ接続により、ある程度の位置情報を保持することが出来る。
スキップ接続により、細部の特徴を捉えることが出来る。
他のアーキテクチャよりもパラメータ数が少なく、高速である。
効率的な学習が可能で、他のモデルよりも少量のトレーニングデータで訓練可能である。

以上の特徴から、データ数が比較的少なく、かつ位置情報や１ピクセルレベルの細部の情報が重要な医療分野の2セグメンテーションタスクで使用される目的で、 U-Netは開発されました。現代では、自然画像処理や自動運転、衛星技術などの画像セグメンテーションで広く用いられ、 Diffusion modelのEncoderとして画像生成での分野でも再注目されています。手が空けば、PytorchでのU-Netの実装について、別記事にまとめようと思います。