Computer Vision Model Fine-Tuning for Custom Tasks

Fine-tuning моделей Computer Vision под кастомные задачи

Взять ImageNet-pretrained модель и дообучить на своих данных — звучит просто. На практике большинство проектов спотыкается на одном и том же: обучение улучшает train mAP до 0.91, а production даёт 0.58. Причина почти всегда не в архитектуре, а в несоответствии распределений: аугментации не покрывают production-условия, train/val split сделан по файлам а не по сценам, leakage между похожими изображениями.

Почему переобучение — главная проблема fine-tuning CV

Типичный кейс: детекция дефектов на производстве. 3200 изображений, YOLOv8m, 100 эпох. val [email protected] = 0.89. Запускаем на новой смене — 0.53. Анализ confusion matrix показывает: модель научилась детектировать дефекты по фону (конкретная линия конвейера), а не по самому дефекту. Решение — аугментации, симулирующие смену условий:

import albumentations as A
from albumentations.pytorch import ToTensorV2

# Аугментации для производственного CV
# Имитируем смену освещения, камеры, угла съёмки
production_augments = A.Compose([
    # Геометрические — небольшой диапазон для детекции
    A.ShiftScaleRotate(
        shift_limit=0.05, scale_limit=0.1,
        rotate_limit=10, p=0.5
    ),
    A.HorizontalFlip(p=0.5),
    A.Perspective(scale=(0.02, 0.05), p=0.3),

    # Освещение — ключевое для производства
    A.OneOf([
        A.RandomBrightnessContrast(
            brightness_limit=0.3, contrast_limit=0.3
        ),
        A.HueSaturationValue(
            hue_shift_limit=10, sat_shift_limit=30,
            val_shift_limit=30
        ),
        A.CLAHE(clip_limit=4.0, tile_grid_size=(8, 8)),
    ], p=0.7),

    # Шум и артефакты камеры
    A.OneOf([
        A.GaussNoise(var_limit=(10, 50)),
        A.ISONoise(color_shift=(0.01, 0.05)),
        A.ImageCompression(quality_lower=75, quality_upper=100),
    ], p=0.4),

    # Имитация загрязнения объектива, запотевания
    A.RandomFog(fog_coef_lower=0.1, fog_coef_upper=0.3, p=0.15),
    A.RandomShadow(num_shadows_lower=1, num_shadows_upper=2, p=0.2),

    A.Normalize(mean=(0.485, 0.456, 0.406),
                std=(0.229, 0.224, 0.225)),
    ToTensorV2()
], bbox_params=A.BboxParams(
    format='yolo', label_fields=['class_labels'],
    min_visibility=0.3   # удаляем bbox, если <30% видно после crop
))

Стратегия разбивки данных при fine-tuning

Стратифицированный split по файлам — ошибка, если изображения сняты сериями. Правильно: split по уникальным сценам/объектам/сессиям.

from sklearn.model_selection import GroupShuffleSplit
import pandas as pd

df = pd.read_csv('annotations.csv')
# scene_id — уникальный идентификатор сцены/объекта/сессии
gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=42)

train_idx, val_idx = next(
    gss.split(df, df['label'], groups=df['scene_id'])
)

train_df = df.iloc[train_idx]
val_df   = df.iloc[val_idx]

# Проверка: нет пересечения scene_id между split'ами
assert len(
    set(train_df['scene_id']) & set(val_df['scene_id'])
) == 0, "Data leakage detected!"

Selection backbone и learning rate schedule

Главная ошибка с ViT при малом датасете — обучать все слои сразу. Правильный подход: сначала замораживаем transformer blocks, обучаем только classifier head 10–15 эпох, потом постепенно размораживаем с LR в 10x меньше базового.

import timm
import torch

model = timm.create_model(
    'vit_base_patch16_224',
    pretrained=True,
    num_classes=num_classes
)

# Этап 1: только head
for name, param in model.named_parameters():
    if 'head' not in name:
        param.requires_grad = False

optimizer_stage1 = torch.optim.AdamW(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=1e-3, weight_decay=0.01
)

# После 15 эпох — этап 2: размораживаем последние 4 блока
for name, param in model.named_parameters():
    if any(f'blocks.{i}' in name for i in range(8, 12)):
        param.requires_grad = True

optimizer_stage2 = torch.optim.AdamW(
    [
        {'params': model.head.parameters(), 'lr': 1e-4},
        {'params': [p for n, p in model.named_parameters()
                    if 'blocks' in n and p.requires_grad],
         'lr': 1e-5}
    ],
    weight_decay=0.01
)

Дисбаланс классов

Precision 0.73 при recall 0.91 на редком классе — типичная картина при дисбалансе 1:50. Решения в порядке эффективности:

1. Focal Loss (γ=2.0) — снижает вес лёгких примеров в функции потерь
2. WeightedRandomSampler — oversample редких классов в DataLoader
3. Class-aware augmentation — агрессивнее аугментировать редкие классы

from torch.utils.data import WeightedRandomSampler
import torch

# class_counts: [n_class0, n_class1, ...]
class_weights = 1.0 / torch.tensor(class_counts, dtype=torch.float)
sample_weights = class_weights[targets]  # targets: метки всего датасета

sampler = WeightedRandomSampler(
    weights=sample_weights,
    num_samples=len(sample_weights),
    replacement=True
)

Трекинг экспериментов

MLflow или Weights & Biases обязательны — без трекинга невозможно воспроизвести лучший результат:

import mlflow

mlflow.set_experiment('defect_detection_v3')
with mlflow.start_run(run_name='yolov8m_focal_weighted_sampler'):
    mlflow.log_params({
        'model': 'yolov8m',
        'img_size': 640,
        'epochs': 100,
        'batch_size': 16,
        'lr0': 0.01,
        'loss': 'focal',
        'augment_strategy': 'production_v2'
    })
    # ... обучение ...
    mlflow.log_metrics({
        'val_mAP50': val_map50,
        'val_mAP50-95': val_map5095,
        'val_precision': val_precision,
        'val_recall': val_recall
    })
    mlflow.pytorch.log_model(model, 'model')

Сроки