YOLO26如何训练自己的obb数据集（原理介绍+代码详见+结构框图） | （无人机航拍船舶为案列）

重复模块m (n次迭代):
┌─────────────────────────────────────────────────────────┐
│                                                         │
│  如果 attn=True:                                        │
│    Sequential(                                          │
│        Bottleneck(self.c, self.c),                     │  ←─ 先特征提取
│        PSABlock(self.c, attn_ratio=0.5, num_heads=...) │  ←─ 后注意力增强
│    )                                                    │
│                                                         │
│  否则如果 c3k=True:                                     │
│    C3k(self.c, self.c, 2)                             │  ←─ 同YOLOv11
│                                                         │
│  否则:                                                  │
│    Bottleneck(self.c, self.c)                         │  ←─ 同YOLOv11
│                                                         │
└─────────────────────────────────────────────────────────┘

class BboxLoss(nn.Module):
    """Criterion class for computing training losses for bounding boxes."""

    def __init__(self, reg_max: int = 16):
        """Initialize the BboxLoss module with regularization maximum and DFL settings."""
        super().__init__()
        self.dfl_loss = DFLoss(reg_max) if reg_max > 1 else None

import torch
import torch.nn as nn
import math

class AnglePredictor(nn.Module):
    """角度预测模块"""
    def __init__(self, in_channels):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, in_channels//2, 1),
            nn.BatchNorm2d(in_channels//2),
            nn.SiLU(),
            nn.Conv2d(in_channels//2, 1, 1)  # 输出角度参数
        )
        
    def forward(self, x):
        angle = self.conv(x)  # [B, 1, H, W]
        angle = torch.tanh(angle) * (math.pi / 2)  # 限制在[-π/2, π/2]
        return angle

class YOLO26_OBB_Head(nn.Module):
    """YOLO26 OBB检测头"""
    def __init__(self, num_classes=1, num_anchors=3):
        super().__init__()
        self.num_classes = num_classes
        self.num_anchors = num_anchors
        self.out_channels = (5 + num_classes) * num_anchors
        
        # 位置预测分支 (x, y, w, h)
        self.reg_conv = nn.Sequential(
            nn.Conv2d(256, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.SiLU(),
            nn.Conv2d(256, 4 * num_anchors, 1)
        )
        
        # 角度预测分支
        self.angle_predictor = AnglePredictor(256)
        
        # 分类分支
        self.cls_conv = nn.Sequential(
            nn.Conv2d(256, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.SiLU(),
            nn.Conv2d(256, num_classes * num_anchors, 1)
        )
        
        # 置信度分支
        self.obj_conv = nn.Sequential(
            nn.Conv2d(256, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.SiLU(),
            nn.Conv2d(256, 1 * num_anchors, 1)
        )
        
    def forward(self, x):
        """
        输入: [B, C, H, W]
        输出: [B, anchors, H, W, (5+num_classes)]
        """
        B, C, H, W = x.shape
        
        # 位置预测
        reg_pred = self.reg_conv(x)  # [B, 4*anchors, H, W]
        reg_pred = reg_pred.view(B, self.num_anchors, 4, H, W)
        
        # 角度预测
        angle_pred = self.angle_predictor(x)  # [B, 1, H, W]
        angle_pred = angle_pred.unsqueeze(1).repeat(1, self.num_anchors, 1, 1, 1)
        
        # 分类预测
        cls_pred = self.cls_conv(x)  # [B, cls*anchors, H, W]
        cls_pred = cls_pred.view(B, self.num_anchors, self.num_classes, H, W)
        
        # 置信度预测
        obj_pred = self.obj_conv(x)  # [B, anchors, H, W]
        obj_pred = obj_pred.view(B, self.num_anchors, 1, H, W)
        
        # 合并所有预测
        pred = torch.cat([
            reg_pred,           # [B, anchors, 4, H, W]
            angle_pred,         # [B, anchors, 1, H, W]
            obj_pred,           # [B, anchors, 1, H, W]
            cls_pred            # [B, anchors, num_classes, H, W]
        ], dim=2)  # [B, anchors, 5+num_classes, H, W]
        
        # 调整维度: [B, anchors, H, W, 5+num_classes]
        pred = pred.permute(0, 1, 3, 4, 2).contiguous()
        
        return pred

class YOLO26_OBB(nn.Module):
    """完整的YOLO26 OBB模型"""
    def __init__(self, num_classes=1):
        super().__init__()
        
        # Backbone
        self.backbone = CSPDarknet26()
        
        # Neck
        self.neck = PANet_BiFPN()
        
        # 三个检测头 (不同尺度)
        self.head_small = YOLO26_OBB_Head(num_classes, num_anchors=3)
        self.head_medium = YOLO26_OBB_Head(num_classes, num_anchors=3)
        self.head_large = YOLO26_OBB_Head(num_classes, num_anchors=3)
        
        # 锚框配置 (针对船舶检测优化)
        self.anchors = {
            'small': torch.tensor([[10, 13], [16, 30], [33, 23]]),
            'medium': torch.tensor([[30, 61], [62, 45], [59, 119]]),
            'large': torch.tensor([[116, 90], [156, 198], [373, 326]])
        }
        
    def forward(self, x):
        # Backbone提取特征
        features = self.backbone(x)
        
        # Neck多尺度融合
        small, medium, large = self.neck(features)
        
        # 各尺度检测
        pred_small = self.head_small(small)
        pred_medium = self.head_medium(medium)
        pred_large = self.head_large(large)
        
        return {
            'small': pred_small,
            'medium': pred_medium,
            'large': pred_large
        }

import warnings

warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO('ultralytics/cfg/models/26/yolo26-obb.yaml')
    #model.load('yolo26n.pt') # loading pretrain weights
    model.train(data='data/ship_detection.yaml',
                cache=False,
                imgsz=640,
                epochs=200,
                batch=32,
                close_mosaic=10,
                workers=2,
                device='0',
                # optimizer='SGD',   # using SGD,auto
                project='runs/train',
                name='exp',
                )

YOLO26-obb summary (fused): 132 layers, 2,446,602 parameters, 0 gradients, 5.4 GFLOPs
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100% ━━━━━━━━━━━━ 11/11 3.8it/s 2.9s
                   all        343        338      0.717      0.657      0.756      0.449