- 이번 포스팅에서는 실시간 객체 탐지 모델 중 가장 실용성, 정확성, 속도가 우수한 v11 을 정리하겠습니다.
직전에 정리한 v12 가 정확도와 속도는 v11 보다 올라갔지만, 아직 실험적인 단계의 모델인 것으로 보였으며,
특히나 v11 보다 우수한 성능을 내기 위해선 전용의 고성능 GPU 가 준비되어 있어야 한다는 것을 보고 제 기술 스택에는 v11 을 추가하기로 결정하였습니다.(경량 + 실시간성 + 디바이스 범용성이 충족되지 못해도 된다면 차라리 무거운 트랜스포머 계열의 모델을 사용하는 것이 나을 테니까요...)
(YOLO v11 설명)
- 논문 : YOLOv11: An Overview of the Key Architectural Enhancements 2024
v11 역시 작년에 출시된 최신 모델입니다.
논문 내용이 짧은 편이네요.
- YOLO 모델의 역사
논문 가장 처음에는 위와 같은 자료가 실려있네요.
간단히 살펴보겠습니다.
최초의 Single Stage Object Detector 모델(You Only Look Once)인 YOLO v1 은 2015 년에서 시작해서, 2024년 까지 11버전까지 나올 정도이므로 거의 1년에 1개 이상의 속도로 발전해온 모델입니다.
v2에서 v4 까지는 이미 정리했으므로 넘어가고,
이전 버전들은 자세히 설명글을 올리지 않았기에 이참에 간단하게 설명드리자면,
v5 : C/C++ 기반의 다크넷 프레임워크에서 파이썬 파이토치 기반으로 변경되어 학습과 적용이 유연해졌습니다.
기존의 ReLU 보다 부드러운 Swish 활성 함수를 사용하여 학습 성능을 높였고,
기존 YOLO 에서 사용하던 앵커 기반 탐지를 on/off 할 수 있는 실험적 기능을 도입하였습니다.
Neck 부분에는 피쳐 피라미드 구조인 FPN 의 발전형인 PANet (Path Aggregation Network) 를 처음으로 도입하여,
Top-Down + Bottom-Up 경로의 의미 정보를 결합(v12 정리글 참고)하였으며,
B-Box 기반 객체탐지뿐 아니라 Head 만 바꾼다면 Instance Segmentation 이 가능하도록 개선되었습니다.
v6 : 앵커 프리가 기본적으로 적용되어서 이제 앵커 계산에 연산량을 낭비하지 않습니다.
전통적인 CNN 모델은 국소 정보 분석에는 강하지만 전역 컨텍스트 정보 해석에는 한계가 있었는데,
Head 부위에 경량 Self-Attention 을 적용하여 이를 보완하여 탐지 정확도를 향상시켰습니다.
CNN 블록에는 ELAN 을 적용하였으며, ELAN 은 v12 정리글에서 정리했듯, 여러 단계의 레이어에서 반환하는 특징맵들을 concat 하여 정보를 보다 풍부하게 만드는 기법입니다.
v6 에서 적용된 앵커프리는, 모델이 앵커 없이 각 픽셀 위치에서 직접 중심점과 크기를 예측하는 방식으로, 앵커 설정 및 앵커 개수만큼 반복되는 연산이 절약됩니다.
여러모로 수동으로 탐지하던 영역을 자동으로 바꾸고 효율을 높인 혁신성 높은 모델이네요.
v7 : v6 에서 사라졌던 앵커가 다시 도입되었습니다.
Self-Attention 을 더욱 적극적으로 도입하여 Head 뿐 아니라 백본과 넥에도 적용하였습니다.
이로인해 기존 모델보다는 무거워졌지만, 더욱 정교한 경계 인식이 가능해짐과 동시에 전역 문맥 정보 파악에도 좋아졌습니다.
ELAN 도 개선되어 Efficient ELAN 을 적용하였으며,
객체 탐지뿐 아니라, Instance Segmentation, Object Tracking 까지 지원합니다.
v8 : v7 에 비해 백본을 더 간결하게 정리했으며,
앵커 프리 모델입니다.
GAN 기반 강화 학습을 적용하여 학습을 했다고 하며, ONNX, CoreML 을 지원 강화하여 추론 속도 최적화를 하였습니다.
v8 모델은 특히 객체 탐지 뿐만 아니라
Instance Segmentaion, 사물(Thing) 과 배경(Stuff) 를 분리하는 Panoptic Segmentation 이나,
사람, 동물 등의 특정 부위(관절, 눈, 손 등)의 좌표를 추정하는 KeyPoint Estimation 까지 지원하는 다재다능한 모델입니다.
v9 : Programmable Gradient Information (PGI) 을 적용하였습니다.
심층 네트워크에서 발생하는 정보 병목과 gradient 흐름을 개선하는 기술로,
주요 분기와 보조 가역 분기를 나누어서, 역전파시 gradient 를 안정시킨다고 합니다.
자세히는 설명하지 않지만 residual 과 비슷한것같네요.
또, Generalized Efficient Layer Aggregation Network (GELAN) 를 도입하였으며,
이는 CSPNet 과 ELAN 을 결합하여 확장한 기술입니다.
CSP, ResBlock, DarkBlock 등을 선택 가능하며, 모듈화가 잘 되어있어서 경량 -> 대형 모델로 확장 가능하면서 성능과 연산 효율을 유지한다고 합니다.
무슨이유인지 v9 부터는 v8 이 지원하는 panoptic segmentation 이나 keypoint estimation 은 지원하지 않는다고 하는군요.
아무래도 객체 탐지 성능 향상에 집중하기 위해서인 것 같습니다.
그래도 본질은 변하지 않으므로 커스텀을 하면 사용이 가능할 것 같네요.
v10 : v9 대비 Neck, SPP 등이 개선되었습니다.
가장 큰 특징으로는, NMS 를 없앤 모델이라는 것입니다.
자세한 내용은 생략합니다.
v11 : 이번에 리뷰할 v11 은 객체 탐지부터 하여 인스턴스 세분화, 포즈 추정까지 지원하는 모델인데,
자세한 내용을 아래서 알아보겠습니다.
YOLO 모델의 구조는 비슷하기 때문에 부위별로 알아보겠습니다.
- Backbone
이미지 특징을 추출하는 백본은 2가지 특징적인 구조로 구성됩니다.
1. C3K2
기존의 C2f 블록을 대신하여 사용한 C3K2 블록은,
큰 커널 하나를 사용하는 대신, 작은 커널 두개를 사용하는 방식을 채택하여, C3 구조에서 커널 2개를 추가하여 K2 입니다.
class ConvBNAct(nn.Module):
def __init__(self, in_ch, out_ch, kernel, stride=1, padding=None, activation=True):
super().__init__()
padding = kernel // 2 if padding is None else padding
layers = [
nn.Conv2d(in_ch, out_ch, kernel, stride, padding, bias=False),
nn.BatchNorm2d(out_ch)
]
if activation:
layers.append(nn.SiLU(inplace=True))
self.block = nn.Sequential(*layers)
def forward(self, x):
return self.block(x)
class BottleneckK2(nn.Module):
def __init__(self, ch, shortcut=True):
super().__init__()
mid = ch // 2
self.conv1 = ConvBNAct(ch, mid, 1)
self.conv2 = ConvBNAct(mid, ch, 3)
self.use_shortcut = shortcut
def forward(self, x):
y = self.conv2(self.conv1(x))
return x + y if self.use_shortcut else y
class C3K2(nn.Module):
def __init__(self, in_ch, out_ch, num_blocks=3, shortcut=True):
super().__init__()
mid = in_ch // 2
self.cv1 = ConvBNAct(in_ch, mid, 1)
self.cv2 = ConvBNAct(in_ch, mid, 1)
self.blocks = nn.Sequential(*[BottleneckK2(mid, shortcut) for _ in range(num_blocks)])
self.cv3 = ConvBNAct(mid * 2, out_ch, 1)
def forward(self, x):
y1 = self.blocks(self.cv1(x))
y2 = self.cv2(x)
return self.cv3(torch.cat([y1, y2], dim=1))
이미지 특징 추출의 핵심 블록이라 생각하시면 됩니다.
Stem 이후에 이것으로 특징 추출 -> stride 2 conv 로 다운 스케일링 을 반복하고, 마지막에 SPPF 를 적용하는 것으로 YOLO v11 백본이 구성됩니다.
2. SPPF(Spatial Pyramid Pooling - Fast)
이건 지난 정리글인 v12 에서도 다룬 적이 있는 기술이네요.
v11 이 이전 모델이므로 여기서도 정리하겠습니다.
SPPF 는 YOLO v5 에서 처음 등장한 구조로, SPP(Spatial Pyramid Pooling) 구조를 간소화하여 연산 속도와 효율을 높인 모듈입니다.
SPPF 의 입력은 최종 특징맵 (B, C, H, W) 입니다.
그리고 출력도 (B, C, H, W) 입니다.
Input → 1×1 Conv → [X, MP(X), MP(MP(X)), MP(MP(MP(X)))] → Concat → 1×1 Conv → Output
구조는 위와 같습니다.
1x1 conv 로 채널 수를 줄여 연산량을 절약하고,
이에 5x5 max pool 을 3번 직렬로 수행하며 중간 결과들을 저장합니다.
그러면 원본인 X 와 Max Pool 이 적용된 결과가 3개 준비되어 총 4개의 결과를 concat 하여 4배의 채널을 가진 특징맵에 1x1 conv 로 다시 처음의 채널 수로 압축하는 프로세스입니다.
class SPPF(nn.Module):
def __init__(self, ch, k=5):
super().__init__()
mid = ch // 2
self.cv1 = ConvBNAct(ch, mid, 1)
self.pool = nn.MaxPool2d(k, stride=1, padding=k//2)
self.cv2 = ConvBNAct(mid * 4, ch, 1)
def forward(self, x):
x_ = self.cv1(x)
y1 = self.pool(x_)
y2 = self.pool(y1)
y3 = self.pool(y2)
return self.cv2(torch.cat([x_, y1, y2, y3], dim=1))
이 블록의 역할을 간단히 설명드리자면,
서로 다른 크기의 풀링을 병렬로 적용하여 멀티 스케일 정보를 융합하는 것입니다.
4. 백본 전체 구조
위에서 정리한 YOLO v11 백본의 핵심 기술인 C3K2, C2PSA, SPPF 를 사용하여 백본 전체 구조를 확인하겠습니다.
class YOLOv11Backbone(nn.Module):
def __init__(self, ch=3):
super().__init__()
c1 = YOLO_CONFIG['base_channels']
n1 = YOLO_CONFIG['base_blocks']
self.stem = ConvBNAct(ch, c1, YOLO_CONFIG['stem_kernel'], YOLO_CONFIG['stem_stride'])
self.stage1 = nn.Sequential(
C3K2(c1, c1 * 2, n1),
ConvBNAct(c1 * 2, c1 * 2, 3, 2)
)
self.stage2 = nn.Sequential(
C3K2(c1 * 2, c1 * 4, n1 * 2),
ConvBNAct(c1 * 4, c1 * 4, 3, 2)
)
self.stage3 = nn.Sequential(
C3K2(c1 * 4, c1 * 8, n1),
SPPF(c1 * 8, k=5)
)
def forward(self, x):
x1 = self.stem(x)
p3 = self.stage1(x1)
p4 = self.stage2(p3)
p5 = self.stage3(p4)
return p3, p4, p5
앞서 설명드렸듯, Stem 이후, C3K2 로 특징 추출, 2 stride conv 로 다운 스케일을 반복하다가 마지막에 SPPF 를 한 구조입니다.
- Neck
Neck 의 전체 구조는 백본에서 반환하는 p3, p4, p5 를 받아들이고,
FPN -> PAN -> C2PSA 를 적용하여 최종 피쳐인 Enhanced p5 를 반환하는 것입니다.
1. FPN(Feature Pyramid Network)
고해상도 피쳐에 상위 계층의 정보를 전달하기 위해 사용하는 기법입니다.
YOLO v5 부터 사용되었던 기술이죠.
class FPN(nn.Module):
def __init__(self, in3, in4, in5, out_ch):
super().__init__()
self.l5 = ConvBNAct(in5, out_ch, 1)
self.l4 = ConvBNAct(in4, out_ch, 1)
self.l3 = ConvBNAct(in3, out_ch, 1)
self.n4 = ConvBNAct(out_ch * 2, out_ch, 3)
self.n3 = ConvBNAct(out_ch * 2, out_ch, 3)
def forward(self, p3, p4, p5):
p5_l = self.l5(p5)
p4_l = self.l4(p4)
p4_td = self.n4(torch.cat([
p4_l,
F.interpolate(p5_l, scale_factor=2, mode='nearest')
], dim=1))
p3_l = self.l3(p3)
p3_td = self.n3(torch.cat([
p3_l,
F.interpolate(p4_td, scale_factor=2, mode='nearest')
], dim=1))
return p3_td, p4_td, p5_l
서로 해상도가 다른 Feature Pyramid 가 있을 때,
위와 같이 채널 수를 1x1 conv 로 맞추고, 업샘플링을 하며 위로 전달해 나갑니다.
이로써 가장 정보가 응축된 하위 피쳐 정보를 상위 피쳐들로 전달합니다.
2. PAN(Path Aggregation Network)
PAN 은 top-down FPN 출력인 p3_td, p4_td, p5_td 를 이용해 이번에는 반대로 bottom-up 경로로 정보를 융합하는 기법입니다.
class PAN(nn.Module):
def __init__(self, ch, num_blocks=2):
super().__init__()
self.down4 = ConvBNAct(ch * 2, ch, 3, 2)
self.c4 = C3K2(ch, ch, num_blocks)
self.down5 = ConvBNAct(ch * 2, ch, 3, 2)
self.c5 = C3K2(ch, ch, num_blocks)
def forward(self, p3_td, p4_td, p5_l):
p3_out = p3_td
p4_bu = self.c4(self.down4(torch.cat([p4_td, p3_out], dim=1)))
p5_bu = self.c5(self.down5(torch.cat([p5_l, p4_bu], dim=1)))
return p3_out, p4_bu, p5_bu
기존 PAN 과 다른 점으로는, 앞서 살펴본 C3K2 를 사용했다는 것입니다.
3. C2PSA(Cross Stage Partial + Parellel Spatial Attention)
C2PSA 는 주요 두가지 아이디어가 합쳐진 모듈입니다.
_Cross Stage Partial(CSP) 구조 :
입력 피쳐를 두갈래로 나누어 일부만 복잡한 블록을 통과시켜서 연산 효율을 높이는 구조입니다.
입력 → 분기 →
(1) CSP 경로: Bottleneck 블록 여러 개 + PSA 적용 → y1
(2) 단순 Conv 경로 → y2
y1, y2 concat → 1x1 Conv → 출력
위와 같이 복잡한 구조와 단순한 구조로 통과시켜 나온 y1, y2 를 concat -> 1x1 conv 채널 압축 을 한 것입니다.
_Parellel Spatial Attention(PSA) :
Spatial Attention 은 공간 정보에서 중요한 영역에 가중치를 주는 Attention 매커니즘입니다.
여러 Spatial Attention 을 병렬로 적용하여 다양한 공간 특징을 동시에 포착하는 것이 Parellel spatial Attention 입니다.
무거운 Self-Attention 을 사용하지 않고도 효율적으로 공간적 측면에서의 중요 정보를 강조할 수 있도록 설계된 기술로,
여러 크기의 커널(1x1, 3x3, 5x5)로 구성된 병렬 Conv Layer를 통해 다양한 사이지의 영역에서 spatial attention map을 계산하고, 이를 통합하여 특징 맵에 곱하는 방식의 attention입니다.
채널 단위로 중요도를 표시하는 SE 를 공간 단위로 표시하도록 바꿨다고 생각하시면 개념은 이해가 가실 것입니다.
class PSA(nn.Module):
def __init__(self, ch):
super().__init__()
self.conv1 = nn.Conv2d(ch, 1, 1, bias=False)
self.conv3 = nn.Conv2d(ch, 1, 3, padding=1, bias=False)
self.conv5 = nn.Conv2d(ch, 1, 5, padding=2, bias=False)
self.sig = nn.Sigmoid()
def forward(self, x):
m1 = self.sig(self.conv1(x))
m3 = self.sig(self.conv3(x))
m5 = self.sig(self.conv5(x))
m = (m1 + m3 + m5) / 3
return x * m
위와 같이 NxN 커널 연산으로 해당 사이즈 영역을 연산하여 sigmoid 로 중요도 개념으로 변환하고,
이러한 값을 평균내어 어텐션 스코어로 만들어서 x 에 곱해주면, 각각 1x1, 3x3, 5x5 사이즈로 파악한 어텐션 중요도가 값에 반영됩니다.
전체 코드는 아래와 같습니다.
class C2PSA(nn.Module):
def __init__(self, in_ch, out_ch, num_blocks=2):
super().__init__()
mid = in_ch // 2
self.branch1 = nn.Sequential(
ConvBNAct(in_ch, mid, 1),
*[BottleneckK2(mid) for _ in range(num_blocks)],
PSA(mid)
)
self.branch2 = ConvBNAct(in_ch, mid, 1)
self.cv = ConvBNAct(mid * 2, out_ch, 1)
def forward(self, x):
b1 = self.branch1(x)
b2 = self.branch2(x)
return self.cv(torch.cat([b1, b2], dim=1))
이를 기반으로 Neck 의 진행 과정을 정리하자면,
Backbone 에서 나온 특징맵 3개를 FPN 에서 입력받아 top-down,
PAN 에서는 FPN 반환값 3개를 bottom-up,
C2PSA 에서는 모든 스케일의 feature 에 각각 attention 적용,
Head 에서는 Neck 최종적으로 어텐션이 반영된 특징 3개를 받아들여서 각각 작은 객체(p3), 중간 객체(p4), 큰 객체(p5)를 탐지.
위와 같이 진행하게 됩니다.
- Head
Head는 객체 검출 결과(Bounding Box, Class, Objectness)를 예측하는 최종 모듈입니다.
YOLO 시리즈에서 Head 는 대부분 구조적으로는 비슷하며, Neck 에서 반환된 3개의 특징을 각각 분석합니다.
class YOLOHead(nn.Module):
def __init__(self, in_ch, num_classes, anchors, stride):
super().__init__()
self.num_classes = num_classes
self.num_anchors = len(anchors)
self.stride = stride
# register anchors as buffer for device mismatch safety
self.register_buffer('anchors', torch.tensor(anchors, dtype=torch.float32).view(-1, 2))
self.conv = nn.Conv2d(in_ch, self.num_anchors * (num_classes + 5), 1)
def forward(self, x):
bs, _, h, w = x.size()
pred = self.conv(x).view(bs, self.num_anchors, self.num_classes + 5, h, w)
pred = pred.permute(0, 1, 3, 4, 2).contiguous()
# create grid with explicit indexing='xy'
grid_x = torch.arange(w, device=x.device)
grid_y = torch.arange(h, device=x.device)
grid_x, grid_y = torch.meshgrid(grid_x, grid_y, indexing='xy')
grid = torch.stack((grid_x, grid_y), dim=-1).view(1, 1, h, w, 2).float()
xy = (torch.sigmoid(pred[..., :2]) + grid) * self.stride
wh = torch.exp(pred[..., 2:4]) * self.anchors.view(1, -1, 1, 1, 2) * self.stride
obj = torch.sigmoid(pred[..., 4:5])
cls = torch.sigmoid(pred[..., 5:])
return torch.cat([xy, wh, obj, cls], dim=-1)
- YOLO v11 전체 코드 구현
import torch
import torch.nn as nn
import torch.nn.functional as F
# --------------------------------------------------------------------
# Configuration
# --------------------------------------------------------------------
YOLO_CONFIG = {
"stem_kernel": 6,
"stem_stride": 2,
"base_channels": 64,
"base_blocks": 3,
"neck_channels": 256,
"neck_blocks": 2,
"anchors": [
[(10, 13), (16, 30), (33, 23)],
[(30, 61), (62, 45), (59, 119)],
[(116, 90), (156, 198), (373, 326)],
],
"strides": [8, 16, 32],
}
# --------------------------------------------------------------------
# Basic Blocks
# --------------------------------------------------------------------
class ConvBNAct(nn.Module):
def __init__(self, in_ch, out_ch, kernel, stride=1, padding=None, activation=True):
super().__init__()
padding = kernel // 2 if padding is None else padding
layers = [
nn.Conv2d(in_ch, out_ch, kernel, stride, padding, bias=False),
nn.BatchNorm2d(out_ch)
]
if activation:
layers.append(nn.SiLU(inplace=True))
self.block = nn.Sequential(*layers)
def forward(self, x):
return self.block(x)
class BottleneckK2(nn.Module):
def __init__(self, ch, shortcut=True):
super().__init__()
mid = ch // 2
self.conv1 = ConvBNAct(ch, mid, 1)
self.conv2 = ConvBNAct(mid, ch, 3)
self.use_shortcut = shortcut
def forward(self, x):
y = self.conv2(self.conv1(x))
return x + y if self.use_shortcut else y
class C3K2(nn.Module):
def __init__(self, in_ch, out_ch, num_blocks=3, shortcut=True):
super().__init__()
mid = in_ch // 2
self.cv1 = ConvBNAct(in_ch, mid, 1)
self.cv2 = ConvBNAct(in_ch, mid, 1)
self.blocks = nn.Sequential(*[BottleneckK2(mid, shortcut) for _ in range(num_blocks)])
self.cv3 = ConvBNAct(mid * 2, out_ch, 1)
def forward(self, x):
y1 = self.blocks(self.cv1(x))
y2 = self.cv2(x)
return self.cv3(torch.cat([y1, y2], dim=1))
class SPPF(nn.Module):
def __init__(self, ch, k=5):
super().__init__()
mid = ch // 2
self.cv1 = ConvBNAct(ch, mid, 1)
self.pool = nn.MaxPool2d(k, stride=1, padding=k//2)
self.cv2 = ConvBNAct(mid * 4, ch, 1)
def forward(self, x):
x_ = self.cv1(x)
y1 = self.pool(x_)
y2 = self.pool(y1)
y3 = self.pool(y2)
return self.cv2(torch.cat([x_, y1, y2, y3], dim=1))
class PSA(nn.Module):
def __init__(self, ch):
super().__init__()
self.conv1 = nn.Conv2d(ch, 1, 1, bias=False)
self.conv3 = nn.Conv2d(ch, 1, 3, padding=1, bias=False)
self.conv5 = nn.Conv2d(ch, 1, 5, padding=2, bias=False)
self.sig = nn.Sigmoid()
def forward(self, x):
m1 = self.sig(self.conv1(x))
m3 = self.sig(self.conv3(x))
m5 = self.sig(self.conv5(x))
m = (m1 + m3 + m5) / 3
return x * m
class C2PSA(nn.Module):
def __init__(self, in_ch, out_ch, num_blocks=2):
super().__init__()
mid = in_ch // 2
self.branch1 = nn.Sequential(
ConvBNAct(in_ch, mid, 1),
*[BottleneckK2(mid) for _ in range(num_blocks)],
PSA(mid)
)
self.branch2 = ConvBNAct(in_ch, mid, 1)
self.cv = ConvBNAct(mid * 2, out_ch, 1)
def forward(self, x):
b1 = self.branch1(x)
b2 = self.branch2(x)
return self.cv(torch.cat([b1, b2], dim=1))
# --------------------------------------------------------------------
# Backbone
# --------------------------------------------------------------------
class YOLOv11Backbone(nn.Module):
def __init__(self, ch=3):
super().__init__()
c1 = YOLO_CONFIG['base_channels']
n1 = YOLO_CONFIG['base_blocks']
self.stem = ConvBNAct(ch, c1, YOLO_CONFIG['stem_kernel'], YOLO_CONFIG['stem_stride'])
self.stage1 = nn.Sequential(
C3K2(c1, c1 * 2, n1),
ConvBNAct(c1 * 2, c1 * 2, 3, 2)
)
self.stage2 = nn.Sequential(
C3K2(c1 * 2, c1 * 4, n1 * 2),
ConvBNAct(c1 * 4, c1 * 4, 3, 2)
)
self.stage3 = nn.Sequential(
C3K2(c1 * 4, c1 * 8, n1),
SPPF(c1 * 8, k=5)
)
def forward(self, x):
x1 = self.stem(x)
p3 = self.stage1(x1)
p4 = self.stage2(p3)
p5 = self.stage3(p4)
return p3, p4, p5
# --------------------------------------------------------------------
# Neck (FPN + PAN)
# --------------------------------------------------------------------
class FPN(nn.Module):
def __init__(self, in3, in4, in5, out_ch):
super().__init__()
self.l5 = ConvBNAct(in5, out_ch, 1)
self.l4 = ConvBNAct(in4, out_ch, 1)
self.l3 = ConvBNAct(in3, out_ch, 1)
self.n4 = ConvBNAct(out_ch * 2, out_ch, 3)
self.n3 = ConvBNAct(out_ch * 2, out_ch, 3)
def forward(self, p3, p4, p5):
p5_l = self.l5(p5)
p4_l = self.l4(p4)
p4_td = self.n4(torch.cat([
p4_l,
F.interpolate(p5_l, scale_factor=2, mode='nearest')
], dim=1))
p3_l = self.l3(p3)
p3_td = self.n3(torch.cat([
p3_l,
F.interpolate(p4_td, scale_factor=2, mode='nearest')
], dim=1))
return p3_td, p4_td, p5_l
class PAN(nn.Module):
def __init__(self, ch, num_blocks=2):
super().__init__()
self.down4 = ConvBNAct(ch * 2, ch, 3, 2)
self.c4 = C3K2(ch, ch, num_blocks)
self.down5 = ConvBNAct(ch * 2, ch, 3, 2)
self.c5 = C3K2(ch, ch, num_blocks)
def forward(self, p3_td, p4_td, p5_l):
p3_out = p3_td
p4_bu = self.c4(self.down4(torch.cat([p4_td, p3_out], dim=1)))
p5_bu = self.c5(self.down5(torch.cat([p5_l, p4_bu], dim=1)))
return p3_out, p4_bu, p5_bu
# --------------------------------------------------------------------
# Detection Head
# --------------------------------------------------------------------
class YOLOHead(nn.Module):
def __init__(self, in_ch, num_classes, anchors, stride):
super().__init__()
self.num_classes = num_classes
self.num_anchors = len(anchors)
self.stride = stride
# register anchors as buffer for device mismatch safety
self.register_buffer('anchors', torch.tensor(anchors, dtype=torch.float32).view(-1, 2))
self.conv = nn.Conv2d(in_ch, self.num_anchors * (num_classes + 5), 1)
def forward(self, x):
bs, _, h, w = x.size()
pred = self.conv(x).view(bs, self.num_anchors, self.num_classes + 5, h, w)
pred = pred.permute(0, 1, 3, 4, 2).contiguous()
# create grid with explicit indexing='xy'
grid_x = torch.arange(w, device=x.device)
grid_y = torch.arange(h, device=x.device)
grid_x, grid_y = torch.meshgrid(grid_x, grid_y, indexing='xy')
grid = torch.stack((grid_x, grid_y), dim=-1).view(1, 1, h, w, 2).float()
xy = (torch.sigmoid(pred[..., :2]) + grid) * self.stride
wh = torch.exp(pred[..., 2:4]) * self.anchors.view(1, -1, 1, 1, 2) * self.stride
obj = torch.sigmoid(pred[..., 4:5])
cls = torch.sigmoid(pred[..., 5:])
return torch.cat([xy, wh, obj, cls], dim=-1)
# --------------------------------------------------------------------
# Full Model
# --------------------------------------------------------------------
class YOLOv11(nn.Module):
def __init__(self, num_classes, in_ch=3):
super().__init__()
self.backbone = YOLOv11Backbone(in_ch)
base_ch = YOLO_CONFIG['base_channels']
neck_ch = YOLO_CONFIG['neck_channels']
neck_blk = YOLO_CONFIG['neck_blocks']
self.fpn = FPN(base_ch * 2, base_ch * 4, base_ch * 8, neck_ch)
self.pan = PAN(neck_ch, neck_blk)
self.attn = nn.ModuleList([C2PSA(neck_ch, neck_ch, neck_blk) for _ in range(3)])
self.heads = nn.ModuleList([
YOLOHead(neck_ch, num_classes, YOLO_CONFIG['anchors'][i], YOLO_CONFIG['strides'][i])
for i in range(3)
])
def forward(self, x):
p3, p4, p5 = self.backbone(x)
p3_td, p4_td, p5_l = self.fpn(p3, p4, p5)
p3_out, p4_bu, p5_bu = self.pan(p3_td, p4_td, p5_l)
feats = [p3_out, p4_bu, p5_bu]
enhanced = [m(f) for m, f in zip(self.attn, feats)]
return [h(f) for h, f in zip(self.heads, enhanced)]