train.py是YOLOv5中用于模型训练的脚本文件，其主要功能是读取配置文件、设置训练参数、构建模型结构、加载数据、训练/验证模型、保存模型权重文件、输出日志等

参考笔记：

【YOLOv3】源码（train.py）_yolo原始代码-CSDN博客

【yolov5】 train.py详解_yolov5 train.py-CSDN博客

学习视频：

模型训练之train.py文件_哔哩哔哩_bilibili

---

目录

1. train.py的主要功能

2. 主要模块

2.1 参数解析与初始化

2.2 加载模型与数据

2.3 优化器与学习率更新器设置

2.4 训练循环(核心)

2.5 最终的模型验证

3. 完整train.py代码

---

1. train.py的主要功能

• 读取命令选项、训练参数配置文件：train.py通过argparse库读取指定的命令行参数，例如batch_size、epoch、weights等；读取yaml文件中的各种训练参数，例如learing_rate、momentum、weight_decay、IoU阈值、高宽比阈值anchor_t等

• 构建模型结构：train.py中要么通过命令行参数weights指定权重文件构建模型结构，并加载参数，如果没有使用命令行参数weights，则通过命令行参数cfg指定YOLOv5模型结构构建一个新的初始化模型

• 数据加载和预处理：train.py中定义了create_dataloader函数，用于加载训练数据和验证数据，并对其进行预处理。其中，预处理过程包括：自适应图像缩放、图像增强、标签转换等操作

• 训练和验证过程：train.py的train函数用于进行模型的训练和验证过程。训练过程中，train.py会对训练数据进行多次迭代，在每个epoch结束时，会对模型在验证集上的表现进行评估，记录的指标有：P，R，mAP@.5，mAP@.5-.95，val_loss(box，obj，cls)

• 模型保存和日志输出：train.py在每个epoch结束时，保存当前epoch的模型权重(允许覆盖)last.pt，通过模型在验证集上的(P，R，mAP@.5，mAP@.5-.95)计算出模型的适应度，利用适应度保存训练过程中的最佳模型权重best.pt，并将训练和验证过程中的各种指标输出到日志文件中

注意：模型训练结束后，train.py会利用best.pt在验证集上作最后一次验证，控制台输出验证结果

2. 主要模块

2.1 参数解析与初始化

常用参数说明

• --weights：模型初始权重存放路径
• --cfg：模型结构的YAML配置文件路径，例如models/yolov5l.yaml
• --data：数据集YAML配置文件路径，在YAML文件中中定义训练/验证数据集的存放路径和类别等
• --hyp：训练超参数YAML配置文件路径，控制learning_rate、weight_decay、momentum等训练超参数
• --epochs：训练的总轮数
• --batch-size：批量大小
• --imgsz：将输入图像自适应缩放到imgsz尺寸大小，其作用阶段是在加载数据时，具体的代码实现在utils/agumentations.py下的letterbox函数

•  --resume：断点续训，有时候服务器崩溃了导致训练中断，就可以利用这个参数继续训练。训练过程保存的pt存放在runs/train/expn/weights下，分别是last.pth、best.pth，使用该参数时，指定--resume runs/train/expn/weights/last.pth即可继续上次中断的训练
• --noautoanchor：是否不作自适应锚框计算，该功能默认开启，代码位置如下：

•  --multi-scale：多尺度训练；该功能的执行阶段在自适应图像缩放之后模型训练过程之中，自适应图像缩放之后，在训练过程中要对图片作处理时，再对图像尺寸作随机缩放或扩充，但必须确保仍然是最大下采样倍数(通常是32)的倍数。多尺度训练能够使得模型对不同尺度的目标具有更强的鲁棒性。代码的具体位置如下：

•  --label-smoothing：标签平滑；具体的代码实现是utils/loss.py下的smooth_BCE函数：

eps是在训练超参数YAML配置文件中定义的

• --patience：早停机制；如果模型在指定的epoch之内仍然没有性能提升，则训练提前终止，代码的具体位置如下：

•  --cos-lr：是否通过余弦函数更新学习率，不开启该功能则使用线性函数更新学习率，代码的具体代码位置如下：

• --freeze：冻结某些层的所有参数不参与训练；具体代码位置如下：

 常用参数和其他参数作用可参考代码注释：

def parse_opt(known=False):parser = argparse.ArgumentParser()#权重文件路径parser.add_argument('--weights', type=str, default=ROOT /'weights/yolov5s.pt',help='initial weights path (初始权重文件路径)')#模型结构yaml配置文件路径parser.add_argument('--cfg', type=str, default='',help='model.yaml path (模型结构配置文件路径)')#数据集yaml配置文件路径parser.add_argument('--data', type=str, default=ROOT / 'data/VOC-hat.yaml',help='dataset.yaml path (数据集配置文件路径)')#训练超参数yaml配置文件路径parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml',help='hyperparameters path (超参数配置文件路径)')#训练轮数parser.add_argument('--epochs', type=int, default=300)#批量大小parser.add_argument('--batch-size', type=int, default=16,help='total batch size for all GPUs, -1 for autobatch')#imgsz指定训练和验证时将输入图片自适应缩放到相应的尺寸parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640,help='指定训练、验证时将输入图片自适应缩放到相应的尺寸')#是否使用矩形训练parser.add_argument('--rect', action='store_true',help='rectangular training(是否使用矩形训练)')'''断点续训，有时候服务器崩溃了导致训练中断，就可以利用这个参数继续训练，训练过程保存的pth存放在runs/train/expxx/weights下，分别是last.pth和best.pth，使用该参数时，指定--resume runs/train/expxx/weights/last.pth即可继续上次中断的训练'''parser.add_argument('--resume', nargs='?', const=True, default=False,help='resume most recent training(断点续训)')'''该参数指定到了最后一个epoch才保存模型的last.pt和best.pt，如果你只想获取最后一个epoch的模型权重文件，那可以开启这个功能，默认是不开启的'''parser.add_argument('--nosave', action='store_true',help='only save final checkpoint(只保存最后一个epoch的last.pt和best.pt)')#仅验证最终epoch,该参数指定到了最后一个epoch才去计算模型在验证集上的性能指标,默认不开启parser.add_argument('--noval', action='store_true',help='only validate final epoch(仅验证最终epoch)')#是否禁用自适应锚框策略，该功能默认开启parser.add_argument('--noautoanchor', action='store_true',help='disable AutoAnchor')#不保存训练中生成的图标文件，默认是保存的parser.add_argument('--noplots', action='store_true',help='save no plot files')#使用遗传算法优化超参数，可指定优化代数，默认不开启该功能parser.add_argument('--evolve', type=int, nargs='?', const=300,help='evolve hyperparameters for x generations')##谷歌云盘bucket，一般用不到parser.add_argument('--bucket', type=str, default='',help='gsutil bucket')#是否缓存数据集到RAM或磁盘parser.add_argument('--cache', type=str, nargs='?', const='ram',help='--cache images in "ram" (default) or "disk"(缓存数据集)')#是否使用加权的图像进行训练,对于那些训练不好的图片，会在下一个epoch中增加一些权重parser.add_argument('--image-weights', action='store_true',help='use weighted image selection for training')#指定训练的设备parser.add_argument('--device', default='',help='cuda device, i.e. 0 or 0,1,2,3 or cpu')#是否使用多尺度训练(即随机将自适应缩放之后的尺寸再作随机缩放或增加，但必须确保是最大下采样倍数(通常是32)的倍数)#该功能使模型对不同尺度的目标具有更强的鲁棒性parser.add_argument('--multi-scale', action='store_true',help='vary img-size +/- 50%%(多尺度训练)')#不知有什么用，默认不开启parser.add_argument('--single-cls', action='store_true',help='train multi-class data as single-class')#优化器，默认是SGDparser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'],default='SGD', help='optimizer')#是否启用同步BatchNorm(只在多GPU中使用)，默认不启用parser.add_argument('--sync-bn', action='store_true',help='use SyncBatchNorm, only available in DDP mode')#数据加载器的最大工作线程数parser.add_argument('--workers', type=int, default=8,help='max dataloader workers (per RANK in DDP mode)')#训练结果存放的根目录，默认是runs/trianparser.add_argument('--project', default=ROOT / 'runs/train',help='save to project/name(根目录)')#训练结果存放的子目录，默认是runs/train/expnparser.add_argument('--name', default='exp',help='save to project/name(子目录)')#是否用当前的训练目录覆盖以前的expn训练目录,默认不开启parser.add_argument('--exist-ok', action='store_true',help='existing project/name ok, do not increment(允许覆盖以前的训练结果)')#是否使用四元数据加载器,默认不开启parser.add_argument('--quad', action='store_true', help='quad dataloader')#是否使用余弦函数更新学习率,默认是线性函数更新学习率parser.add_argument('--cos-lr', action='store_true',help='cosine LR scheduler')#是否启用标签平滑parser.add_argument('--label-smoothing', type=float, default=0.0,help='Label smoothing epsilon(标签平滑)')#早停机制，设置早停的epoch数parser.add_argument('--patience', type=int, default=100,help='EarlyStopping patience (epochs without improvement)')#指定冻结不进行训练的层索引parser.add_argument('--freeze', nargs='+', type=int, default=[0],help='Freeze layers: backbone=10, first3=0 1 2')#设置多少个epoch保存模型权重文件，保存路径是runs/train/weights/epochx.pt#该功能模型不开启,一般只保存last.pt和best.ptparser.add_argument('--save-period', type=int, default=-1,help='Save checkpoint every x epochs (disabled if < 1)')#本地进程排名(DDP模式用)parser.add_argument('--local_rank', type=int, default=-1,help='DDP parameter, do not modify(DDP模式的进程排名)')#---------------------------- W&B（Weights & Biases）参数配置 ----------------------------parser.add_argument('--entity', default=None,help='W&B: Entity (W&B 实体名称)')parser.add_argument('--upload_dataset', nargs='?', const=True, default=False,help='W&B: Upload dataset as artifact table (上传数据集到 W&B Artifact Table)')parser.add_argument('--bbox_interval', type=int, default=-1,help='W&B: Set bounding-box image logging interval (设置目标框日志记录间隔)')parser.add_argument('--artifact_alias', type=str, default='latest',help='W&B: Version of dataset artifact to use (使用的数据集版本别名)')#解析参数opt = parser.parse_known_args()[0] if known else parser.parse_args()return opt

2.2 加载模型与数据

加载模型权重和配置文件，设置模型参数，加载训练数据

分析：该部分代码属于yolov5结构中的哪个阶段？

主要发生在训练前的准备工作，即还没有进入模型的前向传播和反向传播阶段

运行逻辑分析

• 模型加载与构建
  • 如果提供了预训练权重，加载模型结构和权重参数
  • 如果没有提供权重，则根据模型结构配置文件yolov5xx.yaml构建新模型
• 冻结层设置
  • 通过命令行的freeze参数指定需要冻结哪些层(例如Backbone层)，以适应迁移学习或微调场景
• 训练数据准备
  • 加载数据集，创建数据迭代器Dataloader

    '-----------------------------模型部分---------------------------'#Model，模型加载check_suffix(weights, '.pt')  #检查权重文件的后缀是否为.ptpretrained = weights.endswith('.pt') #判断权重文件是否为以pt结尾，如果是的话则为true，说明在命令行指定了pt文件路径if pretrained:#如果本地找不到权重文件，则去YOLOv5的官方仓库中下载权重文件with torch_distributed_zero_first(LOCAL_RANK):weights = attempt_download(weights)#加载权重文件ckpt = torch.load(weights, map_location='cpu')#创建YOLOv5模型model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)'''cfg为模型结构的yaml说明书文件，pt权重文件里也是会保存模型结构的yaml说明书文件的，所以这两个任选一个即可ch：为输入通道数（一般为3） nc：数据集类别数 anchors：先验anchor尺寸'''exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else []#定义需要排除的键csd = ckpt['model'].float().state_dict()  # checkpoint state_dict as FP32csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) #交集，获取匹配的参数model.load_state_dict(csd, strict=False)  #给创建的YOLOv5模型加载权重参数LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}')  #输出转移的参数数量else:#如果在命令行的weights参数没有指定使用预训练权重，则使用给定的cfg创建新模型model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create'-----------------------------冻结层部分---------------------------'#Freeze，冻结一些层的参数，训练过程中不更新#保存需要冻结的层，eg:[model.0,model.1,model.6..]freeze = [f'model.{x}.' for x in (freeze if len(freeze)>1 else range(freeze[0]))]for k, v in model.named_parameters():v.requires_grad = True  #默认所有参数参与训练if any(x in k for x in freeze):#如果当前参数所在的层是否在冻结列表中LOGGER.info(f'freezing {k}')v.requires_grad = False#冻结参数，不进行训练'-----------------------------数据加载部分---------------------------'#Trainloader，加载训练集迭代器train_loader, dataset = create_dataloader(path=train_path,#训练集存放路径imgsz=imgsz,#自适应缩放图片大小batch_size=batch_size // WORLD_SIZE,#批量大小stride=gs,#最大下采样倍数single_cls=single_cls,hyp=hyp,#训练超参数存放路径augment=True,#数据增强cache=None if opt.cache == 'val' else opt.cache,rect=opt.rect,#是否使用矩形训练rank=LOCAL_RANK,workers=workers,#线程数image_weights=opt.image_weights,#是否使用加权的图像进行训练quad=opt.quad,#是否启用四元数据加载方式prefix=colorstr('train: '),shuffle=True#是否随机打乱)

2.3 优化器与学习率更新器设置

    '-------------------------权重衰退调整------------------------------'#这个部分看不太懂，应该是某种调整权重衰退的tricknbs = 64accumulate = max(round(nbs / batch_size), 1)  # accumulate loss before optimizinghyp['weight_decay'] *= batch_size * accumulate / nbs  # scale weight_decayLOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}")'-------------------------参数分组------------------------------'g0, g1, g2 = [], [], []#初始化3个数组'''将模型的参数分为三类：g0: BatchNorm 的权重g1: 卷积层或全连接层的权重g2: 偏置（bias）'''#遍历模型的每个模块for v in model.modules():#如果模块有偏置参数，并且是 nn.Parameter 类型，则将其加入 g2if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):g2.append(v.bias)#如果模块是 BatchNorm2d，则将其权重加入 g0if isinstance(v, nn.BatchNorm2d):g0.append(v.weight)#如果模块有权重参数，并且是 nn.Parameter 类型，则将其加入 g1elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):g1.append(v.weight)'-------------------------优化器设置------------------------------'#优化器选择if opt.optimizer == 'Adam':optimizer = Adam(g0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999))  # adjust beta1 to momentumelif opt.optimizer == 'AdamW':optimizer = AdamW(g0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999))  # adjust beta1 to momentumelse:optimizer = SGD(g0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)#为优化器添加参数组#添加g1参数组，并设置权重衰退weight_decay为超参数yaml文件中的值optimizer.add_param_group({'params': g1, 'weight_decay': hyp['weight_decay']})# 添加g2参数组（偏置），不使用权重衰减optimizer.add_param_group({'params': g2})#日志打印相关信息LOGGER.info(f"{colorstr('optimizer:')} {type(optimizer).__name__} with parameter groups "f"{len(g0)} weight (no decay), {len(g1)} weight, {len(g2)} bias")del g0, g1, g2'-------------------------学习率更新器设置------------------------------'#定义学习率调度器的变化方式（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# 创建学习率更新器，基于上述的学习率变化函数lfscheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)  # plot_lr_scheduler(optimizer, scheduler, epochs)

2.4 训练循环(核心)

训练流程

主要阶段总结

• Warmup热身阶段
• 多尺度训练：通过multi-scale参数可以开启模型的多尺度训练功能，即随机将输入图片自适应缩放之后的尺寸再作随机缩放或扩充，但必须确保仍然是最大下采样倍数(通常是32)的倍数
• 前向传播与损失计算
• 参数/学习率更新
• 日志记录
• 模型验证与保存：在每个epoch结束时模型在验证集上作验证，评估其在验证集上的性能指标，分别是(P，R，mAP@.5，mAP@.5-.95，val_loss(box，obj，cls))；并保存last.pt和best.pt
• 早停机制：通过patience参数可以指定早停的epoch数，训练过程中判断是否连续patience轮训练模型仍然没有提升，如果是则训练早停

    '-------------------------模型训练------------------------------'#Start training:开始训练t0 = time.time()#记录开始时间nw = max(round(hyp['warmup_epochs'] * nb), 100)  #number of warmup iterations, max(3 epochs, 100 iterations)last_opt_step = -1#最后一次优化步骤maps = np.zeros(nc)  #每个类别的map初始化results = (0, 0, 0, 0, 0, 0, 0)  # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)scheduler.last_epoch = start_epoch - 1#设置学习率更新器的最后epoch为当前epoch之前scaler = amp.GradScaler(enabled=cuda)#初始化混合精度训练的梯度缩放器stopper = EarlyStopping(patience=opt.patience)#初始化早停机制，设定耐心值compute_loss = ComputeLoss(model)#初始化损失计算类#记录训练信息callbacks.run('on_train_start')LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n'f"Logging results to {colorstr('bold', save_dir)}\n"f'Starting training for {epochs} epochs...')#开始训练循环for epoch in range(start_epoch, epochs):  # epoch ------------------------------------------------------------------callbacks.run('on_train_epoch_start')model.train()#可选：更新图像权重（仅适用于单GPU）if opt.image_weights:#cw为类别权重，model.class_weights表示每个类别的初始权重，对于样本较少的类别就赋予更高的权重cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc#利用类别权重cw和图像的标签信息，计算每张图片的采样权重iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw)#数据集中图像的顺序被重新排列，新的顺序基于权重iw，权重较高的图像(包含弱势类别)被采样的概率更高dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n)#初始化平均损失记录 (box_loss, obj_loss, cls_loss)mloss = torch.zeros(3, device=device)#如果是分布式训练if RANK != -1:train_loader.sampler.set_epoch(epoch)#设置当前epoch,确保分布式训练的数据加载一致#进度条设置pbar = enumerate(train_loader)#遍历训练集迭代器LOGGER.info(('\n' + '%10s' * 7) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'labels', 'img_size'))#日志打印训练信息if RANK in (-1, 0):#如果是主进程pbar = tqdm(pbar, total=nb, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}')#显示训练进度条optimizer.zero_grad()#优化器梯度清0for i, (imgs, targets, paths, _) in pbar:#遍历批次数据callbacks.run('on_train_batch_start')ni = i + nb * epoch#计算全局迭代步数imgs = imgs.to(device, non_blocking=True).float() / 255#将像素值归一化到 [0,1] 并移到设备上'------------------------Warmup热身阶段-------------------------------'if ni <= nw:#如果在热身阶段xi = [0, nw]#热身范围#动态调整累积步数（accumulate）和学习率accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())#计算累积步骤for j, x in enumerate(optimizer.param_groups):  #遍历优化器参数组#动态调整学习率：偏置学习率从0.1降到lr0，其他学习率从0.0升到lr0x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])#动态调整动量if 'momentum' in x:x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])'------------------------多尺度训练-------------------------------'if opt.multi_scale:#启用了多尺度训练，使模型对不同尺度的目标具有更强的鲁棒性sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs#随机生成新的训练图像尺寸，并且保证是gs的倍数sf = sz / max(imgs.shape[2:])#计算缩放因子if sf != 1:#如果需要缩放ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]#计算新的图像尺寸并调整（保证为gs的倍数）imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)#重新调整图像大小'------------------------前向传播-------------------------------'with amp.autocast(enabled=cuda): #混合精度训练pred = model(imgs)  #模型前向传播loss, loss_items = compute_loss(pred, targets.to(device))  #计算损失if RANK != -1:#如果是分布式训练loss *= WORLD_SIZE  #按照分布式规模调整损失if opt.quad:#如果使用四元数据加载方式，则损失乘以4loss *= 4.'------------------------反向传播-------------------------------'scaler.scale(loss).backward()#使用梯度缩放反向传播'------------------------参数更新-------------------------------'if ni - last_opt_step >= accumulate:  #如果达到累计步数条件scaler.step(optimizer)  #更新优化器参数scaler.update()  # 更新梯度缩放比例optimizer.zero_grad()  #梯度清0if ema:  # 如果启用EMA（指数移动平均）ema.update(model)  #更新模型的指数移动平均last_opt_step = ni  #更新最后一次优化步数'------------------------日志记录-------------------------------'if RANK in (-1, 0):#如果是主进程mloss = (mloss * i + loss_items) / (i + 1)#更新平均损失mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G'#GPU内存使用情况、#更新进度条显示内容pbar.set_description(('%10s' * 2 + '%10.4g' * 5) %(f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))callbacks.run('on_train_batch_end', ni, model, imgs, targets, paths, plots)if callbacks.stop_training:return'------------------------学习率更新-------------------------------'lr = [x['lr'] for x in optimizer.param_groups]#获取当前学习率scheduler.step()#更新学习率调度器'------------------------评估与保存-------------------------------'if RANK in (-1, 0):#如果是主进程callbacks.run('on_train_epoch_end', epoch=epoch)ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights'])#更新ema属性final_epoch = (epoch + 1 == epochs) or stopper.possible_stop#检查是否是最后一个epoch#计算mAPif not noval or final_epoch:#利用验证集验证模型，获取验证结果results:(P,R,mAP@.5,mAP@.5-.95,val_loss(box, obj, cls))results, maps, _ = val.run(data_dict,batch_size=batch_size // WORLD_SIZE * 2,#批量imgsz=imgsz,#自适应缩放图像尺寸model=ema.ema,#使用ema模型进行验证single_cls=single_cls,#是否为单类别dataloader=val_loader,#验证集迭代器save_dir=save_dir,#保存路径plots=False,#是否绘制图callbacks=callbacks,compute_loss=compute_loss)'更新模型最佳适应度'#results的存放结果(P,R,mAP@.5,mAP@.5-.95,val_loss(box, obj, cls))#调用fitness计算当前(P,R,mAP@.5,mAP@.5-.95)的最佳适应度fi = fitness(np.array(results).reshape(1, -1))if fi > best_fitness: #更新最佳适应度best_fitness = filog_vals = list(mloss) + list(results) + lr#记录日志值callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)'保存模型训练时的超参数和权重文件'if (not nosave) or (final_epoch and not evolve):#如果需要保存模型ckpt = {'epoch': epoch,#记录当前epoch'best_fitness': best_fitness,#记录最佳适应度'model': deepcopy(de_parallel(model)).half(),#深拷贝模型并转换为半精度'ema': deepcopy(ema.ema).half(),#深拷贝EMA模型并转换为半精度'updates': ema.updates,#EMA更新次数'optimizer': optimizer.state_dict(),#优化器的参数信息'wandb_id': loggers.wandb.wandb_run.id if loggers.wandb else None,'date': datetime.now().isoformat()}#保存日期#保存last.pt和best.pttorch.save(ckpt, last)#保存为最后一次权重文件if best_fitness == fi:#如果当前epoch的适应度最优torch.save(ckpt, best)#保存为最佳权重文件#按周期保存权重文件，这行代码默认是不执行的，YOLOv5默认只保存last.pt和best.ptif (epoch > 0) and (opt.save_period > 0) and (epoch % opt.save_period == 0):torch.save(ckpt, w / f'epoch{epoch}.pt')del ckptcallbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi)'是否早停'#利用当前epoch算出来的适应度fi判断是否进行早停(仅在单GPU下)if RANK == -1 and stopper(epoch=epoch, fitness=fi):break# end epoch ----------------------------------------------------------------------------------------------------# end training -----------------------------------------------------------------------------------------------------

2.5 最终的模型验证

模型训练结束后，利用训练过程中保存的best.pt在验证集上作最后一次验证，控制台输出验证结果

    '---------------------------最终的模型验证---------------------------'#如果是主进程，模型训练结束时，下面代码将利用最佳模型在验证集计算性能指标显示在控制台上if RANK in (-1, 0):# 记录已完成的周期和耗时LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.')#对于最后一个和最佳模型进行处理for f in last, best:if f.exists():#如果文件存在strip_optimizer(f)  #去除优化器状态以减小模型文件大小if f is best:#如果是最佳模型LOGGER.info(f'\nValidating {f}...')#记录验证信息#验证模型，得到验证结果results:(P, R, mAP @ .5, mAP @ .5 - .95,box_loss,obj_loss,cls_loss)results, _, _ = val.run(data_dict,batch_size=batch_size // WORLD_SIZE * 2,#批量imgsz=imgsz,#自适应图片缩放尺寸model=attempt_load(f, device).half(),# 加载模型权重文件并转换为半精度iou_thres=0.65 if is_coco else 0.60,  #设置IoU阈值single_cls=single_cls,  #是否为单类别dataloader=val_loader,  #验证集迭代器save_dir=save_dir,      #保存路径save_json=is_coco,      #是否保存将验证结果保存为JSON文件verbose=True,           #是否输出详细信息plots=plots,            #是否绘制图callbacks=callbacks,compute_loss=compute_loss)  #计算损失if is_coco:callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)callbacks.run('on_train_end', last, best, plots, epoch, results)#运行训练结束的回调

3. 完整train.py代码

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
"""
Train a YOLOv5 model on a custom dataset.Models and datasets download automatically from the latest YOLOv5 release.
Models: https://github.com/ultralytics/yolov5/tree/master/models
Datasets: https://github.com/ultralytics/yolov5/tree/master/data
Tutorial: https://github.com/ultralytics/yolov5/wiki/Train-Custom-DataUsage:$ python path/to/train.py --data coco128.yaml --weights yolov5s.pt --img 640  # from pretrained (RECOMMENDED)$ python path/to/train.py --data coco128.yaml --weights '' --cfg yolov5s.yaml --img 640  # from scratch
"""import argparse
import math
import os
import random
import sys
import time
from copy import deepcopy
from datetime import datetime
from pathlib import Pathimport numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
import yaml
from torch.cuda import amp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim import SGD, Adam, AdamW, lr_scheduler
from tqdm import tqdmFILE = Path(__file__).resolve()#获取当前py文件存放路径，eg:WindowsPath('D:/BaiduNetdiskDownload/yolov5/train.py')
ROOT = FILE.parents[0]#获取根目录，即WindowsPath('D:/BaiduNetdiskDownload/yolov5')
if str(ROOT) not in sys.path:#将路径存放在sys.path中sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  #绝对路径转换成相对路径import val  # for end-of-epoch mAP
from models.experimental import attempt_load
from models.yolo import Model
from utils.autoanchor import check_anchors
from utils.autobatch import check_train_batch_size
from utils.callbacks import Callbacks
from utils.datasets import create_dataloader
from utils.downloads import attempt_download
from utils.general import (LOGGER, check_dataset, check_file, check_git_status, check_img_size, check_requirements,check_suffix, check_yaml, colorstr, get_latest_run, increment_path, init_seeds,intersect_dicts, is_ascii, labels_to_class_weights, labels_to_image_weights, methods,one_cycle, print_args, print_mutation, strip_optimizer)
from utils.loggers import Loggers
from utils.loggers.wandb.wandb_utils import check_wandb_resume
from utils.loss import ComputeLoss
from utils.metrics import fitness
from utils.plots import plot_evolve, plot_labels
from utils.torch_utils import EarlyStopping, ModelEMA, de_parallel, select_device, torch_distributed_zero_firstLOCAL_RANK = int(os.getenv('LOCAL_RANK', -1))  # https://pytorch.org/docs/stable/elastic/run.html
RANK = int(os.getenv('RANK', -1))#如果是单卡训练，则RANK=-1
WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))#训练函数
def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictionary'''参数解释：:param hyp:训练时的超参数yaml配置文件:param opt:命令行超参数:param device:设备:param callbacks:---'''#设置训练相关的目录和参数save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = \Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \opt.resume, opt.noval, opt.nosave, opt.workers, opt.freezecallbacks.run('on_pretrain_routine_start')#Directories，创建保存模型权重的目录，一般是runs/train/expn/weightsw = save_dir / 'weights'# 如果需要演化，则创建父目录，否则只创建权重目录(w.parent if evolve else w).mkdir(parents=True, exist_ok=True)  # make dir#定义最后和最好的模型文件存放路径last, best = w / 'last.pt', w / 'best.pt'#Hyperparameters，加载yaml文件中的训练超参数if isinstance(hyp, str):with open(hyp, errors='ignore') as f:hyp = yaml.safe_load(f)  #从YAML文件中加载超参数字典，以key:value的形式存放超参数#日志打印训练超参数信息LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))#Save run settings，保存训练过程中的超参数和命令行参数，分别是hyp.yaml、opt.yaml，存放路径一般是runs/train/expnif not evolve:with open(save_dir / 'hyp.yaml', 'w') as f:yaml.safe_dump(hyp, f, sort_keys=False)#保存超参数到YAML文件with open(save_dir / 'opt.yaml', 'w') as f:yaml.safe_dump(vars(opt), f, sort_keys=False)#保存训练的选项到YAML文件#Loggers：初始化日志记录器data_dict = None#初始化数据字典if RANK in [-1, 0]:loggers = Loggers(save_dir, weights, opt, hyp, LOGGER)  # loggers instanceif loggers.wandb:data_dict = loggers.wandb.data_dictif resume:weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size# Register actionsfor k in methods(loggers):callbacks.register_action(k, callback=getattr(loggers, k))#Config，配置plots = not evolve and not opt.noplots#是否创建绘图，演示模式下不创建cuda = device.type != 'cpu'#检查是否使用GPUinit_seeds(1 + RANK)#初始化随机种子，确保每个进程的种子不同#在分布式训练的主进程中执行以下操作with torch_distributed_zero_first(LOCAL_RANK):'''check_dataset会根据yaml中配置的path、train、val、test路径检查是否存在数据集，如果不存在则会根据yaml中的download地址去自动帮我们下载数据集，并按设置的path、train、val路径存放数据集'''data_dict = data_dict or check_dataset(data) #检查数据集，如果数据字典为None，则通过传入的data参数从yaml文件中加载数据字典#获取训练集、验证集存放路径train_path, val_path = data_dict['train'], data_dict['val']#获取类别数量，单类别情况下数量为1nc = 1 if single_cls else int(data_dict['nc'])#获取类别名称，如果是单类别且名称列表长度不为1，则设为 ['item']names = ['item'] if single_cls and len(data_dict['names']) != 1 else data_dict['names']#检查类别名称的数量是否与nc匹配assert len(names) == nc, f'{len(names)} names found for nc={nc} dataset in {data}'  # check#检查验证集集是否为COCO数据集is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt')#Model，模型加载check_suffix(weights, '.pt')  #检查权重文件的后缀是否为.ptpretrained = weights.endswith('.pt') #判断权重文件是否为以pt结尾，如果是的话则为true，说明在命令行指定了pt文件路径if pretrained:#如果本地找不到权重文件，则去YOLOv5的官方仓库中下载权重文件with torch_distributed_zero_first(LOCAL_RANK):weights = attempt_download(weights)#加载权重文件ckpt = torch.load(weights, map_location='cpu')#创建YOLOv5模型model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)'''cfg为模型结构的yaml说明书文件，pt权重文件里也是会保存模型结构的yaml说明书文件的，所以这两个任选一个即可ch：为输入通道数（一般为3） nc：数据集类别数 anchors：先验anchor尺寸'''exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else []#定义需要排除的键csd = ckpt['model'].float().state_dict()  # checkpoint state_dict as FP32csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) #交集，获取匹配的参数model.load_state_dict(csd, strict=False)  #给创建的YOLOv5模型加载权重参数LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}')  #输出转移的参数数量else:#如果在命令行的weights参数没有指定使用预训练权重，则使用给定的cfg创建新模型model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create#Freeze，冻结一些层的参数，训练过程中不更新#保存需要冻结的层，eg:[model.0,model.1,model.6..]freeze = [f'model.{x}.' for x in (freeze if len(freeze)>1 else range(freeze[0]))]for k, v in model.named_parameters():v.requires_grad = True  #默认所有参数参与训练if any(x in k for x in freeze):#如果当前参数所在的层是否在冻结列表中LOGGER.info(f'freezing {k}')v.requires_grad = False#冻结参数，不进行训练#Image size，图片大小处理gs = max(int(model.stride.max()), 32)#获取模型的最大步幅作为下采样倍数，确保最大为32imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2)'验证在命令行指定的自适应缩放图像大小是否是32的倍数，且不小于gs的两倍，如果不是则进行调整'#Batch size，如果我们在命令行指定batch_size=-1，则会执行下面代码，#根据GPU情况自动选择一个最好的batch_sizeif RANK == -1 and batch_size == -1:batch_size = check_train_batch_size(model, imgsz)loggers.on_params_update({"batch_size": batch_size})'-------------------------权重衰退调整------------------------------'#这个部分看不太懂，应该是某种调整权重衰退的tricknbs = 64accumulate = max(round(nbs / batch_size), 1)  # accumulate loss before optimizinghyp['weight_decay'] *= batch_size * accumulate / nbs  # scale weight_decayLOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}")'-------------------------参数分组------------------------------'g0, g1, g2 = [], [], []#初始化3个数组'''将模型的参数分为三类：g0: BatchNorm 的权重g1: 卷积层或全连接层的权重g2: 偏置（bias）'''#遍历模型的每个模块for v in model.modules():#如果模块有偏置参数，并且是 nn.Parameter 类型，则将其加入 g2if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):g2.append(v.bias)#如果模块是 BatchNorm2d，则将其权重加入 g0if isinstance(v, nn.BatchNorm2d):g0.append(v.weight)#如果模块有权重参数，并且是 nn.Parameter 类型，则将其加入 g1elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):g1.append(v.weight)'-------------------------优化器设置------------------------------'#优化器选择if opt.optimizer == 'Adam':optimizer = Adam(g0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999))  # adjust beta1 to momentumelif opt.optimizer == 'AdamW':optimizer = AdamW(g0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999))  # adjust beta1 to momentumelse:optimizer = SGD(g0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)#为优化器添加参数组#添加g1参数组，并设置权重衰退weight_decay为超参数yaml文件中的值optimizer.add_param_group({'params': g1, 'weight_decay': hyp['weight_decay']})# 添加g2参数组（偏置），不使用权重衰减optimizer.add_param_group({'params': g2})#日志打印相关信息LOGGER.info(f"{colorstr('optimizer:')} {type(optimizer).__name__} with parameter groups "f"{len(g0)} weight (no decay), {len(g1)} weight, {len(g2)} bias")del g0, g1, g2'-------------------------学习率更新器设置------------------------------'#定义学习率调度器的变化方式（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# 创建学习率更新器，基于上述的学习率变化函数lfscheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)  # plot_lr_scheduler(optimizer, scheduler, epochs)#EMAema = ModelEMA(model) if RANK in [-1, 0] else None#初始化开始轮次和最佳适应度start_epoch, best_fitness = 0, 0.0#如果指定了预训练权重文件if pretrained:#Optimizerif ckpt['optimizer'] is not None:#检查权重文件中是否包含优化器optimizer.load_state_dict(ckpt['optimizer'])#加载优化器best_fitness = ckpt['best_fitness']#更新最佳适应度#EMAif ema and ckpt.get('ema'):ema.ema.load_state_dict(ckpt['ema'].float().state_dict())ema.updates = ckpt['updates']#Epochsstart_epoch = ckpt['epoch'] + 1    #设置训练开始的epoch#如果选择了断点续训if resume:assert start_epoch > 0, f'{weights} training to {epochs} epochs is finished, nothing to resume.'#如果训练次数还未到达命令行设置的epoch数，则继续接着训练if epochs < start_epoch:LOGGER.info(f"{weights} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {epochs} more epochs.")epochs += ckpt['epoch']  # finetune additional epochsdel ckpt, csd#DP modeif cuda and RANK == -1 and torch.cuda.device_count() > 1:#如果在多 GPU 环境下且未使用分布式训练，发出警告LOGGER.warning('WARNING: DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n''See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started.')model = torch.nn.DataParallel(model)#SyncBatchNormif opt.sync_bn and cuda and RANK != -1:#如果启用了同步批归一化且处于分布式训练模式，转换模型为同步批归一化model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)LOGGER.info('Using SyncBatchNorm()')#Trainloader，加载训练集迭代器train_loader, dataset = create_dataloader(path=train_path,#训练集存放路径imgsz=imgsz,#自适应缩放图片大小batch_size=batch_size // WORLD_SIZE,#批量大小stride=gs,#最大下采样倍数single_cls=single_cls,hyp=hyp,#训练超参数存放路径augment=True,#数据增强cache=None if opt.cache == 'val' else opt.cache,rect=opt.rect,#是否使用矩形训练rank=LOCAL_RANK,workers=workers,#线程数image_weights=opt.image_weights,#是否使用加权的图像进行训练quad=opt.quad,#是否启用四元数据加载方式prefix=colorstr('train: '),shuffle=True#是否随机打乱)#从标签中找出最大的类别索引(用来检查最大类别索引是否超过类别总数)mlc = int(np.concatenate(dataset.labels, 0)[:, 0].max())#计算批次的数量(num of batches)nb = len(train_loader)#检查最大类别标签索引是否超过类别总数assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'#处理进程0if RANK in [-1, 0]:#加载验证集的迭代器val_loader = create_dataloader(val_path,imgsz,batch_size // WORLD_SIZE * 2,gs,single_cls,hyp=hyp,cache=None if noval else opt.cache,rect=True,rank=-1,workers=workers * 2,pad=0.5,prefix=colorstr('val: '))[0]#如果不是断点续训if not resume:labels = np.concatenate(dataset.labels, 0)#合并所有标签#是否绘制图if plots:#绘制训练集的标签分布图labels.jpg和labels_correlogram.jpgplot_labels(labels, names, save_dir)#是否作自适应锚框计算,根据训练集计算出更加适用于本任务的先验anchor尺寸if not opt.noautoanchor:check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)model.half().float()#预先减少锚框精度callbacks.run('on_pretrain_routine_end')#DDP modeif cuda and RANK != -1:model = DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)#关于损失函数的超参数调整设置，使其更适应当前任务nl = de_parallel(model).model[-1].nl#nl:检测头数量，通常为3hyp['box'] *= 3 / nl#定位损失项的权重系数hyp['cls'] *= nc / 80 * 3 / nl#类别损失项的权重系数hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl#置信度损失项的权重系数hyp['label_smoothing'] = opt.label_smoothing#设置标签平滑参数#将类别数量附加到模型model.nc = nc#将训练超参数文件存放路径附加到模型model.hyp = hyp#计算并附加类别权重到模型model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc#将类别名称附加到模型model.names = names'-------------------------模型训练------------------------------'#Start training:开始训练t0 = time.time()#记录开始时间nw = max(round(hyp['warmup_epochs'] * nb), 100)  #number of warmup iterations, max(3 epochs, 100 iterations)last_opt_step = -1#最后一次优化步骤maps = np.zeros(nc)  #每个类别的map初始化results = (0, 0, 0, 0, 0, 0, 0)  # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)scheduler.last_epoch = start_epoch - 1#设置学习率更新器的最后epoch为当前epoch之前scaler = amp.GradScaler(enabled=cuda)#初始化混合精度训练的梯度缩放器stopper = EarlyStopping(patience=opt.patience)#初始化早停机制，设定耐心值compute_loss = ComputeLoss(model)#初始化损失计算类#记录训练信息callbacks.run('on_train_start')LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n'f"Logging results to {colorstr('bold', save_dir)}\n"f'Starting training for {epochs} epochs...')#开始训练循环for epoch in range(start_epoch, epochs):  # epoch ------------------------------------------------------------------callbacks.run('on_train_epoch_start')model.train()#可选：更新图像权重（仅适用于单GPU）if opt.image_weights:#cw为类别权重，model.class_weights表示每个类别的初始权重，对于样本较少的类别就赋予更高的权重cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc#利用类别权重cw和图像的标签信息，计算每张图片的采样权重iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw)#数据集中图像的顺序被重新排列，新的顺序基于权重iw，权重较高的图像(包含弱势类别)被采样的概率更高dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n)#初始化平均损失记录 (box_loss, obj_loss, cls_loss)mloss = torch.zeros(3, device=device)#如果是分布式训练if RANK != -1:train_loader.sampler.set_epoch(epoch)#设置当前epoch,确保分布式训练的数据加载一致#进度条设置pbar = enumerate(train_loader)#遍历训练集迭代器LOGGER.info(('\n' + '%10s' * 7) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'labels', 'img_size'))#日志打印训练信息if RANK in (-1, 0):#如果是主进程pbar = tqdm(pbar, total=nb, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}')#显示训练进度条optimizer.zero_grad()#优化器梯度清0for i, (imgs, targets, paths, _) in pbar:#遍历批次数据callbacks.run('on_train_batch_start')ni = i + nb * epoch#计算全局迭代步数imgs = imgs.to(device, non_blocking=True).float() / 255#将像素值归一化到 [0,1] 并移到设备上'------------------------Warmup热身阶段-------------------------------'if ni <= nw:#如果在热身阶段xi = [0, nw]#热身范围#动态调整累积步数（accumulate）和学习率accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())#计算累积步骤for j, x in enumerate(optimizer.param_groups):  #遍历优化器参数组#动态调整学习率：偏置学习率从0.1降到lr0，其他学习率从0.0升到lr0x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])#动态调整动量if 'momentum' in x:x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])'------------------------多尺度训练-------------------------------'if opt.multi_scale:#启用了多尺度训练，使模型对不同尺度的目标具有更强的鲁棒性sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs#随机生成新的训练图像尺寸，并且保证是gs的倍数sf = sz / max(imgs.shape[2:])#计算缩放因子if sf != 1:#如果需要缩放ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]#计算新的图像尺寸并调整（保证为gs的倍数）imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)#重新调整图像大小'------------------------前向传播-------------------------------'with amp.autocast(enabled=cuda): #混合精度训练pred = model(imgs)  #模型前向传播loss, loss_items = compute_loss(pred, targets.to(device))  #计算损失if RANK != -1:#如果是分布式训练loss *= WORLD_SIZE  #按照分布式规模调整损失if opt.quad:#如果使用四元数据加载方式，则损失乘以4loss *= 4.'------------------------反向传播-------------------------------'scaler.scale(loss).backward()#使用梯度缩放反向传播'------------------------参数更新-------------------------------'if ni - last_opt_step >= accumulate:  #如果达到累计步数条件scaler.step(optimizer)  #更新优化器参数scaler.update()  # 更新梯度缩放比例optimizer.zero_grad()  #梯度清0if ema:  # 如果启用EMA（指数移动平均）ema.update(model)  #更新模型的指数移动平均last_opt_step = ni  #更新最后一次优化步数'------------------------日志记录-------------------------------'if RANK in (-1, 0):#如果是主进程mloss = (mloss * i + loss_items) / (i + 1)#更新平均损失mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G'#GPU内存使用情况、#更新进度条显示内容pbar.set_description(('%10s' * 2 + '%10.4g' * 5) %(f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))callbacks.run('on_train_batch_end', ni, model, imgs, targets, paths, plots)if callbacks.stop_training:return'------------------------学习率更新-------------------------------'lr = [x['lr'] for x in optimizer.param_groups]#获取当前学习率scheduler.step()#更新学习率调度器'------------------------评估与保存-------------------------------'if RANK in (-1, 0):#如果是主进程callbacks.run('on_train_epoch_end', epoch=epoch)ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights'])#更新ema属性final_epoch = (epoch + 1 == epochs) or stopper.possible_stop#检查是否是最后一个epoch#计算mAPif not noval or final_epoch:#利用验证集验证模型，获取验证结果results:(P,R,mAP@.5,mAP@.5-.95,val_loss(box, obj, cls))results, maps, _ = val.run(data_dict,batch_size=batch_size // WORLD_SIZE * 2,#批量imgsz=imgsz,#自适应缩放图像尺寸model=ema.ema,#使用ema模型进行验证single_cls=single_cls,#是否为单类别dataloader=val_loader,#验证集迭代器save_dir=save_dir,#保存路径plots=False,#是否绘制图callbacks=callbacks,compute_loss=compute_loss)'更新模型最佳适应度'#results的存放结果(P,R,mAP@.5,mAP@.5-.95,val_loss(box, obj, cls))#调用fitness计算当前(P,R,mAP@.5,mAP@.5-.95)的最佳适应度fi = fitness(np.array(results).reshape(1, -1))if fi > best_fitness: #更新最佳适应度best_fitness = filog_vals = list(mloss) + list(results) + lr#记录日志值callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)'保存模型训练时的超参数和权重文件'if (not nosave) or (final_epoch and not evolve):#如果需要保存模型ckpt = {'epoch': epoch,#记录当前epoch'best_fitness': best_fitness,#记录最佳适应度'model': deepcopy(de_parallel(model)).half(),#深拷贝模型并转换为半精度'ema': deepcopy(ema.ema).half(),#深拷贝EMA模型并转换为半精度'updates': ema.updates,#EMA更新次数'optimizer': optimizer.state_dict(),#优化器的参数信息'wandb_id': loggers.wandb.wandb_run.id if loggers.wandb else None,'date': datetime.now().isoformat()}#保存日期#保存last.pt和best.pttorch.save(ckpt, last)#保存为最后一次权重文件if best_fitness == fi:#如果当前epoch的适应度最优torch.save(ckpt, best)#保存为最佳权重文件#按周期保存权重文件，这行代码默认是不执行的，YOLOv5默认只保存last.pt和best.ptif (epoch > 0) and (opt.save_period > 0) and (epoch % opt.save_period == 0):torch.save(ckpt, w / f'epoch{epoch}.pt')del ckptcallbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi)'是否早停'#利用当前epoch算出来的适应度fi判断是否进行早停(仅在单GPU下)if RANK == -1 and stopper(epoch=epoch, fitness=fi):break# end epoch ----------------------------------------------------------------------------------------------------# end training -----------------------------------------------------------------------------------------------------'---------------------------最终的模型验证---------------------------'#如果是主进程，模型训练结束时，下面代码将利用最佳模型在验证集计算性能指标显示在控制台上if RANK in (-1, 0):# 记录已完成的周期和耗时LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.')#对于最后一个和最佳模型进行处理for f in last, best:if f.exists():#如果文件存在strip_optimizer(f)  #去除优化器状态以减小模型文件大小if f is best:#如果是最佳模型LOGGER.info(f'\nValidating {f}...')#记录验证信息#验证模型，得到验证结果results:(P, R, mAP @ .5, mAP @ .5 - .95,box_loss,obj_loss,cls_loss)results, _, _ = val.run(data_dict,batch_size=batch_size // WORLD_SIZE * 2,#批量imgsz=imgsz,#自适应图片缩放尺寸model=attempt_load(f, device).half(),# 加载模型权重文件并转换为半精度iou_thres=0.65 if is_coco else 0.60,  #设置IoU阈值single_cls=single_cls,  #是否为单类别dataloader=val_loader,  #验证集迭代器save_dir=save_dir,      #保存路径save_json=is_coco,      #是否保存将验证结果保存为JSON文件verbose=True,           #是否输出详细信息plots=plots,            #是否绘制图callbacks=callbacks,compute_loss=compute_loss)  #计算损失if is_coco:callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)callbacks.run('on_train_end', last, best, plots, epoch, results)#运行训练结束的回调torch.cuda.empty_cache()#清空CUDA缓存#返回最佳模型在验证集上的results:(P, R, mAP @ .5, mAP @ .5 - .95,box_loss,obj_loss,cls_loss)return resultsdef parse_opt(known=False):parser = argparse.ArgumentParser()#权重文件路径parser.add_argument('--weights', type=str, default=ROOT /'weights/yolov5s.pt',help='initial weights path (初始权重文件路径)')#模型结构yaml配置文件路径parser.add_argument('--cfg', type=str, default='',help='model.yaml path (模型结构配置文件路径)')#数据集yaml配置文件路径parser.add_argument('--data', type=str, default=ROOT / 'data/VOC-hat.yaml',help='dataset.yaml path (数据集配置文件路径)')#训练超参数yaml配置文件路径parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml',help='hyperparameters path (超参数配置文件路径)')#训练轮数parser.add_argument('--epochs', type=int, default=300)#批量大小parser.add_argument('--batch-size', type=int, default=16,help='total batch size for all GPUs, -1 for autobatch')#imgsz指定训练和验证时将输入图片自适应缩放到相应的尺寸parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640,help='指定训练、验证时将输入图片自适应缩放到相应的尺寸')#是否使用矩形训练parser.add_argument('--rect', action='store_true',help='rectangular training(是否使用矩形训练)')'''断点续训，有时候服务器崩溃了导致训练中断，就可以利用这个参数继续训练，训练过程保存的pth存放在runs/train/expxx/weights下，分别是last.pth和best.pth，使用该参数时，指定--resume runs/train/expxx/weights/last.pth即可继续上次中断的训练'''parser.add_argument('--resume', nargs='?', const=True, default=False,help='resume most recent training(断点续训)')'''该参数指定到了最后一个epoch才保存模型的last.pt和best.pt，如果你只想获取最后一个epoch的模型权重文件，那可以开启这个功能，默认是不开启的'''parser.add_argument('--nosave', action='store_true',help='only save final checkpoint(只保存最后一个epoch的last.pt和best.pt)')#仅验证最终epoch,该参数指定到了最后一个epoch才去计算模型在验证集上的性能指标,默认不开启parser.add_argument('--noval', action='store_true',help='only validate final epoch(仅验证最终epoch)')#是否禁用自适应锚框策略，该功能默认开启parser.add_argument('--noautoanchor', action='store_true',help='disable AutoAnchor')#不保存训练中生成的图标文件，默认是保存的parser.add_argument('--noplots', action='store_true',help='save no plot files')#使用遗传算法优化超参数，可指定优化代数，默认不开启该功能parser.add_argument('--evolve', type=int, nargs='?', const=300,help='evolve hyperparameters for x generations')##谷歌云盘bucket，一般用不到parser.add_argument('--bucket', type=str, default='',help='gsutil bucket')#是否缓存数据集到RAM或磁盘parser.add_argument('--cache', type=str, nargs='?', const='ram',help='--cache images in "ram" (default) or "disk"(缓存数据集)')#是否使用加权的图像进行训练,对于那些训练不好的图片，会在下一个epoch中增加一些权重parser.add_argument('--image-weights', action='store_true',help='use weighted image selection for training')#指定训练的设备parser.add_argument('--device', default='',help='cuda device, i.e. 0 or 0,1,2,3 or cpu')#是否使用多尺度训练(即随机将自适应缩放之后的尺寸再作随机缩放或增加，但必须确保是最大下采样倍数(通常是32)的倍数)#该功能使模型对不同尺度的目标具有更强的鲁棒性parser.add_argument('--multi-scale', action='store_true',help='vary img-size +/- 50%%(多尺度训练)')#不知有什么用，默认不开启parser.add_argument('--single-cls', action='store_true',help='train multi-class data as single-class')#优化器，默认是SGDparser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'],default='SGD', help='optimizer')#是否启用同步BatchNorm(只在多GPU中使用)，默认不启用parser.add_argument('--sync-bn', action='store_true',help='use SyncBatchNorm, only available in DDP mode')#数据加载器的最大工作线程数parser.add_argument('--workers', type=int, default=8,help='max dataloader workers (per RANK in DDP mode)')#训练结果存放的根目录，默认是runs/trianparser.add_argument('--project', default=ROOT / 'runs/train',help='save to project/name(根目录)')#训练结果存放的子目录，默认是runs/train/expnparser.add_argument('--name', default='exp',help='save to project/name(子目录)')#是否用当前的训练目录覆盖以前的expn训练目录,默认不开启parser.add_argument('--exist-ok', action='store_true',help='existing project/name ok, do not increment(允许覆盖以前的训练结果)')#是否使用四元数据加载器,默认不开启parser.add_argument('--quad', action='store_true', help='quad dataloader')#是否使用余弦函数更新学习率,默认是线性函数更新学习率parser.add_argument('--cos-lr', action='store_true',help='cosine LR scheduler')#是否启用标签平滑parser.add_argument('--label-smoothing', type=float, default=0.0,help='Label smoothing epsilon(标签平滑)')#早停机制，设置早停的epoch数parser.add_argument('--patience', type=int, default=100,help='EarlyStopping patience (epochs without improvement)')#指定冻结不进行训练的层索引parser.add_argument('--freeze', nargs='+', type=int, default=[0],help='Freeze layers: backbone=10, first3=0 1 2')#设置多少个epoch保存模型权重文件，保存路径是runs/train/weights/epochx.pt#该功能模型不开启,一般只保存last.pt和best.ptparser.add_argument('--save-period', type=int, default=-1,help='Save checkpoint every x epochs (disabled if < 1)')#本地进程排名(DDP模式用)parser.add_argument('--local_rank', type=int, default=-1,help='DDP parameter, do not modify(DDP模式的进程排名)')#---------------------------- W&B（Weights & Biases）参数配置 ----------------------------parser.add_argument('--entity', default=None,help='W&B: Entity (W&B 实体名称)')parser.add_argument('--upload_dataset', nargs='?', const=True, default=False,help='W&B: Upload dataset as artifact table (上传数据集到 W&B Artifact Table)')parser.add_argument('--bbox_interval', type=int, default=-1,help='W&B: Set bounding-box image logging interval (设置目标框日志记录间隔)')parser.add_argument('--artifact_alias', type=str, default='latest',help='W&B: Version of dataset artifact to use (使用的数据集版本别名)')#解析参数opt = parser.parse_known_args()[0] if known else parser.parse_args()return optdef main(opt, callbacks=Callbacks()):#Checks，主进程检查if RANK in (-1, 0):print_args(vars(opt))#打印所有参数信息check_git_status()#检查Git仓库状态（确保代码是最新版本）#check_requirements(exclude=['thop'])检查有没有按照requirements.txt文件按照依赖包，排除'thop'包#Resume，断点续训if opt.resume and not check_wandb_resume(opt) and not opt.evolve:   #检查是否从中断位置恢复ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run()  #获取指定的pt文件最近的检查点路径assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist'   #检查pt文件是否存在#从指定的pt文件目录下的opt.yaml中加载训练参数设置with open(Path(ckpt).parent.parent / 'opt.yaml', errors='ignore') as f:opt = argparse.Namespace(**yaml.safe_load(f))  #加载训练配置到opt变量opt.cfg, opt.weights, opt.resume = '', ckpt, True  #设置pt权重文件路径LOGGER.info(f'Resuming training from {ckpt}')#如果不是断点续训else:#校验数据集yaml配置文件、模型结构yaml配置文件、超参数yaml配置文件、模型权重pt文件，并校验路径opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = \check_file(opt.data), check_yaml(opt.cfg), check_yaml(opt.hyp), str(opt.weights), str(opt.project)#确保cfg或weights参数至少一个已经由用户指定，否则报错assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'#如果用户选择进行超参数优化，将保存的根目录改为runs/evolveif opt.evolve:if opt.project == str(ROOT / 'runs/train'):  # if default project name, rename to runs/evolveopt.project = str(ROOT / 'runs/evolve')opt.exist_ok, opt.resume = opt.resume, False  # pass resume to exist_ok and disable resumeif opt.name == 'cfg':opt.name = Path(opt.cfg).stem  # use model.yaml as name#设置训练结果保存目录，一般是runs/train/expnopt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))#DDP Mode（分布式数据并行），这里一般也用不到，大多数都是单卡训练device = select_device(opt.device, batch_size=opt.batch_size)if LOCAL_RANK != -1:msg = 'is not compatible with YOLOv5 Multi-GPU DDP training'assert not opt.image_weights, f'--image-weights {msg}'assert not opt.evolve, f'--evolve {msg}'assert opt.batch_size != -1, f'AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size'assert opt.batch_size % WORLD_SIZE == 0, f'--batch-size {opt.batch_size} must be multiple of WORLD_SIZE'assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'torch.cuda.set_device(LOCAL_RANK)device = torch.device('cuda', LOCAL_RANK)dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")#Train，开始训练模型if not opt.evolve:#如果未选择超参数进化，进行正常的模型训练train(opt.hyp, opt, device, callbacks) #调用train函数进行训练if WORLD_SIZE > 1 and RANK == 0:#判断是否是多卡训练，如果是多GPU模式训练，需要销毁进程组LOGGER.info('Destroying process group... ')dist.destroy_process_group()    #销毁DDP进程组#进行超参数演化（可选功能）else:# Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)meta = {'lr0': (1, 1e-5, 1e-1),  # initial learning rate (SGD=1E-2, Adam=1E-3)'lrf': (1, 0.01, 1.0),  # final OneCycleLR learning rate (lr0 * lrf)'momentum': (0.3, 0.6, 0.98),  # SGD momentum/Adam beta1'weight_decay': (1, 0.0, 0.001),  # optimizer weight decay'warmup_epochs': (1, 0.0, 5.0),  # warmup epochs (fractions ok)'warmup_momentum': (1, 0.0, 0.95),  # warmup initial momentum'warmup_bias_lr': (1, 0.0, 0.2),  # warmup initial bias lr'box': (1, 0.02, 0.2),  # box loss gain'cls': (1, 0.2, 4.0),  # cls loss gain'cls_pw': (1, 0.5, 2.0),  # cls BCELoss positive_weight'obj': (1, 0.2, 4.0),  # obj loss gain (scale with pixels)'obj_pw': (1, 0.5, 2.0),  # obj BCELoss positive_weight'iou_t': (0, 0.1, 0.7),  # IoU training threshold'anchor_t': (1, 2.0, 8.0),  # anchor-multiple threshold'anchors': (2, 2.0, 10.0),  # anchors per output grid (0 to ignore)'fl_gamma': (0, 0.0, 2.0),  # focal loss gamma (efficientDet default gamma=1.5)'hsv_h': (1, 0.0, 0.1),  # image HSV-Hue augmentation (fraction)'hsv_s': (1, 0.0, 0.9),  # image HSV-Saturation augmentation (fraction)'hsv_v': (1, 0.0, 0.9),  # image HSV-Value augmentation (fraction)'degrees': (1, 0.0, 45.0),  # image rotation (+/- deg)'translate': (1, 0.0, 0.9),  # image translation (+/- fraction)'scale': (1, 0.0, 0.9),  # image scale (+/- gain)'shear': (1, 0.0, 10.0),  # image shear (+/- deg)'perspective': (0, 0.0, 0.001),  # image perspective (+/- fraction), range 0-0.001'flipud': (1, 0.0, 1.0),  # image flip up-down (probability)'fliplr': (0, 0.0, 1.0),  # image flip left-right (probability)'mosaic': (1, 0.0, 1.0),  # image mixup (probability)'mixup': (1, 0.0, 1.0),  # image mixup (probability)'copy_paste': (1, 0.0, 1.0)}  # segment copy-paste (probability)with open(opt.hyp, errors='ignore') as f:hyp = yaml.safe_load(f)  # load hyps dictif 'anchors' not in hyp:  # anchors commented in hyp.yamlhyp['anchors'] = 3opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir)  # only val/save final epoch# ei = [isinstance(x, (int, float)) for x in hyp.values()]  # evolvable indicesevolve_yaml, evolve_csv = save_dir / 'hyp_evolve.yaml', save_dir / 'evolve.csv'if opt.bucket:os.system(f'gsutil cp gs://{opt.bucket}/evolve.csv {evolve_csv}')  # download evolve.csv if existsfor _ in range(opt.evolve):  # generations to evolveif evolve_csv.exists():  # if evolve.csv exists: select best hyps and mutate# Select parent(s)parent = 'single'  # parent selection method: 'single' or 'weighted'x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1)n = min(5, len(x))  # number of previous results to considerx = x[np.argsort(-fitness(x))][:n]  # top n mutationsw = fitness(x) - fitness(x).min() + 1E-6  # weights (sum > 0)if parent == 'single' or len(x) == 1:# x = x[random.randint(0, n - 1)]  # random selectionx = x[random.choices(range(n), weights=w)[0]]  # weighted selectionelif parent == 'weighted':x = (x * w.reshape(n, 1)).sum(0) / w.sum()  # weighted combination# Mutatemp, s = 0.8, 0.2  # mutation probability, sigmanpr = np.randomnpr.seed(int(time.time()))g = np.array([meta[k][0] for k in hyp.keys()])  # gains 0-1ng = len(meta)v = np.ones(ng)while all(v == 1):  # mutate until a change occurs (prevent duplicates)v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0)for i, k in enumerate(hyp.keys()):  # plt.hist(v.ravel(), 300)hyp[k] = float(x[i + 7] * v[i])  # mutate# Constrain to limitsfor k, v in meta.items():hyp[k] = max(hyp[k], v[1])  # lower limithyp[k] = min(hyp[k], v[2])  # upper limithyp[k] = round(hyp[k], 5)  # significant digits# Train mutationresults = train(hyp.copy(), opt, device, callbacks)callbacks = Callbacks()# Write mutation resultsprint_mutation(results, hyp.copy(), save_dir, opt.bucket)# Plot resultsplot_evolve(evolve_csv)LOGGER.info(f'Hyperparameter evolution finished {opt.evolve} generations\n'f"Results saved to {colorstr('bold', save_dir)}\n"f'Usage example: $ python train.py --hyp {evolve_yaml}')def run(**kwargs):# Usage: import train; train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt')opt = parse_opt(True)for k, v in kwargs.items():setattr(opt, k, v)main(opt)return optif __name__ == "__main__":opt = parse_opt()main(opt)