目录
上一篇博客精读了YOLO v5 6.0版本的detect模块，这篇博客来精读train模块。
导入部分
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
"""
Train a YOLOv5 model on a custom dataset
在数据集上训练 yolo v5 模型
Usage:
$ python path/to/train.py --data coco128.yaml --weights yolov5s.pt --img 640
训练数据为coco128 coco128数据集中有128张图片 80个类别，是规模较小的数据集
"""
import argparse
import logging
import math
import os
import random
import sys
import time
from copy import deepcopy
from pathlib import Path
import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
import yaml
from torch.cuda import amp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim import Adam, SGD, lr_scheduler
from tqdm import tqdm
"""
确保root目录正确，避免导包时出现错误
因为以下导入的是自定义的包，若根目录错误就会导致导入失败，这里不再过多解释
"""
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] # YOLOv5 root directory
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT)) # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
import val # for end-of-epoch mAP
from models.experimental import attempt_load
from models.yolo import Model
from utils.autoanchor import check_anchors
from utils.datasets import create_dataloader
from utils.general import labels_to_class_weights, increment_path, labels_to_image_weights, init_seeds, \
strip_optimizer, get_latest_run, check_dataset, check_git_status, check_img_size, check_requirements, \
check_file, check_yaml, check_suffix, print_args, print_mutation, set_logging, one_cycle, colorstr, methods
from utils.downloads import attempt_download
from utils.loss import ComputeLoss
from utils.plots import plot_labels, plot_evolve
from utils.torch_utils import EarlyStopping, ModelEMA, de_parallel, intersect_dicts, select_device, \
torch_distributed_zero_first
from utils.loggers.wandb.wandb_utils import check_wandb_resume
from utils.metrics import fitness
from utils.loggers import Loggers
from utils.callbacks import Callbacks
导入自定义包时，要先检查根路径是否正确正确，避免导包时出现错误。
"""设置分布式训练的参数"""
LOGGER = logging.getLogger(__name__)
# 设置环境变量
# LOCAL_RANK表示进程内GPU编号，默认设为-1
LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html
# 表示进程编号，默认为-1
RANK = int(os.getenv('RANK', -1))
# WORLD_SIZE表示全局的进程数，默认为1
WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
接下来是设置分布式训练时所需的环境变量。分布式训练指的是多GPU训练，将训练参数分布在多个GPU上进行训练，有利于提升训练效率。
主函数
接下来按照程序的运行顺序，跳过函数定义部分，先看主函数
if __name__ == "__main__":
# 接收命令行参数
opt = parse_opt()
# 将命令行参数传入main函数
main(opt)
主函数中调用了两个函数，第parse_opt()函数用于接收命令行参数，并将包含参数的对象传递给main()函数。
parse_opt()函数
parse_opt()函数用于添加自定义的命令行参数。
def parse_opt(known=False):
"""
weight: 预训练的权重参数文件
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: 单类别 将多个类的数据按照一个类训练
adam: 是否使用adam优化器
sync-bn: 使用多卡归一化
workers: 最大工作数 用于设置分布式训练
project: 保存结果的路径
name: 保存结果的目录名
exist_ok: 是否重新结果目录 默认为False
linear-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('--adam', action='store_true', help='use torch.optim.Adam() 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('--linear-lr', action='store_true', help='linear LR')
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
main()函数
def main(opt, callbacks=Callbacks()):
# opt为命令行参数 callbacks用于存储Loggers日志记录器中的函数，方便在每个训练阶段控制日志的记录情况main()函数的参数中opt为命令行参数对象，callbacks中记录了Loggers类中的所有函数名。例如，可以通过callbacks.run("on_train_epoch_end")调用Loggers类中的on_train_epoch_end()函数。这样做方便在不同训练阶段，对日志的记录做统一管理。
前期准备
首先看main函数接收的参数，opt为命令行参数；callbacks为训练过程中保存的一些参数
# Checks
"""检查分布式训练环境"""
set_logging(RANK)
# 若进程编号为-1或0
if RANK in [-1, 0]:
# 打印参数
print_args(FILE.stem, opt)
# 检测YOLO v5的github仓库是否更新，若已更新，给出提示
check_git_status()
# 检测依赖库
check_requirements(exclude=['thop'])先检查（分布式）训练的环境，若RANK为-1或0，打印参数并检查github仓库和依赖库。
# Resume
"""判断是否从中断中恢复（断点训练）"""
# wandb为可视化初始化工具
if opt.resume and not check_wandb_resume(opt) and not opt.evolve: # resume an interrupted run
# 若resume是字符串类型,否则通过fet_latest_run()函数获取last.pt模型路径，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, opt.weights, opt.resume = '', ckpt, True # reinstate
# 日志输出提示信息
LOGGER.info(f'Resuming training from {ckpt}')
# 不使用断点训练 直接加载加载配置文件
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:
# 设置新的项目输出目录
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.save_dir
opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))接下来是关于断点训练的判断和准备。断点训练是当训练异常终止或想调节超参数时，系统会保留训练异常终止前的超参数与训练参数，当下次训练开始时，并不会从头开始，而是从上次中断的地方继续训练。
超参数进化：超参数进化是一种使用遗传算法，进行超参数优化的方法。遗传算法是用于解决最优化问题的一种搜索算法，它将要解决的问题模拟为一个生物进化的过程，过复制、交叉、突变等操作产生下一代的解，并逐步淘汰掉适应度函数值低的解，增加适应度函数值高的解。YOLO v5的超参数比较多，用传统的网格搜索会让超参数的调整变得很棘手。这时使用遗传算法调整超参数就是一种合适的方式。
# DDP mode
"""DDP训练是一种多GPU训练的方式 目前DDP模式只能在Linux下应用"""
# 选择程序装载的位置
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训练的配置，并设置GPU。DDP（Distributed Data Parallel）用于单机或多机的多GPU分布式训练，但目前DDP只能在Linux下使用。
训练的开始及结束
准备完毕后，接下来就开始训练
# 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()函数。
# 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
若使用超参数进化，以上代码为超参数训练的前期准备，首先指定每个超参数的突变范围、最大值、最小值，再为超参数的结果保存做好准备。
"""
接下来通过 遗传算法调整超参数
遵循优胜劣汰的原则
"""
# 选择超参数的遗传迭代次数 默认为迭代300次
for _ in range(opt.evolve): # generations to evolve
# 如果evolve.csv文件存在
if evolve_csv.exists(): # if evolve.csv exists: select best hyps and mutate
"""
遗传算法
参考博客
https://blog.csdn.net/weixin_44751294/article/details/125163790?ops_request_misc=&request_id=&biz_id=102&utm_term=%E8%B6%85%E5%8F%82%E6%95%B0%E8%BF%9B%E5%8C%96&utm_medium=distribute.pc_search_result.none-task-blog-2~all~sobaiduweb~default-0-125163790.142^v68^control,201^v4^add_ask,213^v2^t3_esquery_v2&spm=1018.2226.3001.4187
"""
# Select parent(s)
# 选择单亲或双亲
# 使用single模式
parent = 'single' # parent selection method: 'single' or 'weighted'
# 加载csv文件 x为异变后的结果
x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1)
# 最多选择5个最好的变异结果来挑选
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只能从小到大排序, 添加负号实现从大到小排序, 算是排序的一个代码技巧
# 挑选出适应度最好的前n个样本数据, 每个样本包含29个超参数变异结构{array:29}
x = x[np.argsort(-fitness(x))][:n] # top n mutations
# 根据(mp, mr, map50, map)的加权和来作为权重
w = fitness(x) - fitness(x).min() + 1E-6 # weights (sum > 0)
# single方式: 根据每个hyp的权重随机选择一个之前的hyp作为base hyp
# weighted方式: 根据每个hyp的权重对之前所有的hyp进行融合获得一个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)
for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300)
# [i+7]是因为x中前7个数字为results的指标(P,R,mAP,F1,test_loss=(box,obj,cls)),之后才是超参数hyp
hyp[k] = float(x[i + 7] * v[i]) # mutate 超参变异
# Constrain to limits
# 遍历元超参数
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
# 将调整后的超参数用于训练
# 7个结果 'metrics/precision', 'metrics/recall',
# 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95','val/box_loss', 'val/obj_loss', 'val/cls_loss'
results = train(hyp.copy(), opt, device, callbacks)
# Write mutation results 将突变结果保存至evolve.csv文件
print_mutation(results, hyp.copy(), save_dir, opt.bucket)
# Plot results
# 绘图: 每个超参数有一个子图, 显示适应度（y 轴）与超参数值（x 轴）.黄色表示更高的浓度
plot_evolve(evolve_csv)
# 将结果输出至控制台
print(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}')
接下来开始超参数进化，我将超参数进化的步骤总结为以下几步：1）若存在训练数据文件，读取文件中的训练数据，选择结果最优的训练数据突变超参数；2）将突变后的超参数限定值阈值范围内；3）使用突变后的超参数进行训练；4）训练结束后，将训练结果保存至evolution.csv，用于下一次的超参数突变。根据生物进化，优胜劣汰，适者生存的原则，每次迭代都会保存更优秀的结果，直至迭代结束。最后的结果即为最优的超参数。
使用超参数进化时至少要经过至少300次迭代，每次迭代都会经过一次完整的训练。因此超参数进化及其耗时，根据自己需求慎用。
run()函数
train模块中还有一个函数 run()函数，run()函数内的内容与主函数差不多，都是调用了parse_opt()函数与main()函数，只不过run()函数是为导入时提供的，别的模块导入了train模块，即可通过调用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)因为是为导入提供的，这里无法使用命令行设置参数。这里通过打包传参的方式，将参数传入run()函数，再用setattr函数，通过设置属性值的方式将参数设置给opt对象。
train()函数
接收参数
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.freezetrain()函数首先做的就是接收参数。hyp为超参数，不使用超参数进化的前提下也可以从opt中获取；opt指的是全部的命令行参数；device指的是装载程序的设备；callbacks指的是训练过程中产生的一些参数。
前期准备
# Directories
# 设置模型的保存路径
w = save_dir / 'weights' # weights dir
# 如果超参数进化，创建父文件夹，否则直接创建路径
(w.parent if evolve else w).mkdir(parents=True, exist_ok=True) # make dir
# 两个模型的路径 last为训练最后一轮产生的模型 best为fitness最好的模型
last, best = w / 'last.pt', w / 'best.pt'
接下来设置模型保存的路径。训练结束后，系统会产生两个模型，一个是last.pt，一个是best.pt。顾名思义，last.pt即为训练最后一轮产生的模型，而best.pt是训练过程中，效果最好的模型。那么何谓训练效果最好，评判效果最好的标准又是什么？
在这里评判best.pt的训练效果好坏的标准，绝大部分依赖的是mAP@0.5:0.95，是根据上文在超参数的突变时，提到的fitness()函数得出，后面遇到了再细讲。
# 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()))检查超参数是字典还是字符串，若为字符串，则认定为.yaml文件路径，再将yaml文件加载为字典。这里导致超参数的数据类型不同的原因是，超参数进化时，传入train()函数的超参数即为字典。而从命令行参数中读取的则为文件路径。
# 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设置训练过程中参数的保存路径
# 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))加载日志记录器，并将日志记录器中的函数记录到callbacks内，方便在训练的不同阶段，利用callbacks.run()函数对日志的记录做统一处理。
# 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
# 检查模型文件格式是否支持
check_suffix(weights, '.pt') # check weights
# 预训练权重参数是否以.pt结尾
pretrained = weights.endswith('.pt')
# 若以.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加载模型，分为使用预训练权重参数文件与不使用预训练权重参数。这里使用预训练权重参数，是类似于迁移学习。预训练的模型是检测coco数据集的模型，数据集中有80个类别，而自己的训练集类别以及类别的数量，并不与coco数据集相同。所以要先加载一个新的模型，把预训练的参数加载至模型作为初始参数，再把识别的类别改成自己的数据集要识别的类别。接下来将预训练参数中与新模型中相同的参数加载至模型。
新模型与预训练模型中为何会出现相同的参数？这里是因为预训练的概念类似于迁移学习，预训练模型中很多负责特征提取的参数，可能会对训练的新模型有所帮助，所以就会把这些参数加载至新的模型中。
# 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):
print(f'freezing {k}')
# 东接层不需要梯度
v.requires_grad = FalseFreeze会冻结模型的某些层，被冻结的层训练时不会更新参数。冻结层的原理是通过设置每个层参数中的requires_grad属性实现的。requires_grad属性在上一篇detect模块的精读中提到过。若require_grad为True，在反向传播时就会求出此tensor的梯度，若require_grad为False，则不会求该tensor的梯度。冻结就是通过对某些层不求梯度实现的。默认不进行参数冻结。
# 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']}")
接下来是优化器的设置。这里的nbs = 64，nbs指的是nominal batch size，名义上的batch_size。这里的nbs跟命令行参数中的batch_size不同，命令行中的batch_size默认为16，nbs设置为64。由于yolo在训练时占用的内存比较多，因此batch_size可设置的值就会比较小，batch_size小时会降低训练的速度，那么就要想办法，用opt.batch_size=16的情况，实现batch_size=64的效果。
accumulate 为累计次数，在这里 nbs/batch_size计算出 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发生了变化，所有权重参数也要做相应的缩放。
# 将模型参数分为三组(批归一化的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)将模型的参数分为三组，g0表示归一化层中的所有权重参数，g1表示卷积层中所有的权重参数，g2表示所有的偏置参数。
接下里循环遍历所有的层，将对应的参数添加至列表中。这里的循环遍历很有意思。model.modules()返回的不是每个layer，而是由大到小依次遍历每个层次。例如，第一次遍历出的v是模型本身；第二次遍历出的是模型内的Sequential；第三次遍历Sequential里的第一个模块；接下来继续往下遍历，直至遍历到卷积层或批归一化层。这有点像二叉树中的深度优先遍历。
# 若使用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)
判断是否使用adam优化器，初始参数为批归一化层中的参数。如果不使用adam优化器，则直接使用SGD随机梯度下降。
# 将卷积层的参数添加至优化器 并做权重衰减
# 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接下里将g1(卷积层中的权重参数)，g2(偏置参数)，添加进优化器。 add_param_group()函数可以为优化器中添加一个参数组。一个优化器可以更新多个参数组，不同的参数组可以使用不同的超参数。
# Scheduler
# linear_lr是一种学习率变化的策略 这里是线性学习率变化
if opt.linear_lr:
# 给定起始factor和最终的factor，LinearLR会在中间阶段做线性插值
lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear
else:
# 使用One Cycle学习率变化策略
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
# 为优化器设置学习率变化策略
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)
设置学习率调整策略。在训练过程中变更学习率可能会让训练效果更好，yolo v5提供了两种学习率变化的策略。一种是linear_lr，是通过线性插值的方式调整学习率；另一种则是One Cycle学习率调整策略，即周期性学习率调整中，周期被设置为1。在一周期策略中，最大学习率被设置为 LR Range test 中可以找到的最高值，最小学习率比最大学习率小几个数量级。这里默认one_cycle。
# EMA
# 全称是指数移动平均，是一种给予近期数据更高权重的平均方法
# 考虑历史值对参数的影响
ema = ModelEMA(model) if RANK in [-1, 0] else None
接下来是EMA，它的内容比较复杂，可以参考这篇博客
# Resume
# 中断训练 中断训练其实就是把上次训练结束的模型作为预训练模型，并从中加载参数
# 开始的epoch定义为0 最佳的fitness定义为0.0
start_epoch, best_fitness = 0, 0.0
# 如果进行预训练
if pretrained:
# Optimizer
# 若优化器不为None
if ckpt['optimizer'] is not None:
# 将预训练模型中的参数加载进优化器
optimizer.load_state_dict(ckpt['optimizer'])
# 获取预训练模型中的最佳fitness
best_fitness = ckpt['best_fitness']
# EMA 加载ema
if ema and ckpt.get('ema'):
ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
ema.updates = ckpt['updates']
# Epochs 加载上次运行到第几轮结束 并+1
start_epoch = ckpt['epoch'] + 1
# 如果使用中断训练
if resume:
# 判断上次中断已经训练结束
assert start_epoch > 0, f'{weights} training to {epochs} epochs is finished, nothing to resume.'
# 如果训练的轮数小于开始的轮数
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
中断训练的恢复。恢复终端训练其实就可以理解为，把上次中断结束时的模型，作为新的预训练模型，然后从中获取上次训练时的参数，并恢复训练状态。
# Image sizes
# 设置grid size大小，最小为32
gs = max(int(model.stride.max()), 32) # grid size (max stride)
# 检测层的个数 (三个感受野)
nl = model.model[-1].nl # number of detection layers (used for scaling hyp['obj'])
# 输入模型的尺寸大小
imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # verify imgsz is gs-multiple保存一些关于图片预测的信息，gs表示将图片划分的grid_size大小，gs最小为32；nl表示图片检测层的个数，这里是三个层，分别对应着不同的图片感受野，用于检测图片中不同大小的物体；imgz表示输入模型图片的尺寸。
# DP mode
# DP模式: 单机多卡模式
if cuda and RANK == -1 and torch.cuda.device_count() > 1:
logging.warning('DP not recommended, instead 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类将数据平均分布在多个GPU上
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()')DP模式为单机多卡模式，通过torch.nn.DataParallel(model)将模型中的数据平均分布在多张GPU上，并返回新的模型。再做多卡归一化，将多张GPU上的数据做归一化。
加载训练数据
# Trainloader
# 加载训练数据
# 返回一个训练数据加载器，一个数据集对象
# 训练数据加载器是一个可迭代的对象 可以通过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)加载训练数据时，通过create_dataloader()函数得到两个对象。一个为train_loader，另一个为dataset。train_loader为训练数据迭代器，可以通过for循环遍历出每个batch的训练数据；dataset为数据集的一些基本信息，包括所有训练图片的路径，所有标签，每张图片的大小，图片的配置，超参数等等。
# Model parameters
# 为了平衡三个损失函数，会在损失函数前乘一个因子 这三个超参数为各自损失函数前的因子
# 当修改网络结构时，缩放这些参数
# 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接下来根据检测层数和类别等，对损失函数因子超参数进行缩放调整。再将各个参数保存至模型。
开始训练
# Start training
"""开始训练"""
# 记录开始训练时间
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, mAP@.5, mAP@.5-.95, box_loss, obj_loss, cls_loss
results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, 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...')训前准备，做一些参数的初始化。这里要提到两个点：
第一个是warmup。warmup是一种学习率的优化方法，最早出现在ResNet的论文中。简单来说，在模型刚开始训练时，使用较小的学习率开始摸索，经过几轮迭代后使用大的学习率加速收敛，在快接近目标时，再使用小学习率，避免错过目标。
第二个是早停机制。当训练一定的轮数后，如果模型效果未提升，就让模型提前停止训练。这里的默认轮数为100轮，判断模型的效果为fitness，fitness为0.1乘mAP@0.5加上0.9乘mAP@0.5:0.95。
训练过程
# 开始训练
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
# 分布式训练的设置
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) # progress bar
# 将优化器中的所有参数梯度设为0
optimizer.zero_grad()定义一些训练中的参数，制作控制台输出的表头以及进度条效果waa清零
"""加载每批数据"""
for i, (imgs, targets, paths, _) in pbar: # batch -------------------------------------------------------------
# 记录批次 计算从第一轮开始一共运算了多少个批次
ni = i + nb * epoch # number integrated batches (since train start)
# 将图片加载至设备 并做归一化
imgs = imgs.to(device, non_blocking=True).float() / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
# Warmup
# warmup为热身训练 是一种学习率的优化方法
# 若总批次再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
# 对于bias参数组的学习率 会从0.1逐渐降低到初始学习率 使用线性插值的方式更新
# 对于其他参数组的学习率，将会从0逐渐升高到初始学习率 使用线性插值的方式更新
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
# 使用多尺度训练
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)
接下来就开始分批加载训练数据，并作图片的归一化。然后进行热身训练（warmup），这里只对训练初期使用较小的学习率。对于bias参数组的学习率策略是从0.1逐渐降低至初始学习率，其余参数组则从0开始逐渐增长至初始学习率。有助于使模型收敛速度变快，效果更佳.
# Forward
"""前向传播"""
with amp.autocast(enabled=cuda):
# 首先将图片喂给模型得到一个预测结果
pred = model(imgs) # forward
# 再将预测的结果与标签值计算损失
# 得到的loss为每项加权后的损失之和 loss_item是每一项的损失
# 每项损失分别为加权后的 预测框损失 置信度损失 分类损失
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
# scale为使用自动混合精度运算
# 这一步将误差反传
scaler.scale(loss).backward()
# Optimize
# 这里会对多批数据进行累积
# 只有达到累计次数的时候才会更新参数
# 再还没有达到累积次数时 loss会不断的叠加 不会被新的反传替代
if ni - last_opt_step >= accumulate:
scaler.step(optimizer) # optimizer.step
# 更新参数
scaler.update()
# 完成一次累积后，再将梯度清零，方便下一次清零
optimizer.zero_grad()
if ema:
ema.update(model)
# 计数
last_opt_step = ni
正向传播、反向传播、以及更新参数。
正向传播即将图片输入模型，并做一次正向传播，最后得到一个结果。这个结果在训练初期的效果可能会比较差，将这个结果与图片的标签值求损失，目的就是让这个损失越来越小。
接下来将这个误差，通过链式求导法则，反向传播回每一层，求出每层的梯度。
在更新参数时这里有一个设计，并不会在每次反向传播时更新残数，而是做一定的累积，反向传播的结果并不会顶替上一次反向传播结果，而是做一个累积。等到到一定的累计次数使，才会更新参数。这样做是为了以更小的batch_size实现更高的batch_size效果。
# 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(('%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乘mAP@0.5 加上 0.9乘mAP@0.5:0.95
# 更加看重mAPA0.5:0.95的作用
fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
# 若当前的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乘mAP@0.5 加上 0.9乘mAP@0.5:0.95，在评判标准中，更加强调mAP@0.5:0.95的作用。mAP@0.5:0.95大代表模型在多个IOU阈值的情况下，都可以较好的识别物体。
# 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最佳的模型保存下来。至此训练结束。
训练结束，输出结果
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 resultstrain()函数的最后，把最佳的模型取出，用这个最佳的模型跑一边验证集。再将结果保存下来，至此训练完成。若使用了超参数进化，还会进行多次训练，来完成超参数的调整。
1587
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
