物体检测-系列教程27：YOLOV5 源码解析17（训练脚本解读：训练函数3）

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

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

22、epoch循环训练------准备工作

    for epoch in range(start_epoch, epochs):  
        model.train()
        if dataset.image_weights:
            if rank in [-1, 0]:
                w = model.class_weights.cpu().numpy() * (1 - maps) ** 2  # class weights
                image_weights = labels_to_image_weights(dataset.labels, nc=nc, class_weights=w)
                dataset.indices = random.choices(range(dataset.n), weights=image_weights, k=dataset.n)  
            if rank != -1:
                indices = torch.zeros([dataset.n], dtype=torch.int)
                if rank == 0:
                    indices[:] = torch.tensor(dataset.indices, dtype=torch.int)
                dist.broadcast(indices, 0)
                if rank != 0:
                    dataset.indices = indices.cpu().numpy()
        mloss = torch.zeros(4, device=device)  # mean losses
	    if rank != -1:
	        dataloader.sampler.set_epoch(epoch)
	    pbar = enumerate(dataloader)
	    logger.info(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'GIoU', 'obj', 'cls', 'total', 'targets', 'img_size'))
	    if rank in [-1, 0]:
	        pbar = tqdm(pbar, total=nb)  # progress bar
	    optimizer.zero_grad()

1. 逐个epoch进行模型训练，epoch 是当前的epoch，start_epoch是开始epoch可能是从checkpoint中恢复训练的，epochs是总epoch数量
2. 开启模型训练模式
3. 是否开启了图像类别权重训练：
4. 如果在分布式训练中的主节点
5. w，计算每个类别的动态权重。这里使用模型的类别权重（考虑了类别不平衡）和每个类别的mAP损失（1 - maps，其中maps是上一轮评估阶段计算得到的每个类别的平均精度）的平方进行加权，从而得到新的权重。这样做的目的是给那些模型当前识别性能较差的类别更高的权重
6. image_weights ，调用辅助函数labels_to_image_weights根据上一步计算得到的类别权重和数据集的标签，计算每个图像的权重，让模型在训练过程中更多地关注难以识别的类别的图像
7. 根据计算得到的图像权重，通过加权随机抽样选择图像进行训练，生成一个新的图像索引列表dataset.indices。这里dataset.n是数据集中图像的总数
8. 如果当前是分布式训练环境中：
9. indices ，创建一个全为0的Tensor，用来存储全局的图像索引
10. 如果当前是主进程：
11. 用dataset.indices的数据对indices 进行填充更新
12. 使用PyTorch的分布式通信包dist中的broadcast函数，将主进程中的indices张量广播给所有进程，确保每个进程都使用相同的数据索引进行训练
13. 如果不是主进程
14. 非主进程将接收到的广播索引张量转换回NumPy数组，并更新其本地数据集的索引。这确保了在分布式训练环境中，所有进程都基于相同的数据索引进行模型的训练
15. mloss，一个长度为4用于存储平均损失值，并且这个张量是放在由变量device指定的设备上（通常是CPU或者某个GPU），这四个损失值分别代表不同类型的损失，如定位损失、分类损失等
16. 如果是在分布式训练中：
17. 调整数据采样器，根据当前的epoch改变采样方式，确保每个epoch中看到不同顺序的数据，有助于提高模型的泛化能力
18. 定义一个进度条，使用enumerate对数据加载器（dataloader）进行迭代，这样不仅可以加载数据，还可以得到每个批次的索引，通常用于在训练过程中记录进度
19. 这行代码使用日志记录器（logger）来打印一条信息，格式化字符串用于创建表格头部，显示训练过程中将要监控的各项指标，如每轮训练的序号、GPU内存使用情况、各种损失值和目标数量等
20. 记录一条日志信息，格式化字符串用于创建表格头部，显示训练过程中将要监控的各项指标，如每轮训练的序号、GPU内存使用情况、各种损失值和目标数量等
21. 如果不是分布式训练或者在分布式训练的主进程：
22. 使用tqdm库来为数据加载器的迭代器pbar添加一个进度条，total=nb参数指定了进度条的总长度，nb通常是数据集的批次数量。这样做可以在训练过程中直观地看到进度和估计的剩余时间
23. 梯度清零

23、批处理循环

使用PyTorch框架进行图像识别或对象检测的训练过程中的一个批处理循环。包含了多个训练阶段的关键操作，例如数据预处理、模型预热、多尺度训练、自动混合精度训练、反向传播和优化器步骤

    for i, (imgs, targets, paths, _) in pbar:  
        ni = i + nb * epoch  # number integrated batches (since train start)
        imgs = imgs.to(device, non_blocking=True).float() / 255.0  # uint8 to float32, 0-255 to 0.0-1.0
        # Warmup
        if ni <= nw:
            xi = [0, nw]  # x interp
            accumulate = max(1, np.interp(ni, xi, [1, nbs / total_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, [0.1 if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
                if 'momentum' in x:
                    x['momentum'] = np.interp(ni, xi, [0.9, hyp['momentum']])
        # Multi-scale
        if opt.multi_scale:
            sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs  # size
            sf = sz / max(imgs.shape[2:])  # scale factor
            if sf != 1:
                ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]  # new shape (stretched to gs-multiple)
                imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
        # Autocast
        with amp.autocast(enabled=cuda):
            pred = model(imgs)
            loss, loss_items = compute_loss(pred, targets.to(device), model)  # scaled by batch_size
            if rank != -1:
                loss *= opt.world_size  # gradient averaged between devices in DDP mode
        # Backward
        scaler.scale(loss).backward()
        # Optimize
        if ni % accumulate == 0:
            scaler.step(optimizer)  # optimizer.step
            scaler.update()
            optimizer.zero_grad()
            if ema is not None:
                ema.update(model)
        # Print
        if rank in [-1, 0]:
            mloss = (mloss * i + loss_items) / (i + 1)  # update mean losses
            mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0)  # (GB)
            s = ('%10s' * 2 + '%10.4g' * 6) % (
                '%g/%g' % (epoch, epochs - 1), mem, *mloss, targets.shape[0], imgs.shape[-1])
            pbar.set_description(s)
            if ni < 3:
                f = str(log_dir / ('train_batch%g.jpg' % ni))  # filename
                result = plot_images(images=imgs, targets=targets, paths=paths, fname=f)
                if tb_writer and result is not None:
                    tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)

1. 遍历训练数据集的单个batch，pbar是进度条对象，imgs是图像数据，targets是目标标签，paths是图像路径
2. ni ，自训练开始以来处理的batch总数
3. img，将图像数据转移到GPU，并将其类型转换为float，同时进行了归一化处理，将像素值从0-255映射到0.0-1.0之间

预热

4. 如果当前的总批次数小于或等于预热期nw的批次数，则执行以下预热逻辑：
5. xi ，定义一个列表用于后续的线性插值，表示预热开始和结束的批次数
6. accumulate ，根据当前批次数ni，使用线性插值计算梯度累积次数accumulate，确保至少为1
7. 遍历优化器的参数组：
8. 对每个参数组，根据批次数调整学习率。如果是偏置项（通常j=2时），学习率从0.1调整到初始学习率；否则从0开始调整
9. 如果参数组中包含动量设置，则根据批次数调整动量值

多尺度输入处理

10. 如果启用了多尺度输入处理，则执行以下逻辑
11. sz，随机选择一个新的输入尺寸，保证是gs（网格大小）的倍数
12. sf，计算缩放因子，即新尺寸与当前最大边长的比值
13. 如果缩放因子不为1：
14. 调整图片大小到新的尺寸

自动混合精度训练

15. 使用自动混合精度进行训练，以减少计算资源消耗并加速训练
16. 通过模型获取预测结果
17. 调用辅助函数compute_loss计算损失值，loss_items通常包含不同类型损失的详细信息
18. 如果在分布式训练环境中：
19. 根据训练设备的数量调整损失值，以便梯度平均

反向传播和优化

20. 使用梯度缩放进行反向传播，减少浮点数精度误差导致的梯度消失
21. 如果达到累积步数
22. 更新模型参数，并重置梯度
23. 如果使用指数移动平均(EMA)更新模型参数

训练监控和日志记录

24. 如果不是分布式训练或者在分布式训练的主线程
25. mloss ，更新训练过程中的平均损失(mloss)，mloss是一个累积变量，用于计算到目前为止所有批次的平均损失，loss_items是当前批次的损失，i是当前批次的索引。通过这种方式，可以平滑地跟踪损失的变化，有助于监控训练过程
26. mem，计算并格式化GPU内存使用情况，单位为Gigabytes(GB)，torch.cuda.memory_reserved()返回当前为PyTorch保留的总GPU内存量，通过除以1E9转换为GB。如果CUDA可用，它会显示当前GPU内存使用情况；如果不可用，则显示为0
27. s，构建一个包含当前训练周期(epoch)、内存使用(mem)、平均损失(mloss)、当前批次的目标数量(targets.shape[0])和图像尺寸(imgs.shape[-1])的字符串。这个字符串将用于更新进度条的描述，以便在训练时实时显示这些关键信息
28. 更新进度条(pbar)的描述为上一步构建的字符串s，训练过程中的关键信息能够实时展示在进度条上，提高用户监控训练进度的便利性
29. 检查ni（自训练开始以来累积的批次数）是否小于3，用于训练初期，保存一些批次的图像，用于后续的分析和监控
30. 生成要保存的训练批次图像的文件名。这里使用了Python的路径操作，log_dir是日志目录的路径，ni是累积的批次数，将它们结合生成一个文件名，用于保存图像
31. 调用辅助函数plot_images，将当前批次的图像(imgs)、目标(targets)和路径(paths)绘制到图像上，并保存到上一步生成的文件名(f)指定的位置，用于可视化训练数据，帮助理解模型是如何看待这些数据的
32. 如果TensorBoard的写入器(tb_writer)存在且result（绘制的图像）非空
33. 将这个图像添加到TensorBoard中，dataformats='HWC’指明图像数据的格式，global_step=epoch用于指定当前的训练周期，以便在TensorBoard中按训练周期组织图像