之前的文章介绍了YOLOv5的网络结构与目录结构源码以及detect.py的详细解读,今天带来的是YOLOv5的 train.py 代码参数逐行解读以及注释,废话不多说,让我们一起学习YOLOv5的 train.py 源码吧!
YOLOv5所使用版本:v6.1(由于YOLOV5版本一直在更新,不同版本代码略有差异,但差别不大,可供用来学习。)
YOLOv5 源码地址:GitHub - ultralytics/yolov5 at v6.1
其他学习YOLOv5直通车:
YOLOv5【使用云GPU进行训练】超详细教程!!!
YOLOv5【使用云GPU连接本地Pycharm进行训练】超详细教程!!!
YOLOv5【网络结构】超详细解读!!✨✨
YOLOv5【目录结构源码】超详细解读!!!
YOLOv5【detect.py源码及参数】超详细注释解读!!
train.py是yolov5中用于训练模型的主要脚本文件,其主要功能是通过读取配置文件,设置训练参数和模型结构,以及进行训练和验证的过程。
(1)、修改完参数直接在pycharm中点击运行按钮。
(2)、在终端输入:
python train.py --data coco.yaml --cfg yolov5s.yaml --weights yolov5s.pt --epoch 150 --batch-size 32 --device 0
此方法为在命令行中设置修改的参数并进行训练。
train.py 执行主要流程如下✨✨:1.导入相关的库及配置文件---->2.导入完包以后执行main()函数---->3.执行main()函数时用到parse_opt()这个函数,它的功能主要是解析参数,在parse_opt()执行完成之后,会将opt传给函数main()。---->4.执行train()函数, 在main()函数训练时会执行train()函数,train()函数主要分为了六个部分:train()函数主要分为了六个部分:(1)传入参数、(2)初始化参数与配置信息、(3)加载模型、(4)加载训练参数并保存至模型、(5)进行训练、(6)训练的结束。---->5.执行run()函数
总体上抓住 参数+数据 + 模型 + 学习率 + 优化器 + 训练 即可!
接下来让我们按照执行流程来依次解析train.py源码把!!!✨✨
我们先从 train.py 导入库以及相关配置文件开始介绍
✨
导入所需要的相关库✨
'''------------------1.1✨导入所需要的相关库✨--------------------'''
import argparse # 解析命令行参数模块
import math # 数学公式模块
import os # 与操作系统进行交互的模块 包含文件路径操作和解析
import random # 生成随机数模块
import sys # sys系统模块 包含了与Python解释器和它的环境有关的函数
import time # 时间模块 更底层
from copy import deepcopy # 深度拷贝模块
from datetime import datetime # datetime模块能以更方便的格式显示日期或对日期进行运算。
from pathlib import Path # Path将str转换为Path对象 使字符串路径易于操作的模块
import numpy as np # numpy数组操作模块
import torch # 引入torch
import torch.distributed as dist # 分布式训练模块
import torch.nn as nn # 对torch.nn.functional的类的封装 有很多和torch.nn.functional相同的函数
import yaml # yaml是一种直观的能够被电脑识别的的数据序列化格式,容易被人类阅读,并且容易和脚本语言交互。一般用于存储配置文件。
from torch.cuda import amp # PyTorch amp自动混合精度训练模块
from torch.nn.parallel import DistributedDataParallel as DDP # 多卡训练模块
from torch.optim import SGD, Adam, lr_scheduler # tensorboard模块
from tqdm import tqdm # 进度条模块
第一部分导入了所需要的相关库。
这段代码主要作用: 获取当前文件的绝对路径,将当前项目添加到系统路径上,并将项目的绝对路径转换成相对路径保存在ROOT中。
'''----------------------1.2✨路径代码✨----------------------'''
FILE = Path(__file__).resolve() # __file__指的是当前文件(即train.py),FILE最终保存着当前文件的绝对路径,比如D://yolov5/train.py
ROOT = FILE.parents[0] # YOLOv5 root directory ROOT保存着当前项目的父目录,比如 D://yolov5
if str(ROOT) not in sys.path: # sys.path即当前python环境可以运行的路径,假如当前项目不在该路径中,就无法运行其中的模块,所以就需要加载路径
sys.path.append(str(ROOT)) # add ROOT to PATH 把ROOT添加到运行路径上
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative ROOT设置为相对路径
这样做原因:在不同的操作系统和环境下,路径的表示方式可能有所不同。将路径表示为相对路径可以确保代码在不同的环境中具有相同的行为。
'''-------------------- 1.3 ✨导入自定义辅助库✨-----------------'''
import val #测试集
from models.experimental import attempt_load # 实验性质的代码,包括MixConv2d、跨层权重Sum等
from models.yolo import Model #yolo的特定模块,包括BaseModel,DetectionModel,ClassificationModel,parse_model等
from utils.autoanchor import check_anchors #定义了自动生成锚框的方法
from utils.autobatch import check_train_batch_size #定义了自动生成批量大小的方法
from utils.callbacks import Callbacks #定义了回调函数,主要为logger服务
from utils.datasets import create_dataloader #dateset和dateloader定义代码
from utils.downloads import attempt_download #谷歌云盘内容下载
from utils.general import (LOGGER, NCOLS, 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, 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 #用于wandb
from utils.loss import ComputeLoss #存放各种损失函数
from utils.metrics import fitness #模型验证指标,包括ap,混淆矩阵等
from utils.plots import plot_evolve, plot_labels ##定义了Annotator类,可以在图像上绘制矩形框和标注信息
from utils.torch_utils import EarlyStopping, ModelEMA, de_parallel, select_device, torch_distributed_zero_first 定义了一些与PyTorch有关的工具函数,比如选择设备、同步时间等
通过导入这些辅助模块,可以更方便地进行目标检测的相关任务,并且减少了代码的复杂度和冗余。
'''-------------------- 1.4 ✨pytorch 分布式训练初始化✨-----------------'''
# https://pytorch.org/docs/stable/elastic/run.html 详细介绍
LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # -本地序号。 Worker 是这台机器上的第几个 Worker
RANK = int(os.getenv('RANK', -1)) # -进程序号。全局第几个 Worker
WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1)) # 总共有几个 Worker
接下来按照程序的运行顺序,跳过函数定义部分,先看主函数。
if __name__ == "__main__":
# 接收命令行参数
opt = parse_opt()
# 将命令行参数传入main函数
main(opt)
主函数中调用了两个函数,第一个parse_opt()函数用于接收命令行参数,并将包含参数的对象传递给main()函数。
parse_opt()函数用于添加自定义的命令行参数。执行main函数时用到并parse_opt()这个函数,它的功能主要是为模型进行训练时提供参数,在parse_opt()执行完成之后,会将opt传给函数main()。
'''--------------------三、parse_opt()函数-----------------'''
def parse_opt(known=False):
"""
weights: 预训练的权重参数文件
cfg: 模型的配置文件 包括类别、anchor、网络结构等
data: 数据集配置文件路径,包括训练集、验证集、测试集路径以及类别数量以及类别名等
hyp: 超参数配置文件
epochs: 训练迭代的轮数
batch-size: 每批输入GPU或内存图片的数量,batch_size越大占用内存越大,训练越快
imgsz: 输入模型的图片大小
rect: 矩形训练,默认关闭
resume: 是否从中断中恢复
nosave: 不保存模型 默认为False即为保存模型
noval: 表示只在最后一轮验证
noautoanchor: 不使用自动的anchor
evolve: 超参数进化的训练模式,默认为300
bucket: 表示谷歌云盘
cache:是否缓存数据集。
image-weights: 图片采样策略 ,就是根据图片的权重来决定其采样顺序。
device: 参数装载的设备
multi-scale: 是否使用多尺度训练 默认不使用
single-cls: 单类别 将多个类的数据按照一个类训练
optimizer: 是否使用优化器
workers: 最大工作数 用于设置分布式训练
project: 保存结果的路径
name: 保存结果的目录名
exist_ok: 是否重新结果目录 默认为False
quad :是否使用quad dataloader
cos-lr: 训练学习率衰减策略是否使用余弦退火策略。
label-smoothing: 标签平滑度 默认为0
patience: 早停忍耐次数 默认100
freeze: 表示冻结前几层的参数,即不训练前几层的参数 默认为0 即为不冻结
save-period: 多少个epoch保存一下checkpoint
local_rank: 进程内GPU编号 默认为-1
entity: 使用tensorboard 可视化工具
upload_dataset: 是否上传dataset到wandb
bbox_interval: 设置界框图像记录间隔
artifact_alias: 使用数据的版本
"""
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.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')
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 check')
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='--cache images in "ram" (default) or "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='maximum number of dataloader workers')
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', type=int, default=0, help='Number of layers to freeze. backbone=10, all=24')
parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
# Weights & Biases arguments
parser.add_argument('--entity', default=None, help='W&B: Entity')
parser.add_argument('--upload_dataset', action='store_true', help='W&B: Upload dataset as 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')
# known为True时 parse_known_args()用于接收命令行中未设置的参数,known默认为false
opt = parser.parse_known_args()[0] if known else parser.parse_args()
return opt
接下来我们看主要的参数设置吧:
parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
这个参数是:预训练的权重参数文件,默认是yolov5s.pt,可以使用自己训练的权重,也可以使用官网提供的权重,下载好后放在根目录就好。默认官网的权重yolov5s.pt‘
命令行用法:
python train.py --weights yolov5s.pt
(yolov5n.pt/yolov5s.pt/yolov5m.pt/yolov5l.pt/yolov5x.pt/区别在于网络的宽度和深度以此增加)。
下图是官方提供的预训练权重:
注意:
1,若在命令行中使用"–weights" 参数,可指定预训练权重文件(路径);
2,若在命令行不使用"–weights" 参数,则预训练权重文件(路径)为自定义的default默认值;
3,若既使用命令行"–weights" 参数,又自定义了default默认值,则模型使用的是命令行"–weights" 参数指定的预训练权重文件(路径)
4,若不进行预训练,可使用"–weights" 参数指定一个空字符串:“”,或者将default默认值设置为空字符串:“”;
5,若使用yolov5s.pt、yolov5m.pt、yolov5l.pt、yolov5x.pt等yolov5官方预训练权重文件,若没有下载,代码会自动帮你下载,放在ROOT路径下,也就是你yolov5工程项目路径下,但是下载速度一般会很慢,建议先去yolov5官方GitHub中下载好。
parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
这个参数是:模型的配置文件 包括类别、anchor、网络结构等。源码里面提供了下图这5个配置文件:
命令行用法:
python train.py --cfg models/yolov5s.yaml
配置文件里面指定了一些 包括类别、anchor等参数信息以及backbone和head网络结构,如下图所示:
注意:
1,在已经使用"–weights" 参数加载了预训练权重的情况下,可以不使用该参数,模型结构直接使用预训练权重中保存的模型结构。
2,不使用"–weights" 参数且使用"–cfg" 参数,表示模型从头开始训练,不进行预训练。
3,“–weights” 参数和"–cfg" 参数必须要有一个,不然代码会报错。
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
这个参数是:数据集配置文件的路径。包括数据集存放示例、数据集路径(训练集、验证集、测试集的路径)数据集类别信息(类别数量以及类别名)以及数据集下载地址等。如下图所示:
命令行用法:
python train.py --data data/coco128.yaml
我们在训练自己数据集时需要新建.yaml 文件。源码里提供了9种常用数据集的配置文件,如下图所示:
注意:
1,如果没有检查到数据集,代码会自动下载coco128数据集,也可以自己下载;
2,把yolov5官方的数据集配置文件中的数据集下载部分内容给注释掉,代码则不会自动下载。
parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml', help='hyperparameters path')
这个参数是:训练超参数配置文件路径。超参数配置文件里面包含了大量的参数信息,如下图所示:
命令行用法:
python train.py --hyp data/hyps/hyp.scratch-low.yaml
源码里提供了5个超参数配置文件,如下图所示:
注意:
1.hyps.scratch-low.yaml :YOLOv5 COCO训练从头优化,数据增强低,适用于较小型号,比如:v5n、v5s。
2.hyps.scratch-med.yaml :数据增强中,适用于中型型号。比如:v5m。
3. hyps.scratch-high.yaml :数据增强高,适用于大型型号,比如:v3、v3-spp、v5l、v5x。
parser.add_argument('--epochs', type=int, default=300)
这个参数是:训练迭代的轮数。
命令行用法:
python train.py --epochs n
注意:
1,epochs表示训练整个训练集的次数,epoch为n表示将整个训练集训练n次。
parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs, -1 for autobatch')
这个参数是:每批输入GPU或内存图片的数量。batch_size越大占用内存越大,训练越快。
命令行用法:
python train.py --batch-size n
注意:
1.训练批量大小表示每个 mini-batch 中的样本数,batch-size设置为n表示一次性从训练集中获取n张图片送入模型进行训练;
2. batch-size大小需要根据自己设备GPU的资源合理设置。3. default=-1将时自动调节batchsize大小。
补充:epoch、batchsize、iteration三者之间的联系
1、batchsize:批次大小。假如取batchsize=16,则表示每次训练时在训练集中取16个训练样本进行训练。
2、iteration:迭代次数。1个iteration就等于一次使用24(batchsize大小)个样本进行训练。
3、epoch:训练1个epoch就等于训练集中全部样本训练1次。
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
这个参数是:输入模型的图片大小。
命令行用法:
python train.py --imgsz/img/–img-size 640
parser.add_argument('--rect', action='store_true', help='rectangular training')
这个参数是:是否使用矩形训练,默认关闭。
命令行用法:
python train.py --rect
注意:
1.矩形训练过程中会对输入的矩形图片进行预处理,通过保持原图高宽比进行resize后,对resize后的图片进行填充,填充到32的最小整数倍,然后进行矩形训练,减少训练时间。2.矩形推理的大致流程如下:
step1:对原始图片进行等比例缩放。
step2:对短边进行像素填充,填充到32的最小整数倍。
parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
这个参数是:是否断点续训,默认关闭。
命令行用法:
python train.py ----weights /path/last.pt --rect
注意:
1,断点续训就是从上一个训练任务中断的地方继续训练,直至训练完成;
2,当模型按指定的epoch训练完成后,则无法进行断点续训;
3,需要搭配"–weights" 参数使用,指定训练中断保存的最后一次模型权重文件。
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
这个参数是:只保留最后一次训练的权重,默认关闭。
命令行用法:
python train.py --nosave
正常是保存 best.pt 和 last.pt 两个权重文件。
parser.add_argument('--noval', action='store_true', help='only validate final epoch')
这个参数是:只对最后一次训练进行验证,默认关闭。
命令行用法:
python train.py --noval
parser.add_argument('--noautoanchor', action='store_true', help='disable AutoAnchor')
这个参数是:关闭自动计算锚框功能,默认关闭。即默认使用自动计算锚框功能。可以简化训练过程。
命令行用法:
python train.py --noautoanchor
注意:
yolov5采用的是kmeans聚类算法来计算anchor box的大小和比例,最终自动计算出一组最合适训练的锚框。
parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
这个参数是:使用遗传算法对超参数进行优化,进行自动调参,默认关闭。超参数进化会耗费大量的资源和时间,建议不要开启。
命令行用法:
python train.py --evolve n
注意:
1.yolov5采用遗传算法对超参数进行优化,寻找一组最优的训练超参数;
2.开启后传入参数n,训练每迭代n次进行一次超参数进化;
3.开启后不传入参数,则默认为const=300。
parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
这个参数是:从谷歌云盘下载或上传数据。
命令行用法:
python train.py --bucket gsutil bucket
注意:
1.该参数用于指定 gsutil bucket 的名称,其中 gsutil 是 Google 提供的一个命令行工具,用于访问 Google Cloud Storage(GCS)服务;
2.GCS 是 Google 提供的一种对象存储服务,用户可以将任意数量和类型的数据存储在其中。用户可以通过 gsutil 命令行工具上传、下载、复制、删除等操作 GCS 中的数据。在训练模型时,如果需要使用 GCS 中的数据集,就需要指定 bucket 的名称。
parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"')
这个参数是:是否缓存数据集,默认关闭。
命令行用法:
python train.py --cache
注意:
1.缓存数据集图片到内存中,训练时模型直接从内存中读取,加快数据加载和训练速度
2.若"–cache"参数指定值,可以指定的;值:ram/disk;
3.若"–cache"参数不指定值,则默认为const=‘ram’。
parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
这个参数是:图片采样策略 ,就是根据图片的权重来决定其采样顺序。对数据集图片进行加权训练,默认关闭.
命令行用法:
python train.py --image-weights --rect
注意:需要搭配"–rect"参数一起使用。
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
这个参数是:选择训练使用的设备处理器,CPU或者GPU,默认为空。
命令行用法:
python train.py --device 0/0,1,2,3/cpu
注意:
1.默认为空时,代码会进行自动选择,若检查到设备有GPU则优先使用GPU进行训练,若没有则使用CPU进行训练;
2.使用GPU训练时,0,1,2,3分别表示第1,2,3,4张GPU;
3.设备没有GPU,使用CPU训练:
python train.py --device cpu
4.设备有单个GPU,使用单个GPU训练:
python train.py --device 0
5.设备有多个GPU,使用单个GPU训练:
python train.py --device 0 (使用第1张GPU训练);
python train.py --device 2 (使用第3张GPU训练);
6.设备有多个GPU,使用多个GPU训练:
python train.py --device 0,1,2(使用第1,2,3张GPU训练训练)。
parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
这个参数是:是否使用多尺度训练,默认不使用。
命令行用法:
python train.py --multi-scale
注意:
1.多尺度训练是指设置几种不同的图片输入尺度,训练时每隔一定iterations随机选取一种尺度训练,这样训练出来的模型鲁棒性更强。
2.开启多尺度训练,训练过程中每次输入图片会放大或缩小50%。
parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
这个参数是:是否设定训练数据集是单类别,将多个类的数据按照一个类训练。默认为 false多类别。
命令行用法:
python train.py --single-cls
parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer')
这个参数是:选择训练使用的优化器,默认使用SGD.
命令行用法:
python train.py --optimizer SGD
注意:
choices=[‘SGD’, ‘Adam’, ‘AdamW’] 表示只能选择’SGD’, ‘Adam’, 'AdamW’这三种优化器,当然也可以添加自定义的优化器,但代码中其他地方也要做相应的更改。
parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
这个参数是:是否使用SyncBatchNorm (Synchronized Batch Normalization:同步批量归一化),只有在使用DDP模式(分布式训练)时有效,默认关闭。
命令行用法:
python train.py --sync-bn
注意:
1,关闭时,训练使用传统的批量归一化;
2,在传统的批归一化(Batch Normalization,简称 BN)中,每个 GPU 会对数据的均值和方差进行单独计算,因此在多 GPU 训练时,每个 GPU 计算的均值和方差可能会不同,导致模型训练不稳定。为了解决这个问题,SyncBN 技术将 BN 的计算放在了整个分布式训练过程中进行,确保所有 GPU 上计算的均值和方差是一致的,从而提高模型训练的稳定性和效果,但同时也会增加训练时间和硬件要求,因此需要根据具体的训练数据和硬件资源来决定是否使用 SyncBN。
parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
这个参数是:设置Dataloader使用的最大numworkers,默认设置为8.
命令行用法:
python train.py --workers 8
注意:
1,Dataloader中numworkers表示加载处理数据使用的线程数,使用多线程加载数据时,每个线程会负责加载和处理一批数据,数据加载处理完成后,会送入相应的队列中,最后主线程会从队列中读取数据,并送入GPU中进行模型计算;
2,numworkers为0表示不使用多线程,仅使用主线程进行数据加载和处理。
parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name')
这个参数是:设置每次训练结果保存的主路径名称。
命令行用法:
python train.py --project ‘runs/train’
注意:这里主路径的意思是,你每次训练会生成一个单独的子文件夹,主路径就是存放你这些单独子文件夹的地方,可以自己命名,例如’runs/train’。比如说你第一次训练保存结果的文件夹是exp1,第二次是exp2,第三次是exp3,则这些子文件夹都会放在主路径’runs/train’下面。
parser.add_argument('--name', default='exp', help='save to project/name')
这个参数是:设置每次训练结果保存的 子路径(目录)名称。
命令行用法:
python train.py --name exp
注意:这里子路径的意思是就是上面在’–project’中提到的每次训练生成的单独的子文件夹,可以自己命名,例如’exp’。你每次训练生成的模型权重文件、可视化结果以及其它结果文件保存的地方。
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
这个参数是:否覆盖同名的训练结果保存路径,默认关闭,表示不覆盖。
命令行用法:
python train.py --exist-ok
注意:
1.不使用’–exist-ok’参数时,如果’–name’指定的名称不变,比如’exp’,每次训练会按顺序新建文件夹,例如exp1、exp2、exp3、… 、expn;
2.使用’–exist-ok’参数时,如果’–name’指定的名称不变,比如’exp’,每次训练则不会新建文件夹,训练结果会覆盖原先文件夹中保存的所有结果。
parser.add_argument('--quad', action='store_true', help='quad dataloader')
这个参数是:是否使用quad dataloader,默认关闭
命令行用法:
python train.py --quad
注意:quad dataloader 是一种数据加载器,它可以并行地从磁盘读取和处理多个图像,并将它们打包成四张图像,从而减少了数据读取和预处理的时间,并提高了数据加载的效率。
parser.add_argument('--cos-lr', action='store_true', help='cosine LR scheduler')
这个参数是:训练学习率衰减策略是否使用余弦退火策略,默认关闭。
命令行用法:
python train.py --cos-lr
注意:余弦退火策略在训练初期加快学习速度,训练后期减小学习率,从而更好地学习数据的分布,避免模型陷入局部最优。
parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
这个参数是:训练使用标签平滑策略,防止过拟合.
命令行用法:
python train.py --label-smoothing
注意:
1.默认为0.0,即标签平滑策略使用的epsilon为0.0;
2.将标签平滑策略使用的epsilon设置为0.1:python train.py --label-smoothing 0.1
表示在每个标签的真实概率上添加一个 epsilon=0.1的噪声,从而使模型对标签的波动更加鲁棒;
3.–label-smoothing 参数的值应该根据具体的数据集和模型来调整,以达到最优的训练效果。
parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')
这个参数是:训练使用EarlyStopping策略,防止过拟合.
命令行用法:
python train.py --patience 100
注意:‘–patience’参数指定为整数n时,表示模型在训练时,若连续n个epoch验证精度都没有提升,则认为训练已经过拟合,停止训练。’–patience’可根据具体的模型和数据集进行调整。
parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone=10, first3=0 1 2')
这个参数是:表示冻结前几层的参数,即不训练前几层的参数 默认为0 即为不冻结.
命令行用法:
python train.py --freeze
注意:
1.冻结训练是指在训练过程中冻结模型中的某些层,冻结的层不进行权重参数更新.
2.指定’0’或’-1’,不冻结任何层,更新所有层的权重参数:python train.py --freeze 0/-1
3.指定n,冻结前n(0python train.py --freeze n
parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
这个参数是:每训练n个epoch保存一次训练权重,默认关闭.
命令行用法:
python train.py --save-period n
注意:
1.n>0,每训练n个epoch保存一次训练权重;
2.n<=0,关闭save-period,只保存best和last权重。
parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
这个参数是:每训练n个epoch保存一次训练权重,默认关闭.
命令行用法:
python train.py --local_rank 0,1,2,3
注意:如果你有 4 个 GPU,想要使用第 2 号和第 3 号 GPU 进行训练,那么可以在启动训练脚本时设置如下参数:python train.py --local_rank 1,2
这样,第一个进程将使用第 2 号 GPU,第二个进程将使用第 3 号 GPU。注意,如果使用了 --local_rank 参数,那么在启动训练脚本时需要使用 PyTorch 的分布式训练工具,例如 torch.distributed.launch。
parser.add_argument('--entity', default=None, help='W&B: Entity')
这个参数是:用于指定模型实体的参数。
命令行用法:
python train.py --entity None
注意:
1.模型实体可以是一个实体名称或实体 ID,通常用于在实体存储库中管理模型的版本控制和记录。
2.在使用实体存储库时,你需要创建一个实体来存储模型,并在训练时指定该实体,这样训练结果就可以与该实体相关联并保存到实体存储库中。
3.该参数默认值为 None,如果未指定实体,则训练结果将不会与任何实体相关联。
parser.add_argument('--upload_dataset', nargs='?', const=True, default=False, help='W&B: Upload data, "val" option')
这个参数是:用于上传数据集到wandb,默认关闭.
命令行用法:
python train.py --upload_dataset False
注意:
1.如果命令行未使用’–upload_dataset’参数,则默认值为default=False,表示不上传数据集。
2.如果命令行使用’–upload_dataset’参数,但没有传递参数,则默认值为const=True,表示上传数据集。
3.如果命令行使用’–upload_dataset’参数,并且传递了参数’val’,则默认为True,表示要上传val数据集。
parser.add_argument('--bbox_interval', type=int, default=-1, help='W&B: Set bounding-box image logging interval')
这个参数是:指定在训练过程中每隔多少个epoch记录一次带有边界框的图片,默认关闭.
命令行用法:
python train.py --bbox_interval n
注意:
1.n>0,每隔n个epoch记录一次带有边界框的图片;
2.n<=0,关闭–bbox_interval。
parser.add_argument('--artifact_alias', type=str, default='latest', help='W&B: Version of dataset artifact to use')
这个参数是:指定在训练过程中每隔多少个epoch记录一次带有边界框的图片,默认关闭.
命令行用法:
python train.py --artifact_alias latest
注意:在使用MLFlow等工具跟踪模型训练和数据集版本时,会给每个版本分配唯一的别名。通过指定此参数,可以使用特定版本的数据集工件。默认情况下,使用最新版本的数据集工件。
呼~参数逐行解读终于结束啦~~缓一缓接下来让我们学习main函数吧~~
✨
检查训练环境✨
main()函数的参数中opt为命令行参数对象,callbacks中记录了Loggers类中的所有函数名。例如,可以通过callbacks.run("on_train_epoch_end")调用Loggers类中的on_train_epoch_end()函数。这样做方便在不同训练阶段,对日志的记录做统一管理。
首先看main函数接收的参数,opt为命令行参数;callbacks为训练过程中保存的一些参数。
def main(opt, callbacks=Callbacks()):
'''-------------------- 4.1 ✨检查训练环境✨-----------------'''
if RANK in [-1, 0]: # 若进程编号为-1或0
print_args(FILE.stem, opt) # 输出所有训练参数 / 参数以彩色的方式表现
check_git_status() # 检测YOLOv5的github代码是否是最新的
check_requirements(exclude=['thop']) # 检查requirements.txt所需包是否都满足
这段代码的作用是:先检查(分布式)训练的环境,若RANK为-1或0,打印参数并检查github仓库和依赖库。
若RANK为-1或0,会执行下面三行代码:
1. print_args(FILE.stem, opt) :输出所有训练参。
2.check_git_status() :检测YOLOv5的github代码是否是最新的。
3.check_requirements(exclude=['thop']) :检查requirements.txt所需包是否都满足。
✨
判断是否使用断点续训resume✨
这段代码的作用是:判断是否使用断点训练。断点训练是当训练异常终止或想调节超参数时,系统会保留训练异常终止前的超参数与训练参数,当下次训练开始时,并不会从头开始,而是从上次中断的地方继续训练。 若使用断点续训,就从last.pt中读取相关参数。
'''-------------------- 4.2 ✨判断是否使用断点续训resume✨-----------------'''
if opt.resume and not check_wandb_resume(opt) and not opt.evolve: #判断是否使用断点续训resume, 读取参数
ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run() # specified or most recent path
# 使用断点续训 就从last.pt中读取相关参数
# 如果resume是str,则表示传入的是模型的路径地址
# 如果resume是True,则通过get_lastest_run()函数找到runs为文件夹中最近的权重文件last.pt
assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist' # 判断是否为文件,若不是文件抛出异常
with open(Path(ckpt).parent.parent / 'opt.yaml', errors='ignore') as f: # opt.yaml是训练时的命令行参数文件
opt = argparse.Namespace(**yaml.safe_load(f)) # 超参数替换,将训练时的命令行参数加载进opt参数对象中
opt.cfg, opt.weights, opt.resume = '', ckpt, True # reinstate ## opt.cfg设置为'' 对应着train函数里面的操作(加载权重时是否加载权重里的anchor)
LOGGER.info(f'Resuming training from {ckpt}') # 打印从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') # 设置新的项目输出目录
opt.exist_ok, opt.resume = opt.resume, False # pass resume to exist_ok and disable resume # 将resume传递给exist_ok
opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # 根据opt.project生成目录,并赋值给opt.save_dir 如: runs/train/exp1
✨
DDP(分布式)mode训练配置检测✨
这段代码的作用是:检查DDP(分布式)训练的配置,并设置GPU。这段代码会选择使用cpu还是gpu。若采用的是分布式训练,会额外执行下面的一些操作,我们一般不会用到分布式,所以也就没有执行什么东西。
'''--------------------4.3 ✨DDP(分布式)mode训练配置检测✨-----------------'''
# 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")
DDP(Distributed Data Parallel):用于单机或多机的多GPU分布式训练,目前只能在Linux下使用。
✨选择正常训练模式(不使用超参数进化-遗传算法)
✨
✨
选择正常训练模式(不使用超参数进化-遗传算法)✨
准备完毕后,接下来就开始训练:
'''---------------4.4✨选择正常训练模式(不使用超参数进化-遗传算法)✨------------'''
# 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()
这段代码的作用是:判断是否使用超参数进化(遗传算法),边进化边训练。这段代码判断的是不使用遗传进化算法,即正常训练,不使用超参数进化。正常训练时直接把命令行参数传入train函数,训练完成后销毁所有进程。train函数暂时跳过,先分析main()函数。
超参数进化:超参数进化是一种使用遗传算法,进行超参数优化的方法。遗传算法是用于解决最优化问题的一种搜索算法,它将要解决的问题模拟为一个生物进化的过程,过复制、交叉、突变等操作产生下一代的解,并逐步淘汰掉适应度函数值低的解,增加适应度函数值高的解。YOLO v5的超参数比较多,用传统的网格搜索会让超参数的调整变得很棘手。这时使用遗传算法调整超参数就是一种合适的方式。
✨
使用超参数进化进行训练✨
1.✨
前期准备 ✨
若使用超参数进化,以上代码为超参数训练的前期准备,首先指定每个超参数的突变范围、最大值、最小值,再为超参数的结果保存做好准备。
'''-------------------- 4.4.2 ✨使用超参数进化进行训练✨-----------------'''
# 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)
# 打开超参数配置文件 errors表示编码错误的处理方式
with open(opt.hyp, errors='ignore') as f:
# 加载超参数配置文件
hyp = yaml.safe_load(f) # load hyps dict
# 如果anchors不在超参数中
if 'anchors' not in hyp: # anchors commented in hyp.yaml
# 将anchors设置为3
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 {save_dir}') # download evolve.csv if exists
2.✨
超参数训练 ✨
这段代码代表的开始超参数进化:
注意: 使用超参数进化时至少要经过至少300次迭代,每次迭代都会经过一次完整的训练。因此超参数进化及其耗时,根据自己需求慎用。
超参数进化的步骤:
(1)若存在训练数据文件,读取文件中的训练数据,选择结果最优的训练数据突变超参数;(2)将突变后的超参数限定值阈值范围内;
(3)使用突变后的超参数进行训练;
(4)训练结束后,将训练结果保存至evolution.csv,用于下一次的超参数突变。
根据生物进化,优胜劣汰,适者生存的原则,每次迭代都会保存更优秀的结果,直至迭代结束。最后的结果即为最优的超参数。
'''------------------------ 使用超参数进化--------------------'''
# 选择超参数的遗传迭代次数 默认为迭代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, [email protected], [email protected]: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}')
呼~到这里main()函数解读终于结束啦~~缓一缓接下来让我们学习train()函数吧~~
在main()函数训练时会执行train()函数,train()函数主要分为了六个部分:(1)传入参数、(2)初始化参数与配置信息、(3)加载模型、(4)加载训练参数并保存至模型、(5)进行训练、(6)训练的结束。
✨
传入参数✨
train() 函数首先做的就是接收参数。这段函数就是来获取传入的参数。
hyp:超参数,不使用超参数进化的前提下也可以从opt中获取。
opt:指的是全部的命令行参数。
device:指的是装载程序的设备。
callbacks:指的是训练过程中产生的一些参数。
'''---------------------- 5.1 ✨传入参数✨-------------------'''
def train(hyp, # 超参数 可以是超参数配置文件的路径 或 超参数字典 path/to/hyp.yaml or hyp
# dictionary
opt, # 命令行参数
device, # 装载程序的设备
callbacks # 用于存储Loggers日志记录器中的函数,方便在每个训练阶段控制日志的记录情况
):
# 接收命令行参数
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参数 + 保存opt + 加载数据配置信息 + 打印日志信息(logger + wandb) + 其他参数(plots、cuda、nc、names、is_coco)
✨
设置模型(last.py与best.py)的保存路径✨
这段代码主要是设置模型保存的路径,权重名字和训练日志txt文件
每次训练结束后,会产生两个模型,一个是last.pt,一个是best.pt。顾名思义,last.pt即为训练最后一轮产生的模型,而best.pt是训练过程中,效果最好的模型。
评判效果最好的标准又是什么?
在这里评判best.pt的训练效果好坏的标准,绝大部分依赖的是[email protected]:0.95,是根据上文在超参数的突变时,提到的fitness()函数得出,后面遇到了再细讲。
然后创建文件夹,保存训练结果的模型文件路径 以及验证集输出结果的txt文件路径,包含迭代的次数,占用显存大小,图片尺寸,精确率,召回率,位置损失,类别损失,置信度损失和map等。
'''--------------5.2 ✨设置模型(last.py与best.py)的保存路径✨--------------'''
# Directories 获取记录训练日志的保存路径
# 设置保存权重路径 如runs/train/exp1/weights
w = save_dir / 'weights' # weights dir
# 新建文件夹 weights train evolve
(w.parent if evolve else w).mkdir(parents=True, exist_ok=True) # make dir
# 保存训练结果的目录,如last.pt和best.pt
last, best = w / 'last.pt', w / 'best.pt'
✨
判断超参数✨
这段代码是:检查超参数是字典还是字符串,并通过日志打印出来。
若为字符串,则认定为.yaml文件路径,再将yaml文件加载为字典。这里导致超参数的数据类型不同的原因是,超参数进化时,传入train()函数的超参数即为字典。而从命令行参数中读取的则为文件路径。
'''------------------5.2.2 ✨判断超参数✨-----------------'''
# Hyperparameters
# 判断hyp是字典还是字符串
if isinstance(hyp, str):
# 若hyp是字符串,即认定为路径,则加载超参数为字典
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()))
✨
保存训练参数✨
这段代码是:保存训练过程中的参数(包括超参数和命令行参数)。
'''------------------5.2.3 ✨保存训练过程中的参数✨-----------------'''
# Save run settings
# 保存训练过程中的参数
with open(save_dir / 'hyp.yaml', 'w') as f:
# 保存超参数为yaml配置文件
yaml.safe_dump(hyp, f, sort_keys=False)
with open(save_dir / 'opt.yaml', 'w') as f:
# 保存命令行参数为yaml配置文件
yaml.safe_dump(vars(opt), f, sort_keys=False)
# 定义数据集字典
data_dict = None
✨
加载日志信息✨
这段代码是:加载日志记录器,并将日志记录器中的函数记录到callbacks内,方便在训练的不同阶段,利用callbacks.run()函数对日志的记录做统一处理。
'''----------------------5.2.4 ✨加载日志信息✨-----------------'''
# Loggers
# 加载日志信息
# 如果进程编号为-1或0
if RANK in [-1, 0]:
# 初始化日志记录器实例
loggers = Loggers(save_dir, weights, opt, hyp, LOGGER) # loggers instance
# wandb为可视化参数工具
if loggers.wandb:
data_dict = loggers.wandb.data_dict
# 如果使用中断训练 再读取一次参数
if resume:
weights, epochs, hyp = opt.weights, opt.epochs, opt.hyp
# Register actions
# 遍历日志记录器中的所有方法
for k in methods(loggers):
# 将日志记录器中的方法与字符串进行绑定
callbacks.register_action(k, callback=getattr(loggers, k))
✨
设置配置信息✨
这段代码是:做一些变量的配置,包括是否绘图,设置随机数种子,保存训练集、验证集路径,保存类别数量以及类别名,并完成检查。
'''----------------------5.2.5 ✨设置配置信息✨-----------------'''
# Config
# 如果不使用超参数进化,训练结束时将训练数据绘图
plots = not evolve # create plots
# 是否使用cuda
cuda = device.type != 'cpu'
# 不同库中使用同样的随机数种子
init_seeds(1 + RANK)
# 设置分布式训练 存在子进程-分布式训练
# torch_distributed_zero_first(LOCAL_RANK): 用于同步不同进程对数据读取的上下文管理器
with torch_distributed_zero_first(LOCAL_RANK):
# 检查数据集是否存在 若不存在则下载
# data_dict为读取数据集yaml文件后的字典
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 = ['item'] if single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names
# 检查类别名的数量是否与类别数对应
assert len(names) == nc, f'{len(names)} names found for nc={nc} dataset in {data}' # check
# 判断数据集是否为coco数据集
is_coco = data.endswith('coco.yaml') and nc == 80 # COCO dataset
✨
加载模型(Model)✨
✨判断是否预训练
✨
这段代码是:加载模型,分为使用预训练权重参数文件与不使用预训练权重参数文件。
这里使用预训练权重参数,是类似于迁移学习。预训练的模型是检测coco数据集的模型,数据集中有80个类别,而自己的训练集类别以及类别的数量,并不与coco数据集相同。所以要先加载一个新的模型,把预训练的参数加载至模型作为初始参数,再把识别的类别改成自己的数据集要识别的类别。接下来将预训练参数中与新模型中相同的参数加载至模型。
'''---------------------- 5.3 ✨加载模型(Model)✨-----------------'''
# 检查模型文件格式是否支持
check_suffix(weights, '.pt') # check weights
# 预训练权重参数是否以.pt结尾
pretrained = weights.endswith('.pt')
'''----------------------5.3.1 ✨判断是否预训练✨-----------------'''
# 若以.pt结尾
if pretrained:
# torch_distributed_zero_first(LOCAL_RANK): 用于同步不同进程对数据读取的上下文管理器
with torch_distributed_zero_first(LOCAL_RANK):
# 检查本地是否存在预训练权重参数文件 若不存在则从官网下载
weights = attempt_download(weights) # download if not found locally
# 加载预训练权重参数文件
ckpt = torch.load(weights, map_location=device) # load checkpoint
# 创建模型
# 若cfg为空 则从预训练权重参数文件中加载模型 ch为通道数 nc为检测的类别数 achors为超参数,从超参数字典中加载
# to(device)将程序装载至对应的位置
# 这里预训练模型是coco数据集,检测80个类别,这里新建模型检测类别修改为自己的类别数量
model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
# 若cfg 或 hyp.get('anchors')不为空 且 不使用中断训练 exclude=['anchor'] 否则 exclude=[]
exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else [] # exclude keys
# 将预训练模型中的所有参数保存下来 赋值给csd
csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32
# 这一部分类似迁移学习
# 判断预训练参数和新创建的模型参数有多少是相同的
# 筛选字典中的键值对 把exclude删除
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
✨
冻结层✨
这段代码是:Freeze会冻结模型的某些层,被冻结的层训练时不会更新参数。在模型训练过程中不会变化,只训练冻结层以外的权重参数。
冻结层的原理是通过设置每个层参数中的requires_grad属性实现的。requires_grad属性在上一篇detect模块的精读中提到过。若require_grad为True,在反向传播时就会求出此tensor的梯度,若require_grad为False,则不会求该tensor的梯度。冻结就是通过对某些层不求梯度实现的。默认不进行参数冻结。
'''---------------------- 5.3.2 ✨冻结层(Freeze)✨-----------------'''
# Freeze
"""Freeze为冻结某些层数"""
# freeze 为命令行参数 默认为0 表示不冻结
# 若 freeze为10,则代表将网络的前10层冻结 不训练前10层的参数
freeze = [f'model.{x}.' for x in range(freeze)] # layers to freeze
# 首先遍历所有层
for k, v in model.named_parameters():
# 为所有层的参数设置梯度
v.requires_grad = True # train all layers
# 判断是否需要冻结
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) # grid size (max stride)
# 检查输入图片分辨率是否能被32整除
imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # verify imgsz is gs-multiple
# Batch size 设置一次训练所选取的样本数
if RANK == -1 and batch_size == -1: # single-GPU only, estimate best batch size
# 确保batch size满足要求
batch_size = check_train_batch_size(model, imgsz)
1. ✨设置优化器参数✨
这段代码是参数设置:nbs、accumulate、hyp[‘weight_decay’]。
1.nbs(nominal batch size):名义上的batch_size。这里的nbs跟命令行参数中的batch_size不同,命令行中的batch_size默认为16,nbs设置为64。
2.accumulate :累计次数,在这里 nbs/batch_size(64/16)计算出 opt.batch_size输入多少批才达到nbs的水平。简单来说,nbs为64,代表想要达到的batch_size,这里的数值是64;
3.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发生了变化,所有权重参数也要做相应的缩放。
# Optimizer
'''---------------------- 5.3.3 ✨优化器(Optimizer)✨-----------------'''
# 名义上的batch_size 命令行参数的batch_size默认为16
# 这里是为了实现用更小的内存或显存实现batch_size=64的效果
nbs = 64 # nominal batch size
# accumulate 为累计次数 不能小于1
# accumulate 将梯度累积起来 实现batch_size=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']}")
2. ✨优化器参数分组优化✨
这段代码是:参数分组优化。g0表示归一化层中的所有权重参数,g1表示卷积层中所有的权重参数,g2表示所有的偏置参数。
接下里循环遍历所有的层,将对应的参数添加至列表中。这里的循环遍历很有意思。model.modules()返回的不是每个layer(层),而是由大到小依次遍历每个层次。
例如,第一次遍历出的v是模型本身;第二次遍历出的是模型内的Sequential;第三次遍历Sequential里的第一个模块;接下来继续往下遍历,直至遍历到卷积层或批归一化层。这有点像二叉树中的深度优先遍历。
# 将模型参数分为三组(批归一化的weight、卷积层weights、biases)来进行分组优化
g0, g1, g2 = [], [], [] # optimizer parameter groups
# 遍历网络中的所有层
# 这里的循环是一个层次一个层次的进行遍历
# 首先是最外层的Model 然后是Model内的Sequential 接下来是Sequential内的Conv模块、Bottleneck模块等
# 每次遍历完一层就向更深层遍历
for v in model.modules():
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): # bias
# 将层的bias添加至g2
g2.append(v.bias)
# YOLO v5的模型架构中只有卷积层和BatchNorm层
if isinstance(v, nn.BatchNorm2d): # weight (no decay)
# 将批归一化层的权重添加至g0 未经过权重衰减
g0.append(v.weight)
elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter): # weight (with decay)
# 将层的weight添加至g1 经过了权重衰减
# 这里指的是卷积层的weight
g1.append(v.weight)
3. ✨判断是否使用adam优化器✨
这段代码是:判断是否使用adam优化器,若使用adam优化器,初始参数为批归一化层中的参数。如果不使用adam优化器,则直接使用SGD随机梯度下降。
# 若使用adam优化器
if opt.adam:
# adam优化器 lr表示初始的学习率 betas表示两个平滑常数
# 将批归一化层的参数放入优化器
optimizer = Adam(g0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum
else:
# 若不使用Adam优化器 则直接使用随机梯度下降 momentum 动量:结合当前梯度与上一次更新信息,用于当前更新 nesterov为是否采用NAG(动量SGD优化算法)
# 将批归一化层的参数放入优化器
optimizer = SGD(g0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
4. ✨将卷积层的参数添加至优化器✨
接下来将g1(卷积层中的权重参数),g2(偏置参数),添加进优化器。 add_param_group()函数可以为优化器中添加一个参数组。一个优化器可以更新多个参数组,不同的参数组可以使用不同的超参数。
# 将卷积层的参数添加至优化器 并做权重衰减
# add_param_group()函数为添加一个参数组 同一个优化器可以更新很多个参数组 不同的参数组可以设置不同的超参数
optimizer.add_param_group({'params': g1, 'weight_decay': hyp['weight_decay']}) # add g1 with weight_decay
# 将所有的bias添加至优化器
optimizer.add_param_group({'params': g2}) # add g2 (biases)
# 输出优化器信息
LOGGER.info(f"{colorstr('optimizer:')} {type(optimizer).__name__} with parameter groups "
f"{len(g0)} weight, {len(g1)} weight (no decay), {len(g2)} bias")
# 在内存中删除g0 g1 g2 节省空间
del g0, g1, g2
这段代码是:设置学习率衰减方式。
在训练过程中变更学习率可能会让训练效果更好,YOLOv5提供了两种学习率变化的策略:
1.linear_lr(线性学习率),是通过线性插值的方式调整学习率
2.One Cycle(余弦退火学习率),即周期性学习率调整中,周期被设置为1。在一周期策略中,最大学习率被设置为 LR Range test 中可以找到的最高值,最小学习率比最大学习率小几个数量级。这里默认one_cycle。
'''---------------------- 5.3.4 ✨学习率✨-----------------'''
# Scheduler 设置学习率策略:两者可供选择,线性学习率和余弦退火学习率
if opt.linear_lr:
# 使用线性学习率
lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear
else:
# 使用余弦退火学习率
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
# 可视化 scheduler
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)
这段代码是:EMA与中断训练(Resume)。
- EMA(指数加权平均或滑动平均):其将前面模型训练权重,偏差进行保存,在本次训练过程中,假设为第n次,将第一次到第n-1次以指数权重进行加和,再加上本次的结果,且越远离第n次,指数系数越大,其所占的比重越小。
- 中断训练:可以理解为把上次中断结束时的模型,作为新的预训练模型,然后从中获取上次训练时的参数,并恢复训练状态。
- epoch(迭代次数):1个epoch等于使用训练集中的全部样本训练一次,epoch的大小跟迭代次数有着密切的关系,通常在迭代次数处于2000-3000之间损失已经处于平稳。
'''---------------------- 5.3.5 ✨EMA与中断训练(Resume)✨-----------------'''
# EMA 设置ema(指数移动平均),考虑历史值对参数的影响,目的是为了收敛的曲线更加平滑
ema = ModelEMA(model) if RANK in [-1, 0] else None # 为模型创建EMA指数滑动平均,如果GPU进程数大于1,则不创建
# Resume 断点续训
# 断点续训其实就是把上次训练结束的模型作为预训练模型,并从中加载参数
start_epoch, best_fitness = 0, 0.0
if pretrained:# 如果有预训练
# Optimizer 加载优化器与best_fitness
if ckpt['optimizer'] is not None:
# 将预训练模型中的参数加载进优化器
optimizer.load_state_dict(ckpt['optimizer'])
# best_fitness是以[0.0, 0.0, 0.1, 0.9]为系数并乘以[精确度, 召回率, [email protected], [email protected]:0.95]再求和所得
# 获取预训练模型中的最佳fitness,保存为best.pt
best_fitness = ckpt['best_fitness']
# EMA
# 加载ema模型和updates参数,保持ema的平滑性,现在yolov5是ema和model都保存了
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.'
"""
如果新设置epochs小于加载的epoch,
则视新设置的epochs为需要再训练的轮次数而不再是总的轮次数
"""
# 如果训练的轮数小于开始的轮数
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
这段代码是:单机多卡模式与多卡归一化的设置。
1.单机多卡模式(DP mode):DataParallel单机多卡模式自动将数据切分 load 到相应 GPU,将模型复制到相应 GPU,进行正向传播计算梯度并汇总。值得注意的是,模型和数据都需要先导入进 GPU 中,DataParallel 的 module 才能对其进行处理,否则会报错。
SyncBatchNorm。2.多卡归一化(SyncBatchNorm):主要用于解决多卡归一化同步问题,每张卡单独计算均值,然后同步,得到全局均值。用全局均值计算每张卡的方差,然后同步即可得到全局方差,但两次会消耗时间挺长。
'''---------------------- 5.3.6 ✨单机多卡模式与多卡归一化✨-----------------'''
# DP mode 使用单机多卡模式训练,目前一般不使用
# rank为进程编号。如果rank=-1且gpu数量>1则使用DataParallel单机多卡模式,效果并不好(分布不平均)
# rank=-1且gpu数量=1时,不会进行分布式
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)
# SyncBatchNorm 多卡归一化
if opt.sync_bn and cuda and RANK != -1:# 多卡训练,把不同卡的数据做个同步
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
LOGGER.info('Using SyncBatchNorm()')
✨加载训练数据并保存至模型
✨
✨
加载训练数据✨ 这段代码是:加载训练数据。
加载训练数据时,通过create_dataloader()函数得到两个对象。一个为train_loader,另一个为dataset。train_loader为训练数据迭代器,可以通过for循环遍历出每个batch的训练数据;dataset为数据集的一些基本信息,包括所有训练图片的路径,所有标签,每张图片的大小,图片的配置,超参数等等。
'''---------------------- 5.4.2 ✨将训练数据保存至模型✨-----------------'''
'''加载训练数据
返回一个训练数据加载器,一个数据集对象
训练数据加载器是一个可迭代的对象 可以通过for循环加载1个batch_size的数据
数据集对象包括数据集的一些参数 包括所有标签值、所有的训练数据路径、每张图片的尺寸等'''
train_loader, dataset = create_dataloader(train_path, imgsz, batch_size // WORLD_SIZE, gs, single_cls,
hyp=hyp, augment=True, cache=opt.cache, rect=opt.rect, rank=LOCAL_RANK,
workers=workers, image_weights=opt.image_weights, quad=opt.quad,
prefix=colorstr('train: '))
# 最大标签的编号值 例如coco有80个类别 编号从0开始,那么最大的编号就是79
mlc = int(np.concatenate(dataset.labels, 0)[:, 0].max()) # max label class
# nb表示数据的批数 例如coco128有128张图片 batch_size为16 128/16=8
nb = len(train_loader) # number of batches
# 检查类别数是否正确
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, pad=0.5,
prefix=colorstr('val: '))[0]
# 如果不使用断点训练
if not resume:
# 将dataset.labels中的每一项以y轴方向叠在一起
labels = np.concatenate(dataset.labels, 0)
# c = torch.tensor(labels[:, 0]) # classes
# cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency
# model._initialize_biases(cf.to(device))
# 如果绘图
if plots:
# 将标签分布图画出
plot_labels(labels, names, save_dir)
# Anchors
# 如果使用Anchor
if not opt.noautoanchor:
# 检查超参数中定义的anchor是否与训练集中的标注框契合
# 若anchor不合适则调整anchor
check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
# 半精度
model.half().float() # pre-reduce anchor precision
# 预训练常规配置结束 输出相应的日志
callbacks.run('on_pretrain_routine_end')
# DDP mode
if cuda and RANK != -1:
# 分布式多卡训练的设置
model = DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
这段代码是:根据检测层数和类别等,对损失函数因子超参数进行缩放调整。再将各个参数保存至模型。
'''---------------------- 5.4.2 ✨将训练数据保存至模型✨-----------------'''
# 为了平衡三个损失函数,会在损失函数前乘一个因子 这三个超参数为各自损失函数前的因子
# 当修改网络结构时,缩放这些参数
# box为预测框的损失
hyp['box'] *= 3. / nl # scale to layers
# cls为分类的损失
hyp['cls'] *= nc / 80. * 3. / nl # scale to classes and layers
# obj为置信度损失
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
这里要提到两个点:
1.warmup:warmup是一种学习率的优化方法,最早出现在ResNet的论文中。简单来说,在模型刚开始训练时,使用较小的学习率开始摸索,经过几轮迭代后使用大的学习率加速收敛,在快接近目标时,再使用小学习率,避免错过目标。
2.早停机制:当训练一定的轮数后,如果模型效果未提升,就让模型提前停止训练。这里的默认轮数为100轮,判断模型的效果为fitness,fitness为0.1乘[email protected]加上0.9乘[email protected]:0.95。
'''---------------------- 5.5 ✨进行训练✨ -----------------'''
'''----------------------5.5.1 ✨训练热身✨ -----------------'''
# 记录开始训练时间
t0 = time.time()
# 这里nw指的是warmup的迭代次数
# warmup指训练初期使用较小的学习率 是一种学习率的优化方式
nw = max(round(hyp['warmup_epochs'] * nb), 1000) # number of warmup iterations, max(3 epochs, 1k iterations)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
last_opt_step = -1
# 初始化每个类别的mAP
maps = np.zeros(nc) # mAP per class
# 初始化训练的返回结果 P, R, [email protected], [email protected], box_loss, obj_loss, cls_loss
results = (0, 0, 0, 0, 0, 0, 0) # P, R, [email protected], [email protected], val_loss(box, obj, cls)
# 设置学习率衰减所进行到的轮次,即使打断训练,使用resume接着训练也能正常衔接之前的训练进行学习率衰减
scheduler.last_epoch = start_epoch - 1 # do not move
# 使用自动混合精度运算
scaler = amp.GradScaler(enabled=cuda)
# 设置早停机制
# 训练了一定epoch,如果模型效果未提升,就让模型提前停止训练
stopper = EarlyStopping(patience=opt.patience)
# 初始化损失计算类
compute_loss = ComputeLoss(model) # init loss class
# 日志输出信息
LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'
f'Using {train_loader.num_workers} dataloader workers\n'
f"Logging results to {colorstr('bold', save_dir)}\n"
f'Starting training for {epochs} epochs...')
'''----------------------5.5.2 ✨开始训练✨ -----------------'''
# 开始训练
for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------
# 告诉模型现在是训练阶段 因为BN层、DropOut层、两阶段目标检测模型等
# 训练阶段阶段和预测阶段进行的运算是不同的,所以要将二者分开
# model.eval()指的是预测推断阶段
model.train()
# Update image weights (optional, single-GPU only)
# 更新图片的权重
if opt.image_weights:
# (1-maps)计算出来的是不精确度
# 经过一轮训练,若哪一类的不精确度高,那么这个类就会被分配一个较高的权重 来增加它被采样的概率
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
这段代码是:释放训练开始命令和更新权重。
训练命令释放,首先通过model.train()函数告诉模型,现在是训练阶段。因为有些层或模型在训练阶段与预测阶段进行的操作是不一样的,所以要通过model.train()函数用来声明,接下来是训练。若是预测阶段,则可以用model.eval().
更新图片的权重。训练时有些类的准确率可能比较难以识别,准确率并不会很高。在更新图片权重时就会把这些难以识别的类挑出来,并为这个类产生一些权重高的图片,以这种方式来增加识别率低的类别的数据量。提高准确率。
mloss = torch.zeros(3, device=device) # mean losses
# 分布式训练的设置
# DDP模式打乱数据,并且dpp.sampler的随机采样数据是基于epoch+seed作为随机种子,每次epoch不同,随机种子不同
if RANK != -1:
train_loader.sampler.set_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]:
# 通过tqdm创建进度条,方便训练信息的展示
pbar = tqdm(pbar, total=nb, ncols=NCOLS, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar
# 将优化器中的所有参数梯度设为0
optimizer.zero_grad()
这段代码是:分布式训练的设置,以及训练时控制台的显示。
首先DDP模式打乱数据,并进行随机采样。然后设置训练时控制台的显示。LOGGER.info是输出的表头, tqdm 显示进度条效果最后将优化器中所有的参数梯度设为0。
for i, (imgs, targets, paths, _) in pbar: # batch -------------------------------------------------------------
# ni: 计算当前迭代次数 iteration
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 热身训练
'''
热身训练(前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']])
# 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)
这段代码是:分批加载数据和热身训练以及多尺度训练。
- 分批加载训练数据:用ni计算当前迭代的次数,并作图片的归一化。
- 热身训练(warmup):这里只对训练初期使用较小的学习率。对于bias参数组的学习率策略是从0.1逐渐降低至初始学习率,其余参数组则从0开始逐渐增长至初始学习率。
- 多尺度训练: 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.
# Backward 反向传播 scale为使用自动混合精度运算
scaler.scale(loss).backward()
# 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
这段代码是:正向传播、反向传播、以及更新参数。
1.正向传播:将图片输入模型,并做一次正向传播,最后得到一个结果。这个结果在训练初期的效果可能会比较差,将这个结果与图片的标签值求损失,目的就是让这个损失越来越小。
2.反向传播:将这个误差,通过链式求导法则,反向传播回每一层,求出每层的梯度。
3.更新参数:利用optimizer.step更新参数。
注意
在更新参数时这里有一个不一样的地方,并不会在每次反向传播时更新参数,而是做一定的累积,反向传播的结果并不会顶替上一次反向传播结果,而是做一个累积。完成一次积累后,再将梯度清零,方便下一次清零。这样做是为了以更小的batch_size实现更高的batch_size效果。
# Log 打印Print一些信息 包括当前epoch、显存、损失(box、obj、cls、total)、当前batch的target的数量和图片的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)
# end batch ------------------------------------------------------------------------------------------------
# Scheduler 进行学习率衰减
lr = [x['lr'] for x in optimizer.param_groups] # for loggers
# 根据前面设置的学习率更新策略更新学习率
scheduler.step()
这段代码是:打印训练相关信息首先将每批最后的数据输出至控制台。到此每批循环体结束。然后在每所有批训练结束时,做权重衰减,进入下一轮的训练。
if RANK in [-1, 0]:
# mAP
# 调用Loggers中的on_train_epoch_end方法 做一次epoch自增的记录
callbacks.run('on_train_epoch_end', epoch=epoch)
# 更新ema的后面几个属性
ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights'])
# 结束时的epoch
final_epoch = (epoch + 1 == epochs) or stopper.possible_stop
# 如果使用每轮的验证 或 已经是训练的最后一轮
if not noval or final_epoch: # Calculate mAP
# 验证模型效果 返回验证结果以及 mAP
results, maps, _ = val.run(data_dict,
batch_size=batch_size // WORLD_SIZE * 2,
imgsz=imgsz,
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
# 计算fitness
# fitness是 0.1乘[email protected] 加上 0.9乘[email protected]:0.95
# 更加看重mAPA0.5:0.95的作用
fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, [email protected], [email protected]]
# 若当前的fitness大于最佳的fitness
if fi > best_fitness:
# 将最佳fitness更新为当前fitness
best_fitness = fi
# 保存验证结果
log_vals = list(mloss) + list(results) + lr
# 记录验证数据
callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)
判断是否应当结束训练,若选择每轮验证或当前已是最后一轮的情况下,做一次验证。并计算出最好的模型。这里“最好”的评判标准即为fitness。fitness是 0.1乘[email protected] 加上 0.9乘[email protected]:0.95,
在评判标准中,更加强调[email protected]:0.95的作用。[email protected]:0.95大代表模型在多个IOU阈值的情况下,都可以较好的识别物体。
'''----------------------5.5.3 ✨保存训练模型✨ -----------------'''
# Save model
# 如果保存模型
if (not nosave) or (final_epoch and not evolve): # if save
# 将当前训练过程中的所有参数赋值给ckpt
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(),
'wandb_id': loggers.wandb.wandb_run.id if loggers.wandb else None}
# Save last, best and delete
# 保存每轮的模型
torch.save(ckpt, last)
# 如果这个模型的fitness是最佳的
if best_fitness == fi:
# 保存这个最佳的模型
torch.save(ckpt, best)
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)
# Stop Single-GPU
# 停止单卡训练
if RANK == -1 and stopper(epoch=epoch, fitness=fi):
break
这段代码是:将最后一轮的模型,以及fitness最佳的模型保存下来。至此训练结束。
'''----------------------5.6 ✨结束训练✨ -----------------'''
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) # strip optimizers
if f is best:
# 把最好的模型在验证集上跑一边 并绘图
LOGGER.info(f'\nValidating {f}...')
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, # best pycocotools results at 0.65
single_cls=single_cls,
dataloader=val_loader,
save_dir=save_dir,
save_json=is_coco,
verbose=True,
plots=True,
callbacks=callbacks,
compute_loss=compute_loss) # val best model with plots
# 记录训练终止时的日志
callbacks.run('on_train_end', last, best, plots, epoch)
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")
torch.cuda.empty_cache()
return results
这段代码是:模型训练的结束。
train.py中还有一个函数 run()函数,run()函数内的内容与主函数差不多,都是调用了parse_opt()函数与main()函数,只不过run()函数是为导入时提供的,别的模块导入了train模块,即可通过调用run()函数执行训练过程。
'''----------------------六、run()函数-----------------'''
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)
这篇YOLOv5的train.py学习和总结到这里就结束啦,如果有什么问题可以在评论区留言呀~
如果帮助到大家,可以一键三连支持下~
其他学习YOLOv5直通车:
YOLOv5【使用云GPU进行训练】超详细教程!!!
YOLOv5【使用云GPU连接本地Pycharm进行训练】超详细教程!!!
YOLOv5【网络结构】超详细解读!!✨✨
YOLOv5【目录结构源码】超详细解读!!!
YOLOv5【detect.py源码及参数】超详细注释解读!!