U-Net

U-Net — это свёрточная нейронная сеть (CNN) для семантической сегментации, то есть задачи присвоения класса каждому пикселю изображения. Олаф Роннебергер, Филипп Фишер и Томас Брокс из Фрайбургского университета представили эту архитектуру в 2015 году в статье «U-Net: свёрточные сети для биомедицинской сегментации изображений». Энкодер уменьшает изображение, чтобы выделить его содержимое, а декодер восстанавливает результат до исходного разрешения. Пропускающие соединения передают декодеру мелкие детали энкодера на каждом масштабе, благодаря чему границы объектов получаются чёткими. U-Net хорошо обучалась на нескольких десятках размеченных микроскопических изображений, с большим отрывом победила в двух категориях конкурса по отслеживанию клеток ISBI 2015 и до сих пор остаётся стандартной моделью для анализа медицинских изображений. Позднее её архитектура стала основой сети шумоподавления во многих диффузионных моделях.

Как работает U-Net#

У сети две симметричные половины, образующие два рукава буквы U:

• Сжимающий путь (энкодер): Повторяющиеся пары свёрток 3×3 с активациями ReLU; после каждой пары применяется максимальная субдискретизация 2×2, уменьшающая разрешение вдвое. На каждом уровне количество каналов удваивается — с 64 до 1024. Так сеть обменивает пространственные детали на более широкое рецептивное поле и более богатые признаки.
• Расширяющий путь (декодер): На каждом уровне повышающая свёртка 2×2 удваивает разрешение и вдвое уменьшает количество каналов. Результат конкатенируется с картой признаков энкодера того же уровня, после чего две свёртки 3×3 объединяют их. Последняя свёртка 1×1 преобразует признаки каждого пикселя в оценки классов.

В исходной сети 23 свёрточных слоя и четыре этапа субдискретизации. В ней использовались свёртки без дополнения, поэтому на входной тайл 572×572 пикселя получалась карта сегментации 388×388 пикселей, а большие изображения сегментировались как перекрывающиеся тайлы. В большинстве современных реализаций для свёрток используют дополнение, чтобы размер выхода совпадал с размером входа.

Пропускающие соединения#

При субдискретизации теряется точное местоположение каждого признака. На уровне узкого места сеть знает, что в области есть клеточная мембрана, но не знает, через какие именно пиксели она проходит. Пропускающие соединения передают высокоразмерные карты признаков каждого уровня энкодера через U декодеру, который объединяет информацию о том, «что» содержится в области, из пути повышающей дискретизации с информацией о том, «где» это находится, из энкодера. U-Net конкатенирует эти карты по размерности каналов, а остаточные соединения ResNet складывают признаки внутри блока. Этот код на PyTorch создаёт один уровень энкодера, узкое место и один уровень декодера:

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 входных пикселей получаются оценки двух классов.

Обучение на небольшом наборе данных#

Исходный набор данных для распознавания нейронных структур на электронных микрофотографиях состоял из 30 изображений размером 512×512 пикселей. U-Net компенсировала нехватку данных за счёт интенсивной аугментации данных, особенно случайных эластичных деформаций, имитирующих изгибы тканей, а также взвешенной функции потерь кросс-энтропии, которая усиливает вклад узких промежутков фона между соприкасающимися клетками, чтобы соседние клетки выделялись как отдельные области. Сегодня при обучении моделей для медицинской сегментации часто используют функцию потерь Dice: она напрямую измеряет перекрытие масок и хорошо работает, даже когда на небольшую структуру приходится лишь несколько процентов пикселей.

Варианты и производные архитектуры#

• 3D U-Net и V-Net: Заменяют двумерные операции трёхмерными, чтобы сегментировать объёмные данные, например результаты КТ и МРТ.
• Attention U-Net и UNet++: Добавляют вентили внимания или вложенные плотные пропускающие соединения, чтобы точнее отбирать признаки энкодера для передачи декодеру.
• nnU-Net: Среда автоматически настраивает предобработку, архитектуру и обучение под новый набор данных. Исследование MICCAI 2024 показало, что тщательно настроенные U-Net на основе CNN в этой среде при строгой валидации по-прежнему превосходят более новые архитектуры Transformer и Mamba в трёхмерной медицинской сегментации.
• Диффузионные U-Net: В DDPM и Stable Diffusion 1.x U-Net предсказывает шум в изображении или латентном представлении. Слои внимания позволяют задавать генерацию с помощью текстового промпта. В U-Net модели Stable Diffusion 1.x около 860 миллионов параметров, а ControlNet подключается к энкодеру этой U-Net. В новых генераторах U-Net заменяют на диффузионные трансформеры.

U-Net и родственные модели#

• Полносвёрточные сети (FCN): В 2014 году FCN превратили классификаторы в модели попиксельного предсказания и объединили грубые и точные предсказания, складывая карты оценок классов. U-Net построила полноценный декодер с конкатенацией признаков на каждом уровне.
• Автоэнкодеры: Также состоят из энкодера и декодера, но обучаются восстанавливать входные данные через узкое место. Пропускающие соединения позволили бы автоэнкодеру обойти узкое место, тогда как U-Net использует их для формирования выровненных попиксельных меток.
• Модели сегментации экземпляров: U-Net выдаёт один класс для каждого пикселя, поэтому два соприкасающихся объекта одного класса получают общую метку, если обучение или постобработка не разделяют их.

Модели Ultralytics YOLO26 для семантической сегментации, обозначенные суффиксом -sem, используют упрощённый декодер U-Net. Их шея дважды повышает разрешение выхода бэкбона, изначально работающего на масштабе 1/32, по пути конкатенируя признаки бэкбона на масштабах 1/16 и 1/8. Затем модель классифицирует пиксели на масштабе 1/8 и повышает разрешение оценок классов в 8 раз, до размера входа.

Примеры применения в реальных условиях#

• Анализ спутниковых изображений: Карты земного покрова размечают каждый пиксель как воду, лес, пашню или застройку, чтобы отслеживать вырубку лесов и рост городов.
• Промышленный контроль: Маски царапин, трещин и коррозии позволяют точно измерить размер и форму дефекта — то, чего не может дать ограничивающая рамка.
• Медицинская визуализация: U-Net выделяет органы, опухоли и очаги поражения на КТ- и МРТ-снимках для диагностики и планирования лучевой терапии, а также разделяет отдельные клетки на микроскопических изображениях. В конкурсе по отслеживанию клеток ISBI 2015 U-Net, обученная на 20 частично размеченных изображениях клеток HeLa, достигла среднего IoU 77,5% против 46% у метода, занявшего второе место.