[YOLO专题-24]：YOLO V5 - ultralytics代码解析-train.py训练代码的详细执行流程

作者主页(​​文火冰糖的硅基工坊​​​)：​​文火冰糖（王文兵）的博客_文火冰糖的硅基工坊​

​目录​

​​第1步 执行程序​​

​​第2步 训练前准备​​

​​第3步 准备数据集​​

​​第4步 准备网络模型​​

​​第5步 准备训练参数​​

​​第6步 开始训练​​

​​第7步 训练结束，进行后期处理​​

第1步 执行程序

python train --img 640 --batch 1 --epochs 5 --data coco128.yaml --weights yolov5s.pt

第2步 训练前准备

（1）主函数：main()

（2）检查条件是否满足

• 打印输入命令函数参数：print_args(FILE.stem, opt)
• 检查本地文件与github服务器文件的status
• 检查依赖条件是否满足

（3）是否是前一次训练处开始训练：opt.resume

• 如果是，则从上次训练的ckpt文件处开始训练
• 否则重头开始训练

（4）是否需要在多平台并行训练

• 如果是，则初始化多GPU环境：dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")

（5）train函数：train(opt.hyp, opt, device, callbacks)

（6）获取上次训练的权重参数： last, best = w / 'last.pt', w / 'best.pt'

（7）加载网络训练的超参数：hyp = yaml.safe_load(f)

（8）启动一个loggers，把log写到文件中，而不影响GPU训练性能：loggers = Loggers(save_dir, weights, opt, hyp, LOGGER)

该log实际上是tensorboard。

TensorBoard: Start with 'tensorboard --logdir runs\train', view at http://localhost:6006/

第3步 准备数据集

（1）检查并准备好训练数据：data_dict = data_dict or check_dataset(data)

（2）检查分类数：nc = 1 if single_cls else int(data_dict['nc']) # number of classes

第4步 准备网络模型

（1）检查本地是否有预训练模型权重

• 如果没有，则先下载预训练模型：weights = attempt_download(weights)
• 如果有，则直接加载预训练模型：ckpt = torch.load(weights, map_location=device) # load checkpoint

（2）创建模型：model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create

• 用预训练模型的权重初始：model.load_state_dict(csd, strict=False) # load
• 初始化先验框：hyp.get('anchors')
• 确定网络需要迭代的参数：v.requires_grad = True

（3）确保模型的输入图片尺寸满足要求：

• 确保图片本身的尺寸满足要求：imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2)
• 确保batch size满足要求：check_train_batch_size(model, imgsz)

第5步 准备训练参数

（1）确定优化器更新的频率：不是每个batch都更新，如果是batch size =64时更新一次权重

nbs = 64 # nominal batch size

accumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizing

（2) 对待更新参数进行分组：

• group0：所有不待衰减的权重参数
• group1：所有待衰减的权重参数
• group2：所有的偏置参数b

（3）创建优化器：

• optimizer = Adam(group0）或optimizer = SGD（group0）
• 把group1和group2添加到优化器中

（4）创建动态权重优化的调度器

• 创建调度器：scheduler = lr_scheduler.LambdaLR(）

（5）指数平均初始化模型：ema = ModelEMA(model)

（6）恢复上次训练前的上下文

• 恢复上次优化器的状态：optimizer.load_state_dict(ckpt['optimizer'])
• 恢复上次滑动平均的状态：ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
• 恢复上次的epoch计数：epochs += ckpt['epoch'] # finetune additional epochs

（7）初始化单机多GPU训练：model = torch.nn.DataParallel(model)

（8）同步batch nomralization：model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)

（9）创建train dataloader和dataset：train_loader, dataset = create_dataloader()

（10）创建验证 dataloader和dataset：val_loader = create_dataloader()

（11）初始化多机多GPU的训练：# DDP mode：model = DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)

（12）初始化模型的超参数等模型属性

（13）启动训练：初始化过程中使用到的全局变量

• 记录训练开始的时间
• 其他全局变量：warm up的计数、分类数、epoch、输出结构等等。
• 打印log

第6步 开始训练

（1）开始epoch迭代

        （1）设置在训练模式下：model.train()

        （2） Update image weights (optional, single-GPU only)

        （3）# Update mosaic border (optional)：数据增强

        （4）从train loader中读取样本：pbar = enumerate(train_loader)

        （5）迭代前信息提示：

                - LOGGER.info(('\n' + '%10s' * 7) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'labels', 'img_size'))

                - 打印进度条：pbar = tqdm(pbar, total=nb, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar

        （6）清除梯度：optimizer.zero_grad()

        （7）开始批次内迭代

                （1）把image分配到GPU上

                （2）warm up迭代

                （3）多次度训练准备

                （4）前向预测计算：pred = model(imgs) # forward

                （5）反向梯度计算：scaler.scale(loss).backward()

                （6）优化器梯度迭代

                        - 梯度迭代：scaler.step(optimizer) # optimizer.step

                        - 更新学习率：scaler.update()

                        - 清除梯度：optimizer.zero_grad()

                        - 更新模型的滑动平均：ema.update(model)

                （7）打印训练log

                        - 更新进度

                                -- GPU内存使用情况

                                -- 检测到box中对象的置信度

                                -- 检测到box中对象的分类可能性

                                -- 分类标签

                                -- image size

        （8）每个批次迭代后，更新学习率调度器：scheduler.step()

        （9）对模型进行评估：results, maps, _ = val.run()， 详见模型评估过程

                - class：所有分类

                - images：128张图片

                - Labels：929个标签对象

                - P： 准确率为0.648

                - R:   召回率为0.596

                - mAP: mean Average Precision， 0.413.

          (10) 更新最佳mAP: fi = fitness(np.array(results).reshape(1, -1))  

        （11）存储最佳和最新模型

                - torch.save(ckpt, last)

                - torch.save(ckpt, best) 

        （12）关闭部分GPU

第7步 训练结束，进行后期处理

（1）打印训练时间

        5 epochs completed in 0.402 hours.

（2）分别用最新模型和最好模型对训练结果进行评估：for f in last, best:

        Validating runs\train\exp24\weights\best.pt...

（3）打印训练结果

        - class：所有分类

        - images：128张图片

        - Labels：929个标签对象， 254个personal标签对象，6个bicycle对象

        - P：    全部对象的平均准确率为0.693， personal对象的准确率0.831, bicycle的准确率1.

        - R:      全部对象的平均召回率为0.606，  personal对象的准确率0.654, bicycle的准确率0.45.

        - mAP: 全部对象的平均召回率为0.683，  personal对象的准确率0.795, bicycle的准确率0.667

(4) 存储log：LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")

(5) 清空GPU cache：torch.cuda.empty_cache()

作者主页(​​文火冰糖的硅基工坊​​​)：​​文火冰糖（王文兵）的博客_文火冰糖的硅基工坊​