Что такое U-Net?

Что такое U-Net?

U-Net — это популярная архитектура глубокой нейронной сети, специально разработанная для задач семантической сегментации изображений. Она была введена в 2015 году в работе "U-Net: Convolutional Networks for Biomedical Image Segmentation" и с тех пор широко используется в медицинской визуализации и компьютерном зрении.

Структура U-Net

Архитектура имеет форму буквы "U", откуда она и получила своё название:

Encoder (сжатие) → Bottleneck → Decoder (расширение)

Левая часть (Encoder) — сужающийся путь:

• Несколько блоков свёртки (convolution) + pooling
• Уменьшает размер изображения вдвое на каждом шаге
• Извлекает контекстные признаки

Узкое место (Bottleneck) — самая узкая часть:

• Максимально сжатое представление

Правая часть (Decoder) — расширяющийся путь:

• Несколько блоков транспонированной свёртки (upsampling)
• Увеличивает размер вдвое
• Восстанавливает пространственную информацию

Skip Connections (пропускающие соединения)

Ключевое инновационное свойство U-Net — skip connections. Признаки из Encoder соединяются с соответствующими слоями Decoder:

# Упрощённо:
# На каждом уровне:
# decoder_output = concatenate(upsampled_decoder, encoder_feature)
# Это помогает восстановить детали изображения

Пример реализации на PyTorch

import torch
import torch.nn as nn

class UNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1):
        super().__init__()
        
        # Encoder
        self.enc1 = self._conv_block(in_channels, 64)
        self.pool1 = nn.MaxPool2d(2)
        
        self.enc2 = self._conv_block(64, 128)
        self.pool2 = nn.MaxPool2d(2)
        
        # Bottleneck
        self.bottleneck = self._conv_block(128, 256)
        
        # Decoder
        self.upconv2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.dec2 = self._conv_block(256, 128)  # 256 = 128 + 128 (concat)
        
        self.upconv1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.dec1 = self._conv_block(128, 64)
        
        self.final = nn.Conv2d(64, out_channels, 1)
    
    def _conv_block(self, in_c, out_c):
        return nn.Sequential(
            nn.Conv2d(in_c, out_c, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_c, out_c, 3, padding=1),
            nn.ReLU(inplace=True)
        )
    
    def forward(self, x):
        # Encoder с сохранением признаков
        enc1 = self.enc1(x)
        x = self.pool1(enc1)
        
        enc2 = self.enc2(x)
        x = self.pool2(enc2)
        
        # Bottleneck
        x = self.bottleneck(x)
        
        # Decoder с skip connections
        x = self.upconv2(x)
        x = torch.cat([x, enc2], dim=1)  # Skip connection
        x = self.dec2(x)
        
        x = self.upconv1(x)
        x = torch.cat([x, enc1], dim=1)  # Skip connection
        x = self.dec1(x)
        
        x = self.final(x)
        return x

Преимущества U-Net

• Сохранение пространственной информации: Skip connections помогают сохранить детали
• Эффективность: Работает хорошо даже с малым набором данных
• Универсальность: Применима к разным типам изображений
• Простота: Относительно простая архитектура для обучения
• Скорость: Быстрая на GPU

Применение U-Net

Медицинская визуализация:

• Сегментация опухолей
• Сегментация органов (печень, сердце, мозг)
• Анализ снимков МРТ и КТ

Компьютерное зрение:

• Удаление фона из изображений
• Сегментация дорог в спутниковых снимках
• Восстановление изображений (image restoration)

Другие области:

• Сегментация клеток в микроскопии
• Выделение объектов на фото

U-Net vs FCN

# FCN (Fully Convolutional Network)
# - Только encoder-decoder
# - Теряет локальные детали

# U-Net
# - Encoder + skip connections + decoder
# - Сохраняет детали благодаря skip connections
# - Более точна для мелких объектов

Вариации U-Net

• 3D U-Net — для объёмных данных (CT, MRI)
• V-Net — аналог для 3D с модификациями
• Attention U-Net — добавляет механизм внимания
• ResU-Net — комбинирует U-Net с residual connections

U-Net остаётся одной из самых популярных и эффективных архитектур для задач сегментации благодаря своей простоте и мощности.