解读YOLO v7的代码(三)损失函数

原创

已于 2023-04-06 16:45:42 修改 · 1.3w 阅读

文章标签：

#深度学习 #人工智能

于 2023-02-09 08:58:12 首次发布

该文详细解析了YOLOv7训练过程中损失函数的计算，重点介绍了ComputeLossOTA类的__call__函数，包括如何从原始目标数据生成正样本，计算预测值与目标的IOU，以及正样本的选择策略。通过这些步骤，确定用于损失计算的正样本，并计算IOU损失、对象存在性损失和分类损失，最终组合成整体的损失值。

在前两篇博客中我分析了YOLO v7的模型结构以及训练数据的准备。这里将对损失函数的代码进行分析。

在train.py中，我们可以看到以下的代码是进行损失值计算的

                if 'loss_ota' not in hyp or hyp['loss_ota'] == 1:
                    loss, loss_items = compute_loss_ota(pred, targets.to(device), imgs)  # loss scaled by batch_size
                    break
                else:
                    loss, loss_items = compute_loss(pred, targets.to(device))  # loss scaled by batch_size

因为在超参数中定义了loss_ota，因此是用compute_loss_ota来进行计算。这个对应的是loss.py里面定义的ComputeLossOTA的类。

查看这个类里面的__call__函数，其输入变量是模型输出的预测值p, 训练图像包含的targets, 以及训练图像imgs。

例如以下是输入的一幅训练图片：

图片包含的targets如下：

tensor([[0.0000e+00, 4.5000e+01, 5.1015e-01, 8.6860e-02, 9.6169e-01, 1.7279e-01],
        [0.0000e+00, 5.0000e+01, 8.6156e-01, 8.8347e-02, 2.5677e-01, 1.7359e-01],
        [0.0000e+00, 1.5000e+01, 9.0563e-01, 8.2660e-01, 1.8675e-01, 3.4358e-01],
        [0.0000e+00, 3.9000e+01, 7.5925e-02, 1.2316e-01, 1.5037e-01, 2.3427e-01],
        [0.0000e+00, 8.0000e+00, 5.1819e-01, 1.0259e-01, 4.6240e-01, 2.0505e-01],
        [0.0000e+00, 7.0000e+00, 7.8777e-01, 4.4500e-02, 2.7215e-02, 2.3584e-02],
        [0.0000e+00, 6.0000e+00, 5.3230e-01, 4.1727e-01, 5.3604e-01, 3.4526e-01],
        [0.0000e+00, 1.6000e+01, 5.0707e-01, 3.8664e-01, 3.7291e-02, 3.2220e-02]],
       device='cuda:0')

这个targets数组共包括了8个target，每个target对应一个6维的变量，第一维代表Batch的第几个图片，第2维代表物体的类别，后面4维对应中心点的xy以及wh，这里xywh是表示target的大小与图片的大小比例。

模型输出的预测值P是一个列表，包含大中小三个维度的预测结果，其维度分别为(batch_size, 3, 80, 80, 85), (batch_size, 3, 40, 40, 85), (batch_size, 3, 20, 20, 85)。这里面3表示有3个anchor，80*80代表最细粒度的检测，把图片分为80*80个格子，85表示每个格子的一个预测变量，其中包括了预测物体存在概率，物体的xywh以及80个分类的概率，总共是85维。

我们把targets数组的物体在图像中标识出来看一下：