train.py是YOLOv5中用于模型训练的脚本文件，其主要功能是读取配置文件、设置训练参数、构建模型结构、加载数据、训练/验证模型、保存模型权重文件、输出日志等
参考笔记：
【YOLOv3】源码（train.py）_yolo原始代码-CSDN博客
【yolov5】 train.py详解_yolov5 train.py-CSDN博客
学习视频：
目录
[1. train.py的主要功能](#1. train.py的主要功能)
[2. 主要模块](#2. 主要模块)
[2.1 参数解析与初始化](#2.1 参数解析与初始化)
[2.2 加载模型与数据](#2.2 加载模型与数据)
[2.3 优化器与学习率更新器设置](#2.3 优化器与学习率更新器设置)
[2.4 训练循环(核心)](#2.4 训练循环(核心))
[2.5 最终的模型验证](#2.5 最终的模型验证)
[3. 完整train.py代码](#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：**冻结某些层的所有参数不参与训练；具体代码位置如下：
常用参数和其他参数作用可参考代码注释：
python
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')
#优化器，默认是SGD
parser.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/trian
parser.add_argument('--project', default=ROOT / 'runs/train',
help='save to project/name(根目录)')
#训练结果存放的子目录，默认是runs/train/expn
parser.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.pt
parser.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 opt2.2 加载模型与数据
加载模型权重和配置文件，设置模型参数，加载训练数据
分析：该部分代码属于yolov5结构中的哪个阶段？
主要发生在训练前的准备工作，即还没有进入模型的前向传播和反向传播阶段
运行逻辑分析
- 模型加载与构建
- 如果提供了预训练权重，加载模型结构和权重参数
- 如果没有提供权重，则根据模型结构配置文件yolov5xx.yaml构建新模型
- 冻结层设置
- 通过命令行的freeze参数指定需要冻结哪些层(例如Backbone层)，以适应迁移学习或微调场景
- 训练数据准备
- 加载数据集，创建数据迭代器Dataloader
python
'-----------------------------模型部分---------------------------'
#Model，模型加载
check_suffix(weights, '.pt') #检查权重文件的后缀是否为.pt
pretrained = 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 FP32
csd = 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 优化器与学习率更新器设置
python
'-------------------------权重衰退调整------------------------------'
#这个部分看不太懂，应该是某种调整权重衰退的trick
nbs = 64
accumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizing
hyp['weight_decay'] *= batch_size * accumulate / nbs # scale weight_decay
LOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}")
'-------------------------参数分组------------------------------'
g0, g1, g2 = [], [], []#初始化3个数组
'''
将模型的参数分为三类：
g0: BatchNorm 的权重
g1: 卷积层或全连接层的权重
g2: 偏置（bias）
'''
#遍历模型的每个模块
for v in model.modules():
#如果模块有偏置参数，并且是 nn.Parameter 类型，则将其加入 g2
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
g2.append(v.bias)
#如果模块是 BatchNorm2d，则将其权重加入 g0
if isinstance(v, nn.BatchNorm2d):
g0.append(v.weight)
#如果模块有权重参数，并且是 nn.Parameter 类型，则将其加入 g1
elif 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 momentum
elif opt.optimizer == 'AdamW':
optimizer = AdamW(g0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum
else:
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
# 创建学习率更新器，基于上述的学习率变化函数lf
scheduler = 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轮训练模型仍然没有提升，如果是则训练早停
python
'-------------------------模型训练------------------------------'
#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()#优化器梯度清0
for 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升到lr0
x['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:#如果使用四元数据加载方式，则损失乘以4
loss *= 4.
'------------------------反向传播-------------------------------'
scaler.scale(loss).backward()#使用梯度缩放反向传播
'------------------------参数更新-------------------------------'
if ni - last_opt_step >= accumulate: #如果达到累计步数条件
scaler.step(optimizer) #更新优化器参数
scaler.update() # 更新梯度缩放比例
optimizer.zero_grad() #梯度清0
if 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
#计算mAP
if 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 = fi
log_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.pt
torch.save(ckpt, last)#保存为最后一次权重文件
if best_fitness == fi:#如果当前epoch的适应度最优
torch.save(ckpt, best)#保存为最佳权重文件
#按周期保存权重文件，这行代码默认是不执行的，YOLOv5默认只保存last.pt和best.pt
if (epoch > 0) and (opt.save_period > 0) and (epoch % opt.save_period == 0):
torch.save(ckpt, w / f'epoch{epoch}.pt')
del ckpt
callbacks.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在验证集上作最后一次验证，控制台输出验证结果
python
'---------------------------最终的模型验证---------------------------'
#如果是主进程，模型训练结束时，下面代码将利用最佳模型在验证集计算性能指标显示在控制台上
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代码
python
# 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-Data
Usage:
$ 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 Path
import 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 tqdm
FILE = 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_first
LOCAL_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.freeze
callbacks.run('on_pretrain_routine_start')
#Directories，创建保存模型权重的目录，一般是runs/train/expn/weights
w = 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/expn
if 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 instance
if loggers.wandb:
data_dict = loggers.wandb.data_dict
if resume:
weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size
# Register actions
for k in methods(loggers):
callbacks.register_action(k, callback=getattr(loggers, k))
#Config，配置
plots = not evolve and not opt.noplots#是否创建绘图，演示模式下不创建
cuda = device.type != 'cpu'#检查是否使用GPU
init_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']
#获取类别数量，单类别情况下数量为1
nc = 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') #检查权重文件的后缀是否为.pt
pretrained = 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 FP32
csd = 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)#获取模型的最大步幅作为下采样倍数，确保最大为32
imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2)
'验证在命令行指定的自适应缩放图像大小是否是32的倍数，且不小于gs的两倍，如果不是则进行调整'
#Batch size，如果我们在命令行指定batch_size=-1，则会执行下面代码，
#根据GPU情况自动选择一个最好的batch_size
if RANK == -1 and batch_size == -1:
batch_size = check_train_batch_size(model, imgsz)
loggers.on_params_update({"batch_size": batch_size})
'-------------------------权重衰退调整------------------------------'
#这个部分看不太懂，应该是某种调整权重衰退的trick
nbs = 64
accumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizing
hyp['weight_decay'] *= batch_size * accumulate / nbs # scale weight_decay
LOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}")
'-------------------------参数分组------------------------------'
g0, g1, g2 = [], [], []#初始化3个数组
'''
将模型的参数分为三类：
g0: BatchNorm 的权重
g1: 卷积层或全连接层的权重
g2: 偏置（bias）
'''
#遍历模型的每个模块
for v in model.modules():
#如果模块有偏置参数，并且是 nn.Parameter 类型，则将其加入 g2
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
g2.append(v.bias)
#如果模块是 BatchNorm2d，则将其权重加入 g0
if isinstance(v, nn.BatchNorm2d):
g0.append(v.weight)
#如果模块有权重参数，并且是 nn.Parameter 类型，则将其加入 g1
elif 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 momentum
elif opt.optimizer == 'AdamW':
optimizer = AdamW(g0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum
else:
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
# 创建学习率更新器，基于上述的学习率变化函数lf
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)
#EMA
ema = ModelEMA(model) if RANK in [-1, 0] else None
#初始化开始轮次和最佳适应度
start_epoch, best_fitness = 0, 0.0
#如果指定了预训练权重文件
if pretrained:
#Optimizer
if ckpt['optimizer'] is not None:#检查权重文件中是否包含优化器
optimizer.load_state_dict(ckpt['optimizer'])#加载优化器
best_fitness = ckpt['best_fitness']#更新最佳适应度
#EMA
if ema and ckpt.get('ema'):
ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
ema.updates = ckpt['updates']
#Epochs
start_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 epochs
del ckpt, csd
#DP mode
if 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)
#SyncBatchNorm
if 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}'
#处理进程0
if 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.jpg
plot_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 mode
if cuda and RANK != -1:
model = DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
#关于损失函数的超参数调整设置，使其更适应当前任务
nl = de_parallel(model).model[-1].nl#nl:检测头数量，通常为3
hyp['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()#优化器梯度清0
for 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升到lr0
x['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:#如果使用四元数据加载方式，则损失乘以4
loss *= 4.
'------------------------反向传播-------------------------------'
scaler.scale(loss).backward()#使用梯度缩放反向传播
'------------------------参数更新-------------------------------'
if ni - last_opt_step >= accumulate: #如果达到累计步数条件
scaler.step(optimizer) #更新优化器参数
scaler.update() # 更新梯度缩放比例
optimizer.zero_grad() #梯度清0
if 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
#计算mAP
if 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 = fi
log_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.pt
torch.save(ckpt, last)#保存为最后一次权重文件
if best_fitness == fi:#如果当前epoch的适应度最优
torch.save(ckpt, best)#保存为最佳权重文件
#按周期保存权重文件，这行代码默认是不执行的，YOLOv5默认只保存last.pt和best.pt
if (epoch > 0) and (opt.save_period > 0) and (epoch % opt.save_period == 0):
torch.save(ckpt, w / f'epoch{epoch}.pt')
del ckpt
callbacks.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 results
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')
#优化器，默认是SGD
parser.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/trian
parser.add_argument('--project', default=ROOT / 'runs/train',
help='save to project/name(根目录)')
#训练结果存放的子目录，默认是runs/train/expn
parser.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.pt
parser.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
def 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/evolve
if opt.evolve:
if opt.project == str(ROOT / 'runs/train'): # if default project name, rename to runs/evolve
opt.project = str(ROOT / 'runs/evolve')
opt.exist_ok, opt.resume = opt.resume, False # pass resume to exist_ok and disable resume
if opt.name == 'cfg':
opt.name = Path(opt.cfg).stem # use model.yaml as name
#设置训练结果保存目录，一般是runs/train/expn
opt.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 dict
if 'anchors' not in hyp: # anchors commented in hyp.yaml
hyp['anchors'] = 3
opt.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 indices
evolve_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 exists
for _ in range(opt.evolve): # generations to evolve
if 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 consider
x = x[np.argsort(-fitness(x))][:n] # top n mutations
w = 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 selection
x = x[random.choices(range(n), weights=w)[0]] # weighted selection
elif parent == 'weighted':
x = (x * w.reshape(n, 1)).sum(0) / w.sum() # weighted combination
# Mutate
mp, s = 0.8, 0.2 # mutation probability, sigma
npr = np.random
npr.seed(int(time.time()))
g = np.array([meta[k][0] for k in hyp.keys()]) # gains 0-1
ng = 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 limits
for k, v in meta.items():
hyp[k] = max(hyp[k], v[1]) # lower limit
hyp[k] = min(hyp[k], v[2]) # upper limit
hyp[k] = round(hyp[k], 5) # significant digits
# Train mutation
results = train(hyp.copy(), opt, device, callbacks)
callbacks = Callbacks()
# Write mutation results
print_mutation(results, hyp.copy(), save_dir, opt.bucket)
# Plot results
plot_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 opt
if __name__ == "__main__":
opt = parse_opt()
main(opt)