添加损失函数代码:
class DFLoss(nn.Module):"""Criterion class for computing Distribution Focal Loss (DFL)."""def __init__(self, reg_max: int = 16) -> None:"""Initialize the DFL module with regularization maximum."""super().__init__()self.reg_max = reg_maxdef __call__(self, pred_dist: torch.Tensor, target: torch.Tensor) -> torch.Tensor:"""Return sum of left and right DFL losses from https://ieeexplore.ieee.org/document/9792391."""target = target.clamp_(0, self.reg_max - 1 - 0.01)tl = target.long() # target lefttr = tl + 1 # target rightwl = tr - target # weight leftwr = 1 - wl # weight rightreturn (F.cross_entropy(pred_dist, tl.view(-1), reduction="none").view(tl.shape) * wl+ F.cross_entropy(pred_dist, tr.view(-1), reduction="none").view(tl.shape) * wr).mean(-1, keepdim=True)
还需要在BboxLoss 里调用 Alpha-IoU:
class BboxLoss(nn.Module):"""Criterion class for computing training losses for bounding boxes."""def __init__(self, reg_max: int = 16, alpha: float = 0.5):"""Initialize the BboxLoss module with regularization maximum and DFL settings."""super().__init__()self.dfl_loss = DFLoss(reg_max) if reg_max > 1 else Noneself.alpha_iou = AlphaIouLoss(alpha=alpha, reduction="none") # 用 Alpha-IoUdef forward(self,pred_dist: torch.Tensor,pred_bboxes: torch.Tensor,anchor_points: torch.Tensor,target_bboxes: torch.Tensor,target_scores: torch.Tensor,target_scores_sum: torch.Tensor,fg_mask: torch.Tensor,) -> Tuple[torch.Tensor, torch.Tensor]:"""Compute IoU (Alpha-IoU) and DFL losses for bounding boxes."""weight = target_scores.sum(-1)[fg_mask].unsqueeze(-1)# 用 Alpha-IoU 代替原来的 CIoUloss_iou = (self.alpha_iou(pred_bboxes[fg_mask], target_bboxes[fg_mask]) * weight).sum() / target_scores_sum# DFL lossif self.dfl_loss: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]) * weightloss_dfl = loss_dfl.sum() / target_scores_sumelse:loss_dfl = torch.tensor(0.0).to(pred_dist.device)return loss_iou, loss_dfl
整合数据集配置 +训练超参数配置 == 训练配置文件:
# COCO 2017 dataset http://cocodataset.org# download command/URL (optional) #download: bash ./scripts/get_coco.sh# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/] # 数据集配置文件 train: ./datasets/wzry700/train_list.txt # 118287 images val: ./datasets/wzry700/val_list.txt # 5000 images #test: ./coco/test-dev2017.txt # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794# number of classes nc: 2# class names names: [ 'locust', 'longicorn' ]# hyp.yaml # 训练超参数 lr0: 0.01 # 初始学习率 lrf: 0.1 # 最终学习率比例 momentum: 0.937 weight_decay: 0.0005 warmup_epochs: 3 warmup_momentum: 0.8 warmup_bias_lr: 0.1 # 损失函数相关 box: 0.05 # box 损失权重 cls: 0.3 # 分类损失权重 dfl: 1.0 # DFL (Distribution Focal Loss) 权重 iou_loss: alpha_iou alpha: 0.5 cls_loss: varifocal obj_gamma: 2.0 # 其他 fl_gamma: 2.0 hsv_h: 0.015 hsv_s: 0.7 hsv_v: 0.4 degrees: 2.0 translate: 0.1 scale: 0.5 shear: 0.0 perspective: 0.0 mosaic: 0.7 mixup: 0.15
