YOLO v8n 踩的坑

from ultralytics import YOLO

model = YOLO('yolov8n-oiv7.pt')
model.info()
model.predict(source='bus.jpg', save=True)

from ultralytics import YOLO
# Start YOLO v8 training with yaml config file.
print("Start training with yaml config file.")
yolo_yaml = YOLO('./yolov8.yaml')
yolo_yaml.train(epochs=5, data='./TT100K.yaml', resume=True)

yolo_coco.train(name='YOLOv8-COCO', epochs=300, data='./TT100K.yaml', 
                save_period=5, patience=70, batch=16, workers=2, 
                close_mosaic=50, hsv_h=0.02, hsv_s=0.72, hsv_v=0.5, cutmix=0.2,
                device=0,
                amp=False)

try:
    from ultralytics import YOLO

    assert amp_allclose(YOLO("yolo26n.pt"), im)
    LOGGER.info(f"{prefix}checks passed ✅")
except ConnectionError:
    LOGGER.warning(f"{prefix}checks skipped. Offline and unable to download YOLO26n for AMP checks. {warning_msg}")
except (AttributeError, ModuleNotFoundError):
    LOGGER.warning(
        f"{prefix}checks skipped. "
        f"Unable to load YOLO26n for AMP checks due to possible Ultralytics package modifications. {warning_msg}"
    )
except AssertionError:
    LOGGER.error(
        f"{prefix}checks failed. Anomalies were detected with AMP on your system that may lead to "
        f"NaN losses or zero-mAP results, so AMP will be disabled during training."
    )
    return False

"""
CBAM 模块 + 猴子补丁（惰性初始化，适配任意缩放因子）
======================================================
基于 CBAM: Convolutional Block Attention Module (Woo et al., ECCV 2018)。
模块参数在第一次 forward 时根据输入通道数自动构建，无需在 YAML 中指定 c1。

YAML 用法：
    [-1, 1, CBAM, [16, 7]]      # reduction=16, kernel_size=7

导入即完成注册：
    import cbam_patch
"""

import torch
import torch.nn as nn


# ═══════════════════════════════════════════════════════════════════
# 子模块（支持惰性构建）
# ═══════════════════════════════════════════════════════════════════

class ChannelAttention(nn.Module):
    """通道注意力：AvgPool + MaxPool → 共享 MLP → Sigmoid（与原论文一致）"""
    def __init__(self, in_channels: int, reduction: int = 16):
        super().__init__()
        mid_channels = max(1, in_channels // reduction)
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.mlp = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, 1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid_channels, in_channels, 1, bias=False),
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        avg_out = self.mlp(self.avg_pool(x))
        max_out = self.mlp(self.max_pool(x))
        return x * self.sigmoid(avg_out + max_out)


class SpatialAttention(nn.Module):
    """空间注意力：沿通道 mean+max → 拼接 → 7×7 卷积 → Sigmoid"""
    def __init__(self, kernel_size: int = 7):
        super().__init__()
        padding = kernel_size // 2
        self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        avg_map = torch.mean(x, dim=1, keepdim=True)
        max_map, _ = torch.max(x, dim=1, keepdim=True)
        combined = torch.cat([avg_map, max_map], dim=1)
        return x * self.sigmoid(self.conv(combined))


# ═══════════════════════════════════════════════════════════════════
# CBAM 主模块（惰性初始化）
# ═══════════════════════════════════════════════════════════════════

class CBAM(nn.Module):
    """
    CBAM 模块（惰性初始化，自动适配输入通道数）。

    Args:
        reduction : 通道注意力压缩比，默认 16
        kernel_size: 空间注意力卷积核大小，默认 7
    """
    def __init__(self, reduction: int = 16, kernel_size: int = 7):
        super().__init__()
        self.reduction = reduction
        self.kernel_size = kernel_size
        self.channel_att = None
        self.spatial_att = None
        self._built = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not self._built:
            in_channels = x.shape[1]
            self.channel_att = ChannelAttention(in_channels, self.reduction)
            self.spatial_att = SpatialAttention(self.kernel_size)
            self.channel_att.to(device=x.device, dtype=x.dtype)
            self.spatial_att.to(device=x.device, dtype=x.dtype)
            self._built = True
        x = self.channel_att(x)
        x = self.spatial_att(x)
        return x


# ═══════════════════════════════════════════════════════════════════
# 猴子补丁：注册到 ultralytics 命名空间
# ═══════════════════════════════════════════════════════════════════

def _patch():
    try:
        from ultralytics.nn.modules import block
        from ultralytics.nn import tasks

        if not hasattr(block, 'CBAM'):
            setattr(block, 'CBAM', CBAM)    # 猴子补丁写法
        if not hasattr(tasks, 'CBAM'):
            setattr(tasks, 'CBAM', CBAM)    # 猴子补丁写法

        print("✓ [CBAM] success")
    except Exception as e:
        print(f"✗ [CBAM] 注册失败：{e}")

_patch()

cious = iou - (u + alpha * ar)

pb = pred_bboxes[fg_mask]   # [N_fg, 4]
tb = target_bboxes[fg_mask] # [N_fg, 4]

iou = self._iou(pb, tb)     # [N_fg]
L_IoU = 1.0 - iou           # [N_fg]

with torch.no_grad():
    if self.training and L_IoU.numel() > 0:
        mean_L = L_IoU.detach().mean()
        self.running_mean_L_IoU = (
            self.momentum * self.running_mean_L_IoU +
            (1 - self.momentum) * mean_L
        )
    beta = L_IoU / (self.running_mean_L_IoU + 1e-7)
    r = beta / (self.delta * self.alpha ** (beta - self.delta))

    pb_cx = (pb[:, 0] + pb[:, 2]) / 2.0
    pb_cy = (pb[:, 1] + pb[:, 3]) / 2.0
    tb_cx = (tb[:, 0] + tb[:, 2]) / 2.0
    tb_cy = (tb[:, 1] + tb[:, 3]) / 2.0
    gt_w = (tb[:, 2] - tb[:, 0]).detach()
    gt_h = (tb[:, 3] - tb[:, 1]).detach()
    norm = gt_w.pow(2) + gt_h.pow(2) + 1e-7
    dist_sq = (pb_cx - tb_cx).pow(2) + (pb_cy - tb_cy).pow(2)
    R_WIoU = torch.exp(dist_sq / norm)

wiou_per_sample = r * R_WIoU * L_IoU
loss_iou = (wiou_per_sample * weight.squeeze(-1)).sum() / target_scores_sum

"""
WIoU v3 损失函数 + 猴子补丁
===========================
基于 Wise-IoU v3 (Tong et al., 2023)。用离群度驱动的非单调聚焦机制替换 CIoU。

用法（导入即生效）：
    import wiou_patch           # 替换 BboxLoss 为 WIoUBboxLoss
    from ultralytics import YOLO
    model = YOLO('yolov8n.yaml')
    model.train(...)

恢复原始 CIoU：
    wiou_patch.restore()
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

_original_BboxLoss = None


class WIoUBboxLoss(nn.Module):
    # ═══════════ 类属性（供补丁打印等使用）═══════════
    alpha: float = 1.9
    delta: float = 2.9
    momentum: float = 0.99

    def __init__(self, reg_max: int = 16,
                 alpha: float = 1.9,
                 delta: float = 2.9,
                 momentum: float = 0.99):
        super().__init__()
        self.reg_max = reg_max
        self.alpha = alpha
        self.delta = delta
        self.momentum = momentum

        # DFL（与原版一致）
        if reg_max > 1:
            from ultralytics.utils.loss import DFLoss
            self.dfl_loss = DFLoss(reg_max)
        else:
            self.dfl_loss = None

        self.register_buffer('running_mean_L_IoU', torch.tensor(1.0))

    def forward(self, pred_dist, pred_bboxes, anchor_points, target_bboxes,
                target_scores, target_scores_sum, fg_mask, imgsz=None, stride=None):
        weight = target_scores.sum(-1)[fg_mask].unsqueeze(-1)   # [N_fg, 1]

        # ── WIoU v3 核心 ──
        pb = pred_bboxes[fg_mask]   # [N_fg, 4]
        tb = target_bboxes[fg_mask] # [N_fg, 4]

        iou = self._iou(pb, tb)     # [N_fg]
        L_IoU = 1.0 - iou           # [N_fg]

        with torch.no_grad():
            if self.training and L_IoU.numel() > 0:
                mean_L = L_IoU.detach().mean()
                self.running_mean_L_IoU = (
                    self.momentum * self.running_mean_L_IoU +
                    (1 - self.momentum) * mean_L
                )
            beta = L_IoU / (self.running_mean_L_IoU + 1e-7)
            r = beta / (self.delta * self.alpha ** (beta - self.delta))

            # 距离注意力
            pb_cx = (pb[:, 0] + pb[:, 2]) / 2.0
            pb_cy = (pb[:, 1] + pb[:, 3]) / 2.0
            tb_cx = (tb[:, 0] + tb[:, 2]) / 2.0
            tb_cy = (tb[:, 1] + tb[:, 3]) / 2.0
            gt_w = (tb[:, 2] - tb[:, 0]).detach()
            gt_h = (tb[:, 3] - tb[:, 1]).detach()
            norm = gt_w.pow(2) + gt_h.pow(2) + 1e-7
            dist_sq = (pb_cx - tb_cx).pow(2) + (pb_cy - tb_cy).pow(2)
            R_WIoU = torch.exp(dist_sq / norm)

        wiou_per_sample = r * R_WIoU * L_IoU
        loss_iou = (wiou_per_sample * weight.squeeze(-1)).sum() / target_scores_sum

        # ── DFL（与原版完全一致）──
        if self.dfl_loss is not None:
            from ultralytics.utils.tal import bbox2dist
            target_ltrb = bbox2dist(anchor_points, target_bboxes,
                                    self.dfl_loss.reg_max - 1)
            loss_dfl = (
                self.dfl_loss(
                    pred_dist[fg_mask].view(-1, self.dfl_loss.reg_max),
                    target_ltrb[fg_mask],
                ) * weight
            ).sum() / target_scores_sum
        else:
            from ultralytics.utils.tal import bbox2dist
            target_ltrb = bbox2dist(anchor_points, target_bboxes)
            t = target_ltrb * (stride if stride is not None else 1.0)
            if imgsz is not None:
                t[..., 0::2] /= imgsz[1]
                t[..., 1::2] /= imgsz[0]
            p = pred_dist * (stride if stride is not None else 1.0)
            if imgsz is not None:
                p[..., 0::2] /= imgsz[1]
                p[..., 1::2] /= imgsz[0]
            loss_dfl = (
                F.l1_loss(p[fg_mask], t[fg_mask], reduction='none')
                .mean(-1, keepdim=True) * weight
            ).sum() / target_scores_sum

        return loss_iou, loss_dfl

    @staticmethod
    def _iou(box1, box2, eps=1e-7):
        inter_w = (torch.min(box1[:, 2], box2[:, 2])
                   - torch.max(box1[:, 0], box2[:, 0])).clamp(min=0)
        inter_h = (torch.min(box1[:, 3], box2[:, 3])
                   - torch.max(box1[:, 1], box2[:, 1])).clamp(min=0)
        inter = inter_w * inter_h
        area1 = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1])
        area2 = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1])
        union = area1 + area2 - inter + eps
        return inter / union


def _patch():
    global _original_BboxLoss
    try:
        import ultralytics.utils.loss as loss_module
        if _original_BboxLoss is None:
            _original_BboxLoss = loss_module.BboxLoss
        loss_module.BboxLoss = WIoUBboxLoss
        print(f"✓ [WIoU v3] BboxLoss → WIoUBboxLoss (α={WIoUBboxLoss.alpha}, δ={WIoUBboxLoss.delta})")
    except Exception as e:
        print(f"✗ [WIoU v3] 补丁失败：{e}")


def restore():
    global _original_BboxLoss
    if _original_BboxLoss is not None:
        import ultralytics.utils.loss as loss_module
        loss_module.BboxLoss = _original_BboxLoss
        print("✓ 已恢复原始 BboxLoss (CIoU)")

_patch()