之前,我们依据《YOLO目标检测》(ISBN:9787115627094)
一书,提出了新的YOLOV1架构,继续按照此书进行YOLOV1的复现。
经典目标检测YOLO系列(一)YOLOV1的复现(1)总体架构
经典目标检测YOLO系列(一)复现YOLOV1(2)反解边界框及后处理
经典目标检测YOLO系列(一)复现YOLOV1(3)正样本的匹配及损失函数的实现
经典目标检测YOLO系列(一)复现YOLOV1(4)VOC2007数据集的读取及预处理
# 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 = args
self.epoch = 0
self.best_map = -1.
self.device = device
self.criterion = criterion
self.world_size = world_size
self.heavy_eval = False
self.last_opt_step = 0
self.clip_grad = 10
# weak augmentatino stage
self.second_stage = False
self.third_stage = False
self.second_stage_epoch = args.no_aug_epoch
self.third_stage_epoch = args.no_aug_epoch // 2
# path to save model
self.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_cfg
self.model_cfg = model_cfg
self.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 / 64
self.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 move
if 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
训练分为三个阶段:
关闭【马赛克增强】和【混合增强】
,因为这两个数据增强会产生很多几乎不符合真实图像的数据,模型可能会从中学习到一些不合理的特征。
# RT-ODLab/engine.py
def 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 stage
if 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 epoch
weight_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 stage
if 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 epoch
weight_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 epoch
self.epoch = epoch
self.train_one_epoch(model)
# eval one epoch
if self.heavy_eval:
model_eval = model.module if self.args.distributed else model
self.eval(model_eval)
else:
model_eval = model.module if self.args.distributed else model
if (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
# RT-ODLab/engine.py
def train_one_epoch(self, model):
# basic parameters
epoch_size = len(self.train_loader)
img_size = self.args.img_size
t0 = time.time()
nw = epoch_size * self.args.wp_epoch
accumulate = max(1, round(64 / self.args.batch_size))
# train one epoch
for 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 interp
accumulate = 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 lr0
x['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 targets
if self.args.vis_tgt:
vis_data(images*255, targets)
# 然后,我们就可以将这片数据去交给模型做训练阶段的推理(没有后处理阶段),并做标签分配和损失计算,
# 当迭代次数达到了梯度累加指定的次数后,就可以反向传播梯度和更新模型的参数。
# inference
with 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 mode
losses *= distributed_utils.get_world_size()
# backward
self.scaler.scale(losses).backward()
# Optimize
if ni - self.last_opt_step >= accumulate:
if self.clip_grad > 0:
# unscale gradients
self.scaler.unscale_(self.optimizer)
# clip gradients
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=self.clip_grad)
# optimizer.step
self.scaler.step(self.optimizer)
self.scaler.update()
self.optimizer.zero_grad()
# ema
if self.model_ema is not None:
self.model_ema.update(model)
self.last_opt_step = ni
# display
if 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 infor
log = '[Epoch: {}/{}]'.format(self.epoch, self.args.max_epoch)
log += '[Iter: {}/{}]'.format(iter_i, epoch_size)
log += '[lr: {:.6f}]'.format(cur_lr[2])
# loss infor
for k in loss_dict_reduced.keys():
log += '[{}: {:.2f}]'.format(k, loss_dict_reduced[k])
# other infor
log += '[time: {:.2f}]'.format(t1 - t0)
log += '[size: {}]'.format(img_size)
# print log infor
print(log, flush=True)
t0 = time.time()
if self.args.debug:
print("For debug mode, we only train 1 iteration")
break
self.lr_scheduler.step()
训练的主程序代码在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
查看训练过程中的日志。可以利用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 |
|=============================================================================|
+-----------------------------------------------------------------------------+
第一部分参数解释
第二部分参数解释
1、“GPU 型号及序号”
2、“Temp” 温度,这个很好理解就是目前显卡的温度是67 ℃
3、“Perf” 表示显卡的性能状态,有 p0 ~ p12,性能状态从高到底
4、“Persistence-M” 表示持续模式状态,这个参数默认是关闭的,若打开能够让 GPU 更加快速地响应任务,但相应待机功耗也会增加;
5、“Pwr: Usage/Cap” 表示显卡的功耗,可以看到我这里 2080TI 的显卡功耗是 250W,而前面的 212W 是根据执行的任务动态变化的,是卡上所有执行进程所占显存的和;
6、“Bus-Id” 表示 GPU 的总线地址;
7、“Disp.A” 全称是 “Display Active” 表示显卡的显示功能是否初始化,这其实也比较好理解。
8、Memory-Usage
表示显存占用情况
9、Volatile GPU-UTil
表示 GPU 的使用率。
GPU使用率 和 显存使用率
】 是两个比较重要的指标,一般来说,GPU 使用率高的时候显存使用率也会比较高10、“ECC” 全称为 Error Checking and Correcting,表示是否开启错误检查和纠错技术,“0” 代表关闭,“1” 代表开启。
11、“Compute M” 表示计算模式,NVIDIA GPU 有 Default、Exclusive Process、Exclusive Thread、Prohibited 四种计算模式,
12、最后一个是 “MIG M”,这是安培架构新提出的功能,可以用于创建 MIG 实例,在用户层实现 GPU 的物理隔离。
第三部分参数解释
正在执行的 CUDA 程序的一些状态信息,特别要关注的是 “PID” 和 “GPU Memory Usage” 这两个参数,其他就不多说了
我这里在AutoDL的云GPU上进行训练的,平台上也监控了GPU相关性能信息:
如果自己电脑没有显卡,可以去colab白嫖,也可以租一个GPU,例如AutoDL平台,可以参考下面博客租借
训练结束后,模型默认保存在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表示展示检测结果的可视化图片
作者训练好的模型,在VOC2007测试集测试指标如下:
从表格中可以看到,在同样的输入尺寸448×448下,改进后的 YOLOv1 实现了更高的性能:73.2% 的 mAP,超过了官方的 YOLOv1。
模型 | 输入尺寸 | mAP(%) |
---|---|---|
YOLOv1*(官方) | 448 | 63.4 |
YOLOv1 | 416 | 71.9 |
YOLOv1 | 448 | 73.2 |
YOLOv1 | 512 | 74.9 |
YOLOv1 | 640 | 76.7 |