【YOLO】yolov5的训练策略3 -- 图像加权image weights

原创

已于 2023-07-14 16:02:42 修改 · 2.9k 阅读

标签

#YOLO

收录于

于 2023-07-14 15:57:33 首次发布

文章介绍了如何通过训练图像加权来解决样本分布不均的问题。首先计算类别权重，根据每个类别的目标框数的倒数进行赋值，然后利用这些权重筛选训练数据。接着，计算图片权重，结合类别权重和图片中类别数量。最后，在每个训练epoch中更新数据集的图片索引，以增加高权重图片的采样概率，从而达到样本均衡的效果。

一、什么是训练图像加权

根据样本种类分布使用图像调用频率不同的方法解决。
1、读取训练样本中的GT，保存为一个列表；
2、计算训练样本列表中不同类别个数，然后给每个类别按相应目标框数的倒数赋值，数目越多的种类权重越小，形成按种类的分布直方图；
3、对于训练数据列表，训练时按照类别权重筛选出每类的图像作为训练数据。使用random.choice(population, weights=None, *, cum_weights=None, k=1)更改训练图像索引，可达到样本均衡的效果。

二、代码

计算类别的权重：

def labels_to_class_weights(labels, nc=80):
    """
    计算类别权重
    Get class weights (inverse frequency) from training labels
    输出入: 
        labels -- 真实的标签列表 [class xywh]
        nc -- 类别数
    """
    if labels[0] is None:  # no labels loaded
        return torch.Tensor()

    labels = np.concatenate(labels, 0)  # labels.shape = (866643, 5) for COCO
    classes = labels[:,<