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

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

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

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 = 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

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.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

1.1.2 Yolov8Trainer的train_one_epoch函数

 # 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()

1.1.3 开始训练

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

python train.py \
--cuda -d voc \
--root path/to/voc -m yolov1 -bs 16 -size 640 \
--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表示开启多尺度训练策略；

训练的40轮结果如下：

.....
[Epoch: 40/150][Iter: 250/280][lr: 0.007426][loss_obj: 2.09][loss_cls: 0.70][loss_box: 0.35][losses: 146.06][time: 39.83][size: 384]
[Epoch: 40/150][Iter: 260/280][lr: 0.007426][loss_obj: 3.03][loss_cls: 0.72][loss_box: 0.26][losses: 161.56][time: 39.80][size: 672]
[Epoch: 40/150][Iter: 270/280][lr: 0.007426][loss_obj: 2.54][loss_cls: 0.69][loss_box: 0.28][losses: 147.48][time: 33.97][size: 704]
eval ...
im_detect: 1/997 0.006s
im_detect: 501/997 0.008s
Evaluating detections
VOC07 metric? Yes
AP for aeroplane = 0.5796
AP for bicycle = 0.7892
AP for bird = 0.6054
AP for boat = 0.5226
AP for bottle = 0.3170
AP for bus = 0.7942
AP for car = 0.7585
AP for cat = 0.6780
AP for chair = 0.4601
AP for cow = 0.6875
AP for diningtable = 0.6312
AP for dog = 0.7017
AP for horse = 0.7940
AP for motorbike = 0.6346
AP for person = 0.7010
AP for pottedplant = 0.4667
AP for sheep = 0.6005
AP for sofa = 0.5468
AP for train = 0.7788
AP for tvmonitor = 0.5637
Mean AP = 0.6306
Mean AP:  0.630562412972362

二、模型的验证

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

  作者训练好的模型，在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