Обучение моделей сегментации: SAM, U-Net, SegFormer

Выбор архитектуры сегментации изображений — всегда компромисс между точностью, скоростью и объёмом размеченных данных. Часто клиенты приходят с датасетом в 200 снимков КТ и хотят mIoU >0.9, а на деле U-Net справляется, а SegFormer не хватает данных. Или наоборот: промышленный поток с 5000 изображений, где SegFormer даёт +8% mIoU, но latency критична. Мы накопили опыт обучения моделей сегментации для медицинских (сегментация органов, сосудов, опухолей) и промышленных задач (дефекты на конвейере, трещины, контроль качества). За годы работы мы выполнили более 50 проектов по компьютерному зрению, и в каждом случае выбор архитектуры определялся тремя факторами: объёмом данных, требуемой точностью контуров и допустимой задержкой инференса. Средняя экономия на инференсе после оптимизации — до 60% операционных затрат.

U-Net и SegFormer — наиболее популярные архитектуры, каждая со своими сильными сторонами. Первый устойчив к малым данным, второй — лидер по mIoU на больших датасетах. SAM2 же добавляет интерактивность и fine-tuning с минимальной разметкой.

Какую архитектуру выбрать: U-Net, SegFormer или SAM?

U-Net и его вариации — медицина и промышленность

U-Net остаётся стандартом для медицинской сегментации не из-за качества (SegFormer лучше), а из-за устойчивости при малом датасете (100–500 изображений) и интерпретируемости. Мы рекомендуем U-Net++ для задач с тонкими структурами (сосуды, трещины).

import torch import torch.nn as nn import segmentation_models_pytorch as smp def build_unet_model( architecture: str = 'Unet', # 'Unet', 'UnetPlusPlus', 'MAnet' encoder: str = 'efficientnet-b4', # backbone encoder_weights: str = 'imagenet', num_classes: int = 1, # 1 для бинарной сегментации in_channels: int = 3 ) -> nn.Module: model = getattr(smp, architecture)( encoder_name=encoder, encoder_weights=encoder_weights, in_channels=in_channels, classes=num_classes, activation=None # применяем sigmoid/softmax отдельно ) return model # Loss для медицинской сегментации с малым датасетом class CombinedLoss(nn.Module): def __init__(self, dice_weight: float = 0.5, bce_weight: float = 0.5): super().__init__() self.dice_weight = dice_weight self.bce_weight = bce_weight self.bce = nn.BCEWithLogitsLoss() self.dice = smp.losses.DiceLoss(mode='binary', from_logits=True) def forward(self, preds: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: return (self.bce_weight * self.bce(preds, targets) + self.dice_weight * self.dice(preds, targets)) 

SegFormer — semantic segmentation с акцентом на точность

SegFormer-B4 на большинстве задач semantic segmentation обходит U-Net при датасете >1000 изображений. На ADE20K он показывает mIoU 50.3%, что на 8 пунктов выше U-Net с EfficientNetB4. Для задач с высокими требованиями к детализации контуров мы используем SegFormer-B5 с полиномиальным затуханием learning rate.

from transformers import SegformerForSemanticSegmentation, SegformerConfig import torch import torch.nn.functional as F def train_segformer( num_labels: int, id2label: dict, label2id: dict, pretrained_model: str = 'nvidia/mit-b4', learning_rate: float = 6e-5, num_epochs: int = 50 ) -> SegformerForSemanticSegmentation: model = SegformerForSemanticSegmentation.from_pretrained( pretrained_model, num_labels=num_labels, id2label=id2label, label2id=label2id, ignore_mismatched_sizes=True ) optimizer = torch.optim.AdamW( model.parameters(), lr=learning_rate, weight_decay=0.01 ) # Poly LR decay — стандарт для semantic segmentation scheduler = torch.optim.lr_scheduler.PolynomialLR( optimizer, total_iters=num_epochs, power=0.9 ) return model, optimizer, scheduler def segformer_inference( model: SegformerForSemanticSegmentation, pixel_values: torch.Tensor, # (B, 3, H, W) target_size: tuple = None # (H, W) оригинального изображения ) -> torch.Tensor: """ SegFormer выдаёт логиты в 4x уменьшенном разрешении. Нужен билинейный апскейл до оригинального размера. """ outputs = model(pixel_values=pixel_values) logits = outputs.logits # (B, num_labels, H/4, W/4) if target_size is not None: logits = F.interpolate( logits, size=target_size, mode='bilinear', align_corners=False ) return logits 

SAM2 fine-tuning для кастомных доменов

SAM2 из коробки не знает специфические классы (микроскопия, промышленные дефекты). Fine-tuning только decoder mask head — эффективный подход. При замороженном encoder обучается всего ~4M параметров из 224M, что позволяет обучить модель на одном GPU за сутки.

from sam2.build_sam import build_sam2 import torch def finetune_sam2_decoder( checkpoint_path: str, num_epochs: int = 30, learning_rate: float = 1e-4, freeze_image_encoder: bool = True, # encoder тяжёлый — замораживаем freeze_prompt_encoder: bool = True ) -> torch.nn.Module: sam2 = build_sam2( 'sam2_hiera_large.yaml', checkpoint_path, device='cuda' ) # Замораживаем всё кроме mask decoder for param in sam2.image_encoder.parameters(): param.requires_grad = not freeze_image_encoder for param in sam2.prompt_encoder.parameters(): param.requires_grad = not freeze_prompt_encoder # Только mask decoder обучаем for param in sam2.mask_decoder.parameters(): param.requires_grad = True trainable_params = sum( p.numel() for p in sam2.parameters() if p.requires_grad ) print(f'Trainable parameters: {trainable_params:,}') # Для SAM2-Large с frozen encoder: ~4M из 224M optimizer = torch.optim.AdamW( filter(lambda p: p.requires_grad, sam2.parameters()), lr=learning_rate, weight_decay=1e-4 ) return sam2, optimizer 

Почему fine-tuning SAM эффективнее обучения с нуля?

Предобученный image encoder SAM2 содержит обобщённые признаки форм и границ, которые переносятся на новый домен. Это снижает потребность в данных в 5–10 раз по сравнению с обучением U-Net с нуля. Кроме того, SAM поддерживает интерактивную сегментацию — можно дообучить модель под конкретные промпты (точки, боксы).

Процесс работы: пошагово

1. Анализ данных и прототипирование. Мы оцениваем качество разметки, распределение классов, возможные артефакты. Выбираем архитектуру и формат аугментации. Результат — техническое задание с прогнозом метрик.
2. Подготовка пайплайна обучения. Сборка препроцессинга, аугментация (Albumentations, torchvision), логирование в MLflow. Для медицинских задач используем комбинированный loss (Dice + Focal).
3. Обучение и эксперименты. Проводим grid search по learning rate, weight decay, размеру батча. Фиксируем лучшие чекпоинты. Для SegFormer применяем polynomial LR scheduler, для U-Net — ReduceLROnPlateau.
4. Оптимизация инференса. Экспорт в ONNX, квантизация INT8, профилирование на целевой видеокарте. Добиваемся latency p99 не выше порога заказчика.
5. Интеграция и документация. Разворачиваем модель в Docker-контейнере, готовим REST/gRPC API. Пишем model card: описание архитектуры, метрики на тесте, ограничения.

Метрики сегментации

Сравнение архитектур

Что входит в работу

• Аналитика и прототип: аудит данных, выбор архитектуры, оценка достижимой точности.
• Разработка пайплайна: препроцессинг, аугментация, обучение с логированием в MLflow.
• Оптимизация инференса: экспорт в ONNX/TensorRT, квантизация INT8, снижение latency до p99.
• Интеграция API: REST/gRPC эндпоинт, контейнеризация Docker, деплой в Kubernetes.
• Документация: model card, инструкция по повторению эксперимента, описание метрик.
• Обучение команды: воркшоп по работе с моделью и поддержке.

Сроки

Получите консультацию по выбору архитектуры — это бесплатно. Если у вас есть датасет, закажите аудит данных: мы оценим достижимую точность и подберём архитектуру.

Image segmentation