【目标检测】YOLOv5算法实现(七)：模型训练

??本系列文章记录本人硕士阶段YOLO系列目标检测算法自学及其代码实现的过程。其中算法具体实现借鉴于ultralytics YOLO源码Github，删减了源码中部分内容，满足个人科研需求。
??本系列文章主要以YOLOv5为例完成算法的实现，后续修改、增加相关模块即可实现其他版本的YOLO算法。

文章地址：
YOLOv5算法实现(一)：算法框架概述
YOLOv5算法实现(二)：模型加载
YOLOv5算法实现(三)：数据集加载
YOLOv5算法实现(四)：损失计算
YOLOv5算法实现(五)：预测结果后处理
YOLOv5算法实现(六)：评价指标及实现
YOLOv5算法实现(七)：模型训练
YOLOv5算法实现(八)：模型验证
YOLOv5算法实现(九)：模型预测(编辑中…)

1 引言

??本篇文章综合之前文章中的功能，实现模型的训练。模型训练的逻辑如图1所示。

2 超参数文件

??YOLOv5中超参数主要包括学习率、优化器、置信度以及数据增强，源码中某一超参数文件及各参数含义如下所示：

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
# Hyperparameters for VOC training

# 学习率
lr0: 0.00334  # 初始学习率
lrf: 0.15135  # 最终学习率下降比例(lr0 * lrf)
# 优化器(SGD、Adam、AdamW)
momentum: 0.74832 # SGD momentum/Adam beta1
weight_decay: 0.00025 # optimizer weight decay 5e-4 ,权重衰变系数(防止过拟合)
# 热身训练(Warmup)
warmup_epochs: 3.3835 # 学习率热身epoch
warmup_momentum: 0.59462 # 学习率热身初始动量
warmup_bias_lr: 0.18657 # 学习率热身偏置学习率
# 损失增益
box: 0.02 # box loss gain
cls: 0.21638 # cls loss gain
cls_pw: 0.5 # cls BCELoss positive_weight
obj: 0.51728 # obj loss gain (scale with pixels)
obj_pw: 0.67198 # obj BCELoss positive_weight
fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5)
# 正样本匹配阈值
iou_t: 0.2 # IoU training threshold
anchor_t: 3.3744 # anchor-multiple threshold
# 数据增强
## HSV色彩空间增强
hsv_h: 0.01041
hsv_s: 0.54703
hsv_v: 0.27739
## 仿射变换
degrees: 0.0 #图像旋转
translate: 0.04591 #图像平移
scale: 0.75544 #图像仿射变换的缩放比例
shear: 0.0 #设置裁剪的仿射矩阵系数

3 模型训练(train.py)

def train(hyp):
    # ----------------------------------------------------------------------------------
    # hyp超参数解析
    # ----------------------------------------------------------------------------------
    # 训练设备
    device = torch.device(opt.device if torch.cuda.is_available() else "cpu")
    print("Using {} device training.".format(device.type))

    # 权重保存文件
    wdir = "weights" + os.sep  # weights dir

    # 损失和学习率记录txt, 损失和学习率变化曲线, 权重文件保存路径
    results_file, save_file_path, save_path \
        = opt.results_file, opt.save_file_path, opt.save_path

    # 模型结构文件
    cfg = opt.yaml
    # data文件路径, 其中存储了训练和验证数据集的路径
    data = opt.data
    # 训练批次
    epochs = opt.epochs
    batch_size = opt.batch_size
    # 初始化权重路径
    weights = opt.weights  # initial training weights
    # 训练和测试的图片大小
    imgsz_train = opt.img_size
    imgsz_test = opt.img_size
    # 图像要设置成32的倍数
    gs = 32  # (pixels) grid size
    assert math.fmod(imgsz_test, gs) == 0, "--img-size %g must be a %g-multiple" % (imgsz_test, gs)
    # ----------------------------------------------------------------------------------
    # ----------------------------------------------------------------------------------

    # 数据字典 "classes": 类别数, "train":训练数据集路径 "valid":验证数据集路径 "names": 类别名
    data_dict = parse_data_cfg(data)
    train_path = data_dict['train']
    val_path = data_dict['valid']
    # 类别数
    nc = opt.nc
    # Remove previous results
    for f in glob.glob(results_file):
        os.remove(f)

    # ----------------------------------------------------------------------------------
    # 模型初始化(加载预训练权重)
    # ----------------------------------------------------------------------------------
    # 初始化模型
    model = Model(cfg=cfg, ch=3, nc=nc).to(device)
    # 开始训练的epoch
    start_epoch = 0
    best_map = 0
    pretrain = False
    if weights.endswith(".pt") or weights.endswith(".pth"):
        pretrain = True
        ckpt = torch.load(weights, map_location=device)

        # load model
        try:
            ckpt["model"] = {k: v for k, v in ckpt["model"].items() if model.state_dict()[k].numel() == v.numel()}
            model.load_state_dict(ckpt["model"], strict=False)
        except KeyError as e:
            s = "%s is not compatible with %s. Specify --weights '' or specify a --cfg compatible with %s. " \
                "See https://github.com/ultralytics/yolov3/issues/657" % (opt.weights, opt.cfg, opt.weights)
            raise KeyError(s) from e


        # load results
        if ckpt.get("training_results") is not None:
            with open(results_file, "w") as file:
                file.write(ckpt["training_results"])  # write results.txt
        # 加载最好的map
        if "best_map" in ckpt.keys():
            best_map = ckpt['best_map']

        # epochs
        start_epoch = ckpt["epoch"] + 1
        if epochs < start_epoch:
            print('%s has been trained for %g epochs. Fine-tuning for %g additional epochs.' %
                  (opt.weights, ckpt['epoch'], epochs))
            epochs += ckpt['epoch']  # finetune additional epochs

        del ckpt

        hyp['lr0'] = hyp['lr0'] * hyp['lrf']

        print(colorstr('Pretrain') + ': Successful load pretrained weights.')
    # ----------------------------------------------------------------------------------
    # ----------------------------------------------------------------------------------

    # ----------------------------------------------------------------------------------
    # 定义优化器和学习率策略
    # ----------------------------------------------------------------------------------
    nbs = 64  # 训练多少图片进行一次反向传播
    accumulate = max(round(nbs / batch_size), 1)  # accumulate n times before optimizer update (bs 64)
    hyp['weight_decay'] *= batch_size * accumulate / nbs  # scale weight_decay
    # 定义优化器
    optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['weight_decay'])
    # Scheduler
    if opt.cos_lr:
        lf = one_cycle(1, hyp['lrf'], epochs)  # cosine 1->hyp['lrf']
    else:
        lf = lambda x: (1 - x / epochs) * (1.0 - hyp['lrf']) + hyp['lrf']  # linear
    scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
    scheduler.last_epoch = start_epoch  # 指定从哪个epoch开始
    # ----------------------------------------------------------------------------------
    # ----------------------------------------------------------------------------------

    # ----------------------------------------------------------------------------------
    # 加载数据集
    # ----------------------------------------------------------------------------------
    train_dataset = LoadImagesAndLabels(train_path, imgsz_train, batch_size,
                                        augment=True,
                                        hyp=hyp,  # augmentation hyperparameters
                                        rect=False,  # rectangular training
                                        cache_images=opt.cache_images,)
    # dataloader
    num_workers = 0  # number of workers
    train_dataloader = torch.utils.data.DataLoader(train_dataset,
                                                   batch_size=batch_size,
                                                   num_workers=num_workers,
                                                   # Shuffle=True unless rectangular training is used
                                                   shuffle=not opt.rect,
                                                   pin_memory=True,
                                                   collate_fn=train_dataset.collate_fn)
    # 当需要根据map保存最好的权值时, 加载验证数据集
    val_dataset = LoadImagesAndLabels(val_path, imgsz_test, batch_size,
                                      hyp=hyp,
                                      rect=True,
                                      cache_images=opt.cache_images)
    val_dataloader = torch.utils.data.DataLoader(val_dataset,
                                                 batch_size=batch_size,
                                                 num_workers=num_workers,
                                                 pin_memory=True,
                                                 collate_fn=val_dataset.collate_fn)

    # train_dataset.labels (1203, ) 1203张图片中的所有标签(nt,5)
    # np.concatenate 在维度0上对所有标签进行拼接 (1880, 5)
    labels = np.concatenate(train_dataset.labels, 0)
    mlc = int(labels[:, 0].max())  # 最大的类别标签
    assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'
    # ----------------------------------------------------------------------------------
    # ----------------------------------------------------------------------------------

    # ----------------------------------------------------------------------------------
    # 模型参数定义
    # ----------------------------------------------------------------------------------
    nl = model.model[-1].nl  # 输出特征图数量
    hyp['box'] *= 3 / nl  # scale to layers
    hyp['cls'] *= nc / 80 * 3 / nl  # scale to classes and layers
    hyp['obj'] *= (imgsz_test / 640) ** 2 * 3 / nl  # scale to image size and layers
    hyp['label_smoothing'] = opt.label_smoothing
    model.nc = nc  # attach number of classes to model
    model.hyp = hyp  # attach hyperparameters to model
    # ----------------------------------------------------------------------------------
    # ----------------------------------------------------------------------------------
    compute_loss = ComputeLoss(model)

    # ----------------------------------------------------------------------------------
    # 开始训练
    # ----------------------------------------------------------------------------------
    nb = len(train_dataloader)  # batches的数量, dataloader已经将data安装batch_size进行了打包
    nw = max(round(hyp['warmup_epochs'] * nb), 100)  # 热身训练的迭代次数, max(3 epochs, 100 iterations)
    last_opt_step = -1  # 最后一次更新参数的步数
    train_loss, train_box_loss, train_obj_loss, train_cls_loss = [], [], [], []
    learning_rate = []

    if opt.savebest:
        best_map = ComputeAP(model, val_dataloader, device=device)
        print(f"{colorstr('Initialize best_map')}: mAP@0.50 = {best_map: .5f}")

    print(f'Image sizes {imgsz_train} train, {imgsz_test} val\n'
                f'Using {num_workers} dataloader workers\n'
                f'Starting training for {epochs} epochs...')
    for epoch in range(start_epoch, epochs):
        model.train()

        # 训练过程中的信息打印
        metric_logger = utils.MetricLogger(delimiter="  ")
        metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
        header = 'Epoch: [{}]'.format(epoch)

        # 当前训练批次的平均损失
        mloss = torch.zeros(4, device=device)
        now_lr = 0.

        optimizer.zero_grad()

        # imgs: [batch_size, 3, img_size, img_size]
        # targets: [num_obj, 6] , that number 6 means -> (img_index, obj_index, x, y, w, h)
        # 其中(x, y, w, h)绝对作了缩放处理后的相对坐标
        # paths: list of img path(文件路径)
        for i, (imgs, targets, paths, _, _) in enumerate(metric_logger.log_every(train_dataloader, 50, header)):
            ni = i + nb * epoch  # number integrated batches (since train start)
            imgs = imgs.to(device).float() / 255.0  # uint8 to float32, 0 - 255 to 0.0 - 1.0

            # ----------------------------------------------------------------------------------
            # Warmup 热身训练
            # ----------------------------------------------------------------------------------
            if ni <= nw:
                xi = [0, nw]  # x interp
                # compute_loss.gr = np.interp(ni, xi, [0.0, 1.0])  # iou loss ratio (obj_loss = 1.0 or iou)
                accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
                for j, x in enumerate(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, [hyp['warmup_bias_lr'] if j == 0 else 0.0, x['initial_lr'] * lf(epoch)])
                    if 'momentum' in x:
                        x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
            # ----------------------------------------------------------------------------------
            # ----------------------------------------------------------------------------------

            # 前向传播
            pred = model(imgs)
            # 损失计算
            loss_dict = compute_loss(pred, targets.to(device))
            losses = sum(loss for loss in loss_dict.values())

            loss_items = torch.cat((loss_dict["box_loss"],
                                    loss_dict["obj_loss"],
                                    loss_dict["class_loss"],
                                    losses)).detach()
            mloss = (mloss * i + loss_items) / (i + 1)  # update mean losses

            # 反向传播
            losses *= batch_size
            losses.backward()

            # 每训练accumulate次进行参数更新
            if ni - last_opt_step >= accumulate:
                optimizer.step()
                optimizer.zero_grad()

            metric_logger.update(loss=losses, **loss_dict)
            now_lr = optimizer.param_groups[0]["lr"]
            metric_logger.update(lr=now_lr)
            # end batch ----------------------------------------------------------------------
        scheduler.step()


        train_loss.append(mloss.tolist()[-1])
        learning_rate.append(now_lr)
        result_mAP = ComputeAP(model, val_dataloader, device=device)
        voc_mAP = result_mAP[1]  # @0.50
        # write into txt
        with open(results_file, "a") as f:
            # box_loss, obj_clss, cls_loss, train_loss, lr
            result_info = [str(round(i, 4)) for i in [mloss.tolist()[0]]] + \
                          [str(round(i, 4)) for i in [mloss.tolist()[1]]] + \
                          [str(round(i, 4)) for i in [mloss.tolist()[2]]] + \
                          [str(round(i, 4)) for i in [mloss.tolist()[-1]]] + \
                          [str(round(now_lr, 6))] + [str(round(voc_mAP, 6))]
            txt = "epoch:{} {}".format(epoch, '  '.join(result_info))
            f.write(txt + "\n")
        if opt.savebest:
            if voc_mAP > best_map:
                print(f"{colorstr('Save best_map Weight')}: update mAP@0.50 from {best_map: .5f} to {voc_mAP: .5f}")
                best_map = voc_mAP
                with open(results_file, 'r') as f:
                    save_files = {
                        'model': model.state_dict(),
                        'optimizer': optimizer.state_dict(),
                        'training_results': f.read(),
                        'epoch': epoch,
                        'best_map': best_map
                    }
                    torch.save(save_files, save_path.format('best_map'))
        else:
            if (epoch + 1) % 20 == 0 or epoch == epochs - 1:
                with open(results_file, 'r') as f:
                    save_files = {
                        'model': model.state_dict(),
                        'training_results': f.read(),
                        'epoch': epoch}
                    torch.save(save_files, save_path.format(epoch))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # -----------------------------------------
    file = "yolov5s"
    weight_file = f"weights/{file}"  # 权重存储文件
    result_file = f'results/{file}'  # 训练损失、mAP等保存文件
    if not os.path.exists(weight_file):
        os.makedirs(weight_file)
    if not os.path.exists(result_file):
        os.makedirs(result_file)
    # -----------------------------------------
    parser.add_argument('--epochs', type=int, default=300)
    parser.add_argument('--batch-size', type=int, default=4)
    parser.add_argument('--nc', type=int, default=3)
    yaml_path = f'cfg/models/{file}.yaml'
    parser.add_argument('--yaml', type=str, default=yaml_path, help="model.yaml path")
    parser.add_argument('--data', type=str, default='data/my_data.data', help='*.data path')
    parser.add_argument('--hyp', type=str, default='cfg/hyps/hyp.scratch-med.yaml', help='hyperparameters path')
    parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer')
    parser.add_argument('--cos-lr', type=bool, default=True, help='cosine LR scheduler')
    parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
    parser.add_argument('--img-size', type=int, default=640, help='test size')
    parser.add_argument('--rect', action='store_true', help='rectangular training')
    parser.add_argument('--savebest', type=bool, default=False, help='only save best checkpoint')
    # 当内存足够时, 设置为True, 将数据集加载到内存中, 在训练时不用从磁盘中读取图片可以加快训练速度
    parser.add_argument('--cache-images', default=False, help='cache images for faster training')
    # 预训练权重
    weight = f'weights/{file}/{file}.pt'

    parser.add_argument('--weights', type=str, default=weight if os.path.exists(weight) else "", help='initial weights path')

    parser.add_argument('--device', default='cuda:0', help='device id (i.e. 0 or 0,1 or cpu)')


    # 结果保存路径

    time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    results_file = f"./results/{file}/results{time}.txt"
    parser.add_argument("--results_file", default=results_file, help="save results files")
    save_file_path = f'./results/{file}/loss_and_lr{time}.png'
    parser.add_argument('--save_file_path', default=save_file_path, help="save loss and lr fig")
    save_path = f"./weights/{file}/{file}-" + "{}.pt"
    parser.add_argument('--save_path', default=save_path, help="weight save path")
    opt = parser.parse_args()

    # 检查文件是否存在
    opt.cfg = check_file(opt.yaml)
    opt.data = check_file(opt.data)
    opt.hyp = check_file(opt.hyp)
    print(opt)

    with open(opt.hyp) as f:
        hyp = yaml.load(f, Loader=yaml.FullLoader)

    train(hyp)