U-Net

U-Netは、画像の各ピクセルにクラスを割り当てるタスクであるセマンティックセグメンテーション向けの畳み込みニューラルネットワーク（CNN）です。フライブルク大学のOlaf Ronneberger、Philipp Fischer、Thomas Broxは、2015年にU-Net: Convolutional Networks for Biomedical Image Segmentationでこのモデルを発表しました。エンコーダーは画像を縮小して画像内の内容を捉え、デコーダーは結果を元の解像度まで拡大します。また、スキップ接続によって各スケールの細部をエンコーダーからデコーダーに渡し、物体の境界を鮮明にします。U-Netは数十枚の注釈付き顕微鏡画像で良好に学習し、ISBI 2015細胞追跡チャレンジの2部門で大差をつけて優勝しました。現在も医用画像解析の標準モデルとして使われています。その設計は後に、多くの拡散モデルでノイズ除去ネットワークとして採用されました。

U-Netの仕組み#

このネットワークは、U字型の2つのアームとして描かれる、対称な2つの部分から構成されます。

• 縮小パス（エンコーダー）： 3×3の畳み込みとReLU活性化を2回適用する処理を繰り返し、各処理の後に2×2の最大プーリングを適用して解像度を半分にします。チャネル数は各段階で64から1,024まで倍増するため、空間的な詳細を減らす代わりに、より広い受容野と豊かな特徴を得られます。
• 拡張パス（デコーダー）： 各段階で2×2のアップコンボリューションを適用して解像度を2倍、チャネル数を半分にし、その結果を同じ段階のエンコーダーの特徴マップと連結して、3×3畳み込みを2回適用して統合します。最後に1×1の畳み込みを適用し、各ピクセルの特徴をクラススコアに変換します。

元のネットワークには23の畳み込み層と4回のプーリング処理があります。パディングなしの畳み込みを使うため、572×572の入力タイルから得られるセグメンテーションマップは388×388となり、大きな画像は重複するタイルに分割してセグメンテーションしていました。現在の実装の多くは、畳み込みにパディングを使い、出力サイズを入力サイズに合わせます。

スキップ接続#

プーリングでは各特徴の正確な位置情報が失われます。ボトルネックまで処理が進むと、領域に細胞膜が含まれていることは分かっても、その膜がどのピクセルを通るかは分かりません。スキップ接続は、エンコーダー各段階の高解像度特徴マップをU字型の反対側にあるデコーダーへコピーし、アップサンプリングされたパスからの「何が」とエンコーダーからの「どこに」を組み合わせます。U-Netはチャネル方向に2つを連結しますが、ResNetの残差接続はブロック内で特徴を加算します。次のPyTorchコードでは、エンコーダー段階1つ、ボトルネック、デコーダー段階1つを構築します。

import torch
from torch import nn

# One encoder level, a bottleneck, and one decoder level of a U-Net
encoder = nn.Sequential(nn.Conv2d(3, 16, 3, padding=1), nn.ReLU())
bottleneck = nn.Sequential(nn.MaxPool2d(2), nn.Conv2d(16, 32, 3, padding=1), nn.ReLU())
upsample = nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2)
decoder = nn.Sequential(nn.Conv2d(32, 16, 3, padding=1), nn.ReLU())
head = nn.Conv2d(16, 2, kernel_size=1)  # 2 classes, e.g. cell and background

image = torch.randn(1, 3, 128, 128)
skip = encoder(image)  # full-resolution features kept for the decoder
context = bottleneck(skip)  # half resolution, twice the channels, wider context

# Skip connection: concatenate upsampled context with the saved high-resolution features
merged = torch.cat([upsample(context), skip], dim=1)
logits = head(decoder(merged))
print(f"context {tuple(context.shape)} -> merged {tuple(merged.shape)} -> logits {tuple(logits.shape)}")

context (1, 32, 64, 64) -> merged (1, 32, 128, 128) -> logits (1, 2, 128, 128)を出力します。アップサンプリングされた16個のコンテキストチャネルと16個のスキップチャネルが連結されて32個になり、出力には入力の128×128ピクセルそれぞれについて2つのクラススコアが含まれます。

少量データでの学習#

電子顕微鏡画像の神経構造を対象とした元の学習セットは、512×512ピクセルの画像30枚で構成されていました。U-Netは、組織の曲がり方を模倣するランダムな弾性変形を特に多く含む強力なデータ拡張と、隣接する細胞の間にある細い背景部分を重視する重み付き交差エントロピー損失によって、少量のデータを補いました。これにより隣接する細胞を別々の領域として識別できます。現在の医用セグメンテーションモデルでは、マスクの重なりを直接測定し、ピクセルの数パーセントしか占めない小さな構造にも対応できるDice損失を使って学習することがよくあります。

派生モデルと後継モデル#

• 3D U-NetとV-Net： 2D処理を3D処理に置き換え、CTやMRIスキャンなどのボリュームをセグメンテーションします。
• Attention U-NetとUNet++： 注意ゲートや、密に接続されたネスト型のスキップパスを追加し、エンコーダーのどの特徴をデコーダーに渡すかを精緻化します。
• nnU-Net： 新しいデータセット向けに、前処理、アーキテクチャ、学習設定を自動で構成するフレームワークです。MICCAI 2024の研究では、厳密に検証した場合、このフレームワークで適切に調整したCNNベースのU-Netが、3D医用セグメンテーションで新しいTransformerやMambaのアーキテクチャを依然として上回ることが示されました。
• 拡散U-Net： DDPMとStable Diffusion 1.xは、U-Netを使って画像または潜在表現のノイズを予測し、プロンプトが生成を制御できるように注意層を追加します。Stable Diffusion 1.xのU-Netは約8億6,000万個のパラメーターを持ち、ControlNetはこのU-Netのエンコーダーに接続されます。新しい生成モデルでは、U-Netを拡散Transformerに置き換えています。

U-Netと関連モデルの比較#

• 全畳み込みネットワーク（FCN）： 2014年のFCNは、分類器をピクセル単位の予測器に変え、クラススコアマップを加算して粗い予測と細かい予測を統合しました。U-Netは、各段階で特徴を連結する完全なデコーダーを構築しました。
• オートエンコーダー： エンコーダーとデコーダーから構成される点は同じですが、狭いボトルネックを通じて入力を再構成するように学習します。スキップ接続を使うとオートエンコーダーはボトルネックを迂回できますが、U-Netでは位置合わせされたピクセル単位のラベルを生成するためにスキップ接続が必要です。
• インスタンスセグメンテーションモデル： U-Netはピクセルごとに1つのクラスを出力するため、学習または後処理で分離しない限り、同じクラスの接触した2つの物体には同じラベルが割り当てられます。

Ultralytics YOLO26のセマンティックセグメンテーションモデルは、-semサフィックスで識別され、U-Netデコーダーを簡略化した構造を採用しています。ネックは1/32スケールのバックボーン出力を2回アップサンプリングし、その過程で1/16スケールと1/8スケールのバックボーン特徴を連結します。その後、1/8スケールでピクセルを分類し、クラススコアを8倍にアップサンプリングして入力サイズに合わせます。

実際のアプリケーション#

• 衛星画像解析： 土地被覆マップでは、森林破壊や都市の拡大を追跡するため、各ピクセルを水域、森林、農地、建物などに分類します。
• 産業検査： 傷、亀裂、腐食のピクセルマスクによって、バウンディングボックスでは表せない欠陥の正確なサイズと形状を測定できます。
• 医用画像： U-NetはCTやMRIスキャン上で臓器、腫瘍、病変を輪郭抽出し、診断や放射線治療計画を支援します。また、顕微鏡画像で個々の細胞を分離します。ISBI 2015細胞追跡チャレンジでは、HeLa細胞の部分的に注釈付けされた画像20枚で学習したU-Netが、平均IoU 77.5%を達成しました。2位の手法は46%でした。