# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
"""
Train a YOLOv5 model on a custom dataset.
Models and datasets download automatically from the latest YOLOv5 release.
Usage - Single-GPU training:
$ python train.py --data coco128.yaml --weights yolov5s.pt --img 640 # from pretrained (recommended)
$ python train.py --data coco128.yaml --weights '' --cfg yolov5s.yaml --img 640 # from scratch
Usage - Multi-GPU DDP training:
$ python -m torch.distributed.run --nproc_per_node 4 --master_port 1 train.py --data coco128.yaml --weights yolov5s.pt --img 640 --device 0,1,2,3
Models: <https://github.com/ultralytics/yolov5/tree/master/models>
Datasets: <https://github.com/ultralytics/yolov5/tree/master/data>
Tutorial: <https://docs.ultralytics.com/yolov5/tutorials/train_custom_data>
"""
import argparse # 解析命令行参数模块
import math # 数学公式模块
import os # 与操作系统进行交互的模块 包含文件路径操作和解析
import random # 生成随机数的模块
import subprocess # 该模块允许在 Python 程序中生成新的进程，连接到它们的输入/输出/错误管道，并获取它们的返回代码。这可以在 Python 程序中运行其他程序或命令。
import sys # sys系统模块 包含了与Python解释器和它的环境有关的函数
import time # 时间模块 更底层
from copy import deepcopy # 深拷贝模块
from datetime import datetime # 基本日期和时间类型模块
from pathlib import Path # Path模块将str转换为Path对象 使字符串路径易于操作
try:
import comet_ml # must be imported before torch (if installed) 必须在torch之前导入（如果安装）
except ImportError: # 用于导入Comet ML库，启用实验管理和机器学习项目跟踪的功能。
comet_ml = None
import numpy as np # numpy数组操作模块
import torch # torch深度学习框架
import torch.distributed as dist # 分布式训练模块
import torch.nn as nn # 对torch.nn.functional的类的封装 有很多和oneflow.nn.functional相同的函数
import yaml # 操作yaml文件模块
from torch.optim import lr_scheduler # 学习率模块
from tqdm import tqdm # 进度条模块
FILE = Path(__file__).resolve() # 获取当前脚本文件的绝对路径，并将其存储在 FILE 变量中，方便在代码中引用和操作当前脚本文件。
ROOT = FILE.parents[0] # YOLOv5 root directory (YOLOv5根目录)
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 as validate # for end-of-epoch mAP (导入val.py)
from models.experimental import attempt_load # 导入在线下载模块
from models.yolo import Model # 导入YOLOv5的模型定义
from utils.autoanchor import check_anchors # 导入检查anchors合法性的函数
from utils.autobatch import check_train_batch_size # 用于在YOLOv5训练过程中检查和验证批处理大小，以确保与GPU内存兼容。
from utils.callbacks import Callbacks # 和日志相关的回调函数
from utils.dataloaders import create_dataloader # 加载数据集的函数
from utils.downloads import attempt_download, is_url # 尝试下载指定的文件，并提供断点续传功能，确保下载的可靠性和完整性//判断当前字符串是否是链接
from utils.general import (LOGGER, TQDM_BAR_FORMAT, check_amp, check_dataset, check_file, check_git_info,
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,
yaml_save) # 导入utils.general文件中各种函数
from utils.loggers import Loggers # 导入日志管理模块
from utils.loggers.comet.comet_utils import check_comet_resume # 用于检查是否可以从Comet ML平台恢复实验，以便继续记录和分析。
from utils.loss import ComputeLoss # 导入计算Loss的模块
from utils.metrics import fitness # 在YOLOv5中，fitness函数实现对 [P, R, [email protected], [email protected]] 指标进行加权
from utils.plots import plot_evolve # 用于绘制进化算法的进化图，可视化进化过程和结果，帮助用户理解算法效果和变化趋势。
from utils.torch_utils import (EarlyStopping, ModelEMA, de_parallel, select_device, smart_DDP, smart_optimizer,
smart_resume, torch_distributed_zero_first) # 导入utils.torch_utils文件中各种函数
# LOCAL_RANK：当前进程对应的GPU号。
LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # <https://pytorch.org/docs/stable/elastic/run.html>
RANK = int(os.getenv('RANK', -1)) # RANK：当前进程的序号，用于进程间通讯，rank = 0 的主机为 master 节点。
WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1)) # WORLD_SIZE：总的进程数量 (原则上第一个process占用一个GPU是较优的)。
GIT_INFO = check_git_info() # 获取当前代码的 Git 信息，并将其存储在 GIT_INFO 变量中以供使用。
def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictionary
# 载入参数
save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = \\
Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \\
opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze
callbacks.run('on_pretrain_routine_start')
# Directories 保存模型的路径
w = save_dir / 'weights' # weights dir
(w.parent if evolve else w).mkdir(parents=True, exist_ok=True) # make dir
last, best = w / 'last.pt', w / 'best.pt'
# Hyperparameters 超参数
if isinstance(hyp, str): # 检查hyp是否为字符串
with open(hyp, errors='ignore') as f:
hyp = yaml.safe_load(f) # load hyps dict 加载超参数字典
# 日志输出超参信息 hyperparameters: ...
LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
opt.hyp = hyp.copy() # for saving hyps to checkpoints 将超参数保存到检查点中
# Save run settings 保存运行时的参数配置
if not evolve:
yaml_save(save_dir / 'hyp.yaml', hyp)
yaml_save(save_dir / 'opt.yaml', vars(opt)) # vars()用于返回一个对象的属性和属性值的字典。
# Loggers 日志
data_dict = None
if RANK in {-1, 0}:
# loggers 是一个 Loggers 类的实例
# 传入了一些参数，如 save_dir（保存目录）、weights（权重）、opt（训练选项）、hyp（超参数）和 LOGGER（全局日志记录器）
loggers = Loggers(save_dir, weights, opt, hyp, LOGGER) # loggers instance 日志记录器实例
# Register actions 注册
for k in methods(loggers): # 通过 methods(loggers) 函数获取 loggers 实例的方法列表
# 使用 getattr(loggers, k) 获取每个方法的引用。
# 通过 callbacks.register_action(k, callback=getattr(loggers, k)) 将每个方法注册为回调函数。这样，在训练过程中，每次执行相应的回调函数时，会触发对应的方法。
callbacks.register_action(k, callback=getattr(loggers, k))
# Process custom dataset artifact link
data_dict = loggers.remote_dataset
if resume: # If resuming runs from remote artifact
weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size
# Config
# 是否需要画图： 所有的labels信息、迭代的epochs、训练结果等
plots = not evolve and not opt.noplots # create plots
cuda = device.type != 'cpu'
# 初始化随机数种子
init_seeds(opt.seed + 1 + RANK, deterministic=True)
with torch_distributed_zero_first(LOCAL_RANK):
data_dict = data_dict or check_dataset(data) # check if None
train_path, val_path = data_dict['train'], data_dict['val']
# nc: 数据集有多少种类别
nc = 1 if single_cls else int(data_dict['nc']) # number of classes
# 如果只有一个类别并且data_dict里没有names这个key的话，我们将names设置为["item"]代表目标
names = {0: 'item'} if single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names
# 当前数据集是否是coco数据集(80个类别)
is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt') # COCO dataset
# Model
check_suffix(weights, '.pt') # 检查权重文件后缀
pretrained = weights.endswith('.pt') # 检查权重文件后缀，返回布尔类型
# 载入模型
if pretrained:
with torch_distributed_zero_first(LOCAL_RANK):
weights = attempt_download(weights) # 如果本地找不到，则下载权重文件
# 使用预训练
# ---------------------------------------------------------#
# 加载模型及参数
ckpt = torch.load(weights, map_location='cpu') # 将检查点加载到 CPU 上，以避免 CUDA 内存泄漏
# 这里加载模型有两种方式，一种是通过opt.cfg 另一种是通过ckpt['model'].yaml
# 区别在于是否使用resume 如果使用resume会将opt.cfg设为空，按照ckpt['model'].yaml来创建模型
# 这也影响了下面是否除去anchor的key(也就是不加载anchor), 如果resume则不加载anchor
# 原因: 保存的模型会保存anchors，有时候用户自定义了anchor之后，再resume，则原来基于coco数据集的anchor会自己覆盖自己设定的anchor
# 详情参考: <https://github.com/ultralytics/yolov5/issues/459>
# 所以下面设置intersect_dicts()就是忽略exclude
model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # 创建模型
exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else [] # 排除的键列表
csd = ckpt['model'].float().state_dict() # 以 FP32 格式保存检查点的 state_dict
# 筛选字典中的键值对 把exclude删除
csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # 取交集
model.load_state_dict(csd, strict=False) # 加载模型参数
LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}') # 输出日志信息
else:
# 不使用预训练下创建模型
model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
# 检查是否可以使用自动混合精度（AMP），若不支持则继续使用默认的浮点型精度进行训练
amp = check_amp(model) # check AMP
# Freeze
# 冻结权重层，冻结模型中的一些层，使其在训练过程中不更新参数梯度。
# 这里只是给了冻结权重层的一个例子, 但是作者并不建议冻结权重层, 训练全部层参数, 可以得到更好的性能, 不过也会更慢
freeze = [f'model.{x}.' for x in (freeze if len(freeze) > 1 else range(freeze[0]))] # # 要冻结的层
for k, v in model.named_parameters():
v.requires_grad = True # 训练所有层
# v.register_hook(lambda x: torch.nan_to_num(x)) # 将 NaN 替换为 0（注释以避免训练结果不稳定）
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) # 网格大小（最大步长）
imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # 验证 imgsz 是否为 gs 的倍数
# Batch size 批大小
# 判断条件 RANK == -1 and batch_size == -1 是否成立。如果成立，说明当前是单 GPU 的情况，并且批大小未指定。
if RANK == -1 and batch_size == -1: # 仅适用于单 GPU，估计最佳批大小
batch_size = check_train_batch_size(model, imgsz, amp) # 估计最佳批处理大小并赋值给batch_size
loggers.on_params_update({'batch_size': batch_size})
# Optimizer 选择优化器
nbs = 64 # 标称批大小，即：在选择优化器和设置优化器参数时所使用的一个固定的批大小。
# 使用了 round(nbs / batch_size) 来计算累积因子，即标称批大小与实际批大小之间的比值，然后取最大值为 1。
accumulate = max(round(nbs / batch_size), 1) # 在优化之前累积损失
hyp['weight_decay'] *= batch_size * accumulate / nbs # 缩放 weight_decay
optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['weight_decay'])
# Scheduler 学习率调度器
if opt.cos_lr:
# 使用 one_cycle 函数创建一个余弦学习率调度器。
# 该调度器的学习率从 1 线性下降到 hyp['lrf']，然后保持不变。
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
else:
# 使用一个匿名函数定义一个线性学习率调度器。
lf = lambda x: (1 - x / epochs) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear
# 使用 lr_scheduler.LambdaLR 函数将优化器 optimizer 和学习率调度函数 lf 组合起来，创建一个学习率调度器 scheduler。
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)
# EMA 指数移动平均
ema = ModelEMA(model) if RANK in {-1, 0} else None
# Resume 断点续训
best_fitness, start_epoch = 0.0, 0
if pretrained:
if resume:
best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume)
del ckpt, csd
# DP mode
if cuda and RANK == -1 and torch.cuda.device_count() > 1:
LOGGER.warning(
'WARNING ⚠️ DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\\n'
'See Multi-GPU Tutorial at <https://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training> to get started.'
)
model = torch.nn.DataParallel(model)
# SyncBatchNorm
if opt.sync_bn and cuda and RANK != -1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
LOGGER.info('Using SyncBatchNorm()')
# Trainloader
# 数据加载
train_loader, dataset = create_dataloader(train_path,
imgsz,
batch_size // WORLD_SIZE,
gs,
single_cls,
hyp=hyp,
augment=True,
cache=None if opt.cache == 'val' else opt.cache,
rect=opt.rect,
rank=LOCAL_RANK,
workers=workers,
image_weights=opt.image_weights,
quad=opt.quad,
prefix=colorstr('train: '),
shuffle=True,
seed=opt.seed)
labels = np.concatenate(dataset.labels, 0)
mlc = int(labels[:, 0].max()) # max label class
assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'
# Process 0
if RANK in {-1, 0}:
val_loader = create_dataloader(val_path,
imgsz,
batch_size // WORLD_SIZE * 2,
gs,
single_cls,
hyp=hyp,
cache=None if noval else opt.cache,
rect=True,
rank=-1,
workers=workers * 2,
pad=0.5,
prefix=colorstr('val: '))[0]
# 如果不使用断点续训
if not resume:
# Anchors
# 计算默认锚框anchor与数据集标签框的高宽比
# 标签的高h宽w与anchor的高h_a宽h_b的比值 即h/h_a, w/w_a都要在(1/hyp['anchor_t'], hyp['anchor_t'])是可以接受的
# 如果bpr小于98%，则根据k-mean算法聚类新的锚框
if not opt.noautoanchor:
# check_anchors : 这个函数是通过计算bpr确定是否需要改变anchors 需要就调用k-means重新计算anchors。
# bpr(best possible recall): 最多能被召回的ground truth框数量 / 所有ground truth框数量 最大值为1 越大越好
# 小于0.98就需要使用k-means + 遗传进化算法选择出与数据集更匹配的anchor boxes框。
check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) # run AutoAnchor
model.half().float() # pre-reduce anchor precision
callbacks.run('on_pretrain_routine_end', labels, names)
# DDP mode
if cuda and RANK != -1:
model = smart_DDP(model)
# 附加Model attributes
# nl：这个属性代表模型中的检测层数量。它用于调整假设的规模。
nl = de_parallel(model).model[-1].nl # 检测层的数量（用于缩放hyps）
# hyp['box']：该属性被缩放了3除以nl的因子。它用于调整边界框的大小。
hyp['box'] *= 3 / nl # # 缩放到层级
# hyp['cls']：该属性被缩放了nc / 80乘以3除以nl的因子。它用于调整类别和层数的大小。
hyp['cls'] *= nc / 80 * 3 / nl # 缩放到类别和层级
# hyp['obj']：该属性被缩放了(imgsz / 640) ** 2乘以3除以nl的因子。它用于根据图像大小和层数调整对象的大小。
hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl # 缩放到图像大小和层级
# hyp['label_smoothing']：该属性被设置为opt.label_smoothing的值。它用于在训练过程中进行标签平滑处理。
hyp['label_smoothing'] = opt.label_smoothing
# model.nc：该属性代表模型中的类别数量。nc被附加到模型上。
model.nc = nc # 将类别数量附加到模型
# model.hyp：该属性代表模型的超参数。hyp被附加到模型上。
model.hyp = hyp # 将超参数附加到模型
# 从训练样本标签得到类别权重（和类别中的目标数即类别频率成反比）
# model.class_weights：该属性代表模型的类别权重。它使用labels_to_class_weights函数和数据集的标签以及nc作为参数进行计算。然后乘以nc并附加到模型上。
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # 附加类别权重
# model.names：该属性代表模型中类别的名称。
model.names = names # 获取类别名
# Start training
t0 = time.time() # 记录当前时间，用于计算训练时间。
nb = len(train_loader) # number of batches 获取训练数据集的批次数量，即每个批次包含多少个样本。
# 获取预热迭代的次数iterations
# 预热迭代是在训练开始时使用较小的学习率进行一些迭代，以帮助模型更快地收敛。这里使用的是将预热的轮数乘以批次数量，但最少为100次。
nw = max(round(hyp['warmup_epochs'] * nb), 100) # number of warmup iterations, max(3 epochs, 100 iterations)
# ↓使用 min() 函数将计算得到的值与 nw 进行比较，选择较小的值作为最终的预热迭代次数。
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # 限制预热迭代的次数不超过总训练次数的一半。
# ↑目的是确保预热迭代不会占据过多的训练时间，而是在训练的前半部分进行。这样可以更好地平衡模型在预热和正式训练阶段的学习率调整和收敛速度。
last_opt_step = -1 # 初始化最后一次优化的步骤为-1，用于后续记录。
# 初始化maps(每个类别的map)和results
maps = np.zeros(nc) # mAP per class
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 = torch.cuda.amp.GradScaler(enabled=amp) 这个是和amp相关的loss缩放模块，后续one-yolv5支持好amp训练后会打开
scaler = torch.cuda.amp.GradScaler(enabled=amp)
stopper, stop = EarlyStopping(patience=opt.patience), False
# 初始化损失函数
# 这里的bbox_iou_optim是one-yolov5扩展的一个参数，可以启用更快的bbox_iou函数，模型训练速度比PyTorch更快。
compute_loss = ComputeLoss(model) # init loss class
callbacks.run('on_train_start')
# 打印日志信息
LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\\n'
f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\\n'
f"Logging results to {colorstr('bold', save_dir)}\\n"
f'Starting training for {epochs} epochs...')
for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------
callbacks.run('on_train_epoch_start')
model.train()
# Update image weights (optional, single-GPU only)
# Update image weights (optional) 并不一定好 默认是False的
# 如果为True 进行图片采样策略(按数据集各类别权重采样)
if opt.image_weights:
# 根据前面初始化的图片采样权重model.class_weights（每个类别的权重 频率高的权重小）以及maps配合每张图片包含的类别数
# 通过rando.choices生成图片索引indices从而进行采用 （作者自己写的采样策略，效果不一定ok）
cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights
# labels_to_image_weights: 这个函数是利用每张图片真实gt框的真实标签labels和开始训练前通过 labels_to_class_weights函数
# 得到的每个类别的权重得到数据集中每张图片对应的权重。
# <https://github.com/Oneflow-Inc/oneflow-yolo-doc/blob/master/docs/source_code_interpretation/utils/general_py.md#192-labels_to_image_weights>
iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights
dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx
# Update mosaic border (optional)
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
# 初始化训练时打印的平均损失信息
mloss = torch.zeros(3, device=device) # mean losses
if RANK != -1:
# DDP模式打乱数据，并且ddp.sampler的随机采样数据是基于epoch+seed作为随机种子，每次epoch不同，随机种子不同
train_loader.sampler.set_epoch(epoch)
# 进度条，方便展示信息
pbar = enumerate(train_loader)
LOGGER.info(('\\n' + '%11s' * 7) % ('Epoch', 'GPU_mem', 'box_loss', 'obj_loss', 'cls_loss', 'Instances', 'Size'))
if RANK in {-1, 0}:
# 创建进度条
pbar = tqdm(pbar, total=nb, bar_format=TQDM_BAR_FORMAT) # progress bar
# 梯度清零
optimizer.zero_grad()
for i, (imgs, targets, paths, _) in pbar: # batch -------------------------------------------------------------
callbacks.run('on_train_batch_start')
# ni: 计算当前迭代次数 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] 选取较小的accumulate，学习率以及momentum,慢慢的训练
if ni <= nw:
xi = [0, nw] # x interp
# compute_loss.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 0 else 0.0, x['initial_lr'] * lf(epoch)])
if 'momentum' in x:
x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
# Multi-scale 默认关闭
# Multi-scale 多尺度训练 从[imgsz*0.5, imgsz*1.5+gs]间随机选取一个尺寸(32的倍数)作为当前batch的尺寸送入模型开始训练
# imgsz: 默认训练尺寸 gs: 模型最大stride=32 [32 16 8]
if opt.multi_scale:
sz = random.randrange(int(imgsz * 0.5), int(imgsz * 1.5) + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple)
# 下采样
imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
# Forward
with torch.cuda.amp.autocast(amp):
pred = model(imgs) # forward
loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size
if RANK != -1:
loss *= WORLD_SIZE # gradient averaged between devices in DDP mode
if opt.quad:
loss *= 4.
# Backward
# scaler.scale(loss).backward()
# Backward 反向传播
scaler.scale(loss).backward()
# Optimize - <https://pytorch.org/docs/master/notes/amp_examples.html>
# 模型反向传播accumulate次（iterations）后再根据累计的梯度更新一次参数
if ni - last_opt_step >= accumulate:
scaler.unscale_(optimizer) # unscale gradients
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0) # clip gradients
scaler.step(optimizer) # optimizer.step 参数更新
scaler.update()
# 梯度清零
optimizer.zero_grad()
if ema:
# 当前epoch训练结束 更新ema
ema.update(model)
last_opt_step = ni
# Log
# 打印Print一些信息 包括当前epoch、显存、损失(box、obj、cls、total)、当前batch的target的数量和图片的size等信息
if RANK in {-1, 0}:
mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G' # (GB)
pbar.set_description(('%11s' * 2 + '%11.4g' * 5) %
(f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))
callbacks.run('on_train_batch_end', model, ni, imgs, targets, paths, list(mloss))
if callbacks.stop_training:
return
# end batch ------------------------------------------------------------------------------------------------
# Scheduler
lr = [x['lr'] for x in optimizer.param_groups] # for loggers
scheduler.step()
if RANK in {-1, 0}:
# mAP
callbacks.run('on_train_epoch_end', epoch=epoch)
ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights'])
final_epoch = (epoch + 1 == epochs) or stopper.possible_stop
if not noval or final_epoch: # Calculate mAP
# 测试使用的是ema（指数移动平均 对模型的参数做平均）的模型
# results: [1] Precision 所有类别的平均precision(最大f1时)
# [1] Recall 所有类别的平均recall
# [1] [email protected] 所有类别的平均[email protected]
# [1] [email protected]:0.95 所有类别的平均[email protected]:0.95
# [1] box_loss 验证集回归损失, obj_loss 验证集置信度损失, cls_loss 验证集分类损失
# maps: [80] 记录每一个类别的ap值
results, maps, _ = validate.run(data_dict,
batch_size=batch_size // WORLD_SIZE * 2,
imgsz=imgsz,
half=amp,
model=ema.ema,
single_cls=single_cls,
dataloader=val_loader,
save_dir=save_dir,
plots=False,
callbacks=callbacks,
compute_loss=compute_loss)
# Update best mAP
# fi 是我们寻求最大化的值。在YOLOv5中，fitness函数实现对 [P, R, [email protected], [email protected]] 指标进行加权。
fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, [email protected], [email protected]]
stop = stopper(epoch=epoch, fitness=fi) # early stop check
if fi > best_fitness:
best_fitness = fi
log_vals = list(mloss) + list(results) + lr
callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)
# Save model
if (not nosave) or (final_epoch and not evolve): # if save
ckpt = {
'epoch': epoch,
'best_fitness': best_fitness,
'model': deepcopy(de_parallel(model)).half(),
'ema': deepcopy(ema.ema).half(),
'updates': ema.updates,
'optimizer': optimizer.state_dict(),
'opt': vars(opt),
'git': GIT_INFO, # {remote, branch, commit} if a git repo
'date': datetime.now().isoformat()}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
if opt.save_period > 0 and epoch % opt.save_period == 0:
torch.save(ckpt, w / f'epoch{epoch}.pt')
del ckpt
# Write 将测试结果写入result.txt中
callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi)
# EarlyStopping
if RANK != -1: # if DDP training
broadcast_list = [stop if RANK == 0 else None]
dist.broadcast_object_list(broadcast_list, 0) # broadcast 'stop' to all ranks
if RANK != 0:
stop = broadcast_list[0]
if stop:
break # must break all DDP ranks
# end epoch ----------------------------------------------------------------------------------------------------
# end training -----------------------------------------------------------------------------------------------------
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, _, _ = validate.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 at iou 0.65
single_cls=single_cls,
dataloader=val_loader,
save_dir=save_dir,
save_json=is_coco,
verbose=True,
plots=plots,
callbacks=callbacks,
compute_loss=compute_loss) # val best model with plots
if is_coco:
callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)
callbacks.run('on_train_end', last, best, epoch, results)
torch.cuda.empty_cache()
return results
def parse_opt(known=False):
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
parser.add_argument('--data', type=str, default=ROOT / 'D:/dataset/Armor128/Armor128.yaml', help='dataset.yaml path')
parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml', help='hyperparameters path')
parser.add_argument('--epochs', type=int, default=100, help='total training epochs')
parser.add_argument('--batch-size', type=int, default=1, help='total batch size for all GPUs, -1 for autobatch')
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=1280, help='train, val image size (pixels)')
parser.add_argument('--rect', action='store_true', help='rectangular training')
parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
parser.add_argument('--noval', action='store_true', help='only validate final epoch')
parser.add_argument('--noautoanchor', action='store_true', help='disable AutoAnchor')
parser.add_argument('--noplots', action='store_true', help='save no plot files')
parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
parser.add_argument('--cache', type=str, nargs='?', const='ram', help='image --cache ram/disk')
parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer')
parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name')
parser.add_argument('--name', default='exp', help='save to project/name')
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
parser.add_argument('--quad', action='store_true', help='quad dataloader')
parser.add_argument('--cos-lr', action='store_true', help='cosine LR scheduler')
parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')
parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone=10, first3=0 1 2')
parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
parser.add_argument('--seed', type=int, default=0, help='Global training seed')
parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')
# Logger arguments
parser.add_argument('--entity', default=None, help='Entity')
parser.add_argument('--upload_dataset', nargs='?', const=True, default=False, help='Upload data, "val" option')
parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval')
parser.add_argument('--artifact_alias', type=str, default='latest', help='Version of dataset artifact to use')
return parser.parse_known_args()[0] if known else parser.parse_args()
def main(opt, callbacks=Callbacks()):
# Checks
if RANK in {-1, 0}:
# 输出所有训练opt参数 train: ...
print_args(vars(opt))
# 检查代码版本是否是最新的 github: ...
check_git_status()
# 检查requirements.txt所需包是否都满足 requirements: ...
check_requirements(ROOT / 'requirements.txt')
# Resume
# 判断是否使用断点续训resume, 读取参数
if opt.resume and not check_comet_resume(opt) and not opt.evolve: # from specified or most recent last.pt 源自特定的或最近的last.pt
# 使用断点续训 就从last模型文件夹中读取相关参数
# 如果resume是str，则表示传入的是模型的路径地址
# 如果resume是True，则通过get_lastest_run()函数找到runs文件夹中最近的权重文件last
last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run())
opt_yaml = last.parent.parent / 'opt.yaml' # train options yaml
opt_data = opt.data # original dataset
if opt_yaml.is_file():
# 相关的opt参数也要替换成last中的opt参数
with open(opt_yaml, errors='ignore') as f:
d = yaml.safe_load(f)
else:
d = torch.load(last, map_location='cpu')['opt']
opt = argparse.Namespace(**d) # replace
opt.cfg, opt.weights, opt.resume = '', str(last), True # reinstate
if is_url(opt_data):
opt.data = check_file(opt_data) # avoid HUB resume auth timeout
else:
# 不使用断点续训 就从文件中读取相关参数
# opt.hyp = opt.hyp or ('hyp.finetune.yaml' if opt.weights else 'hyp.scratch.yaml')
opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = \\
check_file(opt.data), check_yaml(opt.cfg), check_yaml(opt.hyp), str(opt.weights), str(opt.project) # checks
assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
if opt.evolve: # 当使用进化算法时
if opt.project == str(ROOT / 'runs/train'): # if default project name, rename to runs/evolve
opt.project = str(ROOT / 'runs/evolve')
opt.exist_ok, opt.resume = opt.resume, False # pass resume to exist_ok and disable resume
if opt.name == 'cfg':
opt.name = Path(opt.cfg).stem # use model.yaml as name
# 根据opt.project生成目录 如: runs/train/exp18
opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))
# DDP mode DDP mode设置
"""select_device
select_device 函数： 设置当前脚本的device：cpu或者cuda。
并且当且仅当使用cuda时并且有多块gpu时可以使用ddp模式，否则抛出报错信息。batch_size需要整除总的进程数量。
另外DDP模式不支持AutoBatch功能，使用DDP模式必须手动指定batch size。
"""
device = select_device(opt.device, batch_size=opt.batch_size)
if LOCAL_RANK != -1:
msg = 'is not compatible with YOLOv5 Multi-GPU DDP training'
assert not opt.image_weights, f'--image-weights {msg}'
assert not opt.evolve, f'--evolve {msg}'
assert opt.batch_size != -1, f'AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size'
assert opt.batch_size % WORLD_SIZE == 0, f'--batch-size {opt.batch_size} must be multiple of WORLD_SIZE'
assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
torch.cuda.set_device(LOCAL_RANK)
device = torch.device('cuda', LOCAL_RANK)
dist.init_process_group(backend='nccl' if dist.is_nccl_available() else 'gloo')
# Train
# 不使用进化算法 正常Train
if not opt.evolve:
train(opt.hyp, opt, device, callbacks)
# 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) 优化器的初始学习率。对于SGD，默认值为1e-2，对于Adam，默认值为1e-3。
'lrf': (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf) OneCycleLR调度器的最终学习率。计算方式为lr0 * lrf。
'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1 SGD优化器的动量值或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 训练中的IoU（交并比）阈值。它确定预测的边界框是被视为真正例还是假正例。
'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold 锚框倍数的阈值。根据其大小确定用于训练的锚框。
'anchors': (2, 2.0, 10.0), # anchors per output grid (0 to ignore) 每个输出网格的锚框数量。值为0表示忽略锚框。
'fl_gamma': (0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5) Focal Loss的gamma值。它控制在训练过程中对难例的关注程度。
'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction) 训练中应用的图像HSV-Hue增强的比例。
'hsv_s': (1, 0.0, 0.9), # image HSV-Saturation augmentation (fraction) 训练中应用的图像HSV-Saturation增强的比例。
'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction) 训练中应用的图像HSV-Value增强的比例。
'degrees': (1, 0.0, 45.0), # image rotation (+/- deg) 训练中应用的图像旋转范围（正负角度）。
'translate': (1, 0.0, 0.9), # image translation (+/- fraction) 训练中应用的图像平移范围。
'scale': (1, 0.0, 0.9), # image scale (+/- gain) 训练中应用的图像缩放范围。
'shear': (1, 0.0, 10.0), # image shear (+/- deg) 训练中应用的图像剪切范围（正负角度）。
'perspective': (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001 训练中应用的图像透视范围。
'flipud': (1, 0.0, 1.0), # image flip up-down (probability) 训练中垂直翻转图像的概率。
'fliplr': (0, 0.0, 1.0), # image flip left-right (probability) 训练中水平翻转图像的概率。
'mosaic': (1, 0.0, 1.0), # image mixup (probability) 使用马赛克技术应用图像混合的概率。
'mixup': (1, 0.0, 1.0), # image mixup (probability) 训练中应用图像混合的概率。
'copy_paste': (1, 0.0, 1.0)} # segment copy-paste (probability) 训练中应用分割复制粘贴的概率。
# 载入初始超参
with open(opt.hyp, errors='ignore') as f:
hyp = yaml.safe_load(f) # load hyps dict 使用yaml.safe_load函数加载该文件中的内容为一个字典对象hyp
if 'anchors' not in hyp: # anchors commented in hyp.yaml 通过检查hyp字典中是否存在anchors键来确定是否在hyp.yaml文件中注释了锚框的设定。
hyp['anchors'] = 3 # 如果anchors键不存在，说明锚框的设定被注释掉了，则将anchors设定为默认值3。
if opt.noautoanchor: # 如果在命令行参数opt中设置了noautoanchor选项，则删除hyp字典中的anchors键，并删除meta字典中的anchors键。
del hyp['anchors'], meta['anchors']
# 设置opt.noval为True，表示只在最后一个epoch进行验证（不进行中间epoch的验证）。
# 同时，设置opt.nosave为True，表示只保存最后一个epoch的模型（不保存中间epoch的模型）。
# 最后，将opt.save_dir转换为Path对象，并赋值给save_dir变量。
# 这样设置可以节省存储空间，只保存最终结果，而不保存中间结果。
opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir) # only val/save final epoch 只验证/保存最后一个epoch
# ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices
# evolve_yaml 超参进化后文件保存地址
evolve_yaml, evolve_csv = save_dir / 'hyp_evolve.yaml', save_dir / 'evolve.csv'
if opt.bucket:
# download evolve.csv if exists 下载 evolve.csv（如果存在）
subprocess.run([
'gsutil',
'cp',
f'gs://{opt.bucket}/evolve.csv',
str(evolve_csv),])
"""
使用遗传算法进行参数进化 默认是进化300代
这里的进化算法原理为：根据之前训练时的hyp来确定一个base hyp再进行突变，具体是通过之前每次进化得到的results来确定之前每个hyp的权重，有了每个hyp和每个hyp的权重之后有两种进化方式；
1.根据每个hyp的权重随机选择一个之前的hyp作为base hyp，random.choices(range(n), weights=w)
2.根据每个hyp的权重对之前所有的hyp进行融合获得一个base hyp，(x * w.reshape(n, 1)).sum(0) / w.sum()
evolve.txt会记录每次进化之后的results+hyp
每次进化时，hyp会根据之前的results进行从大到小的排序；
再根据fitness函数计算之前每次进化得到的hyp的权重
(其中fitness是我们寻求最大化的值。在YOLOv5中，fitness函数实现对 [P, R, [email protected], [email protected]] 指标进行加权。)
再确定哪一种进化方式，从而进行进化。
这部分代码其实不是很重要并且也比较难理解，大家如果没有特殊必要的话可以忽略，因为正常训练也不会用到超参数进化。
"""
for _ in range(opt.evolve): # 进化的迭代次数
if evolve_csv.exists(): # 如果存在evolve.csv文件：选择最佳的超参数并进行变异
# 选择父代
parent = 'single' # 父代选择方法：'single'（随机选择）或'weighted'（加权选择）
x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1)
n = min(5, len(x)) # 考虑的先前结果数量（最多为5个）
# fitness是我们寻求最大化的值。在YOLOv5中，fitness函数实现对 [P, R, [email protected], [email protected]] 指标进行加权
x = x[np.argsort(-fitness(x))][:n] # 选择前n个适应度最高的变异配置
w = fitness(x) - fitness(x).min() + 1E-6 # 权重（总和大于0）
if parent == 'single' or len(x) == 1:
# x = x[random.randint(0, n - 1)] # 随机选择
x = x[random.choices(range(n), weights=w)[0]] # 加权选择
elif parent == 'weighted':
x = (x * w.reshape(n, 1)).sum(0) / w.sum() # 加权组合
# 变异
mp, s = 0.8, 0.2 # 变异概率，标准差
npr = np.random
npr.seed(int(time.time()))
g = np.array([meta[k][0] for k in hyp.keys()]) # 增益 0-1
ng = len(meta)
v = np.ones(ng)
while all(v == 1): # 变异直到发生变化（防止重复）
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()):
hyp[k] = float(x[i + 7] * v[i]) # 变异
# 限制在范围内
for k, v in meta.items():
hyp[k] = max(hyp[k], v[1]) # 下限
hyp[k] = min(hyp[k], v[2]) # 上限
hyp[k] = round(hyp[k], 5) # 有效数字
# 训练变异
results = train(hyp.copy(), opt, device, callbacks)
callbacks = Callbacks()
# 写入变异结果
keys = ('metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', 'val/box_loss',
'val/obj_loss', 'val/cls_loss')
print_mutation(keys, results, hyp.copy(), save_dir, opt.bucket)
# 绘制结果
plot_evolve(evolve_csv)
LOGGER.info(f'Hyperparameter evolution finished {opt.evolve} generations\\n'
f"Results saved to {colorstr('bold', save_dir)}\\n"
f'Usage example: $ python train.py --hyp {evolve_yaml}')
def run(**kwargs):
# 封装train接口 支持函数调用执行这个train.py脚本
# Usage: import train; train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt')
opt = parse_opt(True)
for k, v in kwargs.items():
setattr(opt, k, v)
main(opt)
return opt
if __name__ == '__main__':
opt = parse_opt()
main(opt)