0. 前言
虽然v8支持直接训练“无标签的负样本”(v8中叫做background),但默认方式也只是混进正样本中一起练,这会导致新的误检图得不到充分训练
使用本文的“负样本马赛克”功能,可以保证每一轮、每一张图里全都有误检图。能够在较少的训练轮次中,解决漏检问题
使用例:
比如说你发现你模型把摩托车识别成人,那就把那几张摩托车的图放进neg_dir,练个30轮
出来就不会有摩托车误识别的问题了
【但要注意,放进neg_dir里的图 绝 对 不 能 有正样本(例如把摩托车误检成人的图里不能出现真的人)。不然模型会学到错误的负样本,反而会出大问题】
1. 加入新超参
在ultralytics/cfg/default.yaml(默认超参文件)中加入新超参"neg_dir"。
v8的结构是直接读取所有超参数,解析交给后面的函数
neg_dir: ' ' # 负样本文件夹(str),其中只有负样本图片,没有标签(默认为空)
2. 读取新超参
超参在ultralytics/engine/Trainer.py中的BaseTrainer()初始化,
在ultralytics/engine/Trainer.py的108行左右,即BaseTrainer的__init__()方法中,加入下述内容。
目的是读取args里的neg_dir(108行左右改为下述内容),之后超参数会随着trainer的属性hyp一层一层的传下去
# 负样本文件夹——初始化
self.neg_dir = self.args.neg_dir
3. 识别超参
在ultralytics/data/augment.py里的v8_transforms()里的Mosaic()。
增加一个输入neg_dir = hyp.neg_dir。
def v8_transforms(dataset, imgsz, hyp, stretch=False):
"""Convert images to a size suitable for YOLOv8 training."""
pre_transform = Compose([
Mosaic(dataset, imgsz=imgsz, p=hyp.mosaic, neg_dir=hyp.neg_dir, neg_num = hyp.neg_num),
CopyPaste(p=hyp.copy_paste),
RandomPerspective(
degrees=hyp.degrees,
translate=hyp.translate,
scale=hyp.scale,
shear=hyp.shear,
perspective=hyp.perspective,
pre_transform=None if stretch else LetterBox(new_shape=(imgsz, imgsz)),
)])
再到Mosaic()类中的__init__中,添加识别此参数的代码
def __init__(self, dataset, imgsz=640, p=1.0, n=4, neg_dir='', neg_num=''):
"""Initializes the object with a dataset, image size, probability, and border."""
assert 0 <= p <= 1.0, f'The probability should be in range [0, 1], but got {p}.'
assert n in (4, 9), 'grid must be equal to 4 or 9.'
super().__init__(dataset=dataset, p=p)
self.dataset = dataset
self.imgsz = imgsz
self.border = (-imgsz // 2, -imgsz // 2) # width, height
self.n = n
# # 更改此处, 读取负样本加入数(默认为2)
self.neg_num = int(neg_num)
self.img_neg_files = [] # 负样本路径列表
# “默认只有负样文件夹才有无标签负样本”
# 读取负样本文件夹
# additional feature
if os.path.isdir(neg_dir):
# 负样本路径
self.img_neg_files = [os.path.join(neg_dir, i) for i in os.listdir(neg_dir)]
logging.info(
colorstr("Negative dir: ")
+ f"'{neg_dir}', using {len(self.img_neg_files)} picture