经典目标检测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文件中，进入项目文件，可以执行下面的脚本进行训练：

nohup python -u train.py        \--cuda -d voc      \-m yolov1          \-bs 32             \-size 640          \--wp_epoch 3       \--max_epoch 150    \--eval_epoch 10    \--no_aug_epoch 10  \--ema              \--fp16             \--multi_scale      \--num_workers 8 1>./logs/yolo_v1_train_log.txt 2>./logs/yolo_v1_warning_log.txt &

命令行参数的意义：

--cuda表示我们调用GPU来训练模型，切记，一定要输入“–cuda”,否则程序只调用CPU来训练，那将是个极其漫长的训练过程；
-d是数据集类别，比如voc表示使用VOC数据集，coco表示使用COCO数据集；
--root是存放数据集的路径；
-bs 32表示训练期间的batch size为32；读者可以根据自己的GPU显存来修改，如-bs 8或者-bs 16等；
--size表示输入图像的最长边的尺寸；
--wp_epoch是warmup阶段的时长；
--max_epoch是训练的总epoch数量；
--eval_epoch是在训练阶段中去测试模型的性能的节点；
--no_aug_epoch表示最后的多少epoch会关闭数据增强；
--ema表示启动Model EMA策略；
--fp16表示开启混合精度策略；
--multi_scale表示开启多尺度训练策略；
加上nohup让程序可以后台运行，同时保存训练日志到./logs/yolo_v1_train_log.txt文件，报错日志到./logs/yolo_v1_warning_log.txt文件中，我们可以通过tail -f ./logs/yolo_v1_train_log.txt查看训练过程中的日志。
用2080Ti 训练YOLOv1花费时间1h~2h

可以利用nvidia-smi查看GPU利用率等相关信息:

(base) root@autodl-container-adbc11ae52-f2ebff02:~/data/VOCdevkit/VOC2007/ImageSets/Main# nvidia-smi 
Sat Jan 13 08:51:05 2024       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.89.02    Driver Version: 525.89.02    CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  On   | 00000000:B4:00.0 Off |                  N/A |
| 39%   67C    P2   212W / 250W |   9230MiB / 11264MiB |     93%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------++-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

第一部分参数解释
- 都是版本号，有 NVIDIA-SMI 版本号、CUDA 驱动版本号和 CUDA 版本号。
- 这里需要注意的是，这里的 CUDA 版本号并不一定是你自己安装的 CUDA 的版本号，而是这个 CUDA 驱动能够支持的最高的 CUDA 版本。如果要查看你自己安装的 CUDA 版本号，可以使用 nvcc -V 进行查看。
第二部分参数解释
- 1、“GPU 型号及序号”
  - 比如上图中我这里就是只有一块卡，所以序号是 “0”，然后型号是英伟达 GeForce；
  - “Fan” 风扇表示你的 GPU 是主动散热还是被动散热的，比如像 RTX 那些消费级显卡一般都是主动散热的，从物理形态上来看就是带风扇的；A10 是推理计算卡，是被动散热的，所以风扇一直其实都会是 0% (相当于 N/A)；
- 2、“Temp” 温度，这个很好理解就是目前显卡的温度是67 ℃
  - “温度” 这个指标是需要在执行程序的时候经常关注的，特别是对于被动散热的卡，要是机器的散热做不好，就容易使卡达到降频保护的温度，从而卡的性能会下降，跑出来的性能就不准确了；
- 3、“Perf” 表示显卡的性能状态，有 p0 ~ p12，性能状态从高到底
  - 也就是 p0 表示性能最大，p12 表示性能最小
  - 在 linux 下可以使用 nvidia-settings 工具对这个参数进行配置；
- 4、“Persistence-M” 表示持续模式状态，这个参数默认是关闭的，若打开能够让 GPU 更加快速地响应任务，但相应待机功耗也会增加；
- 5、“Pwr: Usage/Cap” 表示显卡的功耗，可以看到我这里 2080TI 的显卡功耗是 250W，而前面的 212W 是根据执行的任务动态变化的，是卡上所有执行进程所占显存的和；
- 6、“Bus-Id” 表示 GPU 的总线地址；
- 7、“Disp.A” 全称是 “Display Active” 表示显卡的显示功能是否初始化，这其实也比较好理解。
  - 如果当前的 GPU 正驱动一个活动的显示器，那么这个参数就会是 On，反之为 Off。
  - 我这里的是不带显示功能的，所以这个参数就是 Off 了；
- 8、Memory-Usage 表示显存占用情况
  - 可以看到我这里的显存是 11264 MiB，前面的 9230 MiB 也是随着执行的程序动态变化的
  - 但是不能超过 11264 MiB，不然就会报 CUDA 显存不足的错误；
- 9、Volatile GPU-UTil 表示 GPU 的使用率。
  - 【GPU使用率和显存使用率】是两个比较重要的指标，一般来说，GPU 使用率高的时候显存使用率也会比较高
  - 但反过来，显存使用率高的时候 GPU 使用率并不一定会高。
    比如在大量数据传输的时候，显存是被占用的，但是 GPU 本身还没开始计算，这个时候的显存占用是高的，但 GPU 使用率并不一定高，这种情况会经常出现在访存密集型的计算中；
- 10、“ECC” 全称为 Error Checking and Correcting，表示是否开启错误检查和纠错技术，“0” 代表关闭，“1” 代表开启。
  - ECC 开启后功能可以用于保护 GPU 硬件，帮助其免受硬件故障和问题的影响。因为它还可以检查内存模块中的错误，并帮助 GPU 从错误中恢复。
  - 此外，ECC 还有助于降低电力消耗，因为它只需要较少的功耗就可以处理大量的内存纠错；
- 11、“Compute M” 表示计算模式，NVIDIA GPU 有 Default、Exclusive Process、Exclusive Thread、Prohibited 四种计算模式，
  - 一般就用默认的 Default 就可以了。
  - 在 Default 模式下，多个进程可以同时共享不同的 GPU；
- 12、最后一个是 “MIG M”，这是安培架构新提出的功能，可以用于创建 MIG 实例，在用户层实现 GPU 的物理隔离。
  - 一看我这 “MIG M” 是 “N/A”，这说明在 2080TI 上是不支持这个功能的呀。
第三部分参数解释
正在执行的 CUDA 程序的一些状态信息，特别要关注的是 “PID” 和 “GPU Memory Usage” 这两个参数，其他就不多说了

我这里在AutoDL的云GPU上进行训练的，平台上也监控了GPU相关性能信息:
在这里插入图片描述
如果自己电脑没有显卡，可以去colab白嫖，也可以租一个GPU，例如AutoDL平台，可以参考下面博客租借

AutoDL平台租借GPU详解

二、模型的验证

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表示展示检测结果的可视化图片