经典目标检测YOLO系列(一)复现YOLOV1(5)模型的训练及验证

之前，我们依据《YOLO目标检测》(ISBN:9787115627094)一书，提出了新的YOLOV1架构，继续按照此书进行YOLOV1的复现。
经典目标检测YOLO系列(一)YOLOV1的复现(1)总体架构

经典目标检测YOLO系列(一)复现YOLOV1(2)反解边界框及后处理

经典目标检测YOLO系列(一)复现YOLOV1(3)正样本的匹配及损失函数的实现

经典目标检测YOLO系列(一)复现YOLOV1(4)VOC2007数据集的读取及预处理

1 YOLOV1模型的训练

1.1 Yolov8Trainer类

作者为了代码具有更好的复用性，参考一些开源的工作单独实现了若干个训练所用的Trainer类。YOLOv1训练使用的是engine.py下的Yolov8Trainer类。初始化代码如下所示：
初始化代码中，关键的超参数配置都参考了YOLOv5和YOLOv8的项目代码：
- 优化器使用SGD，其中的momentum参数设置为0.937，weight decay设置为5e-4；
- 使用梯度累加的策略。当设定的batch size小于64时，会自动计算梯度累加的次数，比如batch size=16时，会累加4次梯度来近似batch size=64的效果；
- 初始学习率设置为0.01，并采用线性梯度衰减策略；
- 使用当前YOLO系列很流行的Model EMA(指数移动平均)策略。
  - 指数移动平均(exponential moving average)，也叫做权重移动平均(weighted moving average)，
    可以用来估计变量的局部均值，使得变量的更新与一段时间内的历史取值有关。
  - 在采用 SGD 或者其他的一些优化算法 (Adam, Momentum) 训练神经网络时，通常会使用EMA的方法。它的意义在于利用滑动平均的参数来提高模型在测试数据上的健壮性。
  - 简单来讲就是，通过在模型学习过程中融合更多的历史状态，从而达到更好的优化效果。
  - 【炼丹技巧】指数移动平均（EMA）的原理及PyTorch实现 - 知乎 (zhihu.com)

# RT-ODLab/engine.py# YOLOv8 Trainer
class Yolov8Trainer(object):def __init__(self, args, data_cfg, model_cfg, trans_cfg, device, model, criterion, world_size):# ------------------- basic parameters -------------------self.args = argsself.epoch = 0self.best_map = -1.self.device = deviceself.criterion = criterionself.world_size = world_sizeself.heavy_eval = Falseself.last_opt_step = 0self.clip_grad = 10# weak augmentatino stageself.second_stage = Falseself.third_stage = Falseself.second_stage_epoch = args.no_aug_epochself.third_stage_epoch = args.no_aug_epoch // 2# path to save modelself.path_to_save = os.path.join(args.save_folder, args.dataset, args.model)os.makedirs(self.path_to_save, exist_ok=True)# ---------------------------- Hyperparameters refer to YOLOv8 ----------------------------self.optimizer_dict = {'optimizer': 'sgd', 'momentum': 0.937, 'weight_decay': 5e-4, 'lr0': 0.01}self.ema_dict = {'ema_decay': 0.9999, 'ema_tau': 2000}self.lr_schedule_dict = {'scheduler': 'linear', 'lrf': 0.01}self.warmup_dict = {'warmup_momentum': 0.8, 'warmup_bias_lr': 0.1}        # ---------------------------- Build Dataset & Model & Trans. Config ----------------------------self.data_cfg = data_cfgself.model_cfg = model_cfgself.trans_cfg = trans_cfg# ---------------------------- Build Transform ----------------------------self.train_transform, self.trans_cfg = build_transform(args=args, trans_config=self.trans_cfg, max_stride=model_cfg['max_stride'], is_train=True)self.val_transform, _ = build_transform(args=args, trans_config=self.trans_cfg, max_stride=model_cfg['max_stride'], is_train=False)# ---------------------------- Build Dataset & Dataloader ----------------------------self.dataset, self.dataset_info = build_dataset(self.args, self.data_cfg, self.trans_cfg, self.train_transform, is_train=True)self.train_loader = build_dataloader(self.args, self.dataset, self.args.batch_size // self.world_size, CollateFunc())# ---------------------------- Build Evaluator ----------------------------self.evaluator = build_evluator(self.args, self.data_cfg, self.val_transform, self.device)# ---------------------------- Build Grad. Scaler ----------------------------self.scaler = torch.cuda.amp.GradScaler(enabled=self.args.fp16)# ---------------------------- Build Optimizer ----------------------------accumulate = max(1, round(64 / self.args.batch_size))print('Grad Accumulate: {}'.format(accumulate))self.optimizer_dict['weight_decay'] *= self.args.batch_size * accumulate / 64self.optimizer, self.start_epoch = build_yolo_optimizer(self.optimizer_dict, model, self.args.resume)# ---------------------------- Build LR Scheduler ----------------------------self.lr_scheduler, self.lf = build_lr_scheduler(self.lr_schedule_dict, self.optimizer, self.args.max_epoch)self.lr_scheduler.last_epoch = self.start_epoch - 1  # do not moveif self.args.resume and self.args.resume != 'None':self.lr_scheduler.step()# ---------------------------- Build Model-EMA ----------------------------if self.args.ema and distributed_utils.get_rank() in [-1, 0]:print('Build ModelEMA ...')self.model_ema = ModelEMA(self.ema_dict, model, self.start_epoch * len(self.train_loader))else:self.model_ema = None

1.1.1 Yolov8Trainer的train函数

训练分为三个阶段：

在第一阶段，我们采用设定好的数据增强手段来处理输出，并去训练模型，这一阶段会很长，几乎占据了训练的绝大部分的时间开销。
- 比如设定总的epoch=300，从0～280 epoch都是第一训练阶段；
在第二阶段，我们会关闭【马赛克增强】和【混合增强】，因为这两个数据增强会产生很多几乎不符合真实图像的数据，模型可能会从中学习到一些不合理的特征。
- 这一阶段会很短，比如只有10个epoch，仍以epoch=300的设定为例，从280～290epoch是第二阶段
- 这一策略是借鉴于YOLOX工作，不过不同于YOLOX的是，在第二阶段，我们虽然关闭了两个强大的增强，但仍旧【保留了较弱的空间扰动和颜色扰动的两个增强】；
在第三阶段，即从第二阶段结束开始，一直到整个训练结束
- 这个阶段也很短，仍以epoch=300的设定为例，从290～300 epoch是第三阶段
- 在这一阶段中，我们【关闭所有的数据增强】，只使用Resize和Normalize操作。
- Note:
  不过，就我们所训练的YOLOv1而言，我们不会用到马赛克增强和混合增强，所以第二阶段和第一阶段是一样的，只有到了第三阶段去关闭所有的数据增强。

# RT-ODLab/engine.pydef train(self, model):for epoch in range(self.start_epoch, self.args.max_epoch):if self.args.distributed:self.train_loader.batch_sampler.sampler.set_epoch(epoch)# check second stageif epoch >= (self.args.max_epoch - self.second_stage_epoch - 1) and not self.second_stage:self.check_second_stage()  # 关闭【马赛克增强】和【混合增强】# save model of the last mosaic epochweight_name = '{}_last_mosaic_epoch.pth'.format(self.args.model)checkpoint_path = os.path.join(self.path_to_save, weight_name)print('Saving state of the last Mosaic epoch-{}.'.format(self.epoch))torch.save({'model': model.state_dict(),'mAP': round(self.evaluator.map*100, 1),'optimizer': self.optimizer.state_dict(),'epoch': self.epoch,'args': self.args}, checkpoint_path)# check third stageif epoch >= (self.args.max_epoch - self.third_stage_epoch - 1) and not self.third_stage:self.check_third_stage()   # 【关闭所有的数据增强】# save model of the last mosaic epochweight_name = '{}_last_weak_augment_epoch.pth'.format(self.args.model)checkpoint_path = os.path.join(self.path_to_save, weight_name)print('Saving state of the last weak augment epoch-{}.'.format(self.epoch))torch.save({'model': model.state_dict(),'mAP': round(self.evaluator.map*100, 1),'optimizer': self.optimizer.state_dict(),'epoch': self.epoch,'args': self.args}, checkpoint_path)# train one epochself.epoch = epochself.train_one_epoch(model)# eval one epochif self.heavy_eval:model_eval = model.module if self.args.distributed else modelself.eval(model_eval)else:model_eval = model.module if self.args.distributed else modelif (epoch % self.args.eval_epoch) == 0 or (epoch == self.args.max_epoch - 1):self.eval(model_eval)if self.args.debug:print("For debug mode, we only train 1 epoch")break

1.1.2 Yolov8Trainer的train_one_epoch函数

 # RT-ODLab/engine.pydef train_one_epoch(self, model):# basic parametersepoch_size = len(self.train_loader)img_size = self.args.img_sizet0 = time.time()nw = epoch_size * self.args.wp_epochaccumulate = max(1, round(64 / self.args.batch_size))# train one epochfor iter_i, (images, targets) in enumerate(self.train_loader):ni = iter_i + self.epoch * epoch_size# Warmup"""在第一个epoch阶段，会有一段Warmup的过程，即缓慢的将所有的weight的学习率从0线性增加至设定的初始学习率0.01，而所有的bias的学习率则从0.1线性降低至设定的初始学习率0.01，并且，SGD优化器的中的momentum参数也会从默认的warmup_momentum线性增加至设定的0.937。这一warmup策略是遵循了YOLOv5的设定，默认会持续3个epoch。"""if ni <= nw:xi = [0, nw]  # x interpaccumulate = max(1, np.interp(ni, xi, [1, 64 / self.args.batch_size]).round())for j, x in enumerate(self.optimizer.param_groups):# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0x['lr'] = np.interp(ni, xi, [self.warmup_dict['warmup_bias_lr'] if j == 0 else 0.0, x['initial_lr'] * self.lf(self.epoch)])if 'momentum' in x:x['momentum'] = np.interp(ni, xi, [self.warmup_dict['warmup_momentum'], self.optimizer_dict['momentum']])# to device# 随后，我们将读取进来的数据放置制定的设备上，比如GPU，并对输入的图像做简单的归一化（除以255）# 注意，这种归一化等价于设定均值和方差为别为0和255。# 之所以这么做是因为均值和方差往往依赖于任务场景的数据，通常不具备泛化性，所以我们不采用统计均值和方差的方法。images = images.to(self.device, non_blocking=True).float() / 255.# Multi scale# 完成这一步后，我们就可以调用类内的rescale_image_targets函数对这批图像做多尺度处理的操作# 即随机将这批图像Resize成其他的尺寸，比如都从640x640调整成512x512# 最后用类内的refine_targets 函数滤出掉无效的目标if self.args.multi_scale:images, targets, img_size = self.rescale_image_targets(images, targets, self.model_cfg['stride'], self.args.min_box_size, self.model_cfg['multi_scale'])else:targets = self.refine_targets(targets, self.args.min_box_size)# visualize train targetsif self.args.vis_tgt:vis_data(images*255, targets)# 然后，我们就可以将这片数据去交给模型做训练阶段的推理（没有后处理阶段），并做标签分配和损失计算，# 当迭代次数达到了梯度累加指定的次数后，就可以反向传播梯度和更新模型的参数。# inferencewith torch.cuda.amp.autocast(enabled=self.args.fp16):# 调用模型获取训练阶段的推理结果outputs = model(images)# loss 标签分配及计算损失loss_dict = self.criterion(outputs=outputs, targets=targets, epoch=self.epoch)losses = loss_dict['losses']losses *= images.shape[0]  # loss * bs# reduce            loss_dict_reduced = distributed_utils.reduce_dict(loss_dict)# gradient averaged between devices in DDP modelosses *= distributed_utils.get_world_size()# backwardself.scaler.scale(losses).backward()# Optimizeif ni - self.last_opt_step >= accumulate:if self.clip_grad > 0:# unscale gradientsself.scaler.unscale_(self.optimizer)# clip gradientstorch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=self.clip_grad)# optimizer.stepself.scaler.step(self.optimizer)self.scaler.update()self.optimizer.zero_grad()# emaif self.model_ema is not None:self.model_ema.update(model)self.last_opt_step = ni# displayif distributed_utils.is_main_process() and iter_i % 10 == 0:t1 = time.time()cur_lr = [param_group['lr']  for param_group in self.optimizer.param_groups]# basic inforlog =  '[Epoch: {}/{}]'.format(self.epoch, self.args.max_epoch)log += '[Iter: {}/{}]'.format(iter_i, epoch_size)log += '[lr: {:.6f}]'.format(cur_lr[2])# loss inforfor k in loss_dict_reduced.keys():log += '[{}: {:.2f}]'.format(k, loss_dict_reduced[k])# other inforlog += '[time: {:.2f}]'.format(t1 - t0)log += '[size: {}]'.format(img_size)# print log inforprint(log, flush=True)t0 = time.time()if self.args.debug:print("For debug mode, we only train 1 iteration")breakself.lr_scheduler.step()

1.1.3 开始训练

训练的主程序代码在train.py文件中，进入项目文件，可以执行下面的脚本进行训练：

python train.py --cuda -d voc \-m yolov1  \-bs 16\-size 416  \--wp_epoch 3  \--max_epoch 150  \--eval_epoch 10 \--no_aug_epoch 10  \--ema  \--fp16  \--multi_scale

命令行参数的意义：

--cuda表示我们调用GPU来训练模型，切记，一定要输入“–cuda”,否则程序只调用CPU来训练，那将是个极其漫长的训练过程；
-d是数据集类别，比如voc表示使用VOC数据集，coco表示使用COCO数据集；
--root是存放数据集的路径；
-bs 16表示训练期间的batch size为16；读者可以根据自己的GPU显存来修改，如-bs 8或者-bs 32等；
--size表示输入图像的最长边的尺寸；
--wp_epoch是warmup阶段的时长；
--max_epoch是训练的总epoch数量；
--eval_epoch是在训练阶段中去测试模型的性能的节点；
--no_aug_epoch表示最后的多少epoch会关闭数据增强；
--ema表示启动Model EMA策略；
--fp16表示开启混合精度策略；
--multi_scale表示开启多尺度训练策略；

二、模型的验证

2.1 mAP指标计算

训练结束后，模型默认保存在weights/voc/yolov1/文件夹下，名为yolov1_voc_best.pth，保存了训练阶段在测试集上mAP指标最高的模型。
运行项目中所提供的eval.py文件可以验证模型的性能，具体命令如下行所示

可以给定不同的图像尺寸来测试实现的YOLOv1在不同输入尺寸下的性能

python eval.py \
--cuda -d voc \
--root path/to/voc -m yolov1 \
--weight path/to/yolov1_voc_best.pth \
-size 416

也可以可视化训练好的模型

python test.py \
--cuda -d voc \
--root path/to/voc -m yolov1 
--weight path/to/yolov1_voc_best.pth \
-size 416 -vt 0.3 \
--show# -size表示输入图像的最大边尺寸
# -vt是可视化的置信度阈值，只有高于此值的才会被可视化出来
# --show表示展示检测结果的可视化图片