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