Bootstrap

物体检测-系列教程26:YOLOV5 源码解析16(训练脚本解读:训练函数2)

😎😎😎物体检测-系列教程 总目录

有任何问题欢迎在下面留言
本篇文章的代码运行界面均在Pycharm中进行
本篇文章配套的代码资源已经上传
点我下载源码

21、训练准备

21.1 模型训练前的配置与优化

图像尺寸的调整、模型并行化、批量标准化同步以及指数移动平均的设置

    gs = int(max(model.stride))  # grid size (max stride)
    imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size]  
    if cuda and rank == -1 and torch.cuda.device_count() > 1:
        model = torch.nn.DataParallel(model)
    if opt.sync_bn and cuda and rank != -1:
        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
        logger.info('Using SyncBatchNorm()')
    ema = ModelEMA(model) if rank in [-1, 0] else None
    if cuda and rank != -1:
        model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank)
  1. gs,计算模型中最大的步长(stride),并将其设置为网格大小(gs)。步长是模型在处理图像时减少空间维度的因子
  2. imgsz, imgsz_test,调用辅助函数check_img_size验证输入图像尺寸是否为网格大小的倍数。这是因为某些模型架构要求输入尺寸满足特定条件,以避免在下采样过程中产生尺寸不匹配的问题
  3. 检查是否在使用CUDA(即GPU训练)、是否不在分布式训练环境中(rank == -1),并且是否有多个GPU可用
  4. 如果条件满足,使用DataParallel来包装模型,实现模型在多个GPU上的数据并行训练
  5. 如果启用同步批量标准化(sync_bn)、在使用CUDA并且处于分布式训练环境中(rank != -1)
  6. 将模型中的所有批量标准化层转换为同步批量标准化层,并将模型移至指定设备。这有助于在多GPU训练时保持各设备间批量标准化层的一致性
  7. 记录一条日志,表明正在使用同步批量标准化
  8. 如果在主进程或非分布式训练环境中,初始化指数移动平均(EMA)对象。EMA有助于平滑模型在训练过程中的参数变化,通常能提升最终模型的稳定性和性能
  9. 再次检查是否在使用CUDA并且处于分布式训练环境
  10. 使用DistributedDataParallel(DDP)包装模型,实现模型在多GPU环境下的分布式训练。device_ids和output_device指定了模型和数据应该在哪个GPU上进行计算

通过一系列配置确保了模型训练能在不同环境下正确运行,包括处理单GPU、多GPU以及分布式训练场景。通过调整图像尺寸、应用数据并行、同步批量标准化和使用指数移动平均等策略,提升了模型训练的效率和效果

21.2 配置训练和测试数据加载器

	# Trainloader
    dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt, hyp=hyp, augment=True,
                                            cache=opt.cache_images, rect=opt.rect, rank=rank,
                                            world_size=opt.world_size, workers=opt.workers)
    mlc = np.concatenate(dataset.labels, 0)[:, 0].max()  # max label class
    nb = len(dataloader)  # number of batches
    assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)

    # Testloader
    if rank in [-1, 0]:
        # local_rank is set to -1. Because only the first process is expected to do evaluation.
        testloader = create_dataloader(test_path, imgsz_test, total_batch_size, gs, opt, hyp=hyp, augment=False,
                                       cache=opt.cache_images, rect=True, rank=-1, world_size=opt.world_size,
                                       workers=opt.workers)[0]
  1. 创建训练数据加载器:使用create_dataloader函数根据训练数据路径、图像尺寸、批量大小等参数创建数据加载器。这包括可选的数据增强、缓存机制、是否使用矩形训练等选项
  2. mlc ,最大标签类别,计算数据集中所有标签的最大类别值mlc。这是为了确保训练数据中的标签类别不会超过模型配置的类别数nc
  3. nb ,批次数量,确定数据加载器中的批次数量nb,这对于遍历训练数据和计算训练进度至关重要
  4. 确保数据中的最大标签类别没有超过模型所配置的类别数nc
  5. 在非分布式训练环境或分布式训练的主进程(rank为-1或0)中
  6. 创建测试数据加载器

21.3 模型训练前的参数配置与初始化

首先调整模型超参数以适应当前数据集,然后为模型设置类别数、超参数、类别权重、类别名称等。接着,对于特定的训练节点(或单机训练),会计算类别频率,初始化模型偏差,并可能进行锚点检查。最后,设置了训练的一些基本参数,如预热迭代次数和日志记录

    hyp['cls'] *= nc / 80.  # scale coco-tuned hyp['cls'] to current dataset
    model.nc = nc  # attach number of classes to model
    model.hyp = hyp  # attach hyperparameters to model
    model.gr = 1.0  # giou loss ratio (obj_loss = 1.0 or giou)
    model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device)  # attach class weights
    model.names = names
    if rank in [-1, 0]:
        labels = np.concatenate(dataset.labels, 0)
        c = torch.tensor(labels[:, 0])  # classes
        plot_labels(labels, save_dir=log_dir)
        if tb_writer:
            tb_writer.add_histogram('classes', c, 0)
        if not opt.noautoanchor:
            check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
    t0 = time.time()
    nw = max(3 * nb, 1e3)  # number of warmup iterations, max(3 epochs, 1k iterations)
    maps = np.zeros(nc)  # mAP per class
    results = (0, 0, 0, 0, 0, 0, 0)  # 'P', 'R', 'mAP', 'F1', 'val GIoU', 'val Objectness', 'val Classification'
    scheduler.last_epoch = start_epoch - 1  # do not move
    scaler = amp.GradScaler(enabled=cuda)
    logger.info('Image sizes %g train, %g test' % (imgsz, imgsz_test))
    logger.info('Using %g dataloader workers' % dataloader.num_workers)
    logger.info('Starting training for %g epochs...' % epochs)
  1. 将预先调整好的超参数hyp[‘cls’]根据当前数据集的类别数nc进行缩放,这样做是为了让超参数更适应于当前的数据集
  2. 将当前数据集的类别数nc赋值给模型的nc属性,用于后续训练和推理过程中识别类别数量
  3. 将超参数字典hyp赋值给模型的hyp属性,使得训练过程可以根据这些超参数进行
  4. 设置模型的gr属性(GIoU loss比例)为1.0。这意味着在计算目标损失时,GIoU损失将被完全考虑
  5. 根据数据集的标签计算类别权重,并将计算结果分配给模型的class_weights属性,这有助于在训练过程中平衡类别不平衡问题
  6. 将数据集中的类别名称赋值给模型的names属性,用于标识每个类别
  7. 如果在分布式训练中的主节点
  8. labels ,将所有数据的标签合并成一个数组,用于后续的类别频率计算和可视化
  9. c,提取合并后标签数组中的类别信息,以便进行进一步的处理
  10. 调用辅助函数 plot_labels,绘制标签的分布图,并保存到日志目录。这有助于了解数据集的类别分布情况
  11. 如果启用了TensorBoard
  12. 将类别分布信息添加到TensorBoard的直方图中,以便于观察
  13. 如果不禁用自动锚点检查
  14. 对锚点进行检查,以确保它们适合当前的数据集和模型配置
  15. t0,记录训练开始的时间
  16. nw,计算预热阶段的迭代次数,取三个epoch内batch的三倍和1000之间的最大值
  17. maps ,始化每个类别的平均精度(mAP)为0,用于后续的评估
  18. results ,初始化训练结果的各项指标,包括精确度、召回率、mAP、F1分数等
  19. 设置学习率调度器的最后一个epoch,确保学习率调整能够正确进行
  20. scaler ,初始化梯度缩放器,用于混合精度训练,以提高训练速度和减少内存使用
  21. 记录几条日志信息,记录训练相关的基础信息,如图像大小、使用的数据加载器工作数、训练周期数等
;