Detection Head

검출 헤드는 객체 검출 네트워크의 마지막 계층 집합입니다. 앞선 계층이 생성한 특징 맵을 읽고 각 후보 위치에서 바운딩 박스와 클래스별 점수를 출력합니다. 헤드 설계에 따라 박스를 앵커 박스에 대한 오프셋으로 인코딩할지, 한 점으로부터의 거리로 인코딩할지, 모델이 후보 박스를 몇 개 생성할지, 비최대 억제(NMS)와 같은 후처리 단계에서 중복을 제거해야 할지가 결정됩니다.

헤드의 위치#

최신 검출기는 세 단계로 구성됩니다. 일반적으로 컨볼루션 네트워크 또는 비전 트랜스포머인 백본은 픽셀을 특징 맵으로 변환하며, 깊어질수록 해상도는 낮아지고 의미 정보는 풍부해집니다. 특징 피라미드 네트워크(FPN)와 같은 넥은 깊은 특징 맵과 얕은 특징 맵을 결합해 각 출력 스케일에 세부 정보와 문맥을 모두 담습니다. 그런 다음 헤드는 일반적으로 스트라이드 8, 16 및 32의 피라미드 각 단계에 작은 컨볼루션 분기를 적용해, 고해상도 맵은 작은 객체를 찾고 저해상도 맵은 큰 객체를 찾도록 합니다.

1단계 검출기는 모든 그리드 셀에서 예측하는 단일 밀집 헤드를 사용합니다. Faster R-CNN과 같은 2단계 검출기는 두 개의 헤드를 사용합니다. 먼저 영역 제안 네트워크가 모든 위치에서 앵커의 점수를 매긴 다음, 관심 영역 헤드가 네트워크에서 제안한 수백 개의 영역을 분류하고 정밀화합니다.

헤드의 박스 및 클래스 예측 방식#

밀집 헤드는 각 그리드 셀에서 세 종류의 값을 예측합니다:

• 클래스 점수: 클래스마다 하나씩 출력하고, 각 위치에서 클래스를 독립적으로 채점할 수 있도록 시그모이드를 적용합니다.
• 박스 회귀: 박스의 위치를 정하는 숫자 네 개입니다. 앵커 박스를 이동하고 크기를 조정하는 오프셋을 사용하거나, 셀 중심에서 박스의 네 모서리까지의 거리를 사용합니다.
• 품질 및 작업별 출력: 박스 순위를 매기는 객체성 점수(YOLOv1~YOLOv5) 또는 중심성 점수(FCOS)와, 인스턴스 세그멘테이션용 마스크 계수나 자세 추정용 키포인트와 같은 작업별 값입니다.

박스 인코딩 방식에 따라 앵커 기반 검출기와 앵커 프리 검출기가 구분됩니다. 각 분기는 몇 개의 컨볼루션으로 구성되며, 마지막에는 예측값마다 출력 채널 하나를 갖는 1 × 1 컨볼루션이 있습니다. RetinaNet과 FCOS는 이미 분류 분기와 회귀 분기를 분리해 사용했으며, YOLOX는 2021년에 이 디커플드 헤드를 YOLO에 도입하고 YOLOv3의 단일 공유 출력 계층보다 학습 수렴이 훨씬 빠르다고 보고했습니다. 아래 도식은 640 × 640 이미지의 세 피라미드 단계에 사용할 마지막 1 × 1 계층을 구성하고 출력을 평탄화해 후보 목록 하나로 만듭니다:

import torch
from torch import nn

# Neck outputs for a 640x640 image at strides 8, 16, and 32 (80x80, 40x40, and 20x20 grids)
channels = (64, 128, 256)
features = [torch.randn(1, c, 640 // s, 640 // s) for c, s in zip(channels, (8, 16, 32))]

# Decoupled head: a box branch (4 values) and a class branch (80 scores) for every pyramid level
box_branches = nn.ModuleList(nn.Conv2d(c, 4, 1) for c in channels)
cls_branches = nn.ModuleList(nn.Conv2d(c, 80, 1) for c in channels)

# Flatten each level's grid and concatenate, giving one candidate per grid cell
boxes = torch.cat([branch(f).flatten(2) for branch, f in zip(box_branches, features)], dim=2)
scores = torch.cat([branch(f).flatten(2) for branch, f in zip(cls_branches, features)], dim=2).sigmoid()
print(f"boxes {tuple(boxes.shape)}, class scores {tuple(scores.shape)}")

boxes (1, 4, 8400), class scores (1, 80, 8400)을 출력합니다. 그리드 셀 6,400개, 1,600개 및 400개가 각각 클래스 점수 80개와 박스 하나를 제안합니다.

원시 출력에서 검출 결과까지#

학습 과정에서 각 객체에 할당되는 예측 수가 결정됩니다. 일대다 할당은 객체 근처의 여러 셀을 양성으로 표시해 헤드에 더 많은 학습 신호를 제공하고, 추론 시 객체마다 서로 겹치는 박스를 여러 개 생성하므로 NMS가 가장 좋은 박스를 선택합니다. DETR와 YOLOv10의 추론 헤드가 사용하는 일대일 할당은 객체마다 정확히 하나의 예측만 학습하므로 출력에 NMS가 필요하지 않습니다.

YOLO26 헤드는 네 모서리까지의 거리를 직접 예측하며, YOLOv8과 YOLO11이 사용한 모서리당 16개 빈의 분포 초점 손실 출력을 제거합니다. 또한 박스 및 클래스 분기의 일대다 복사본과 일대일 복사본을 학습합니다. YOLO26은 기본적으로 일대다 출력에 NMS를 적용하며, nms=False를 전달하면 일대일 분기로 전환되어 이미지당 최대 300개 박스를 NMS 단계 없이 반환합니다.

실제 헤드 설계#

• 새 클래스에 대한 파인튜닝: 클래스 분기는 클래스마다 출력 채널 하나를 사용하므로 클래스 수가 다른 데이터셋으로 학습하려면 최종 분류 계층을 새로 만들어야 합니다. 사전 학습된 백본 및 넥 가중치는 그대로 사용할 수 있습니다.
• 배포: 헤드의 모든 연산은 대상 장치에서 실행되므로 박스를 간단하게 디코딩하고 NMS 없는 출력을 사용하면 내보내기가 쉬워지고 후처리 시간이 줄어듭니다.
• 작은 객체: 스트라이드 4 피라미드 단계를 추가하면 연산 비용이 늘어나는 대신 헤드에 작은 객체를 위한 더 촘촘한 그리드가 제공됩니다. 640 × 640 입력에서는 셀이 160 × 160 = 25,600개 추가되며, 이는 스트라이드 8 단계의 네 배입니다.