Архитектура U-Net: encoder, decoder и skip-connections

U-Net - это сверточная нейронная сеть, предложенная в 2015 году для сегментации биомедицинских изображений. Её ключевая идея: симметричная U-образная форма, где encoder сжимает входное изображение до компактного представления, а decoder восстанавливает пространственную детализацию через skip-connections - прямые соединения с симметричными уровнями encoder. Сеть обучается на малых выборках и сегодня используется от медицины до спутниковых снимков. Чтобы сразу увидеть, как меняются каналы и разрешение по уровням, покрутите калькулятор ниже - а дальше разберём каждый блок строго.

Encoder: путь сжатия

Encoder U-Net - это последовательность блоков понижения разрешения. Каждый блок состоит из двух свёрточных слоев Conv 3x3 с ReLU и операции max-pooling 2x2, которая вдвое уменьшает пространственное разрешение. При этом число каналов (фильтров) удваивается на каждом уровне.

Если обозначить число базовых фильтров как $base$, то последовательность каналов по уровням encoder выглядит так:

$\text{L1}: 1 \to base, \quad \text{L2}: base \to 2 \cdot base, \quad \text{L3}: 2b \to 4b, \quad \text{L4}: 4b \to 8b$

Для оригинальной статьи $base = 64$, что даёт последовательность $64 \to 128 \to 256 \to 512$ каналов. Число обучаемых параметров одного свёрточного слоя 3x3:

$\text{params} = C_{out} \cdot (C_{in} \cdot 9 + 1)$

где $C_{in}$ и $C_{out}$ - числа входных и выходных каналов, $9 = 3 \times 3$ - размер ядра, $+1$ - смещение (bias) на каждый фильтр.

Рассчитать для своей задачи

Важная особенность: свёртки в U-Net оригинальной статьи - без дополнения (padding=0), что делает выходные карты чуть меньше входных на каждом слое. Это требует кропа при skip-connections. В современных реализациях чаще используют padding=1 (same-padding) - тогда размеры совпадают точно.

Bottleneck: самый глубокий уровень

В основании U-образной архитектуры находится bottleneck - блок с максимальным числом каналов и минимальным пространственным разрешением. Для $base = 64$ bottleneck содержит $base \cdot 16 = 1024$ канала. При входе $572 \times 572$ разрешение в bottleneck составляет около $28 \times 28$ пикселей.

Bottleneck захватывает глобальный контекст: на таком маленьком разрешении каждый «пиксель» карты активаций видит рецептивное поле, охватывающее большую часть входного изображения. Именно это позволяет сети принимать решения о классификации с учётом широкого окружения.

Рассчитать для своей задачи

Skip-connections: зачем соединять encoder и decoder

Skip-connections - главная инновация U-Net. Каждому уровню decoder соответствует симметричный уровень encoder: карты активаций с encoder конкатенируются (по оси каналов) с картами decoder перед свёрткой.

Почему это критически важно? Encoder в процессе сжатия неизбежно теряет точную пространственную информацию - где именно находится граница объекта. Decoder по bottleneck-представлению знает «что» есть на изображении, но не «где» точно. Skip-connections передают детализацию напрямую: высокочастотные признаки с ранних слоев encoder (края, текстуры) соединяются с семантическими картами decoder.

Формально, если decoder на уровне $l$ получает карту $D_l$ размером $H_l \times W_l \times C_l$, а соответствующий уровень encoder дал карту $E_l$ размером $H_l \times W_l \times C_l$, то вход следующего свёрточного блока decoder:

$\text{input} = \text{Concat}(E_l,\, \text{Up}(D_{l+1})) \in \mathbb{R}^{H_l \times W_l \times 2C_l}$

где $\text{Up}$ - операция upsampling (транспонированная свёртка или билинейная интерполяция).

Decoder: восстановление разрешения

Decoder зеркально повторяет encoder. На каждом уровне: upsampling (увеличение разрешения вдвое), конкатенация со skip-connection, два блока Conv 3x3 + ReLU. Число каналов вдвое уменьшается:

$\text{L4}: 1024 + 512 \to 512, \quad \text{L3}: 512 + 256 \to 256, \quad \text{L2}: 256 + 128 \to 128, \quad \text{L1}: 128 + 64 \to 64$

Первый элемент в каждой паре - каналы с decoder предыдущего уровня после upsampling, второй - каналы со skip-connection.

В конце - финальный свёрточный слой Conv 1x1, который проецирует $base$ каналов в $K$ каналов числа классов: именно этот слой выдаёт сегментационную маску. Число его параметров:

$\text{params}_{1\times1} = K \cdot (base + 1)$

Формула числа параметров

Суммарное число параметров U-Net с $L = 4$ уровнями (не считая bottleneck):

$N = \sum_{l=1}^{L} \left[\underbrace{2 \cdot C_{out}^{enc,l}(C_{in}^{enc,l} \cdot 9 + 1)}_{\text{encoder}} + \underbrace{2 \cdot C_{out}^{dec,l}(C_{in}^{dec,l} \cdot 9 + 1)}_{\text{decoder}}\right] + N_{bn} + K(base + 1)$

Для стандартных гиперпараметров ($base = 64$, входной канал = 1, 2 класса) это около 31 млн параметров - число, характерное для классической U-Net.

Рассчитать для своей задачи

Рецептивное поле и глубина архитектуры

Рецептивное поле (receptive field) - область входного изображения, на которую влияет один нейрон в данном слое. Для одного Conv 3x3 рецептивное поле равно $3 \times 3$. После двух последовательных Conv 3x3 оно становится $5 \times 5$. После каждого max-pool 2x2 рецептивное поле в пространстве исходного изображения удваивается.

На уровне bottleneck с 4 шагами max-pool рецептивное поле охватывает участок $124 \times 124$ пикселей при входе $572 \times 572$. Это означает, что сеть «видит» крупные объекты целиком и способна принимать решение о их классе с полным контекстом. При этом skip-connections возвращают детализацию для точной локализации границ.

Именно поэтому U-Net выигрывает у FCN (fully convolutional network) без skip-connections: FCN теряет детализацию при восстановлении разрешения, тогда как U-Net её сохраняет через прямые соединения.

Обучение на малых выборках и аугментация

Одна из ключевых практических причин популярности U-Net в биомедицине - способность обучаться на небольших размеченных наборах данных (десятки-сотни изображений). Авторы оригинальной статьи обучили сеть на 30 изображениях клеток EM-микроскопии и получили state-of-the-art результат.

Это стало возможным благодаря двум факторам. Во-первых, extensive data augmentation: случайные эластические деформации, флипы, повороты, изменение яркости и контраста. Эластические деформации особенно эффективны для биологических структур, форма которых варьирует естественным образом.

Во-вторых, взвешенная функция потерь. Для задач с тонкими границами между объектами (разделение клеток) веса потерь у пограничных пикселей увеличиваются. Карта весов предвычисляется по формуле:

$w(x) = w_c(x) + w_0 \cdot \exp\!\left(-\frac{(d_1(x) + d_2(x))^2}{2\sigma^2}\right)$

где $d_1$ и $d_2$ - расстояния до границ двух ближайших объектов, $w_0 = 10$, $\sigma = 5$ px. Это буквально «штрафует» модель за ошибки на границах между соседними клетками.

Варианты и расширения U-Net

Оригинальная U-Net породила целое семейство архитектур. U-Net++ вводит вложенные skip-connections для более плавного перехода между уровнями - вместо прямой конкатенации encoder-decoder каждый skip-путь проходит через несколько промежуточных узлов. Attention U-Net добавляет механизм внимания к skip-connections, обучая сеть подавлять несущественные пространственные признаки.

ResU-Net заменяет простые conv-блоки на residual-блоки (с остаточным соединением внутри каждого уровня), что упрощает обучение глубоких версий сети. TransUNet и Swin-UNet заменяют часть encoder на transformer-блоки, получая глобальный контекст уже на ранних уровнях.

Выбор числа уровней и $base$ зависит от задачи. Для небольших объектов на высоком разрешении нужно больше уровней encoder (глубже bottleneck, шире рецептивное поле). Для задач с малым числом классов и большими однородными областями достаточно мелкой U-Net с $base = 32$.

Частые ошибки

• Путаница с padding. В оригинальной U-Net padding=0: карты уменьшаются на 2 px на каждый Conv 3x3. При skip-connection нужно кропировать карту encoder до размера decoder. В современных реализациях used padding=1 - проверяйте документацию своего фреймворка.
• Неправильный подсчёт параметров. Входная размерность у первого decoder-слоя на уровне $l$ - сумма каналов skip + upsampled decoder (конкатенация, а не сложение).
• Игнорирование batch normalization. Оригинальная U-Net не использует batch norm (2015), но современные варианты добавляют её между Conv и ReLU - это меняет число параметров и нормы.
• Перепутаны upsample и transpose conv. Билинейная интерполяция не имеет обучаемых параметров; транспонированная свёртка - имеет. Это влияет на общий подсчёт.
• Забытый финальный слой. Conv 1x1 в конце имеет $K \cdot (base + 1)$ параметров - при большом числе классов вклад заметен.

FAQ

Почему U-Net называется U-Net? Из-за формы архитектуры на схеме: encoder идёт «вниз» (уменьшение разрешения), decoder идёт «вверх» (восстановление), и вся схема образует букву U. Skip-connections - это «горизонтальные перемычки» между симметричными уровнями.

Зачем нужен bottleneck, если он сжимает информацию? Bottleneck вынуждает сеть сформировать компактное глобальное представление. Без него decoder не «видел» бы широкого контекста и опирался только на локальные признаки. Именно на уровне bottleneck сеть «понимает», что перед ней клетка, сосуд или опухоль - а не просто граница.

Можно ли использовать U-Net не для медицины? Да. U-Net успешно применяется для сегментации спутниковых снимков (дороги, здания, сельхозугодья), промышленного контроля качества (дефекты материалов), вождения (разметка дороги, пешеходы). Архитектура универсальна: любая задача попиксельной классификации - её область.

Коротко

U-Net - это симметричная энкодер-декодерная сеть с skip-connections. Encoder из $L$ уровней на каждом шаге удваивает число каналов (начиная с $base$) и вдвое уменьшает пространственное разрешение через max-pooling. Bottleneck достигает $base \cdot 16$ каналов и минимального разрешения. Decoder зеркально восстанавливает разрешение через upsampling, конкатенируя карты со skip-connections. Финальный Conv 1x1 выдаёт $K$-канальную сегментационную маску. Число параметров Conv 3x3 слоя: $C_{out}(C_{in} \cdot 9 + 1)$; для классической U-Net с $base = 64$ это около 31 миллиона параметров.