YOLOv5/v7 中的类别不平衡问题解决方案研究

目录

• 类别权重
• 图像权重
• 参考：

 

---

本文定义的类别不平衡问题：在图像分类问题中，不同类别的图像的个数不同，并且差异较大。在目标检测问题中，目标类别的个数区别较大。这都是类别不平衡问题。

原理：

当训练图像的所有 类个数不相同 时, 我们可以更改 类权重 , 即而达到更改 图像权重 的目的。然后根据 图像权重 重新采集数据，这在图像类别不均衡的数据下尤其重要。

使用 YOLOv5/v7 训练自己的数据集时，各类别的标签数量难免存在不平衡的问题，在训练过程中为了就减小类别不平衡问题的影响，YOLOv5/v7 中引入了 类别权重 和 图像权重 的设置。

回到顶部

类别权重

类别权重用于根据每个类别的统计真实框的个数，取倒数，然后做个规范化(每个值除以它们的和)，最后乘以类别数目放大下。

# 获取每个类别的权重，比如COCO 是80类，那么返回的是 shape: torch.Size([80]), 该类别的真实框越多，对应位置上的值越小。该类别的真实框越少，对应位置的值越大。
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc  # attach class weights # 注意：最后乘了 nc=80, 
 
def labels_to_class_weights(labels, nc=80):
	# Get class weights (inverse frequency) from training labels
	if labels[0] is None:  # no labels loaded
		return torch.Tensor()
 
	labels = np.concatenate(labels, 0)  # labels.shape = (866643, 5) for COCO
	classes = labels[:, 0].astype(np.int)  # labels = [class xywh]
	weights = np.bincount(classes, minlength=nc)  # occurrences per class
 
	# Prepend gridpoint count (for uCE training)
	# gpi = ((320 / 32 * np.array([1, 2, 4])) ** 2 * 3).sum()  # gridpoints per image
	# weights = np.hstack([gpi * len(labels)  - weights.sum() * 9, weights * 9]) ** 0.5  # prepend gridpoints to start
 
	weights[weights == 0] = 1  # replace empty bins with 1
	weights = 1 / weights  # number of targets per class
	weights /= weights.sum()  # normalize
	return torch.from_numpy(weights)

回到顶部

图像权重

在 YOLOv5/v7 train.py 中： 默认是 False

parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training') # 如果需要开启，需要在模型训练时候命令行加上  --image-weights

每个 epoch 选择训练的图片是不一样的，每个图片都有个概率，

调用代码：

    # Update image weights (optional)
    if opt.image_weights:
        # Generate indices
        if rank in [-1, 0]:
            cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc  # class weights # maps为80个0，所以这里是求平方再除以 80
            iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw)  # image weights
	     # 利用random.choices根据图像的权重选取图片，每次选取和训练集相同的图片数。# 也就意味着训练集中的图片权重大的有的可能会被选取多次，权重小的可能一次都不会被选中用于训练
            dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n)  # rand weighted idx
        # Broadcast if DDP
        if rank != -1:
            indices = (torch.tensor(dataset.indices) if rank == 0 else torch.zeros(dataset.n)).int()
            dist.broadcast(indices, 0)
            if rank != 0:
                dataset.indices = indices.cpu().numpy()

标签到图像权重：

def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)):
	# Produces image weights based on class_weights and image contents
	class_counts = np.array([np.bincount(x[:, 0].astype(np.int), minlength=nc) for x in labels]) # 统计每张图片中的类别个数， 如 array([0, 0, 0, 0, 1, 0, 2, 0, 1]) 2 表示该图片内有该类别的目标两个，COCO 数据集的话长度为 80
	image_weights = (class_weights.reshape(1, nc) * class_counts).sum(1) # 传入的 class_weights 乘以 类别权重，得到每张图像的权重，也就是共 图像个数 个值。 sum(1) 的意思是 按该图片内的所有某个类别的目标乘以类别权重，再求和，就得到每张图像的权重是多少，权重越大，越可能被采样到。
	# index = random.choices(range(n), weights=image_weights, k=1)  # weight image sample 
	return image_weights

回到顶部

参考：

https://www.cnblogs.com/xiaoheizi-12345/p/14458802.html

https://blog.csdn.net/Graceying/article/details/120288985