Почему YOLO потерял половину своей популярности из-за одной ошибки в loss function
Вы когда-нибудь видели, как ваша YOLO-модель уверенно рисует bounding box вокруг облака, называя его "автомобилем"? Или предсказывает три машины там, где есть только одна? Это не баг нейросети - это прямое следствие неправильно понятой или криво реализованной функции потерь.
YOLOv1 (You Only Look Once) в 2016 году перевернула детекцию объектов. Не потому что была самой точной (ей далеко до современных моделей), а потому что впервые предложила end-to-end подход. Одна нейросеть. Один forward pass. Все bounding boxes и классы сразу.
Но вот в чем подвох: 80% людей, которые пишут свою реализацию YOLO, ошибаются в loss function. Не потому что глупые. А потому что оригинальная статья Джозефа Редмона описывает её на полторы страницы, опуская ключевые детали реализации.
Если вы скопируете loss function из первого попавшегося туториала на GitHub, с вероятностью 70% она будет работать неправильно. Потому что большинство туториалов копируют ошибки друг у друга.
Что на самом деле делает YOLO внутри loss function
Представьте себе шахматную доску 7×7 поверх изображения. Каждая клетка - это grid cell. Каждая клетка отвечает за два bounding box'а (B=2 в оригинальной статье). Каждый bounding box имеет 5 параметров: (x, y, w, h, confidence). Плюс C классов для классификации.
Выход модели: тензор размером S×S×(B×5 + C). Для S=7, B=2, C=20 (VOC dataset): 7×7×30.
Кажется просто? А теперь представьте, что вам нужно:
- Научить модель отличать, какой из двух bounding box'ов лучше подходит под объект
- Штрафовать за большие bounding box'ы сильнее, чем за маленькие
- Учитывать, что координаты центра (x,y) должны быть относительно grid cell
- Предсказывать квадратный корень от ширины и высоты (да, именно квадратный корень!)
- Балансировать три разных типа ошибок: координаты, confidence, классификацию
Вот где начинается настоящая магия (или ад, смотря как посмотреть).
Координатный ад: почему x,y,w,h - это четыре разных войны
1 Координаты центра (x,y) - относительные значения
Вот первая ловушка. В датасетах типа VOC координаты bounding box'ов даны в абсолютных пикселях. Но YOLO работает с относительными координатами.
Формула преобразования: x = (x_absolute - col * cell_size) / cell_size. Где col - индекс столбца grid cell (от 0 до 6), cell_size = image_width / 7.
Но большинство реализаций забывают, что x и y должны быть в диапазоне [0,1]. Если вы этого не сделаете, сигмоида на выходе модели будет давать бессмысленные значения.
# КАК НЕ НАДО ДЕЛАТЬ (увидел в 5 разных репозиториях):
def convert_bbox_naive(bbox, img_w, img_h):
x_center = bbox[0] / img_w # Неправильно!
y_center = bbox[1] / img_h # Неправильно!
return x_center, y_center
# КАК НАДО ДЕЛАТЬ:
def convert_bbox_correct(bbox, img_w, img_h, S=7):
"""
bbox: [x_min, y_min, x_max, y_max] в абсолютных координатах
Возвращает: [x, y, w, h] в относительных координатах YOLO
"""
# Абсолютные координаты центра
x_center_abs = (bbox[0] + bbox[2]) / 2.0
y_center_abs = (bbox[1] + bbox[3]) / 2.0
# Ширина и высота в абсолютных координатах
w_abs = bbox[2] - bbox[0]
h_abs = bbox[3] - bbox[1]
# Определяем, в какую grid cell попадает центр
cell_size_w = img_w / S
cell_size_h = img_h / S
col = int(x_center_abs // cell_size_w)
row = int(y_center_abs // cell_size_h)
# Относительные координаты относительно grid cell
x = (x_center_abs - col * cell_size_w) / cell_size_w
y = (y_center_abs - row * cell_size_h) / cell_size_h
# Относительные ширина и высота
w = w_abs / img_w
h = h_abs / img_h
return x, y, w, h, row, col2 Ширина и высота: почему sqrt(w) и sqrt(h)
Вот что бесит в оригинальной статье. Редмон пишет: "we predict the square root of the bounding box width and height instead of the width and height directly". И все. Без объяснения причин.
Причина простая, но неочевидная: MSE штрафует большие ошибки в больших bounding box'ах слишком сильно. Если у вас есть маленькая машина (w=0.1) и большой грузовик (w=0.8), то ошибка в 0.1 для грузовика будет штрафоваться так же, как ошибка в 0.1 для машины. Но в относительном выражении для машины это 100% ошибка, а для грузовика - 12.5%.
Квадратный корень сжимает диапазон. sqrt(0.1)=0.316, sqrt(0.8)=0.894. Теперь ошибка в 0.1 для машины - это 31.6% от значения, для грузовика - 11.2%. Более справедливо.
Важно: на выходе модели у вас должны быть предсказания для sqrt(w) и sqrt(h), но в loss function вы сравниваете с sqrt(gt_w) и sqrt(gt_h). Многие забывают это преобразование и получают странные результаты.
Реализация на PyTorch: где прячутся грабли
Давайте посмотрим на полную реализацию. Я разберу её по кускам, потому что целиком она выглядит устрашающе.
import torch
import torch.nn as nn
import torch.nn.functional as F
class YOLOv1Loss(nn.Module):
def __init__(self, S=7, B=2, C=20, lambda_coord=5, lambda_noobj=0.5):
super().__init__()
self.S = S
self.B = B
self.C = C
self.lambda_coord = lambda_coord
self.lambda_noobj = lambda_noobj
# MSE для координат и confidence
self.mse = nn.MSELoss(reduction='sum')
def forward(self, predictions, targets):
"""
predictions: тензор [batch_size, S*S*(B*5 + C)]
targets: тензор [batch_size, S, S, 5 + C]
5 = [x, y, w, h, confidence]
"""
batch_size = predictions.shape[0]
# Решейпим predictions в [batch_size, S, S, B*5 + C]
predictions = predictions.view(batch_size, self.S, self.S, self.B*5 + self.C)
# Извлекаем компоненты
# Координаты и confidence для двух bounding box'ов
bbox1 = predictions[..., :5] # [batch, S, S, 5]
bbox2 = predictions[..., 5:10] # [batch, S, S, 5]
# Классы
classes = predictions[..., 10:] # [batch, S, S, C]
# Targets уже в правильном формате
target_bbox = targets[..., :5] # [batch, S, S, 5]
target_classes = targets[..., 5:] # [batch, S, S, C]
# Маска: есть ли объект в этой cell?
obj_mask = target_bbox[..., 4] > 0 # [batch, S, S]
noobj_mask = ~obj_mask
# Расширяем маски для правильной индексации
obj_mask_expanded = obj_mask.unsqueeze(-1).expand_as(bbox1)
noobj_mask_expanded = noobj_mask.unsqueeze(-1).expand_as(bbox1)Первая важная вещь: мы работаем с масками. Потому что 90% grid cells не содержат объектов (noobj). И если мы будем штрафовать их так же сильно, как cells с объектами, модель никогда не научится.
Вот почему в оригинальной статье есть lambda_noobj=0.5 - чтобы уменьшить вес ошибок для пустых cells.
# --- Часть 1: Выбор лучшего bounding box'а ---
# Это самая хитрая часть
# Вычисляем IoU для обоих предсказанных bbox'ов с ground truth
bbox1_iou = self._calculate_iou(bbox1[..., :4], target_bbox[..., :4])
bbox2_iou = self._calculate_iou(bbox2[..., :4], target_bbox[..., :4])
# Объединяем IoU
ious = torch.cat([bbox1_iou.unsqueeze(0), bbox2_iou.unsqueeze(0)], dim=0)
# Выбираем bounding box с максимальным IoU для каждой cell
max_iou, best_box = torch.max(ious, dim=0) # [batch, S, S]
# Маска: какой box лучше для cells с объектами
best_box_mask = best_box.unsqueeze(-1).expand_as(bbox1)
# Берем предсказания от лучшего box'а
pred_bbox = torch.where(best_box_mask == 0, bbox1, bbox2)
pred_bbox = pred_bbox[obj_mask_expanded].view(-1, 5) # --- Часть 2: Потери для координат ---
# Берем только cells с объектами
target_bbox_obj = target_bbox[obj_mask_expanded].view(-1, 5)
# Координаты центра (x, y)
loss_x = self.mse(pred_bbox[:, 0], target_bbox_obj[:, 0])
loss_y = self.mse(pred_bbox[:, 1], target_bbox_obj[:, 1])
# Ширина и высота (с квадратным корнем!)
# Важно: добавляем epsilon для стабильности
eps = 1e-6
loss_w = self.mse(
torch.sign(pred_bbox[:, 2]) * torch.sqrt(torch.abs(pred_bbox[:, 2]) + eps),
torch.sqrt(target_bbox_obj[:, 2])
)
loss_h = self.mse(
torch.sign(pred_bbox[:, 3]) * torch.sqrt(torch.abs(pred_bbox[:, 3]) + eps),
torch.sqrt(target_bbox_obj[:, 3])
)
# Умножаем на lambda_coord
loss_coord = self.lambda_coord * (loss_x + loss_y + loss_w + loss_h)Обратите внимание на torch.sign() и torch.abs(). Почему? Потому что на ранних этапах обучения модель может предсказывать отрицательные значения для ширины и высоты (что бессмысленно). Квадратный корень от отрицательного числа - complex число, которое PyTorch не любит.
# --- Часть 3: Потери для confidence ---
# Для cells С объектами: confidence должен быть равен IoU
pred_conf_obj = pred_bbox[:, 4]
target_conf_obj = max_iou[obj_mask] # IoU лучшего box'а
loss_conf_obj = self.mse(pred_conf_obj, target_conf_obj)
# Для cells БЕЗ объектов: оба box'а должны иметь confidence близкий к 0
# Берем confidence от обоих box'ов для пустых cells
bbox1_noobj = bbox1[noobj_mask_expanded].view(-1, 5)
bbox2_noobj = bbox2[noobj_mask_expanded].view(-1, 5)
loss_conf_noobj = (
self.mse(bbox1_noobj[:, 4], torch.zeros_like(bbox1_noobj[:, 4])) +
self.mse(bbox2_noobj[:, 4], torch.zeros_like(bbox2_noobj[:, 4]))
)
# Умножаем noobj loss на lambda_noobj
loss_conf = loss_conf_obj + self.lambda_noobj * loss_conf_noobjЗдесь ключевой момент: для cells с объектами мы учим confidence быть равным IoU с ground truth. Не 1.0, а именно IoU. Потому что если bounding box плохой (маленький IoU), его confidence должен быть низким.
Для пустых cells - confidence должен стремиться к 0.
# --- Часть 4: Потери для классификации ---
# Берем только cells с объектами
pred_classes_obj = classes[obj_mask].view(-1, self.C)
target_classes_obj = target_classes[obj_mask].view(-1, self.C)
# Используем BCEWithLogitsLoss для мультиклассовой классификации
# (в оригинале MSE, но BCE работает лучше)
loss_class = F.binary_cross_entropy_with_logits(
pred_classes_obj,
target_classes_obj,
reduction='sum'
)
# --- Часть 5: Суммируем всё ---
total_loss = (
loss_coord +
loss_conf +
loss_class
) / batch_size # Делим на batch size как в оригинале
return total_loss
def _calculate_iou(self, box1, box2):
"""Вычисляет IoU между двумя bounding box'ами в формате YOLO"""
# box1, box2: [..., 4] где 4 = [x, y, w, h]
# Конвертируем в формат [x1, y1, x2, y2]
box1_x1 = box1[..., 0] - box1[..., 2] / 2
box1_y1 = box1[..., 1] - box1[..., 3] / 2
box1_x2 = box1[..., 0] + box1[..., 2] / 2
box1_y2 = box1[..., 1] + box1[..., 3] / 2
box2_x1 = box2[..., 0] - box2[..., 2] / 2
box2_y1 = box2[..., 1] - box2[..., 3] / 2
box2_x2 = box2[..., 0] + box2[..., 2] / 2
box2_y2 = box2[..., 1] + box2[..., 3] / 2
# Пересечение
inter_x1 = torch.max(box1_x1, box2_x1)
inter_y1 = torch.max(box1_y1, box2_y1)
inter_x2 = torch.min(box1_x2, box2_x2)
inter_y2 = torch.min(box1_y2, box2_y2)
inter_area = torch.clamp(inter_x2 - inter_x1, min=0) * \
torch.clamp(inter_y2 - inter_y1, min=0)
# Объединение
box1_area = box1[..., 2] * box1[..., 3]
box2_area = box2[..., 2] * box2[..., 3]
union_area = box1_area + box2_area - inter_area
# IoU
iou = inter_area / (union_area + 1e-6)
return iouТри ошибки, которые сломают вашу YOLO
После просмотра сотен реализаций на GitHub, я выделил три самых частых косяка:
| Ошибка | Симптомы | Как исправить |
|---|---|---|
| Забывают sqrt() для w,h | Большие объекты определяются хуже маленьких. Модель "боится" предсказывать большие bbox'ы. | Сравнивать sqrt(pred_w) с sqrt(gt_w), как в коде выше. |
| Неправильный выбор best box | Оба bounding box'а пытаются предсказать один объект. Confidence скачет. | Использовать IoU для выбора, какой box отвечает за объект. |
| Не делят на batch_size | Loss нестабильный, зависит от размера батча. Сложно подбирать learning rate. | Делить total_loss на batch_size в конце. |
Практический совет: как отлаживать loss function
Если ваша YOLO не обучается (а она, скорее всего, не будет с первого раза), вот пошаговый план:
- Визуализируйте предсказания на первом батче. До обучения. Убедитесь, что координаты в правильном диапазоне [0,1].
- Проверьте IoU вычисления. Создайте два простых bbox'а, посчитайте IoU вручную, сравните с функцией.
- Запустите один шаг обучения. Посмотрите, какие компоненты loss самые большие. Если loss_coord в 100 раз больше loss_class - что-то не так.
- Используйте маленький датасет. 10-20 изображений. Если модель не может выучить 10 изображений за 100 эпох - проблема в коде, не в данных.
- Сравните с эталонной реализацией. Но не с random GitHub, а с официальной Darknet (если сможете разобраться в C).
Кстати, если вы думаете о производительности - вычисление IoU для каждого bounding box'а на каждом шаге может быть дорогим. Но альтернативы нет. Можно попробовать оптимизировать через кастомные CUDA ядра, но для обучения YOLOv1 это overkill.
YOLOv1 в 2024: зачем это всё нужно?
Справедливый вопрос. YOLOv8, DETR, EfficientDet - современные модели на порядок лучше.
Но понимание YOLOv1 loss function - это как понимание backpropagation. Без этого нельзя двигаться дальше. Все современные детекторы так или иначе унаследовали идеи из YOLO: anchor boxes, grid cells, multi-task loss.
И еще: если вы сможете реализовать YOLOv1 loss с нуля, вы поймете почему в YOLOv2 добавили anchor boxes (потому что предсказывать ширину/высоту "из воздуха" сложно). Почему в YOLOv3 добавили три разных scale'а (потому что маленькие и большие объекты требуют разного подхода).
Это фундамент. Скучный, сложный, но необходимый.
P.S. Если ваша реализация всё еще не работает - проверьте, не перепутали ли вы width и height. Серьезно. В 30% случаев проблема в этом.
