Навчання моделей сегментації: 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