一. opt参数
def parse_opt(known=False):
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
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, help='total training epochs')
parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs, -1 for autobatch')
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
parser.add_argument('--rect', action='store_true', help='rectangular training')
parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
parser.add_argument('--noval', action='store_true', help='only validate final 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')
parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
parser.add_argument('--cache', type=str, nargs='?', const='ram', help='image --cache ram/disk')
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')
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')
parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer')
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)')
parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name')
parser.add_argument('--name', default='exp', help='save to project/name')
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')
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')
parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
parser.add_argument('--seed', type=int, default=0, help='Global training seed')
parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')
# Logger arguments
parser.add_argument('--entity', default=None, help='Entity')
parser.add_argument('--upload_dataset', nargs='?', const=True, default=False, help='Upload data, "val" option')
parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval')
parser.add_argument('--artifact_alias', type=str, default='latest', help='Version of dataset artifact to use')
return parser.parse_known_args()[0] if known else parser.parse_args()
逐步分解代码
-
定义函数
parse_opt
:-
此函数目的是解析命令行参数,以便在运行训练程序时可以通过参数指定相关设置。
-
-
创建参数解析器:
parser = argparse.ArgumentParser()
argparse
是 Python 标准库中用于处理命令行参数的模块。这里创建一个解析器对象,用于添加各种参数。 -
添加参数:
parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
每个参数的功能和含义:
-
每个
add_argument
方法用于定义一个可接受的命令行选项。这些选项包含训练所需的多个配置参数,例如: -
其他参数类似,涉及到训练过程的多个设置。
-
--resume
: 是否恢复最近一次的训练,断点续训:即是否在之前训练的一个模型基础上继续训练,default
值默认是False
;如果想采用断点续训的方式,这里我推荐一种写法,即首先将default=False
改为default=True
。随后在终端中键入如下指令:
python train.py --resume D:\Pycharm_Projects\yolov5-6.1-4_23\runs\train\exp19\weights\last.pt
-
--rect
: 是否使用矩形训练所谓矩阵推理就是不再要求你训练的图片是正方形了;矩阵推理会加速模型的推理过程,减少一些冗余信息。
下图分别是方形推理方式和矩阵推理方式。
-
--imgsz
: 图像的大小(以像素为单位),输入默认640 × 640
。 图片越大,训练速度越慢。 图像在训练时resize
同样花费时间的。 -
--batch-size
:每批次的输入数据量;
这里说一下 epoch、batchsize、iteration三者之间的联系 1、batchsize是批次大小,假如取batchsize=24,则表示每次训练时在训练集中取24个训练样本进行训练。 2、iteration是迭代次数,1个iteration就等于一次使用24(batchsize大小)个样本进行训练。 3、epoch:1个epoch就等于使用训练集中全部样本训练1次。default=-1
将时自动调节batchsize
大小。 -
--epochs
: 总共的训练轮数,默认为300
轮,显示效果是0-299
。 -
--hyp
: 超参数文件的路径,超参数里面包含了大量的参数信息 -
--data
: 数据集的配置文件路径。 -
--cfg
: 模型配置文件的路径。 -
--weights
: 初始权重文件的路径,如果这里设置为空的话,就是自己从头开始进行训练; -
–-nosave:是否只保存最后一轮的
pt
文件;我们默认是保存best.pt
和last.pt
两个的。 -
–-noval:只在最后一轮测试;正常情况下每个
epoch
都会计算mAP
,但如果开启了这个参数,那么就只在最后一轮上进行测试,做科研时不建议开启,因为会看不到曲线变化情况。 -
–-noautoanchor:是否禁用自动锚框;默认是开启的,自动锚点
的好处是可以简化训练过程;
yolov5中预先设定了一下锚定框,这些锚框是针对coco数据集的,其他目标检测也适用,可以在models/yolov5.文件中查看,例如如图所示,这些框针对的图片大小是640640。这是默认的anchor大小。需要注意的是在目标检测任务中,一般使用大特征图上去检测小目标,因为大特征图含有更多小目标信息,因此大特征图上的anchor数值通常设置为小数值,小特征图检测大目标,因此小特征图上anchor数值设置较大。在yolov5 中自动锚定框选项,训练开始前,会自动计算数据集标注信息针对默认锚定框的最佳召回率,当最佳召回率大于等于0.98时,则不需要更新锚定框;如果最佳召回率小于0.98,则需要重新计算符合此数据集的锚定框。 在parse_opt设置了默认自动计算锚框选项,如果不想自动计算,可以设置这个,建议不要改动。
-
–noplots:开启这个参数后将不保存绘图文件。
-
–evolve:遗传超参数进化;yolov5使用遗传超参数进化,提供的默认参数
是通过在COCO数据集上使用超参数进化得来的(也就是下图这些参数)。由于超参数进化会耗费大量的资源和时间,所以建议大家不要动这个参数。
是利用种群搜索技术将种群作为一组问题解,通过对当前种群施加类似生物遗传环境因素的选择、交叉、变异等一系列的遗传操作来产生新一代的种群,并逐步使种群优化到包含近似最优解的状态,遗传算法调优能够求出优化问题的全局最优解,优化结果与初始条件无关,算法独立于求解域,具有较强的鲁棒性,适合于求解复杂的优化问题,应用较为广泛。
-
–bucket:谷歌云盘;通过这个参数可以下载谷歌云盘上的一些东西,但是现在没必要使用了
-
–cache:是否提前缓存图片到内存,以加快训练速度,默认
False
;开启这个参数就会对图片进行缓存,从而更好的训练模型。 -
–image-weights:是否启用加权图像策略,默认是不开启的;主要是为了解决样本不平衡问题;开启后会对于上一轮训练效果不好的图片,在下一轮中增加一些权重;
-
–device:设备选择;cpu还是cuda。
-
–multi-scale是否启用多尺度训练,默认是不开启的;多尺度训练是指设置几种不同的图片输入尺度,训练时每隔一定
iterations
随机选取一种尺度训练,这样训练出来的模型鲁棒性更强。多尺度训练在比赛中经常可以看到他身影,是被证明了有效提高性能的方式。输入图片的尺寸对检测模型的性能影响很大,在基础网络部分常常会生成比原图小数十倍的特征图,导致小物体的特征描述不容易被检测网络捕捉。通过输入更大、更多尺寸的图片进行训练,能够在一定程度上提高检测模型对物体大小的鲁棒性。 -
–single-cls:设定训练数据集是单类别还是多类别;默认为
False
多类别 -
–optimizer:选择优化器;默认为
SGD
,可选SGD
,Adam
,AdamW
-
–sync-bn:是否开启跨卡同步
BN
;开启参数后即可使用SyncBatchNorm
多GPU
进行分布式训练。 -
–workers:最大
worker
数量;这里经常出问题,Windows系统报错时可以设置成 0 0 0。 -
–project:指定训练好的模型的保存路径;默认在
runs/train
。 -
–name:设定保存的模型文件夹名,默认在
exp
-
–exist-ok:每次预测模型的结果是否保存在原来的文件夹;如果指定了这个参数的话,那么本次预测的结果还是保存在上一次保存的文件夹里;如果不指定就是每次预测结果保存一个新的文件夹下。
-
–quad:官方发布的开启这个功能后的实际效果:好处是在比默认
640
大的数据集上训练效果更好。副作用是在640
大小的数据集上训练效果可能会差一些。 -
–cos-lr:是否开启余弦学习率;
这是我不采用
cos-lr
时学习率的曲线:在这里插入图片描述
开启后的学习率应该是这样子:
在这里插入图片描述
-
–label-smoothing
是否对标签进行平滑处理,默认是不启用的;
在训练样本中,我们并不能保证所有sample都标注正确,如果某个样本标注错误,就可能产生负面印象,如果我们有办法“告诉”模型,样本的标签不一定正确,那么训练出来的模型对于少量的样本错误就会有“免疫力”采用随机化的标签作为训练数据时,损失函数
有1-ε的概率与上面的式子相同,比如说告诉模型只有0.95概率是那个标签。
-
–patience早停;如果模型在
default
值轮数里没有提升,则停止训练模型。 -
–freeze
指定冻结层数量;可以在
yolov5s.yaml
中查看主干网络层数。是迁移学习常用的方法,当我们在使用数据量不足的情况下,通常我们会选择公共数据集提供权重作为预训练权重,我们知道网络的backbone主要是用来提取特征用的,一般大型数据集训练好的权重主干特征提取能力是比较强的,这个时候我们只需要冻结主干网络
,fine-tune后面层就可以了,不需要从头开始训练,大大减少了实践而且还提高了性能。
在这里插入图片描述
例如如下指令,代表冻结整个主干,前三层的编号是0,1,2
python train.py --freeze 10
这里分享一个对比冻结效果的项目 这里是项目的部分截图
在这里插入图片描述
-
–save-period用于设置多少个
epoch
保存一下checkpoint
-
–seed这是
v6.2
版本更新的一个非常重要的参数,如果你使用torch>=1.12.0
的单GPU训练, 那你的训练结果完全可再现。 -
–local_rank
DistributedDataParallel
单机多卡训练,单GPU
设备不需要设置; -
–entity在线可视化工具,类似于
tensorboard
。 -
–upload_dataset是否上传
dataset
到wandb tabel
(将数据集作为交互式 dsviz表 在浏览器中查看、查询、筛选和分析数据集) 默认False
。 -
–bbox_interval设置界框图像记录间隔
Set bounding-box image logging interval for W&B
默认-1
。 -
–artifact_alias 这个功能作者还未实现。
-
-
增加特定参数的选项:
parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
--resume
选项可以跟随一个可选的参数,且以布尔值形式表示是否恢复训练。-
一些参数使用
action
或nargs
来定义特殊的解析方式,例如:
-
-
处理日志记录和其他参数:
-
代码还包括一些与日志记录和数据集上传相关的参数,确保在训练过程中能够记录和调试。
-
-
返回解析的结果:
return parser.parse_known_args()[0] if known else parser.parse_args()
如果参数
known
为True
,则调用parse_known_args()
以解析已知参数;否则调用parse_args()
解析所有参数。
总结
这段代码的主要功能是解析命令行参数,用于配置YOLOv5训练的各个方面。它允许用户通过命令行灵活地指定不同的训练设置,例如权重文件、配置文件、超参数、训练轮数、批量大小等选项。这有助于实现高效和可重复的模型训练,使用户能够根据数据集和计算资源的不同需求自定义训练过程。
二 main函数
第一步:检查分布式训练环境
# 检查当前的 RANK 值,并执行相关的检查操作
# 如果 RANK 为 -1 或 0,执行以下检查
if RANK in {-1, 0}:
print_args(vars(opt))
check_git_status()
check_requirements()
这段代码的作用是在特定条件下执行一些检查和打印操作。逐步分解如下:
-
条件检查 (
if RANK in {-1, 0}:
):-
RANK
是一个表示当前进程的排名的变量。通常在分布式训练中使用,多进程环境中每个进程会有一个唯一的排名。 -
{-1, 0}
代表的是两个特殊值:-1
表示是单进程训练,0
表示是主进程。在分布式训练中,通常只有主节点负责打印和代码检查工作,因此这里的条件用于确保只有主进程执行后续操作。
-
-
打印参数 (
print_args(vars(opt))
):-
vars(opt)
是将opt
对象(通常包含命令行参数和配置选项)转换为字典格式,以便后续操作获取其属性。 -
print_args
函数则负责将这些参数以某种格式输出到控制台。这通常用于调试目的,帮助开发者检查和确认当前训练的配置。
-
-
检查 Git 状态 (
check_git_status()
):-
该函数用于检查当前代码库的 Git 状态(例如,是否有未提交的更改、当前分支信息等)。
-
它的目的是确保当前代码的版本是干净的,通常在开始训练之前检查代码状态,以避免意外的代码变动影响实验结果。
-
-
检查依赖 (
check_requirements()
):-
该函数会检查项目所需的 Python 包和库是否正确安装,并确保其版本满足要求。
-
目的是避免因缺少依赖或版本不兼容导致的运行时错误。
-
总结
这段代码主要功能是在训练前进行必要的环境和配置检查。它确保只有主进程执行这些检查,以避免在分布式训练中重复冗余的操作。同时通过打印参数、检查 Git 状态和验证依赖,帮助开发者确认训练环境的正确性,从而提高训练的可靠性和有效性。
第二步:判断是否断点续训
if opt.resume and not check_wandb_resume(opt) and not opt.evolve: # resume an interrupted run
# isinstance()是否是已经知道的类型
# 如果resume是True,则通过get_lastest_run()函数找到runs为文件夹中最近的权重文件last.pt
ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run() # specified or most recent path
# 判断是否为文件,若不是文件抛出异常
assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist'
# opt.yaml是训练时的命令行参数文件
with open(Path(ckpt).parent.parent / 'opt.yaml', errors='ignore') as f:
# 超参数替换,将训练时的命令行参数加载进opt参数对象中
opt = argparse.Namespace(**yaml.safe_load(f)) # replace
# opt.cfg设置为'' 对应着train函数里面的操作(加载权重时是否加载权重里的anchor)
opt.cfg, opt.weights, opt.resume = '', ckpt, True # reinstate
# 打印从ckpt恢复断点训练信息
LOGGER.info(f'Resuming training from {ckpt}')
else:
# 不使用断点续训,就从文件中读取相关参数
# check_file (utils/general.py)的作用为查找/下载文件 并返回该文件的路径。
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) # checks
# 如果模型文件和权重文件为空,弹出警告
assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
# 如果要进行超参数进化,重建保存路径
if opt.evolve:
# 设置新的项目输出目录
opt.project = str(ROOT / 'runs/evolve')
# 将resume传递给exist_ok
opt.exist_ok, opt.resume = opt.resume, False # pass resume to exist_ok and disable resume
# 根据opt.project生成目录,并赋值给opt.save_dir 如: runs/train/exp1
opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))
这段代码主要是关于断点训练的判断和准备。
断点训练是当训练异常终止或想调节超参数时,系统会保留训练异常终止前的超参数与训练参数,当下次训练开始时,并不会从头开始,而是从上次中断的地方继续训练。
- 使用断点续训,就从last.pt中读取相关参数
- 不使用断点续训,就从文件中读取相关参数
逐步分解
-
判断条件:
if opt.resume and not check_comet_resume(opt) and not opt.evolve:
这行代码检查几个条件:
opt.resume
为真,表示用户希望从上次训练的检查点恢复训练。not check_comet_resume(opt)
确保用户不是从 Comet.ml 这样的实验追踪工具中恢复。not opt.evolve
检查用户没有选择进化超参数模式。
-
获取最新检查点路径:
last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run())
如果
opt.resume
是字符串类型,调用check_file(opt.resume)
来检查该路径是否有效;如果不是,则调用get_latest_run()
获取最近一次运行的检查点。 -
获取训练选项配置:
opt_yaml = last.parent.parent / 'opt.yaml' # train options yaml opt_data = opt.data # original dataset
这里设置了
opt_yaml
的路径,该路径指向训练选项的 YAML 配置文件。同时保存了原始数据集的路径。 -
读取配置文件:
if opt_yaml.is_file(): with open(opt_yaml, errors='ignore') as f: d = yaml.safe_load(f) else: d = torch.load(last, map_location='cpu')['opt']
检查
opt_yaml
文件是否存在。如果存在,则以 YAML 格式读取文件内容;如果不存在,则从最近的检查点加载训练选项。 -
更新训练参数:
opt = argparse.Namespace(**d) # replace opt.cfg, opt.weights, opt.resume = '', str(last), True # reinstate
将读取到的训练参数更新到
opt
中,同时将cfg
置为空,将weights
设置为最新检查点的路径,并将resume
设置为True
,表示即将恢复训练。 -
处理数据集路径:
if is_url(opt_data): opt.data = check_file(opt_data) # avoid HUB resume auth timeout
如果
opt.data
是 URL,调用check_file(opt_data)
来确保其有效性。 -
当没有恢复训练时的处理:
else: 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) # checks
如果不进行恢复,则对数据集和模型配置文件等进行有效性检查。
-
确保提供必要的配置或权重:
assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
确保用户至少提供了模型配置或权重文件。
-
处理超参数进化的情况:
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
如果用户选择了超参数进化,则将项目路径更新为
runs/evolve
并设置resume
为False
,以便不在进化过程中恢复训练。 -
设置保存目录和名称:
if opt.name == 'cfg': opt.name = Path(opt.cfg).stem # use model.yaml as name opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))
如果名称是
cfg
,则使用配置文件的基本名称。然后设置保存目录,确保如果目录已经存在则自动递增编号。
总结
这段代码的主要功能是处理模型训练的恢复逻辑和配置文件的加载。当用户选择从上一个训练中恢复时,代码会在配置文件和检查点之间进行选择,并确保所有必要的参数都是有效的。如果不是恢复情境,代码会验证所有输入,并处理超参数的进化同时设置训练的保存路径。最终,它确保训练环境的配置是正确的,从而准备好开始训练过程。
第三步:判断是否分布式训练
# DDP mode --> 支持多机多卡、分布式训练
# 选择程序装载的位置
device = select_device(opt.device, batch_size=opt.batch_size)
# 当进程内的GPU编号不为-1时,才会进入DDP
if LOCAL_RANK != -1:
# 用于DDP训练的GPU数量不足
assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
# WORLD_SIZE表示全局的进程数
assert opt.batch_size % WORLD_SIZE == 0, '--batch-size must be multiple of CUDA device count'
# 不能使用图片采样策略
assert not opt.image_weights, '--image-weights argument is not compatible with DDP training'
# 不能使用超参数进化
assert not opt.evolve, '--evolve argument is not compatible with DDP training'
# 设置装载程序设备
torch.cuda.set_device(LOCAL_RANK)
# 保存装载程序的设备
device = torch.device('cuda', LOCAL_RANK)
# torch.distributed是用于多GPU训练的模块
dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")
这段代码的主要功能是设置和初始化分布式数据并行 (Distributed Data Parallel, DDP) 的训练环境,尤其是在多GPU的情况下。这是深度学习训练中的一种常用方式,目的是提高训练速度和效率。
逐步分解和详细解释:
-
设备选择:
device = select_device(opt.device, batch_size=opt.batch_size)
通过调用
select_device
函数来选择用于训练的计算设备(CPU或GPU)。opt.device
中包含了用户指定的设备信息,例如"0"
(只使用第一个GPU),或者"0,1"
(使用多个GPU)。batch_size
用于在选择设备时考虑内存的可用性。
-
分布式训练条件检查:
if LOCAL_RANK != -1:
LOCAL_RANK
的值用于指示当前进程的排名。如果值为-1
,表示当前进程不是分布式训练中的一部分;如果不等于-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'
这些
assert
语句用于检查一些基本条件,确保在分布式训练中使用的参数有效:- 确保当前设备数量大于进程的
LOCAL_RANK
,以防止在没有足够设备的情况下崩溃。 - 确保批次大小能够被世界大小(
WORLD_SIZE
)整除,即在所有进程间均匀分配。 - 确保批次大小不为
-1
,因为 AutoBatch 需要一个有效的批次大小。 - 确保不会选择超参数进化(
--evolve
)。 - 确保没有启用图像权重(
--image-weights
)选项,因为在DDP训练中不兼容。
- 确保当前设备数量大于进程的
-
设定CUDA设备:
torch.cuda.set_device(LOCAL_RANK) device = torch.device('cuda', LOCAL_RANK)
- 根据
LOCAL_RANK
设置当前CUDA设备,以确保此进程利用正确的GPU进行训练。 torch.device
用于创建一个设备对象,表示将要使用的 GPU。
- 根据
-
初始化进程组:
dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")
CopyInsert
- 这行代码用于初始化分布式训练的进程组。
dist
模块提供了多GPU训练的通信能力。 - 根据
NCCL
是否可用选择后端,如果不可用则使用Gloo
作为后备解决方案。
- 这行代码用于初始化分布式训练的进程组。
总结:
以上代码主要用于初始化分布式数据并行训练的环境。通过选择合适的计算设备、检查输入参数的有效性以及初始化进程组,为后续的深度学习模型训练做好准备。这些步骤确保了模型能够在多GPU设置下高效地进行训练,从而提高训练速度和性能。
第四步:判断是否进化训练
# Train 训练模式: 如果不进行超参数进化,则直接调用train()函数,开始训练
if not opt.evolve:# 如果不使用超参数进化
# 开始训练
train(opt.hyp, opt, device, callbacks)
if WORLD_SIZE > 1 and RANK == 0:
# 如果全局进程数大于1并且RANK等于0
# 日志输出 销毁进程组
LOGGER.info('Destroying process group... ')
# 训练完毕,销毁所有进程
dist.destroy_process_group()
# Evolve hyperparameters (optional) 遗传进化算法,边进化边训练
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
# 如果超参数文件中没有'anchors',则设为3
if 'anchors' not in hyp: # anchors commented in hyp.yaml
hyp['anchors'] = 3
# 使用进化算法时,仅在最后的epoch测试和保存
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 {save_dir}') # download evolve.csv if exists
"""
遗传算法
调参:遵循适者生存、优胜劣汰的法则,即寻优过程中保留有用的,去除无用的。
遗传算法需要提前设置4个参数: 群体大小/进化代数/交叉概率/变异概率
"""
# 选择超参数的遗传迭代次数 默认为迭代300次
for _ in range(opt.evolve): # generations to evolve
# 如果evolve.csv文件存在
if evolve_csv.exists(): # if evolve.csv exists: select best hyps and mutate
# Select parent(s)
# 选择超参进化方式,只用single和weighted两种
parent = 'single' # parent selection method: 'single' or 'weighted'
# 加载evolve.txt
x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1)
# 选取至多前五次进化的结果
n = min(5, len(x)) # number of previous results to consider
# fitness()为x前四项加权 [P, R, mAP@0.5, mAP@0.5:0.95]
# np.argsort只能从小到大排序, 添加负号实现从大到小排序, 算是排序的一个代码技巧
x = x[np.argsort(-fitness(x))][:n] # top n mutations
# 根据(mp, mr, map50, map)的加权和来作为权重计算hyp权重
w = fitness(x) - fitness(x).min() + 1E-6 # weights (sum > 0)
# 根据不同进化方式获得base hyp
if parent == 'single' or len(x) == 1:
# 根据权重的几率随机挑选适应度历史前5的其中一个
# x = x[random.randint(0, n - 1)] # random selection
x = x[random.choices(range(n), weights=w)[0]] # weighted selection
elif parent == 'weighted':
# 对hyp乘上对应的权重融合层一个hpy, 再取平均(除以权重和)
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()))
# 获取突变初始值, 也就是meta三个值的第一个数据
# 三个数值分别对应着: 变异初始概率, 最低限值, 最大限值(mutation scale 0-1, lower_limit, upper_limit)
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)
# 将突变添加到base hyp上
for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300)
hyp[k] = float(x[i + 7] * v[i]) # mutate
# Constrain to limits 限制hyp在规定范围内
for k, v in meta.items():
# 这里的hyp是超参数配置文件对象
# 而这里的k和v是在元超参数中遍历出来的
# hyp的v是一个数,而元超参数的v是一个元组
hyp[k] = max(hyp[k], v[1]) # 先限定最小值,选择二者之间的大值 ,这一步是为了防止hyp中的值过小
hyp[k] = min(hyp[k], v[2]) # 再限定最大值,选择二者之间的小值
hyp[k] = round(hyp[k], 5) # 四舍五入到小数点后五位
# 最后的值应该是 hyp中的值与 meta的最大值之间的较小者
# Train mutation 使用突变后的参超,测试其效果
results = train(hyp.copy(), opt, device, callbacks)
# Write mutation results
# 将结果写入results,并将对应的hyp写到evolve.txt,evolve.txt中每一行为一次进化的结果
# 每行前七个数字 (P, R, mAP, F1, test_losses(GIOU, obj, cls)) 之后为hyp
# 保存hyp到yaml文件
print_mutation(hyp.copy(), results, yaml_file, opt.bucket)
# Plot results 将结果可视化 / 输出保存信息
plot_evolve(evolve_csv)
LOGGER.info(f'Hyperparameter evolution finished\n'
f"Results saved to {colorstr('bold', save_dir)}\n"
f'Use best hyperparameters example: $ python train.py --hyp {evolve_yaml}')
这段代码是一个用于训练YOLOv5模型的部分,主要功能是进行模型训练和超参数的演化(进化)。下面对代码进行逐步分解和详细解释:
代码逐行解释
-
判断是否进行超参数演化:
if not opt.evolve: print(" no evolve") train(opt.hyp, opt, device, callbacks)
- 如果不进行超参数演化,则调用
train
函数进行模型训练,传入超参数、选项、设备和回调。
- 如果不进行超参数演化,则调用
-
超参数演化(可选):
else: meta = { ... } # 定义超参数演化的元数据
- 如果选择进行演化,则定义一个字典
meta
,该字典记录了多个超参数的基本信息,包括:- 超参数的名称
- 初始值
- 最小和最大值
- 如果选择进行演化,则定义一个字典
-
读取初始超参数:
with open(opt.hyp, errors='ignore') as f: hyp = yaml.safe_load(f) if 'anchors' not in hyp: hyp['anchors'] = 3
- 读取位于
opt.hyp
路径的超参数配置文件,并将其解析为字典。如果超参数中没有anchors
,则设置默认值为3。
- 读取位于
-
处理超参数:
if opt.noautoanchor: del hyp['anchors'], meta['anchors']
- 如果不使用自动锚点,则删除锚点相关的超参数。
-
准备演化的结果保存路径:
opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir) 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}')
- 设定了一些选项,比如只进行验证和只保存最终模型。同时定义了演化结果的保存路径,并且如果指定了存储桶,则从云端下载现有的
evolve.csv
文件。
- 设定了一些选项,比如只进行验证和只保存最终模型。同时定义了演化结果的保存路径,并且如果指定了存储桶,则从云端下载现有的
-
超参数演化循环:
for _ in range(opt.evolve): # generations to evolve
- 开始进行多代的超参数演化,循环次数为
opt.evolve
的值。
- 开始进行多代的超参数演化,循环次数为
-
从现有演化结果中选择最佳超参数并进行变异:
- 从
evolve.csv
中加载数据,选择最优秀的超参数进行变异。 - 在选择超参数的过程中,可以选择单一选择或者加权选择,并进行变异处理。
- 从
-
变异操作:
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
- 定义变异概率和变异标准差,并以当前时间作为随机数种子。接着生成变异增益,以便在后续变异中使用。
-
约束超参数值:
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
- 确保变异后的超参数在定义的最小值和最大值之间,并保留一定的有效数字。
-
训练变异模型并保存结果:
results = train(hyp.copy(), opt, device, callbacks) ... print_mutation(keys, results, hyp.copy(), save_dir, opt.bucket)
- 调用
train
函数进行训练,并保存结果。这会逐代训练模型,评估模型性能,同时也会记录变异结果。
- 调用
-
绘制演化成果:
plot_evolve(evolve_csv)
- 将演化过程中记录的超参数和结果图形化,便于可视化分析。
超参数进化的步骤如下:
- 1.若存在evolve.csv文件,读取文件中的训练数据,选择超参进化方式,结果最优的训练数据突变超参数
- 2.限制超参进化参数hyp在规定范围内
- 3.使用突变后的超参数进行训练,测试其效果
- 4.训练结束后,将训练结果可视化,输出保存信息保存至evolution.csv,用于下一次的超参数突变。
原理:根据生物进化,优胜劣汰,适者生存的原则,每次迭代都会保存更优秀的结果,直至迭代结束。最后的结果即为最优的超参数
注意:使用超参数进化时要经过至少300次迭代,每次迭代都会经过一次完整的训练。因此超参数进化及其耗时,大家需要根据自己需求慎用。
总结
这段代码主要是实现YOLOv5模型的训练与超参数的演化。其主要功能包括:
- 在指定条件下进行模型训练,或者基于已有的超参数进行演化。
- 通过定义、变异和选择超参数来优化模型性能,这个过程涵盖了数据的选择、变异概率控制及超参数限制等元素。
- 通过将演化过程可视化,帮助研究者理解超参数对模型性能的影响,从而进行更合理的超参数配置。
四 train函数
第一步:加载参数和初始化配置信息
1)初始化参数
train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictionary
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')
print(save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze)
代码逐步分解:
-
函数定义:
train(hyp, opt, device, callbacks)
这是一个名为
train
的函数,接受四个参数:hyp
:表示超参数,可能是超参数的 YAML 文件路径或者一个字典。opt
:包含训练选项的对象,如保存路径、训练的周期数、批次大小等等。device
:指定用于训练的设备(例如,CPU 或 CUDA 设备)。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
这一行将
opt
对象中的多个属性提取出来,赋值给对应的变量,如下:save_dir
:保存模型的目录,通过Path
模块来处理路径。epochs
:训练的总周期数。batch_size
:每次迭代的批次大小。weights
:预训练模型的权重路径。single_cls
:是否将多类数据当作单类数据进行训练的标志。evolve
:是否进行超参数演化的标志。data
:数据集的路径。cfg
:模型配置文件的路径。resume
:是否从上次保存的状态恢复训练的标志。noval
:是否只在最后一轮进行验证的标志。nosave
:是否只保存最终的检查点的标志。workers
:数据加载的工作线程数。freeze
:需要冻结的层的列表。
-
回调函数:
callbacks.run('on_pretrain_routine_start')
这行代码调用回调对象的
run
方法,传入事件名称'on_pretrain_routine_start'
。这常常用于在训练开始之前执行某些操作,比如记录日志或进行设定。 -
打印参数信息:
print(save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze)
这一行代码将之前提取的参数信息打印出来,有助于调试和追踪训练过程中的参数设置。
总结:
该代码段定义了一个用于训练模型的函数 train
的开头部分,其中主要完成了以下功能:
- 接收并解析训练相关的参数设置,包括模型保存路径、训练周期、批次大小、预训练权重、数据集路径等。
- 调用回调函数以执行预训练例程开始时的操作。
- 打印设置参数,以便于在训练过程中进行调试和验证。
总的来说,这段代码为后续的训练过程做准备,确保所有必要的参数正确加载并便于追踪。
2)创建训练权重目录和保存路径
# Directories
w = save_dir / 'weights' # weights dir
(w.parent if evolve else w).mkdir(parents=True, exist_ok=True) # make dir
last, best = w / 'last.pt', w / 'best.pt'
这段代码主要用于设置保存模型权重的目录,并在必要时创建该目录。具体流程如下:
-
定义权重目录:
w = save_dir / 'weights' # weights dir
- 这里创建了一个名为
w
的变量,它代表了保存模型权重的目录路径。save_dir
是提前定义的路径,'weights'
是子目录的名称,使用/
符号是Path
对象的一种快捷拼接方式。
- 这里创建了一个名为
-
创建目录:
(w.parent if evolve else w).mkdir(parents=True, exist_ok=True) # make dir
- 这里的逻辑判断取决于
evolve
变量:- 如果
evolve
为真,则创建w
的父目录(w.parent
)。 - 如果
evolve
为假,则直接创建w
目录。
- 如果
mkdir(parents=True, exist_ok=True)
的作用是:parents=True
表示如果上层目录不存在,则一并创建。exist_ok=True
表示如果目录已经存在,不会抛出错误,这样可以安全地运行而不担心目录已经存在的问题。
- 这里的逻辑判断取决于
-
定义权重文件路径:
last, best = w / 'last.pt', w / 'best.pt'
- 这里定义了两个文件路径:
last
将会保存为最新的训练权重,文件名叫last.pt
。best
将会保存训练过程中表现最好的权重,文件名叫best.pt
。
- 这两个文件都保存在之前创建的
weights
目录下。
- 这里定义了两个文件路径:
总结
上述代码块的主要功能是创建一个用于存储模型权重的目录,并确保该目录的存在。如果在演化超参数时,权重目录会相应变化。代码还定义了最新权重和最佳权重的保存路径,用于后续的模型训练与验证。通过这些设置,确保训练过程中生成的权重文件能够按预期保存到指定目录,为后续的模型恢复或评估做好准备。
3)读取超参数配置文件
# Hyperparameters
if isinstance(hyp, str):
with open(hyp, errors='ignore') as f:
hyp = yaml.safe_load(f) # load hyps dict
LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
opt.hyp = hyp.copy() # for saving hyps to checkpoints
这段代码主要用于处理超参数(hyperparameters),具体功能和每一行的解释如下:
-
检查超参数的类型:
if isinstance(hyp, str):
这一行的含义是:检查变量
hyp
是否是字符串类型。如果是字符串,通常意味着hyp
是一个文件路径,包含超参数的定义。 -
加载超参数文件:
with open(hyp, errors='ignore') as f: hyp = yaml.safe_load(f) # load hyps dict
如果
hyp
是一个字符串,则打开这个字符串指定的文件(假设该文件是 YAML 格式的),并使用yaml.safe_load
函数将其内容解析为一个字典。这个字典包含了超参数的具体值,例如学习率、权重衰减等。 -
记录超参数信息:
LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
这一行的作用是将所有的超参数信息打印到日志中。它使用
colorstr
函数为输出的文字加上颜色,提升可读性。通过列表推导式', '.join(f'{k}={v}' for k, v in hyp.items())
,将字典中的每个键值对格式化为key=value
的形式,并用逗号连接起来。 -
保存超参数副本:
opt.hyp = hyp.copy() # for saving hyps to checkpoints
最后,将读取到的超参数字典做一个拷贝,保存到
opt
对象的hyp
属性中。这种做法方便后续将这些超参数保存在检查点(checkpoints)中,以便于后续实验或调试使用。
总结
这段代码的主要功能是:
- 根据提供的超参数文件路径加载超参数配置。
- 将加载的超参数信息记录到日志中,以方便观察和调试。
- 将超参数保存在
opt
对象中,以便于在训练过程中使用和保存。
通过这一段代码,模型训练时可以灵活地设置和调整超参数,提升模型性能。
4)设置参数的保存路径
# Save run settings
if not evolve:
yaml_save(save_dir / 'hyp.yaml', hyp)
yaml_save(save_dir / 'opt.yaml', vars(opt))
这段代码的主要功能是保存当前训练运行的配置和超参数设置。现逐步分解并详细解释代码。
逐步分解与解释
-
条件判断:
if not evolve:
- 这行代码检查一个名为
evolve
的变量。如果evolve
为False
,则执行下面的代码。evolve
通常用于指示是否正在进行超参数的进化(即动态调整超参数的训练过程),如果不是在进化阶段,则保存当前的训练配置。
- 这行代码检查一个名为
-
保存超参数:
yaml_save(save_dir / 'hyp.yaml', hyp)
- 这一行使用
yaml_save
函数将超参数保存到一个 YAML 格式的文件中。save_dir
表示保存目录,'hyp.yaml'
是要保存的文件名,hyp
是包含超参数的字典。 - 这样做可以确保用户在下次训练时可以恢复相同的超参数设置,方便追踪和复现实验结果。
- 这一行使用
-
保存其他配置选项:
yaml_save(save_dir / 'opt.yaml', vars(opt))
- 这一行同样使用
yaml_save
函数,但这次保存的是训练选项(选项配置),vars(opt)
将opt
对象转换为一个字典,该对象通常包含与训练过程相关的参数(例如模型配置、数据集路径等)。 - 将选项保存下来,可以使得用户能够通过该文件理解当前训练的具体设置和环境。
- 这一行同样使用
总结
这一代码片段的主要功能是保存训练运行的超参数和选项设置。它通过条件判断确定是否处于超参数进化阶段,仅在非进化模式下将参数保存为 YAML 格式的文件。这种做法确保了当前训练的配置能够被记录和复用,提升了实验的可重复性和可追溯性。
5) 设置日志记录器
# Loggers
data_dict = None
if RANK in {-1, 0}:
loggers = Loggers(save_dir, weights, opt, hyp, LOGGER) # loggers instance
# Register actions
for k in methods(loggers):
callbacks.register_action(k, callback=getattr(loggers, k))
# Process custom dataset artifact link
data_dict = loggers.remote_dataset
if resume: # If resuming runs from remote artifact
weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size
这段代码的主要功能是设置日志记录器,并在训练过程中注册回调操作。以下是对代码的逐步分解和详细解释:
逐步分解
-
初始化
data_dict
data_dict = None
这里初始化了一个变量
data_dict
,用于在后续处理中存储数据字典的引用。起初将其设置为None
。 -
检查
RANK
的值if RANK in {-1, 0}:
仅在分布式训练时的主节点(
RANK
为-1或0)执行下面的操作。这是因为只有主节点负责处理日志记录。 -
创建日志记录器实例
loggers = Loggers(save_dir, weights, opt, hyp, LOGGER) # loggers instance
创建了一个
Loggers
类的实例,传入了一些参数:save_dir
: 保存日志的目录。weights
: 模型权重。opt
: 训练选项。hyp
: 超参数。LOGGER
: 日志记录器实例。
loggers
实例将用于记录训练的相关信息,例如损失、准确率等。 -
注册回调操作
for k in methods(loggers): callbacks.register_action(k, callback=getattr(loggers, k))
通过
methods(loggers)
获取loggers
实例中所有可用的方法,并为每个方法在callbacks
对象中注册一个操作。getattr(loggers, k)
调用相应的方法以便在有事件触发时可以执行。这样可以实现训练过程中的自动日志记录。 -
处理数据集的远程链接
data_dict = loggers.remote_dataset
这里从
loggers
实例中获取远程数据集的链接,并将其赋值给data_dict
。这可能用于保存训练与验证数据的引用。 -
如果在恢复训练,更新相关参数
if resume: # If resuming runs from remote artifact weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size
如果设置了
resume
,说明正在从之前的训练状态中恢复。这时,将当前的weights
、epochs
、hyp
和batch_size
参数更新为训练选项中的值,以确保在恢复训练时使用的参数是最新的。
总结
这段代码的主要功能是设置训练过程中使用的日志记录器。该日志记录器通过记录重要的训练信息(如损失与准确性)来帮助监控模型的性能。代码通过主节点检查(RANK
)来确保只有必要的部分执行,创建日志记录实例,并注册相应的方法以便在训练过程中的特定事件发生时能够记录数据。此外,它还考虑了从远程位置恢复训练时需要更新参数的情况。这些功能为保障训练过程的透明度与可追溯性提供了支持。
6)训练前的配置
# Config
plots = not evolve and not opt.noplots # create plots
cuda = device.type != 'cpu'
init_seeds(opt.seed + 1 + RANK, deterministic=True)
with torch_distributed_zero_first(LOCAL_RANK):
data_dict = data_dict or check_dataset(data) # check if None
train_path, val_path = data_dict['train'], data_dict['val']
nc = 1 if single_cls else int(data_dict['nc']) # number of classes
names = {0: 'item'} if single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names
is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt') # COCO dataset
这段代码主要是在配置模型训练之前的一些设置,下面是对每行代码的逐步分解与详细解释:
-
plots = not evolve and not opt.noplots # create plots
- 该行代码用来判断是否需要生成训练过程的绘图。当不处于超参数进化(
evolve
为True
)且opt.noplots
为False
时,变量plots
将被设为True
,表示生成绘图;反之,则设为False
。
- 该行代码用来判断是否需要生成训练过程的绘图。当不处于超参数进化(
-
cuda = device.type != 'cpu'
- 这里判断当前使用的设备是否为 GPU。如果
device.type
不是'cpu'
,则cuda
变量将被设为True
,表示模型将在 GPU 上训练;否则,为False
,表示使用 CPU。
- 这里判断当前使用的设备是否为 GPU。如果
-
init_seeds(opt.seed + 1 + RANK, deterministic=True)
- 此行调用
init_seeds
函数来初始化随机种子,以确保实验的可重现性。使用的种子是给定的opt.seed
加上1
和当前进程的编号RANK
,并且设置为确定性计算(deterministic=True
),使得每次运行结果一致。
- 此行调用
-
with torch_distributed_zero_first(LOCAL_RANK):
- 这个上下文管理器用于控制在分布式训练中,保证特定任务的执行顺序。
LOCAL_RANK
用于识别当前设备的序号,这样可以避免不同设备间的竞态条件。
- 这个上下文管理器用于控制在分布式训练中,保证特定任务的执行顺序。
-
data_dict = data_dict or check_dataset(data) # check if None
- 这一行代码检查
data_dict
是否为None
。如果是,则调用check_dataset(data)
函数来验证和获取数据集信息,并赋值给data_dict
。这样确保训练时有有效的数据集信息。
- 这一行代码检查
-
train_path, val_path = data_dict['train'], data_dict['val']
- 从
data_dict
中获取训练集和验证集的路径,分别赋值给train_path
和val_path
。这两个路径将用于后续的数据加载。
- 从
-
nc = 1 if single_cls else int(data_dict['nc']) # number of classes
- 根据是否是单类训练(
single_cls
),确定类别数量(nc
)。若是单类,则nc
设置为1
;若为多类,从data_dict
中读取类别数量。
- 根据是否是单类训练(
-
names = {0: 'item'} if single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names
- 此行根据是否是单类训练来定义类别名称。如果是单类并且
data_dict['names']
的长度不为1,则将名称设置为默认的0: 'item'
;否则,直接使用data_dict['names']
来获取所有类别名称。
- 此行根据是否是单类训练来定义类别名称。如果是单类并且
-
is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt') # COCO dataset
- 该行代码检查验证集的路径
val_path
是否为字符串类型,并且是否以'coco/val2017.txt'
结尾,从而判断当前的数据集是否为 COCO 数据集。
- 该行代码检查验证集的路径
总结
总体而言,这段代码负责模型训练前的配置,涉及生成绘图的设置、设备类型判断、随机种子的初始化、数据集的检查与路径设置、类别的确定以及是否使用 COCO 数据集的判断。这些配置将为后续的模型训练奠定基础,确保数据环境的准备和模型参数的准确。
总之,整个初始化步骤主要解析各种yaml的参数+创建训练权重目录和保存路径+ 读取超参数配置文件 + 设置保存参数保存路径 + 加载数据配置信息 + 加载日志信息(logger + wandb) + 加载其他参数(plots、cuda、nc、names、is_coco)
第二步:初始化网络模型
1)创建一个YOLOv5模型实例
# Model
check_suffix(weights, '.pt') # check weights
pretrained = weights.endswith('.pt')
if pretrained:
with torch_distributed_zero_first(LOCAL_RANK):
weights = attempt_download(weights) # download if not found locally
ckpt = torch.load(weights, map_location='cpu') # load checkpoint to CPU to avoid CUDA memory leak
model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else [] # exclude keys
csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32
csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # intersect
model.load_state_dict(csd, strict=False) # load
LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}') # report
else:
model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
amp = check_amp(model) # check AMP
下面是对给定代码的逐步分解和详细解释:
-
检查权重文件后缀:
check_suffix(weights, '.pt') # check weights
这行代码用于检查
weights
变量中指定的文件名是否以.pt
结尾,以确保它是一个有效的PyTorch权重文件格式。 -
确定是否使用预训练权重:
pretrained = weights.endswith('.pt')
这里,
pretrained
变量被设置为一个布尔值,指示权重文件是否为预训练模型的权重。 -
如果是预训练权重:
if pretrained: with torch_distributed_zero_first(LOCAL_RANK): weights = attempt_download(weights) # download if not found locally
如果
pretrained
为True
,那么代码进入此块。torch_distributed_zero_first(LOCAL_RANK)
是一个上下文管理器,确保在分布式训练的某些场景下,只有一个进程会尝试下载权重。attempt_download(weights)
函数尝试从远程位置下载指定的权重文件,如果本地不存在该文件。 -
加载检查点:
ckpt = torch.load(weights, map_location='cpu') # load checkpoint to CPU to avoid CUDA memory leak
使用
torch.load
函数加载权重文件,并将其指定为在CPU上加载,以防止潜在的CUDA内存泄露。 -
创建模型:
model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
根据配置文件
cfg
或从加载的检查点中提取的YAML文件创建一个新的模型实例。ch=3
表示输入通道数,nc
是类别数量,anchors
则由超参数字典hyp
中的锚点获取。模型最后被移动到指定的设备(如CUDA或CPU)。 -
处理从检查点加载的状态字典:
exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else [] # exclude keys csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32 csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # intersect model.load_state_dict(csd, strict=False) # load
exclude
用于指定在加载模型状态字典时需要排除的键。随后,通过将检查点中的状态字典转换为浮点格式并与模型的当前状态字典进行交集检查,以确保只加载兼容的部分。最后,加载这些状态字典到模型中,strict=False
允许不严格地匹配键。 -
记录加载信息:
LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}') # report
这行代码记录了成功加载的参数数量,以便在训练日志中提供反馈。
-
如果不是预训练权重:
else: model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
如果输入的
weights
不是预训练模型权重,则会直接根据指定的配置cfg
创建一个新的模型实例。 -
检查是否支持自动混合精度(AMP):
amp = check_amp(model) # check AMP
最后,调用
check_amp(model)
以检查模型是否支持自动混合精度,这对于加速训练过程和减少内存使用是有帮助的。
总结
这段代码的主要功能是创建一个YOLOv5模型实例,并根据是否有有效的预训练权重来加载相应的状态字典。如果有预训练权重,则下载并加载这些权重;如果没有,则创建一个新的模型。这部分代码还负责确保模型适应分布式训练环境,并支持混合精度训练。整体上,这段代码为后续的训练过程做好了准备。
这里使用预训练权重参数,是类似于迁移学习。
预训练的模型是检测coco数据集的模型,数据集中有80个类别,而自己的训练集类别以及类别的数量,并不与coco数据集相同。所以要先加载一个新的模型,把预训练的参数加载至模型作为初始参数,再把识别的类别改成自己的数据集要识别的类别。接下来将预训练参数中与新模型中相同的参数加载至模型。
2) 冻结模型中的某些层
# Freeze
freeze = [f'model.{x}.' for x in (freeze if len(freeze) > 1 else range(freeze[0]))] # layers to freeze
for k, v in model.named_parameters():
v.requires_grad = True # train all layers
# v.register_hook(lambda x: torch.nan_to_num(x)) # NaN to 0 (commented for erratic training results)
if any(x in k for x in freeze):
LOGGER.info(f'freezing {k}')
v.requires_grad = False
这段代码的目的是在训练过程中冻结某些网络层的参数,以便在训练时不更新这些层的权重。以下是代码的逐步分解和详细解释:
逐步分解
-
定义冻结的层:
freeze = [f'model.{x}.' for x in (freeze if len(freeze) > 1 else range(freeze[0]))] # layers to freeze
- 这里的
freeze
是一个列表推导式,根据原始的freeze
列表(可能是数字或字符串)的内容生成一个新的列表。 - 如果
freeze
的长度大于 1,则遍历这个列表;否则,将其转换为一个范围,从 0 到freeze[0]
,表示需要冻结的层的索引(通常是由外部参数指定的)。 - 结果是
freeze
包含了需要冻结的层标识符,例如model.0.
,model.1.
等。
- 这里的
-
迭代模型中的参数:
for k, v in model.named_parameters():
- 遍历模型的所有参数,
k
是参数的名称,v
是参数的值(即Tensor)。named_parameters
返回一个包含模型所有参数的可迭代对象。
- 遍历模型的所有参数,
-
允许训练所有层:
v.requires_grad = True # train all layers
- 默认情况下,所有层的
requires_grad
属性都设置为True
,这表示在反向传播时都会计算这些参数的梯度。
- 默认情况下,所有层的
-
冻结特定的层:
if any(x in k for x in freeze): LOGGER.info(f'freezing {k}') v.requires_grad = False
- 通过条件判断,检查参数的名称
k
是否在需要冻结的层列表中。 - 如果是,记录日志,显示正在冻结的层的名称
k
,然后将requires_grad
属性设置为False
,从而阻止这些层的权重在训练过程中被更新。
- 通过条件判断,检查参数的名称
总结
这段代码的主要功能是允许用户选择性地冻结模型中的某些层,以便在训练时不对这些层的参数进行更新。这通常用于迁移学习中,用户可能希望保留预训练模型的某些层权重,而只训练特定的层。冻结层通过设置 requires_grad = False
来实现,使得这些层在梯度计算中被忽略,从而提高训练效率并减少过拟合的风险。代码中也包含了日志记录功能,以便跟踪哪些层被冻结。
冻结层的原理是通过设置每个层参数中的requires_grad属性实现的。
- 若require_grad为True,在反向传播时就会求出此tensor的梯度
- 若require_grad为False,则不会求该tensor的梯度。冻结就是通过对某些层不求梯度实现的。默认不进行参数冻结
3)检查图像尺寸
# Image size
gs = max(int(model.stride.max()), 32) # grid size (max stride)
imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # verify imgsz is gs-multiple
这段代码分为两行,主要涉及图像尺寸的处理,具体步骤和功能如下:
gs = max(int(model.stride.max()), 32) # grid size (max stride)
- 功能:计算网格尺寸
gs
,它是模型的最大步幅(stride)和32中的较大值。 - 详细解释:
model.stride.max()
:获取模型检测层的最大步幅。这表示在图像上进行下采样的最大比例,通常影响特征图的尺寸,这对于YOLO(You Only Look Once)模型非常重要,因为它需要在不同的尺度上处理对象。int(...)
:将步幅值转换为整数。max(..., 32)
:确保即使在步幅较小的情况下,gs
也至少为32。这是为了避免在训练过程中出现过小的网格大小,可能会影响模型的性能。
imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # verify imgsz is gs-multiple
- 功能:验证输入的图像大小
imgsz
是否是gs
的倍数,并可能将其调整为适合的值。 - 详细解释:
check_img_size(opt.imgsz, gs, floor=gs * 2)
:此函数会检查传入的opt.imgsz
(用户指定的图像大小)是否能被gs
整除。如果不能,它会根据gs
的值进行调整,以确保图像大小符合网络的要求。floor=gs * 2
的参数意味着,如果需要调整,所获得的图像大小不应小于gs * 2
,即保证图像有足够的分辨率。
这段代码的主要功能是确定合适的网络输入网格大小和图像大小。首先,通过获取模型的最大步幅来计算网格大小,确保其不小于32。然后,通过验证用户指定的图像大小是否符合模型要求,以保证训练过程中的有效性和稳定性。对于YOLO模型而言,正确配置这些参数是确保模型能够正确识别和定位目标的关键。
4)动态设置批处理
# Batch size
if RANK == -1 and batch_size == -1: # single-GPU only, estimate best batch size
batch_size = check_train_batch_size(model, imgsz, amp)
loggers.on_params_update({"batch_size": batch_size})
这段代码的主要功能是根据当前设备的状态动态估算最佳的训练批量大小(batch size),具体过程如下:
-
条件判断(if 语句):
if RANK == -1 and batch_size == -1:
RANK
是一个用于指示当前设备在分布式训练中的排名变量(rank),当RANK
为-1
时,表示当前代码在单个 GPU 上运行。batch_size
是指批量大小。-1
表示用户未指定批量大小,因此代码需要在没有用户输入的情况下做出估计。
-
动态估算最佳批量大小:
batch_size = check_train_batch_size(model, imgsz, amp)
- 这里调用了
check_train_batch_size
函数来估算一个合适的批量大小。 model
参数是当前的 YOLOv5 模型,imgsz
是指定的训练图像尺寸,amp
指是否使用自动混合精度(Automatic Mixed Precision)来加速训练。- 函数会根据模型的结构、GPU的内存、图像尺寸等信息来返回一个适合的批量大小,从而保证在训练时不会超出 GPU 内存限制。
- 这里调用了
-
更新日志:
loggers.on_params_update({"batch_size": batch_size})
- 此行代码调用了
loggers
对象的方法on_params_update
,并将新的batch_size
作为参数传递。 - 这通常用于记录当前的参数设置,便于后续分析和调试。
- 此行代码调用了
这段代码的主要功能是: 在单 GPU 模式下,如果用户没有指定批量大小,算法会动态计算一个最佳的批量大小,以确保训练过程的顺利进行。通过调用一个专门的函数来进行 estimations,并通过日志记录更新当前参数,这有助于更好地管理和监控训练过程。
5) 设置优化器
# Optimizer
nbs = 64 # nominal batch size
accumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizing
hyp['weight_decay'] *= batch_size * accumulate / nbs # scale weight_decay
optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['weight_decay'])
当然可以。让我们逐步分解并详细解释这段代码。
-
nbs = 64
:- 这行代码定义了一个名为
nbs
(nominal batch size,名义批量大小)的变量,并将其设置为64。这通常是在训练过程中希望的标准批量大小。
- 这行代码定义了一个名为
-
accumulate = max(round(nbs / batch_size), 1)
:- 这行代码计算了一个名为
accumulate
的变量。 nbs / batch_size
:这里计算的是“名义批量大小”与实际批量大小的比值。如果实际批量大小batch_size
比较小,这个值就会大于1,因此需要被“累积”多次来达到相同的训练效果。round(...)
函数对计算的比值进行四舍五入,确保得到一个整数。max(..., 1)
确保即使batch_size
非常大(大于nbs
),accumulate
的值至少为1。这意味着如果batch_size
更小,累积次数会增加。
- 这行代码计算了一个名为
-
hyp['weight_decay'] *= batch_size * accumulate / nbs
:- 这行代码调整损失函数中权重衰减(weight decay)的值。
batch_size * accumulate / nbs
计算了一个比例,这个比例用于调整weight_decay
。它的目的是适应当前的批量大小和名义批量大小,使模型在不同的环境下保持训练的稳定性。- 通过*=
来更新
hyp['weight_decay']`的值,这确保了模型在训练中的正则化是适应当前配置的。
-
optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['weight_decay'])
:- 这行代码初始化优化器。
smart_optimizer
是一个函数,其中传入的参数包括:model
: 当前正在训练的模型。opt.optimizer
: 指定的优化器类型(如SGD、Adam等)。hyp['lr0']
: 初始学习率。hyp['momentum']
: 动量值(如果使用的优化器支持)。hyp['weight_decay']
: 调整后的权重衰减值。
- 这个函数将返回一个配置好的优化器。
这段代码的主要目的是为神经网络训练过程设置和初始化优化器。它首先定义了名义批量大小,然后根据实际批量大小计算累积损失的次数,以适应训练的稳定性,并调整权重衰减的值,以便在不同的批量大小下实现一致的正则化效果。最后,它调用smart_optimizer
函数来创建一个优化器,准备在模型训练过程中使用。这一切的目的是优化模型的训练过程,以提高性能和稳定性。
这段代码是参数设置(nbs、accumulate、hyp[‘weight_decay’])
nbs指的是nominal batch size,名义上的batch_size。这里的nbs跟命令行参数中的batch_size不同,命令行中的batch_size默认为16,nbs设置为64。
accumulate 为累计次数,在这里 nbs/batch_size(64/16)计算出 opt.batch_size输入多少批才达到nbs的水平。简单来说,nbs为64,代表想要达到的batch_size,这里的数值是64;batch_size为opt.batch_size,这里的数值是16。64/16等于4,也就是opt.batch_size需要输入4批才能达到nbs,accumulate等于4。(round表示四舍五入取整数,而max表示accumulate不能低于1。)
当给模型喂了4批图片数据后,将四批图片数据得到的梯度值,做累积。当每累积到4批数据时,才会对参数做更新,这样就实现了与batch_size=64时相同的效果。
最后还要做权重参数的缩放,因为batch_size发生了变化,所有权重参数也要做相应的缩放。
6)设置学习率
# 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
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)
这段代码主要涉及学习率调度器(scheduler)的设置,目的是在模型训练过程中动态调整学习率。让我们逐步分解和详细解释这段代码:
-
判断是否使用余弦学习率调度(cos_lr):
if opt.cos_lr:
这一行代码检查用户是否选择了使用余弦学习率调度。如果用户在参数中设置了
--cos-lr
,则opt.cos_lr
为真。 -
定义余弦调度学习率函数:
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
如果使用余弦调度,调用
one_cycle
函数来生成一个学习率变化的函数。这个函数用于在训练过程中根据训练的进度(epoch)调整学习率。1
表示训练开始时的学习率比例。hyp['lrf']
是最终学习率的比例。epochs
是训练的总轮数。
-
定义线性学习率调度函数:
else: lf = lambda x: (1 - x / epochs) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear
如果不使用余弦调度,定义一个线性学习率调整函数
lf
。这个 lambda 函数根据当前的 epoch 的进度线性地调整学习率:x
是当前的 epoch。(1 - x / epochs)
逐渐减少,表示随着训练的进行,学习率开始减小。hyp['lrf']
是设置的最小学习率,确保学习率不会降到某个阈值以下。
-
创建学习率调度器:
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)
使用
LambdaLR
创建一个学习率调度器scheduler
,将定义好的lf
函数作为参数传入。- 这个调度器会根据
lf
的定义动态调整给定优化器optimizer
的学习率。 # plot_lr_scheduler(optimizer, scheduler, epochs)
是注释,暗示可以将调度函数与学习率的变化过程绘制出来,这将在日后分析学习率的效果时非常有用。
- 这个调度器会根据
该段代码设置了一个学习率调度器,用于在模型训练过程中动态调整学习率。根据用户的选项,调度器可以选择使用余弦函数或线性函数的方式来调整学习率。学习率是影响深度学习模型收敛速度和效果的重要超参数,因此这种动态调整机制有助于提升模型的训练效果和稳定性。
在训练过程中变更学习率可能会让训练效果更好,YOLOv5提供了两种学习率变化的策略:
- 一种是linear_lr(线性学习率),是通过线性插值的方式调整学习率
- 另一种则是One Cycle(余弦退火学习率),即周期性学习率调整中,周期被设置为1。在一周期策略中,最大学习率被设置为 LR Range test 中可以找到的最高值,最小学习率比最大学习率小几个数量级。这里默认one_cycle。
7) 设置EMA
# EMA
ema = ModelEMA(model) if RANK in {-1, 0} else None
这段代码的主要功能是初始化一个指数移动平均(Exponential Moving Average, EMA)模型,用于在训练过程中对模型权重的平滑处理。
-
注释部分:
# EMA
这里的注释指出该代码块与“指数移动平均”有关,EMA是一种常用的技巧,用于提高模型的鲁棒性和性能。
-
条件表达式:
ema = ModelEMA(model) if RANK in {-1, 0} else None
ModelEMA(model)
:这是一个调用ModelEMA
类的构造函数,传入当前的模型model
作为参数。ModelEMA
类通常用于创建一个EMA模型,其权重会根据主模型的权重进行更新,以实现更平滑的学习过程。if RANK in {-1, 0}
:这是一个条件判断,用于判断当前训练环境的并行状态。RANK
是分布式训练中的一个参数,它表示当前进程的唯一标识符。在定义分布式训练时,RANK
为-1通常表示单进程(单GPU)状态,而0表示主进程。在多进程环境中,其他进程的RANK
值会大于0。
else None
:如果RANK
不是-1或0(即在非主进程的情况下),将ema
设置为None
,意味着在非主进程中不需要使用EMA模型。
这段代码的主要功能是为了在训练过程中为主进程(单GPU或主进程)初始化一个EMA模型实例,以便利用EMA对训练过程中得到的权重进行平滑处理。EMA模型可以增强模型的泛化能力,对于跟踪模型性能和在验证时选择更好的权重有帮助。而在其他非主进程中,因为无需进行指数移动平均处理,因此ema
被设置为None
。
EMA为指数加权平均或滑动平均。其将前面模型训练权重,偏差进行保存,在本次训练过程中,假设为第n次,将第一次到第n-1次以指数权重进行加和,再加上本次的结果,且越远离第n次,指数系数越大,其所占的比重越小。
8)设置断点续传参数
# Resume
best_fitness, start_epoch = 0.0, 0
if pretrained:
if resume:
best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume)
del ckpt, csd
这段代码的功能是处理模型的恢复(resume)功能,具体来说,是在训练过程中从上一次训练的检查点继续训练。下面是逐步分解和详细解释:
-
初始化变量:
best_fitness, start_epoch = 0.0, 0
这行代码初始化了两个变量:
best_fitness
:用于存储当前最佳的模型表现(fitness)。通常表现为模型在验证集上的指标得分,比如精度(accuracy)或均值平均精度(mAP)。start_epoch
:用于记录训练开始的轮次(epoch),如果我们要恢复训练,这个值将指示从哪个轮次继续训练。
-
检查预训练模型:
if pretrained:
这行代码检查当前是否使用预训练模型(
pretrained
为真时),如果使用,则进入后续处理。 -
检查是否恢复训练:
if resume:
如果
resume
为真,表示用户希望从上一次训练的检查点重新开始训练。 -
调用恢复函数:
best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume)
这一行代码调用一个名为
smart_resume
的函数,并传入以下参数:ckpt
:包含模型状态、优化器状态、历史训练信息等的检查点文件。optimizer
:当前优化器的实例,可能需要恢复优化器的状态。ema
:表示模型的指数移动平均(Exponential Moving Average),用于帮助稳定训练。weights
:当前的权重文件路径。epochs
:总训练轮数。resume
:表示是否恢复训练的标志。
smart_resume
函数返回新的best_fitness
、start_epoch
和epochs
。这些值将根据检查点的状态进行更新。 -
删除不再需要的变量:
del ckpt, csd
这行代码删除了
ckpt
和csd
这两个变量,以释放内存。这些变量在恢复训练之后已经不再需要。
这段代码的主要功能是实现模型的恢复训练功能。当使用预训练模型时,如果用户希望从上次训练中继续,则通过调用smart_resume
函数从检查点中加载模型及其优化器状态,并更新最佳模型表现和起始训练轮次。最后,通过删除不再需要的变量来管理内存。这使得在长时间训练的过程中能够顺利恢复训练,极大地方便了模型训练的灵活性和效率。
9)设置DP模式
# DP mode
if cuda and RANK == -1 and torch.cuda.device_count() > 1:
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)
这段代码的功能是在特定的条件下启用数据并行(Data Parallel)模式。以下是逐步分解和详细解释。
-
条件检查:
if cuda and RANK == -1 and torch.cuda.device_count() > 1:
cuda
:这是一个布尔值,表示是否启用了CUDA(即是否可以使用GPU进行计算)。RANK == -1
:在分布式训练中,RANK
通常表示当前进程的编号。在单GPU运行时,其值为-1,因此这个条件确保我们在单GPU模式下运行(而不是在多GPU分布式模式下)。torch.cuda.device_count() > 1
:检查可用的CUDA设备数量是否大于1,即是否有多个GPU可用。
如果上述条件都满足,意味着当前环境是单个进程环境,并且有多个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.')
- 这行代码记录一条警告信息,提示用户在这种情况下使用数据并行(DataParallel)模式并不是最佳实践。推荐的做法是使用
torch.distributed.run
来进行多GPU的分布式数据并行(Distributed Data Parallel, DDP)。 - 提供了一个指向官方教程的链接,以便用户了解如何设置多GPU训练。
- 这行代码记录一条警告信息,提示用户在这种情况下使用数据并行(DataParallel)模式并不是最佳实践。推荐的做法是使用
-
启用数据并行:
model = torch.nn.DataParallel(model)
- 如果以上条件成立,此行代码会使用
torch.nn.DataParallel
来包装现有的模型,从而启用数据并行功能。这使得输入数据能够在多个GPU之间分配,利用多个GPU加速模型训练。
- 如果以上条件成立,此行代码会使用
这段代码的主要功能是在特定条件下(允许CUDA、多GPU环境且未使用分布式训练的情况下)启用PyTorch的DataParallel
模式。它首先检查CUDA环境及设备数量,并记录一条警告信息,提醒用户数据并行模式可能不是最佳选择,而推荐使用分布式数据并行方式进行训练。
DataParallel单机多卡模式自动将数据切分 load 到相应 GPU,将模型复制到相应 GPU,进行正向传播计算梯度并汇总。值得注意的是,模型和数据都需要先导入进 GPU 中,DataParallel 的 module 才能对其进行处理,否则会报错。
10)设置同步批量归一化(SyncBatchNorm)
# SyncBatchNorm
if opt.sync_bn and cuda and RANK != -1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
LOGGER.info('Using SyncBatchNorm()')
这段代码的主要功能是检查用户是否选择使用同步批量归一化(SyncBatchNorm),并在满足条件时将模型中的批量归一化层转换为同步批量归一化层,适用于分布式训练。
-
if opt.sync_bn and cuda and RANK != -1:
- 这一行是一个条件判断语句。
opt.sync_bn
:检查用户是否在命令行参数中指定了使用同步批量归一化。cuda
:检查当前是否可用CUDA设备,通常用于判断是否使用GPU进行训练。RANK != -1
:检查当前的训练进程是否是分布式训练中的一部分。RANK
一般用于表示进程的编号,-1
表示单进程(非分布式训练)。
-
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
- 如果上述条件成立,就会执行这一行。
torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
:将模型中的所有普通的批量归一化(BatchNorm)层转换为同步批量归一化层。这种修改使得在多GPU的环境下,每个GPU计算的均值和方差能够进行同步,以提高模型在训练时的稳定性和收敛速度。.to(device)
:将转换后的模型移动到指定的设备(如GPU)上,以保证模型在训练时能够使用GPU加速。
-
LOGGER.info('Using SyncBatchNorm()')
- 这一行用于记录日志,表明正在使用同步批量归一化。这对于调试和监控训练过程非常有用,可以帮助用户了解训练设置。
这段代码的主要功能是将模型中的批量归一化层转换为同步批量归一化层,以提升在分布式训练中的性能。在多GPU训练时,使用同步批量归一化能够确保每个GPU上的数据都依据全局的统计信息进行归一化,从而带来更好的训练稳定性和效果。这段代码适用于那些开启了分布式训练并选择使用同步批量归一化的情况。
SyncBatchNorm主要用于解决多卡归一化同步问题,每张卡单独计算均值,然后同步,得到全局均值。用全局均值计算每张卡的方差,然后同步即可得到全局方差,但两次会消耗时间挺长。
加载网络模型小结:
(1)载入模型:载入模型(预训练/不预训练) + 检查数据集 + 设置数据集路径参数(train_path、test_path) + 设置冻结层
(2)优化器:参数设置(nbs、accumulate、hyp[‘weight_decay’]) + 分组优化(pg0、pg1、pg2) + 选择优化器 + 为三个优化器选择优化方式 + 删除变量
(3)学习率:线性学习率 + one cycle学习率 + 实例化 scheduler + 画出学习率变化曲线
(4)训练前最后准备:EMA +断点续训+ 迭代次数的加载 + DP +SyncBatchNorm)
第三步:初始化数据集
# Trainloader
train_loader, dataset = create_dataloader(train_path,
imgsz,
batch_size // WORLD_SIZE,
gs,
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)
labels = np.concatenate(dataset.labels, 0)
mlc = int(labels[:, 0].max()) # max label class
assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'
# Process 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:
if not opt.noautoanchor:
check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) # run AutoAnchor
model.half().float() # pre-reduce anchor precision
callbacks.run('on_pretrain_routine_end', labels, names)
该段代码的主要功能是准备训练和验证数据加载器,并进行一些必要的处理。这段代码在训练YOLOv5模型时被调用,具体的逐步分解和解释如下:
-
创建训练数据加载器:
train_loader, dataset = create_dataloader(train_path, imgsz, batch_size // WORLD_SIZE, gs, 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)
create_dataloader()
函数用于创建一个训练数据加载器和数据集对象,到两个对象。一个为train_loader,另一个为dataset。train_loader为训练数据加载器,可以通过for循环遍历出每个batch的训练数据。dataset为数据集对象,包括所有训练图片的路径,所有标签,每张图片的大小,图片的配置,超参数等等。- 参数说明:
train_path
:训练数据集的路径。imgsz
:训练时图像的大小。batch_size // WORLD_SIZE
:每个GPU的批大小,支持多GPU训练。gs
:网格大小(最大步幅)。single_cls
:是否训练为单类(布尔值)。hyp
:超参数。augment
:是否进行数据增强。cache
:是否缓存图像,取决于给定的选项。rect
:是否进行矩形训练。rank
:当前进程的排名,用于分布式训练。workers
:数据加载的工作线程数量。image_weights
:使用加权图像选择进行训练。quad
:四连数据加载器模式。prefix
:用于显示的前缀字符串。shuffle
:是否随机打乱数据。
-
合并标签:
labels = np.concatenate(dataset.labels, 0)
- 将数据集中所有图片的标签整合成一个数组,方便后续处理。
-
获取最大标签类:
mlc = int(labels[:, 0].max()) # max label class assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'
mlc
是标签中最大的类别索引。assert
语句检查最高标签类是否小于训练任务中的类别数量(nc
),确保没有超出范围。
-
处理验证数据加载器:
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]
- 仅在主进程(RANK为-1或0)创建验证数据加载器。
create_dataloader()
调用类似于训练数据加载器,但参数有所不同,以适应验证集的要求。
-
模型锚框检查与精度预处理:
if not resume: if not opt.noautoanchor: check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) # run AutoAnchor model.half().float() # pre-reduce anchor precision
- 如果不是继续训练,则进行锚框检查(
check_anchors()
),以自动调整锚框的大小。 - 将模型转为半精度(
half()
),以提高性能并节省内存。 -
这段代码主要是计算默认锚点anchor与数据集标签框的长宽比值
check_anchors计算默认锚点anchor与数据集标签框的长宽比值,标签的长h宽w与anchor的长h_a宽w_a的比值, 即h/h_a, w/w_a都要在(1/hyp[‘anchor_t’], hyp[‘anchor_t’])是可以接受的,如果标签框满足上面条件的数量小于总数的99%,则根据k-mean算法聚类新的锚点anchor。
- 如果不是继续训练,则进行锚框检查(
-
调用训练回调:
callbacks.run('on_pretrain_routine_end', labels, names)
CopyInsert
- 调用回调函数,运行训练前例程结束的相关事件。
这段代码的核心功能是设置和预处理YOLOv5模型的训练和验证数据加载器。它创建数据集和加载器、整合标签并确保标签数量的合理性。此外,还检查锚框的质量,以优化检测模型的性能,并在最后执行必要的回调。整个过程为后续模型的训练和验证做准备,确保数据的有效性和模型的性能。
第四步:分布式训练
# DDP mode
if cuda and RANK != -1:
model = smart_DDP(model)
这段代码的作用是针对分布式数据并行(Distributed Data Parallel, DDP)训练模式进行设置。以下是对代码的逐步分解和详细解释:
-
条件判断
if cuda and RANK != -1:
cuda
: 这是一个布尔值,表示当前是否使用GPU进行计算。如果为True
,则说明程序在CUDA设备(GPU)上运行。RANK
: 这是一个整数,表示当前进程在分布式训练中的身份。通常,RANK
为-1表示单个GPU训练,而非负值表示在多GPU环境下的进程编号。这里的判断是确保在使用多GPU的分布式模式下,代码将继续执行。
-
model = smart_DDP(model)
smart_DDP(model)
: 这个函数的作用是将模型包装成一个适合PyTorch的分布式数据并行(DDP)模型。DDP可以有效地把模型的训练过程分发到多个GPU上,进而加速训练。函数smart_DDP
根据当前的环境判断是否需要使用DDP,并处理相应的模型设置。
该代码段的核心功能是检查当前环境是否适合使用分布式数据并行训练,并在确认可以使用的情况下,将模型转换为DDP模式。通过进行这样的设置,代码能够支持在多个GPU上并行训练模型,从而提高训练效率并加快模型的收敛速度。分布式训练是一种常用的深度学习训练策略,特别是在处理大型数据集和复杂模型时,可以显著提高计算速度。
第五步:开始训练
1)初始化训练需要的模型参数
# Model attributes
nl = de_parallel(model).model[-1].nl # number of detection layers (to scale hyps)
hyp['box'] *= 3 / nl # scale to layers
hyp['cls'] *= nc / 80 * 3 / nl # scale to classes and layers
hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl # scale to image size and layers
hyp['label_smoothing'] = opt.label_smoothing
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights
model.names = names
这段代码的主要功能是对YOLOv5模型的某些属性进行初始化和调整,以适应特定的数据集和超参数设置。下面逐步分解并详细解释每一行代码。
-
nl = de_parallel(model).model[-1].nl
- 这一行首先调用
de_parallel(model)
,它将模型从数据并行模式中转换为单个模型对象。接着,.model[-1]
表示获取模型的最后一层(通常是检测层),通过.nl
获取该层的检测层数(number of detection layers,nl)。 - 这里通过
nl
来确定模型中需要调整的层数。
- 这一行首先调用
-
hyp['box'] *= 3 / nl
- 该行代码将超参数
hyp
中的box
(边界框损失权重)按照检测层数nl
进行缩放。原始的box
损失权重值乘以3/nl
,其目的是将损失权重调整到模型中的检测层数,以便适应不同的结构。
- 该行代码将超参数
-
hyp['cls'] *= nc / 80 * 3 / nl
- 此处调整的是类别损失的权重。
nc
表示数据集中类别的数量,80
是COCO数据集中类别的默认数量,这里用nc/80
来进行缩放,再乘以3/nl
,目的是为了平衡不同类别的损失在训练过程中的影响。
- 此处调整的是类别损失的权重。
-
hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl
- 这一行代码则是调整目标损失(object loss)的权重。通过
imgsz / 640
来考虑输入图像的大小变化,640
是默认的输入图像大小。接着对结果进行平方计算,是因为损失与像素数量成比例,因此需要平方比例,同时也进行层数的缩放。
- 这一行代码则是调整目标损失(object loss)的权重。通过
-
hyp['label_smoothing'] = opt.label_smoothing
- 这里将标签平滑的超参数
label smoothing
设置为命令行参数中指定的值。这有助于避免过拟合,让模型对标签的某些变动更具鲁棒性。
- 这里将标签平滑的超参数
-
model.nc = nc
- 将模型的类别数(
nc
)设置为数据集中的类别数,以便模型能够正确地配置输出和计算损失。
- 将模型的类别数(
-
model.hyp = hyp
- 将超参数
hyp
附加到模型上,以便在训练过程中使用这些参数。
- 将超参数
-
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc
- 计算每个类别的权重,并将其存储在模型中,以提高不平衡类别的训练效果。
labels_to_class_weights
是一个函数,它生成基于数据集标签的类别权重。使用.to(device)
将其移动到指定的设备(如GPU)。
- 计算每个类别的权重,并将其存储在模型中,以提高不平衡类别的训练效果。
-
model.names = names
- 将类别的名称(
names
)附加到模型中。这使得在模型推理和结果可视化时更容易理解哪些类别对应于输出。
- 将类别的名称(
这段代码的主要功能是在训练YOLOv5模型之前,根据指定的超参数和数据集特性来初始化模型的各个属性。包括:
- 确定和调整检测层的数量。
- 根据类别的数量和输入图像的尺寸调整损失函数的权重。
- 设置标签平滑、类别数、超参数以及类别权重。
通过这样的初始化,模型能够适当地应对即将开始的训练任务,提高模型的学习效率以及最终的检测性能。
这段代码主要是根据自己数据集的类别数设置分类损失的系数,位置损失的系数。设置类别数,超参数等操作
其中,
- box: 预测框的损失
- cls: 分类的损失
- obj: 置信度损失
- label_smoothing : 标签平滑
2)训练
# Start training
t0 = time.time()
nb = len(train_loader) # number of batches
nw = max(round(hyp['warmup_epochs'] * nb), 100) # number of warmup iterations, max(3 epochs, 100 iterations)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
last_opt_step = -1
maps = np.zeros(nc) # mAP per class
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 # do not move
scaler = torch.cuda.amp.GradScaler(enabled=amp)
stopper, stop = EarlyStopping(patience=opt.patience), False
compute_loss = ComputeLoss(model) # init loss class
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()
# Update image weights (optional, single-GPU only)
if opt.image_weights:
cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights
iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights
dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx
# Update mosaic border (optional)
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
mloss = torch.zeros(3, device=device) # mean losses
if RANK != -1:
train_loader.sampler.set_epoch(epoch)
pbar = enumerate(train_loader)
LOGGER.info(('\n' + '%11s' * 7) % ('Epoch', 'GPU_mem', 'box_loss', 'obj_loss', 'cls_loss', 'Instances', 'Size'))
if RANK in {-1, 0}:
pbar = tqdm(pbar, total=nb, bar_format=TQDM_BAR_FORMAT) # progress bar
optimizer.zero_grad()
for i, (imgs, targets, paths, _) in pbar: # batch -------------------------------------------------------------
callbacks.run('on_train_batch_start')
ni = i + nb * epoch # number integrated batches (since train start)
imgs = imgs.to(device, non_blocking=True).float() / 255 # uint8 to float32, 0-255 to 0.0-1.0
# Warmup
if ni <= nw:
xi = [0, nw] # x interp
# compute_loss.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 0 else 0.0, x['initial_lr'] * lf(epoch)])
if 'momentum' in x:
x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
# Multi-scale
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple)
imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
# Forward
with torch.cuda.amp.autocast(amp):
pred = model(imgs) # forward
loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size
if RANK != -1:
loss *= WORLD_SIZE # gradient averaged between devices in DDP mode
if opt.quad:
loss *= 4.
# Backward
scaler.scale(loss).backward()
# Optimize - https://pytorch.org/docs/master/notes/amp_examples.html
if ni - last_opt_step >= accumulate:
scaler.unscale_(optimizer) # unscale gradients
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0) # clip gradients
scaler.step(optimizer) # optimizer.step
scaler.update()
optimizer.zero_grad()
if ema:
ema.update(model)
last_opt_step = ni
# Log
if RANK in {-1, 0}:
mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G' # (GB)
pbar.set_description(('%11s' * 2 + '%11.4g' * 5) %
(f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))
callbacks.run('on_train_batch_end', model, ni, imgs, targets, paths, list(mloss))
if callbacks.stop_training:
return
# end batch ------------------------------------------------------------------------------------------------
# Scheduler
lr = [x['lr'] for x in optimizer.param_groups] # for loggers
scheduler.step()
if RANK in {-1, 0}:
# mAP
callbacks.run('on_train_epoch_end', epoch=epoch)
ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights'])
final_epoch = (epoch + 1 == epochs) or stopper.possible_stop
if not noval or final_epoch: # Calculate mAP
results, maps, _ = validate.run(data_dict,
batch_size=batch_size // WORLD_SIZE * 2,
imgsz=imgsz,
half=amp,
model=ema.ema,
single_cls=single_cls,
dataloader=val_loader,
save_dir=save_dir,
plots=False,
callbacks=callbacks,
compute_loss=compute_loss)
# Update best mAP
fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
stop = stopper(epoch=epoch, fitness=fi) # early stop check
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)
# Save model
if (not nosave) or (final_epoch and not evolve): # if save
ckpt = {
'epoch': epoch,
'best_fitness': best_fitness,
'model': deepcopy(de_parallel(model)).half(),
'ema': deepcopy(ema.ema).half(),
'updates': ema.updates,
'optimizer': optimizer.state_dict(),
'opt': vars(opt),
'git': GIT_INFO, # {remote, branch, commit} if a git repo
'date': datetime.now().isoformat()}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
if 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)
# EarlyStopping
if RANK != -1: # if DDP training
broadcast_list = [stop if RANK == 0 else None]
dist.broadcast_object_list(broadcast_list, 0) # broadcast 'stop' to all ranks
if RANK != 0:
stop = broadcast_list[0]
if stop:
break # must break all DDP ranks
# end epoch ----------------------------------------------------------------------------------------------------
# end training -----------------------------------------------------------------------------------------------------
这段代码是用于训练YOLOv5模型的核心部分。下面是对这段代码的逐步分解和详细解释:
主要功能
该代码段实现了YOLOv5模型的训练过程,包括数据加载、模型前向传递、损失计算、反向传播、优化步骤以及模型保存等。它支持多种功能,如学习率调度、图像权重更新、早期停止等,以适应不同训练需求。
-
训练开始的初始化:
t0 = time.time() nb = len(train_loader) # batch的数量 nw = max(round(hyp['warmup_epochs'] * nb), 100) # 预热迭代次数 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 # 避免调度器移动 scaler = torch.cuda.amp.GradScaler(enabled=amp) # 自动混合精度 stopper, stop = EarlyStopping(patience=opt.patience), False # 早期停止机制 compute_loss = ComputeLoss(model) # 初始化损失计算类
这段代码是训练前的热身准备,做一些参数的初始化
这里要提到两个点:
第一个是warmup。warmup是一种学习率的优化方法,最早出现在ResNet的论文中。简单来说,在模型刚开始训练时,使用较小的学习率开始摸索,经过几轮迭代后使用大的学习率加速收敛,在快接近目标时,再使用小学习率,避免错过目标。
第二个是早停机制。当训练一定的轮数后,如果模型效果未提升,就让模型提前停止训练。这里的默认轮数为100轮,判断模型的效果为fitness,fitness为0.1乘mAP@0.5加上0.9乘mAP@0.5:0.95。
-
训练开始的日志信息:
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...')
-
训练循环,通过model.train()函数告诉模型已经进入了训练阶段。因为有些层或模型在训练阶段与预测阶段进行的操作是不一样的,所以要通过model.train()函数用来声明,接下来是训练。
for epoch in range(start_epoch, epochs): # 遍历每个epoch callbacks.run('on_train_epoch_start') model.train() # 设置模型为训练模式
-
更新图像权重(如果启用):
if opt.image_weights: cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # 计算类权重 iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # 计算图像权重 dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # 随机选择索引
更新图片的权重。训练时有些类的准确率可能比较难以识别,准确率并不会很高。在更新图片权重时就会把这些难以识别的类挑出来,并为这个类产生一些权重高的图片,以这种方式来增加识别率低的类别的数据量。提高准确率。
-
初始化损失记录:
mloss = torch.zeros(3, device=device) # 初始化损失 if RANK != -1: train_loader.sampler.set_epoch(epoch) # 设置epoch用于数据加载 pbar = enumerate(train_loader) # 遍历训练数据
-
优化步骤:
optimizer.zero_grad() # 清除梯度 for i, (imgs, targets, paths, _) in pbar: # 迭代每个batch ... loss = compute_loss(pred, targets.to(device)) # 计算损失 scaler.scale(loss).backward() # 反向传播 if ni - last_opt_step >= accumulate: scaler.unscale_(optimizer) # 取消梯度缩放 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0) # 梯度修剪 scaler.step(optimizer) # 更新优化器 scaler.update() # 更新缩放器 optimizer.zero_grad() # 清除梯度 if ema: ema.update(model) # 更新EMA模型 last_opt_step = ni
这段代码是YOLOv5模型训练过程中的一个批次处理逻辑,主要用于前向传播、反向传播和优化模型参数。逐步分解如下:
初始化梯度:
optimizer.zero_grad()
在开始新的批次之前,将所有优化器中参数的梯度清零。这是因为PyTorch的梯度是累积的,每次反向传播后不自动清零,必须显式调用。
遍历数据批次:
for i, (imgs, targets, paths, _) in pbar:
通过pbar遍历数据加载器提供的一个批次,包括图像imgs、目标检测标注targets、图像路径paths。
执行回调:
callbacks.run('on_train_batch_start')
执行回调函数,通常用于记录训练状态或执行某些操作。
计算集成批次索引:
ni = i + nb * epoch
计算从开始训练到当前批次的总集成批次索引ni。
数据处理:
imgs = imgs.to(device, non_blocking=True).float() / 255
将图像数据转移到指定的设备(CPU或GPU),同时将数据类型转换为浮点型并进行归一化处理(将值从0-255缩放到0-1)。
首先分批加载训练数据,用ni计算当前迭代的次数,并作图片的归一化。
暖身步骤(Warmup):
# Warmup 热身训练
'''
热身训练(前nw次迭代),热身训练迭代的次数iteration范围[1:nw]
在前nw次迭代中, 根据以下方式选取accumulate和学习率
'''
if ni <= nw:
xi = [0, nw] # x interp
# compute_loss.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
# 遍历优化器中的所有参数组
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
"""
bias的学习率从0.1下降到基准学习率lr*lf(epoch),
其他的参数学习率从0增加到lr*lf(epoch).
lf为上面设置的余弦退火的衰减函数
"""
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']])
在训练的初期进行学习率的暖身调整。根据当前的批次索引ni与预设的暖身批次数nw进行插值计算,调整学习率和动量。然后进行热身训练(warmup),这里只对训练初期使用较小的学习率。对于bias参数组的学习率策略是从0.1逐渐降低至初始学习率,其余参数组则从0开始逐渐增长至初始学习率。
多尺度训练:
Multi-scale 设置多尺度训练,从imgsz * 0.5, imgsz * 1.5 + gs随机选取尺寸
# imgsz: 默认训练尺寸 gs: 模型最大stride=32 [32 16 8]
if opt.multi_scale: # 随机改变图片的尺寸
sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple)
# 下采样
imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
如果开启了多尺度训练,则随机生成一个新的图像大小,调整图像的形状以适应新的大小。
- imgz: 默认训练尺寸
- gs: 模型最大stride=32
前向传播:
# Forward 前向传播
with amp.autocast(enabled=cuda):
# 将图片送入网络得到一个预测结果
pred = model(imgs) # forward
# 计算损失,包括分类损失,objectness损失,框的回归损失
# loss为总损失值,loss_items为一个元组,包含分类损失,objectness损失,框的回归损失和总损失
loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size
if RANK != -1:
# 采用DDP训练,平均不同gpu之间的梯度
loss *= WORLD_SIZE # gradient averaged between devices in DDP mode
if opt.quad:
# 如果采用collate_fn4取出mosaic4数据loss也要翻4倍
loss *= 4.
使用自动混合精度(AMP)库进行前向传播,计算模型的预测结果pred,再通过损失函数计算损失值loss和损失项loss_items。
反向传播:
# Backward 反向传播 scale为使用自动混合精度运算
scaler.scale(loss).backward()
使用scaler进行损失的反向传播,将损失反向传播到模型中并计算每个参数的梯度。
优化步骤:
# Optimize 模型会对多批数据进行累积,只有达到累计次数的时候才会更新参数,再还没有达到累积次数时 loss会不断的叠加 不会被新的反传替代
if ni - last_opt_step >= accumulate:
'''
scaler.step()首先把梯度的值unscale回来,
如果梯度的值不是 infs 或者 NaNs, 那么调用optimizer.step()来更新权重,
否则,忽略step调用,从而保证权重不更新(不被破坏)
'''
scaler.step(optimizer) # optimizer.step 参数更新
# 更新参数
scaler.update()
# 完成一次累积后,再将梯度清零,方便下一次清零
optimizer.zero_grad()
if ema:
ema.update(model)
# 计数
last_opt_step = ni
执行优化,包括反向传播后的参数更新、梯度裁剪、更新last_opt_step等。依据是否属于分布式(DDP)训练模式来调整损失值、参数更新等。
但是要注意,在更新参数时这里有一个不一样的地方,并不会在每次反向传播时更新参数,而是做一定的累积,反向传播的结果并不会顶替上一次反向传播结果,而是做一个累积。完成一次积累后,再将梯度清零,方便下一次清零。这样做是为了以更小的batch_size实现更高的batch_size效果。
日志记录:
if RANK in [-1, 0]:
# 打印显存,进行的轮次,损失,target的数量和图片的size等信息
mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
# 计算显存
mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G' # (GB)
# 进度条显示以上信息
pbar.set_description(('%10s' * 2 + '%10.4g' * 5) % (
f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))
# 调用Loggers中的on_train_batch_end方法,将日志记录并生成一些记录的图片
callbacks.run('on_train_batch_end', ni, model, imgs, targets, paths, plots, opt.sync_bn)
如果是单进程或主进程,则记录日志,包括每个批次的平均损失,并更新进度条的描述。
这段代码是YOLOv5训练过程中处理每个批次的数据流和训练逻辑,主要功能包括:
1)执行梯度清零。
2)加载和处理当前批次的图像数据。
3)进行前向传播和计算损失。
4)进行反向传播和参数优化。
5)支持学习率暖身和多尺度训练。
6)记录和更新训练日志,便于监控训练进度和效果。
7. 更新学习率调度器
# Scheduler
lr = [x['lr'] for x in optimizer.param_groups] # for loggers
scheduler.step()
这段代码的功能是更新学习率调度器(scheduler)以调整训练过程中模型的学习率。下面是对这段代码逐步分解的详细解释:
-
获取学习率:
lr = [x['lr'] for x in optimizer.param_groups]
- 这一行代码遍历优化器(optimizer)的参数组(param_groups)中的每一个参数组,并从中提取当前的学习率('lr')。
optimizer.param_groups
是一个包含所有参数组的列表,每个参数组都是一个字典,包含了学习率、权重衰减等信息。- 通过列表推导式,我们将所有参数组的学习率存储到列表
lr
中。这个步骤通常是为了记录当前的学习率以便于后续的日志记录或调试。
-
更新调度器:
scheduler.step()
- 这一行调用了学习率调度器的
step()
方法,目的是更新学习率。根据设定的调度策略(如预热、余弦退火等),调度器会调整当前使用的学习率。 - 这个方法会根据当前的epoch或batch更新学习率,为后续的训练步骤提供新的学习率。
- 这一行调用了学习率调度器的
这段代码的主要功能是实现学习率的动态调整。在深度学习训练过程中,合理地更新学习率可以帮助模型更好地收敛,提高训练效率。通过调用 scheduler.step()
,模型在每个训练周期或步骤后可以根据预设的调度策略自动调整学习率,从而适应训练过程中的不同阶段。
8. 更新mAP:
# Update best mAP
fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
stop = stopper(epoch=epoch, fitness=fi) # early stop check
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)
这段代码主要用于训练模型的最后一步,负责更新最佳的平均精度(mean Average Precision, mAP),并检查是否满足提前终止训练的条件。我们逐行分析代码的工作原理:
-
计算当前的适应度值(fitness):
fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
- 首先,
results
是一个包含模型在验证集上测试结果的数组,通常包括精确度(Precision)、召回率(Recall)、mAP@0.5 和 mAP@0.5:0.95 等指标。 - 使用
np.array(results).reshape(1, -1)
将结果转换为适合计算适应度函数的二维数组格式。 fitness
函数会根据这些结果计算一个综合适应度值fi
。这个值通常是根据多个指标加权计算的,以便更好地反映模型的整体表现。
- 首先,
-
进行提前终止检查:
stop = stopper(epoch=epoch, fitness=fi) # early stop check
stopper
是一个用于判断是否提前停止训练的机制,该函数接收当前训练的轮次epoch
和计算得到的适应度值fi
。- 如果适应度在一定轮次内没有提高,
stopper
会返回True
,这会导致训练提前终止。
-
更新最佳适应度值:
if fi > best_fitness: best_fitness = fi
- 通过比较当前适应度
fi
与之前记录的最佳适应度best_fitness
,若当前适应度更高,则更新最佳适应度值。 - 这使得模型可以随时保存最佳历史性能,便于后续模型评估与选择。
- 通过比较当前适应度
-
记录训练日志:
log_vals = list(mloss) + list(results) + lr callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)
- 将当前轮次的损失
mloss
和验证结果results
、学习率lr
组合成一个列表log_vals
,便于记录和可视化。 callbacks.run('on_fit_epoch_end', ...)
用于在训练的某一轮结束后执行注册的操作,例如记录日志、更新可视化等等。这是一个设计模式,允许用户在不同训练阶段插入自定义操作。
- 将当前轮次的损失
这段代码主要负责更新和监控训练过程中的模型性能。通过计算当前模型在验证集上的适应度,检查是否需要提前停止训练,并在每个训练轮次结束时记录相关日志信息。其主要功能是确保模型在训练过程中能够自我评估和优化,从而达到最佳的性能。
9. 模型保存
# Save model
if (not nosave) or (final_epoch and not evolve): # if save
ckpt = {
'epoch': epoch,
'best_fitness': best_fitness,
'model': deepcopy(de_parallel(model)).half(),
'ema': deepcopy(ema.ema).half(),
'updates': ema.updates,
'optimizer': optimizer.state_dict(),
'opt': vars(opt),
'git': GIT_INFO, # {remote, branch, commit} if a git repo
'date': datetime.now().isoformat()}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
if 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)
这段代码主要用于保存训练过程中的模型以及相关的状态信息。以下是逐步分解和详细解释:
-
条件检查:
if (not nosave) or (final_epoch and not evolve):
nosave
:如果此标志为True
,则表示不保存模型。final_epoch
:检查当前是否是最后一个训练周期。evolve
:如果处于超参数进化的状态,通常不保存模型。- 本行代码的意思是:如果
nosave
为False
,或者当前是最后一个周期且没有进行超参数进化,则继续执行保存模型的操作。
-
构建检查点:
ckpt = { 'epoch': epoch, 'best_fitness': best_fitness, 'model': deepcopy(de_parallel(model)).half(), 'ema': deepcopy(ema.ema).half(), 'updates': ema.updates, 'optimizer': optimizer.state_dict(), 'opt': vars(opt), 'git': GIT_INFO, 'date': datetime.now().isoformat() }
epoch
:当前训练的周期数。best_fitness
:最佳适应度值,即模型在验证集上表现最好的结果。model
:深拷贝去并行化后的模型,并转换为半精度(half
),以减少内存占用。ema
:用以保存指数移动平均(Exponential Moving Average)模型的状态。updates
:用于记录 EMA 的更新次数。optimizer
:保存优化器的状态字典,以便后续恢复训练。opt
:保存训练时的参数设置。git
:保存与 Git 相关的信息(如远程地址、分支、提交等)。date
:当前时间的 ISO 格式字符串,记录保存操作的时间。
-
保存模型:
torch.save(ckpt, last)
- 使用
torch.save
方法将保存的检查点ckpt
存储到last
文件中,通常这个文件用于保存最新的模型状态。
- 使用
-
保存最佳模型:
if best_fitness == fi: torch.save(ckpt, best)
- 如果当前的适应度(
fi
)等于最佳适应度值,说明当前模型是最好的,那么将它存储到best
文件中。
- 如果当前的适应度(
-
定期保存模型:
if opt.save_period > 0 and epoch % opt.save_period == 0: torch.save(ckpt, w / f'epoch{epoch}.pt')
- 如果设置了保存周期(
save_period
)并且当前周期是保存的周期,便将当前检查点保存为epoch{epoch}.pt
文件。w
是用于保存模型的路径。
- 如果设置了保存周期(
-
删除检查点:
del ckpt
- 删除
ckpt
变量,以释放内存。
- 删除
-
回调函数:
callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi)
- 调用回调函数,传递有关模型保存的信息,这通常用于更新日志或其他后续处理。
这段代码实现了在训练过程中的模型保存功能。它根据指定的条件判断是否应保存模型,并将当前的训练状态(包括模型参数、优化器状态、训练配置等)存储到文件中,以便于后续恢复和分析。同时,它还提供了保存最佳模型和周期性保存模型的功能,从而确保训练过程中的重要状态不会丢失。
10.早期停止检查:
# EarlyStopping
if RANK != -1: # if DDP training
broadcast_list = [stop if RANK == 0 else None]
dist.broadcast_object_list(broadcast_list, 0) # broadcast 'stop' to all ranks
if RANK != 0:
stop = broadcast_list[0]
if stop:
break # must break all DDP ranks
这段代码的目的是在训练过程中实现“提前停止”(Early Stopping)功能,主要用于分布式数据并行(DDP)训练的场景。下面逐步分解并解释这段代码:
-
判断 RANK 是否不等于 -1:
if RANK != -1: # if DDP training
这一行代码检查当前训练是否在分布式模式下运行。
RANK
变量是用于标识当前进程的编号,-1表示单进程模式。 -
创建广播列表:
broadcast_list = [stop if RANK == 0 else None]
这行代码创建一个包含
stop
的列表。如果当前进程是主进程(RANK为0),则将stop
添加到列表中;否则,该列表的内容为None
。 -
广播操作:
dist.broadcast_object_list(broadcast_list, 0) # broadcast 'stop' to all ranks
这行代码通过分布式训练框架的广播机制,将
broadcast_list
列表中的内容广播到所有的进程中。这里,其他的进程会接收到主进程stop
的值。 -
更新停止条件:
if RANK != 0: stop = broadcast_list[0]
如果当前进程不是主进程(RANK不为0),则将
broadcast_list
中的第一个元素(即stop
的值)赋值给当前进程的stop
变量。这意味着所有进程都将拥有相同的stop
值。 -
提前停止逻辑:
if stop: break # must break all DDP ranks
最后,这段代码检查
stop
的值,如果为真,则调用break
退出循环,从而停止训练。注意,所有的 DDP 进程都将会因为这个逻辑而停止,确保训练在所有节点上的一致性。
这段代码实现了分布式训练中的提前停止机制。通过使用 PyTorch 的分布式通信功能(如 broadcast_object_list
),它确保所有参与训练的进程都能依据主进程的停止条件一致地停止训练。这样做的好处是在面对大规模训练时,可以有效避免不必要的资源浪费,尤其是在模型不再提升性能的情况下。
总结
这段代码主要负责YOLOv5模型的训练逻辑,涵盖了数据加载、模型训练、损失计算、梯度更新、模型验证等多个方面。其中,支持的功能如学习率调度、图像权重、早期停止等为训练过程提供了灵活性和高效性。训练完成后,模型也会根据验证后的指标进行保存,以便后续使用。
参考: