YOLOV6之loss代码

最新推荐文章于 2024-10-15 16:32:50 发布

原创

最新推荐文章于 2024-10-15 16:32:50 发布 · 1.9k 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#目标检测

本文详细介绍了YoloV6中损失函数的定义与计算流程。包括损失函数的初始化参数设置，训练过程中损失函数的调用方式，以及损失的具体计算过程。

一、loss定义

位于engine.py文件的train_before_loop()函数下：

from yolov6.models.loss import ComputeLoss

self.compute_loss = ComputeLoss(num_classes=self.data_dict['nc'],
                                ori_img_size=self.img_size,
                                use_dfl=self.cfg.model.head.use_dfl,
                                reg_max=self.cfg.model.head.reg_max,
                                iou_type=self.cfg.model.head.iou_type)

二、loss调用

preds, s_featmaps = self.model(images)
# print(preds[0][i].shape)  # [[bs,64,80,80] [bs,128,40,40] [bs,256,20,20]] --- feats
# print(preds[1].shape)     # [bs,8400,nc] --- pred_scores
# print(preds[2].shape)     # [bs,8400,4] --- pred_distri
# print(targets.shape)      # [n,6]
total_loss, loss_items = self.compute_loss(preds, targets, epoch_num)