二、train.py
参数作用总结与调整建议
1. 数据和模型相关参数
--weights ：初始预训练权重文件的路径。 调整建议：如果在目标任务上数据量较小，可以选用预训练权重；如果数据量充足，甚至可以考虑从头训练。
--cfg ：模型配置文件（例如 YOLOv5m.yaml）的路径，定义模型结构。 调整建议：对于不同模型规模或自定义网络，可修改配置文件；也可根据硬件条件选择轻量级或大模型。
--data ：数据集配置文件的路径，包含训练和验证数据路径、类别数和类别名称等信息。 调整建议：确保数据格式与配置一致；如需要训练自定义数据集，修改此文件中的类别和数据路径。
--hyp ：超参数配置文件的路径，其中包含学习率、动量、权重衰减、数据增强参数等。 调整建议：根据任务特点调整关键超参（例如：lr0、momentum、weight_decay）；如遇训练不稳定可适当调低学习率或调整 warmup 参数。
2. 训练过程相关参数
--epochs ：训练的总轮次。 调整建议：数据量较大或模型较复杂时适当增加；若发现验证指标早已收敛，也可减少轮次以节省时间。
--batch-size ：总的批次大小，若设置为 -1，则会自动调整。 调整建议：根据 GPU 显存大小和模型参数量进行调整；注意分布式训练时需保证 batch size 能被 GPU 数量整除。
--imgsz ：训练和验证图像的输入尺寸。 调整建议：较高的分辨率有助于小目标检测，但会降低训练速度；可根据任务需求与硬件条件做平衡。
--rect ：是否采用矩形训练（保持图像原始长宽比）。 调整建议：对于形状差异较大的数据集，可启用此参数以提高检测精度。
3. 数据增强和采样相关参数
--multi-scale 含义：是否启用多尺度训练，使得输入尺寸在一定范围内随机变化。 调整建议：有助于提高模型鲁棒性，建议在数据量较多时启用；若追求速度可关闭。
--image-weights 含义：是否根据类别权重动态采样训练图像。 调整建议：当类别分布不均时启用，能使模型更关注稀有类别；否则可关闭。
--single-cls 含义：是否将多类别数据当作单类别处理，通常用于单目标检测任务。 调整建议：如果数据集只有一种目标，建议启用该选项。
4. 优化器与学习率调度器相关参数
--optimizer 含义：选择优化器类型，支持 SGD、Adam 和 AdamW。 调整建议：SGD 较为常用，Adam/AdamW 在某些任务中能更快收敛；可根据任务特点选取。
--cos-lr 含义：是否使用余弦学习率调度器。 调整建议：余弦调度器能在训练后期减缓学习率下降，有助于收敛；否则采用线性调度器。
--label-smoothing 含义：标签平滑系数，防止过拟合。 调整建议：当训练数据有噪声或过拟合现象明显时，可适当设置一个非零值（如 0.1）。
--patience 含义：EarlyStopping 的耐心值，即连续多少个 epoch 无改进后停止训练。 调整建议：根据验证集指标波动情况设置，通常 10-20 之间较为合理。
5. 分布式训练及其他高级参数
--device 含义：指定使用的设备（GPU 编号或 cpu）。 调整建议：根据实际硬件资源指定；多 GPU 训练时需要注意 batch size 分布。
--sync-bn 含义：是否在分布式训练下启用同步 BatchNorm。 调整建议：在多 GPU 环境下启用能提高训练稳定性，但可能会稍微降低速度。
--workers 含义：数据加载器的工作线程数。 调整建议：根据 CPU 核数和硬盘速度设置，过高可能导致内存占用过多，过低会影响数据读取速度。
--freeze 含义：指定冻结模型中部分层（例如前几层）的参数，使其不参与训练。 调整建议：如果使用预训练模型并希望保持低层特征不变，可冻结部分 backbone 层；若数据差异较大，则建议全部参与训练。
--save-period 含义：每隔多少个 epoch 保存一次 checkpoint。 调整建议：频繁保存会增加 IO 压力，通常可以设为 10 或 20；如训练较长可适当缩短间隔以防断点丢失。
6. 超参数进化相关参数
--evolve 含义：指定超参数进化的代数（例如 300 代）。 调整建议：仅在希望自动搜索最佳超参数时使用；训练初期建议关闭，待确定基本训练流程后再启用。
超参数元数据（meta 字典） 含义：定义各个超参数的变异尺度、上下限。例如 lr0、lrf、momentum、weight_decay、数据增强参数等。 调整建议：根据任务和数据集的实际表现，可以手动调整某些参数的上下限范围，以便进化过程中搜索到更适合的超参组合。
总结
整个 train.py 文件主要实现了 YOLOv5 的训练流程，涵盖以下关键步骤：
-
环境设置与日志记录：通过读取环境变量、Git 信息以及初始化日志和回调函数确保训练过程有良好的记录与版本控制。
-
超参数加载与保存：支持从 yaml 文件加载超参数，并在训练开始前保存当前配置，便于实验复现和对比。
-
模型加载与参数冻结：加载预训练模型或新建模型，并可根据需求冻结部分层，常用于微调任务。
-
数据加载与增强：创建训练和验证数据加载器，支持矩形训练、多尺度训练、图像权重采样等多种数据增强技术。
-
优化器与学习率调度：根据批次大小自动调整梯度累积步数，并采用预热、调度器（余弦或线性）来控制学习率的变化。
-
训练循环：包括前向传播、损失计算、反向传播、梯度裁剪、梯度累积以及 EMA 更新，同时支持多 GPU 和分布式训练。
-
验证与 EarlyStopping：定期在验证集上评估模型，记录指标（如 mAP），并根据 EarlyStopping 机制决定是否提前终止训练。
-
超参数进化：提供自动超参数进化模块，通过多代变异搜索找到更优的超参数组合，并记录每次变异结果。
如何合理调整参数
-
根据任务数据调整图像尺寸与批次大小：如果目标较小或图像细节要求较高，可提高输入图像分辨率；同时根据显存容量适当调整批次大小。
-
优化器及学习率调度策略：对于大多数目标检测任务，SGD 配合余弦调度器通常效果较好，但如数据噪声较大可尝试 Adam；同时可通过预热和梯度累积使训练更加稳定。
-
超参数细调：初始可采用默认超参数进行训练，观察验证指标变化；如发现过拟合或欠拟合，再调整 lr0、momentum、weight_decay 以及数据增强参数（如 mosaic、mixup、HSV 调整等）。
-
分布式训练设置：在多 GPU 环境下，确保 batch size 能够被 GPU 数量整除，适当启用 SyncBatchNorm；同时注意不要启用与 DDP 不兼容的参数（如 image_weights 和 evolve）。
-
冻结层策略：对于小数据集或目标域差异不大的场景，可以冻结 backbone 前几层，只训练后面的检测头；反之，则建议全部参与训练以充分适应新数据。
通过对这些参数的细致调整，可以使模型在不同场景下达到较好的精度和训练效率。建议在进行大规模实验前，先进行小规模试验，观察指标变化，再逐步调整各项参数以获得最佳效果。
# 从环境变量中获取 RANK（当前进程在分布式训练中的排名），默认为 -1 表示非分布式或单 GPU 训练
RANK = int(os.getenv('RANK', -1))
# WORLD_SIZE 表示全局使用的 GPU 数量，默认为 1
WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
# 获取 Git 仓库信息，用于记录训练时的代码版本（便于重现实验）
GIT_INFO = check_git_info()
def train(hyp, opt, device, callbacks): # hyp 可以是超参数配置文件路径或超参数字典
# 解包训练参数：
# save_dir：训练结果保存目录
# epochs：总训练轮次
# batch_size：批次大小
# weights：预训练权重文件路径
# single_cls：是否将多类别数据当作单类别处理
# evolve：是否进行超参数进化
# data：数据集配置文件路径
# cfg：模型配置文件路径
# resume：是否从之前的训练断点恢复
# noval：是否跳过验证，仅训练最后一次验证
# nosave：是否不保存训练过程中的中间模型
# workers：数据加载时的工作线程数
# freeze：冻结网络中部分层的列表（层号或范围）
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')
# ========= 目录管理 =========
# 定义保存权重文件的目录（如 save_dir/weights）
w = save_dir / 'weights'
# 如果不进行超参数进化，则创建保存权重的父目录，否则直接创建 weights 目录
(w.parent if evolve else w).mkdir(parents=True, exist_ok=True)
# 定义最后一次模型和最佳模型的保存路径
last, best = w / 'last.pt', w / 'best.pt'
# ========= 超参数加载 =========
# 如果 hyp 为字符串，则认为是路径，打开并用 yaml 加载成字典
if isinstance(hyp, str):
with open(hyp, errors='ignore') as f:
hyp = yaml.safe_load(f)
# 将超参数打印到日志中，便于记录本次实验使用的超参配置
LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
# 将超参数复制到 opt 对象中，用于后续保存至 checkpoint
opt.hyp = hyp.copy()
# ========= 保存运行设置 =========
# 如果不是超参数进化过程，则将当前超参数和训练选项保存到对应的 yaml 文件中
if not evolve:
yaml_save(save_dir / 'hyp.yaml', hyp)
yaml_save(save_dir / 'opt.yaml', vars(opt))
# ========= 日志记录器 =========
data_dict = None
# 分布式训练中只有主进程（RANK 为 -1 或 0）进行日志记录
if RANK in {-1, 0}:
loggers = Loggers(save_dir, weights, opt, hyp, LOGGER) # 初始化日志记录器
# 将日志记录器的方法注册到回调系统中（在训练过程中会调用）
for k in methods(loggers):
callbacks.register_action(k, callback=getattr(loggers, k))
# 处理自定义数据集（可能是远程数据集）的链接
data_dict = loggers.remote_dataset
# 如果为断点恢复训练，则重新设置 weights、epochs、hyp、batch_size（从上次记录恢复）
if resume:
weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size
# ========= 配置部分 =========
# 是否生成训练过程中各项指标的图表，默认只有在非进化且未禁用 noplots 时才生成
plots = not evolve and not opt.noplots
# 检查当前设备是否为 CUDA（即 GPU）
cuda = device.type != 'cpu'
# 初始化随机种子（保证每个训练进程有不同但可复现的随机性），并设为确定性模式
init_seeds(opt.seed + 1 + RANK, deterministic=True)
# 通过 torch_distributed_zero_first 保证在多进程中只有一个进程执行数据集检查
with torch_distributed_zero_first(LOCAL_RANK):
data_dict = data_dict or check_dataset(data) # 检查数据集配置（若 data_dict 为 None，则加载数据集配置）
# 从数据字典中获取训练和验证数据集的路径
train_path, val_path = data_dict['train'], data_dict['val']
# 根据 single_cls 参数判断类别数：单类别时设为 1，否则按数据集中的类别数
nc = 1 if single_cls else int(data_dict['nc'])
# 获取类别名称。如果 single_cls 且数据集名称数量不为 1，则统一命名为 'item'
names = {0: 'item'} if single_cls and len(data_dict['names']) != 1 else data_dict['names']
# 检查验证集是否为 COCO 数据集（通过文件名判断）
is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt')
# ========= 模型加载 =========
# 检查预训练权重文件后缀是否为 .pt（PyTorch 格式）
check_suffix(weights, '.pt')
# 判断是否为预训练模型（.pt 文件表示预训练权重）
pretrained = weights.endswith('.pt')
if pretrained:
# 若是预训练权重，首先确保在分布式训练中只由一个进程下载权重
with torch_distributed_zero_first(LOCAL_RANK):
weights = attempt_download(weights)
# 将权重加载到 CPU 中，防止 CUDA 内存泄漏
ckpt = torch.load(weights, map_location='cpu')
# 创建模型实例，使用 cfg 或从 checkpoint 中加载模型配置，同时指定输入通道数（ch=3）、类别数 nc 和 anchors
model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)
# 如果指定了 cfg 或超参数中的 anchors 且不是断点恢复，则需要排除 anchor 参数（防止重复加载）
exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else []
# 将 checkpoint 中的模型参数转换为 FP32
csd = ckpt['model'].float().state_dict()
# 取交集，确保加载到当前模型中时只加载匹配的参数
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)
# 检查是否支持自动混合精度（AMP），返回 amp 标志
amp = check_amp(model)
# ========= 冻结部分层 =========
# freeze 参数用于冻结指定层的参数。若 freeze 参数长度为 1，则认为冻结前 freeze[0] 个层；否则逐个指定
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 # 默认所有层均参与训练
# 如果参数名中包含需要冻结的层，则将其 requires_grad 设为 False（冻结该层参数）
if any(x in k for x in freeze):
LOGGER.info(f'freezing {k}')
v.requires_grad = False
# ========= 图像尺寸设置 =========
# 计算网格尺寸（模型最大 stride 与 32 中的较大值），保证输入尺寸为 stride 的倍数
gs = max(int(model.stride.max()), 32)
# 检查并调整图像尺寸，确保其为 gs 的整数倍，若设置了 floor 限制则下限为 gs * 2
imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2)
# ========= 批次大小调整 =========
# 对于单 GPU 训练（RANK == -1）且 batch_size 为 -1 时，自动估计最优批次大小
if RANK == -1 and batch_size == -1:
batch_size = check_train_batch_size(model, imgsz, amp)
loggers.on_params_update({"batch_size": batch_size})
# ========= 优化器及梯度累积 =========
nbs = 64 # 设定一个名义上的批次大小，用于归一化学习率等计算
# 根据当前批次大小计算需要梯度累积的步数，确保总的样本数达到 nbs
accumulate = max(round(nbs / batch_size), 1)
# 按批次大小和累积步数对 weight_decay 进行缩放
hyp['weight_decay'] *= batch_size * accumulate / nbs
# 根据选定的优化器类型、初始学习率、动量和 weight_decay 初始化优化器
optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['weight_decay'])
# ========= 学习率调度器 =========
if opt.cos_lr:
# 如果采用余弦退火调度，则计算 OneCycle 学习率变化曲线
lf = one_cycle(1, hyp['lrf'], epochs)
else:
# 否则采用线性下降策略：从 1 到 hyp['lrf'] 之间线性变化
lf = lambda x: (1 - x / epochs) * (1.0 - hyp['lrf']) + hyp['lrf']
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# ========= EMA（指数移动平均） =========
# EMA 用于平滑模型权重，通常只有在主进程中（RANK in {-1, 0}）使用
ema = ModelEMA(model) if RANK in {-1, 0} else None
# ========= 断点恢复 =========
best_fitness, start_epoch = 0.0, 0
if pretrained:
if resume:
# 使用 smart_resume 函数从 checkpoint 中恢复训练状态，包括优化器、EMA、当前 epoch 等
best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume)
# 释放 checkpoint 内存
del ckpt, csd
# ========= 多 GPU 并行 =========
# 如果使用 CUDA 且为单 GPU（RANK == -1）但实际可用 CUDA 数量大于1，警告不推荐 DP 模式，提示使用 DDP
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.')
# 使用 DataParallel 封装模型（虽然不推荐，但仍支持）
model = torch.nn.DataParallel(model)
# ========= 同步 BatchNorm =========
if opt.sync_bn and cuda and RANK != -1:
# 如果开启同步 BatchNorm，则将模型中所有 BatchNorm 层转换为同步版，并移动到设备上
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
LOGGER.info('Using SyncBatchNorm()')
# ========= 数据加载 =========
# 创建训练数据加载器，传入训练数据路径、图像尺寸、调整后的批次大小（考虑 WORLD_SIZE 分布式）
# 此外传入一些数据增强参数（如 augment=True）、缓存方式、是否使用矩形训练等
train_loader, dataset = create_dataloader(train_path,
imgsz,
batch_size // WORLD_SIZE,
gs,
single_cls,
hyp=hyp,
augment=True,
cache=None if opt.cache == 'val' else opt.cache,
rect=opt.rect,
rank=LOCAL_RANK,
workers=workers,
image_weights=opt.image_weights,
quad=opt.quad,
prefix=colorstr('train: '),
shuffle=True)
# 将所有样本的标签拼接在一起，用于后续计算类别权重等
labels = np.concatenate(dataset.labels, 0)
# 计算数据集中最大的标签类别索引，确保其小于总类别数 nc
mlc = int(labels[:, 0].max())
assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'
# ========= 验证数据加载 =========
if RANK in {-1, 0}:
# 创建验证数据加载器，注意批次大小翻倍、保持图像比例（rect=True）、多线程加载等参数
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:
# 如果不是断点恢复且没有禁用自动 anchor，则运行 AutoAnchor 检查
if not opt.noautoanchor:
check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
# 将模型先转换为半精度再转换回 float，用于预先降低 anchor 精度
model.half().float()
# 通知回调函数训练前准备工作已结束，同时传入标签和类别名称
callbacks.run('on_pretrain_routine_end', labels, names)
# ========= 分布式数据并行 =========
if cuda and RANK != -1:
model = smart_DDP(model)
# ========= 模型属性与超参数调节 =========
# 获取检测层的数量（最后一层 nl），用于对不同损失权重进行缩放
nl = de_parallel(model).model[-1].nl
# 按检测层数量调整 box、cls、obj 损失的超参数
hyp['box'] *= 3 / nl
hyp['cls'] *= nc / 80 * 3 / nl
hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl
# 添加 label smoothing 参数到超参数中
hyp['label_smoothing'] = opt.label_smoothing
# 将类别数、超参数、类别权重和类别名称附加到模型对象中，方便后续使用
model.nc = nc
model.hyp = hyp
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc
model.names = names
# ========= 开始训练 =========
t0 = time.time()
nb = len(train_loader) # 总共的批次数
# 根据 warmup_epochs 参数和批次数计算需要进行预热的总迭代次数，至少为 100 次
nw = max(round(hyp['warmup_epochs'] * nb), 100)
# last_opt_step 用于记录上次优化器更新的迭代步数
last_opt_step = -1
# 初始化 mAP 数组（每个类别的 mAP），以及训练过程中各项指标的初始值（精度、召回、mAP 等）
maps = np.zeros(nc)
results = (0, 0, 0, 0, 0, 0, 0)
# 设置学习率调度器的初始 epoch（以便断点恢复时保持一致）
scheduler.last_epoch = start_epoch - 1
# 初始化自动混合精度梯度缩放器
scaler = torch.cuda.amp.GradScaler(enabled=amp)
# 初始化 EarlyStopping，用于提前终止训练，patience 表示允许的无改善轮数
stopper, stop = EarlyStopping(patience=opt.patience), False
# 初始化损失计算类
compute_loss = ComputeLoss(model)
callbacks.run('on_train_start')
LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'
f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n'
f"Logging results to {colorstr('bold', save_dir)}\n"
f'Starting training for {epochs} epochs...')
# ===== 开始每个 epoch 的训练循环 =====
for epoch in range(start_epoch, epochs):
callbacks.run('on_train_epoch_start')
model.train() # 设定模型为训练模式
# ===== 可选：更新图像权重 =====
# 若启用 image_weights，则根据类别权重和 mAP 更新每张图像的采样权重
if opt.image_weights:
cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc
iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw)
dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n)
# ===== 可选：更新 Mosaic 边界 =====
# 此处代码被注释，原理为随机调整 Mosaic 拼接图像的边界
mloss = torch.zeros(3, device=device) # 初始化平均损失（box, obj, cls）
if RANK != -1:
train_loader.sampler.set_epoch(epoch)
pbar = enumerate(train_loader)
LOGGER.info(('\n' + '%11s' * 7) % ('Epoch', 'GPU_mem', 'box_loss', 'obj_loss', 'cls_loss', 'Instances', 'Size'))
if RANK in {-1, 0}:
pbar = tqdm(pbar, total=nb, bar_format=TQDM_BAR_FORMAT) # 设置进度条显示
optimizer.zero_grad()
# ===== 开始每个 batch 的训练 =====
for i, (imgs, targets, paths, _) in pbar:
callbacks.run('on_train_batch_start')
# 当前累计迭代步数
ni = i + nb * epoch
# 将图像数据转为 float 并归一化到 [0, 1]
imgs = imgs.to(device, non_blocking=True).float() / 255
# ===== 学习率与梯度累积预热 =====
if ni <= nw:
xi = [0, nw] # 预热区间
# 根据当前迭代步数更新累积步数（保证梯度累积的比例平滑过渡）
accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# 对不同参数组（如 bias 和其他参数）设置不同的预热学习率和动量
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']])
# ===== 多尺度训练 =====
if opt.multi_scale:
# 随机选择一个新的尺寸（在 imgsz 的 50% 到 150% 之间，并保证为 gs 的倍数）
sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs
sf = sz / max(imgs.shape[2:]) # 计算缩放因子
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]
imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
# ===== 前向传播 =====
with torch.cuda.amp.autocast(amp):
pred = model(imgs) # 模型预测输出
loss, loss_items = compute_loss(pred, targets.to(device))
if RANK != -1:
loss *= WORLD_SIZE # 在分布式训练中放大损失以平均梯度
if opt.quad:
loss *= 4.
# ===== 反向传播 =====
scaler.scale(loss).backward()
# ===== 梯度累积与优化器更新 =====
if ni - last_opt_step >= accumulate:
scaler.unscale_(optimizer) # 反缩放梯度
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0) # 梯度裁剪，防止梯度爆炸
scaler.step(optimizer) # 更新参数
scaler.update() # 更新缩放器
optimizer.zero_grad() # 清零梯度
if ema:
ema.update(model) # 更新 EMA 模型
last_opt_step = ni
# ===== 日志记录 =====
if RANK in {-1, 0}:
mloss = (mloss * i + loss_items) / (i + 1) # 更新平均损失
mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G'
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 loop
# ===== 学习率调度器步进 =====
lr = [x['lr'] for x in optimizer.param_groups] # 获取当前各参数组的学习率（用于日志记录）
scheduler.step() # 更新学习率
if RANK in {-1, 0}:
# 触发 epoch 结束回调，可能用于保存中间结果、记录统计数据等
callbacks.run('on_train_epoch_end', epoch=epoch)
# 更新 EMA 模型的其他属性（如 yaml、nc、hyp、names 等）
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: # 如果需要验证
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)
# 根据验证结果计算综合适应度
fi = fitness(np.array(results).reshape(1, -1))
stop = stopper(epoch=epoch, fitness=fi) # 检查是否满足提前停止条件
if fi > best_fitness:
best_fitness = fi
log_vals = list(mloss) + list(results) + lr
callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)
# ===== 保存模型 =====
if (not nosave) or (final_epoch and not evolve): # 如果需要保存模型
ckpt = {
'epoch': epoch,
'best_fitness': best_fitness,
'model': deepcopy(de_parallel(model)).half(), # 保存半精度模型
'ema': deepcopy(ema.ema).half(), # 保存 EMA 模型
'updates': ema.updates,
'optimizer': optimizer.state_dict(),
'opt': vars(opt),
'git': GIT_INFO, # Git 信息
'date': datetime.now().isoformat()}
# 保存 last.pt 和（如果适应度更高）best.pt
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
# 根据 save_period 周期性保存检查点
if opt.save_period > 0 and epoch % opt.save_period == 0:
torch.save(ckpt, w / f'epoch{epoch}.pt')
del ckpt
callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi)
# ===== 分布式训练中广播提前停止信号 =====
if RANK != -1:
broadcast_list = [stop if RANK == 0 else None]
dist.broadcast_object_list(broadcast_list, 0)
if RANK != 0:
stop = broadcast_list[0]
if stop:
break # 满足提前停止条件，则结束训练
# end epoch loop
# ===== 训练结束 =====
if RANK in {-1, 0}:
LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.')
for f in last, best:
if f.exists():
strip_optimizer(f) # 保存时剥离优化器信息以减小模型体积
if f is best:
LOGGER.info(f'\nValidating {f}...')
results, _, _ = 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,
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)
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
# ==================== parse_opt 函数（命令行参数解析） ====================
def parse_opt(known=False):
parser = argparse.ArgumentParser()
# 预训练权重的初始路径
parser.add_argument('--weights', type=str, default=ROOT / 'yolov5m.pt', help='initial weights path')
# 模型配置文件（yaml 格式）的路径
parser.add_argument('--cfg', type=str, default='models/yolov5m.yaml', help='models.yaml path')
# 数据集配置文件路径（yaml 格式）
parser.add_argument('--data', type=str, default=ROOT / 'data/vein.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')
# 批次大小（所有 GPU 上总的 batch size，设为 -1 表示自动调整）
parser.add_argument('--batch-size', type=int, default=5, help='total batch size for all GPUs, -1 for autobatch')
# 训练与验证的图像尺寸（单位：像素）
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
# 是否采用矩形训练（保持图像长宽比例）
parser.add_argument('--rect', action='store_true', help='rectangular training')
# 断点恢复参数，如果给定则恢复最近一次训练
parser.add_argument('--resume', nargs='?', const=False, default=False, help='resume most recent training')
# 如果设置，不保存中间 checkpoint，仅保存最终模型
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
# 如果设置，仅在最后一轮验证
parser.add_argument('--noval', action='store_true', help='only validate final epoch')
# 禁用自动 anchor 检查
parser.add_argument('--noautoanchor', action='store_true', help='disable AutoAnchor')
# 不生成图表文件
parser.add_argument('--noplots', action='store_true', help='save no plot files')
# 超参数进化的代数（默认 300 代）
parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
# 用于 Google Storage 的 bucket 名称
parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
# 图像缓存方式（如：ram/disk）
parser.add_argument('--cache', type=str, nargs='?', const='disk', help='image --cache ram/disk')
# 是否启用基于图像权重的采样
parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
# 训练设备，支持指定 GPU 编号或 cpu
parser.add_argument('--device', default='0', 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')
# 优化器类型，支持 SGD、Adam 和 AdamW
parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer')
# 是否启用分布式训练下的同步 BatchNorm
parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
# 数据加载时使用的工作线程数（每个 GPU 的线程数）
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')
# quad dataloader 开关（对数据加载方式的改进）
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')
# EarlyStopping 容忍的 epoch 数（连续多少轮无改善后提前停止）
parser.add_argument('--patience', type=int, default=20, help='EarlyStopping patience (epochs without improvement)')
# 冻结层参数：可以指定冻结 backbone 的前几层（例如 backbone=10，或具体层号如 0 1 2）
parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone=10, first3=0 1 2')
# 每隔多少个 epoch 保存一次 checkpoint，若设为 <1 则不周期性保存
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 相关参数：例如 wandb entity
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')
# 数据集 artifact 版本标识
parser.add_argument('--artifact_alias', type=str, default='latest', help='Version of dataset artifact to use')
# 根据 known 参数返回已知参数或所有参数
return parser.parse_known_args()[0] if known else parser.parse_args()
# ==================== main 函数 ====================
def main(opt, callbacks=Callbacks()):
# 检查（只有主进程）：
if RANK in {-1, 0}:
print_args(vars(opt))
check_git_status()
check_requirements()
# ===== 断点恢复逻辑 =====
if opt.resume and not check_comet_resume(opt) and not opt.evolve:
# 若 resume 参数有效，则尝试加载最近一次训练的 checkpoint
last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run())
opt_yaml = last.parent.parent / 'opt.yaml' # 获取训练选项文件路径
opt_data = opt.data # 原始数据集路径
if opt_yaml.is_file():
with open(opt_yaml, errors='ignore') as f:
d = yaml.safe_load(f)
else:
d = torch.load(last, map_location='cpu')['opt']
# 用加载的选项替换当前 opt
opt = argparse.Namespace(**d)
opt.cfg, opt.weights, opt.resume = '', str(last), True
if is_url(opt_data):
opt.data = check_file(opt_data)
else:
# 否则对各配置文件进行检查，确保文件存在且格式正确
opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = \
check_file(opt.data), check_yaml(opt.cfg), check_yaml(opt.hyp), str(opt.weights), str(opt.project)
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'):
opt.project = str(ROOT / 'runs/evolve')
opt.exist_ok, opt.resume = opt.resume, False
if opt.name == 'cfg':
opt.name = Path(opt.cfg).stem
# 根据 opt.project 和 opt.name 生成保存目录
opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))
# ===== 分布式训练设置 =====
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")
# ===== 开始训练或超参数进化 =====
if not opt.evolve:
train(opt.hyp, opt, device, callbacks)
else:
# ===== 超参数进化部分 =====
# 定义超参数进化的元数据，每个超参数包括（mutation scale, lower_limit, upper_limit）
meta = {
'lr0': (1, 1e-5, 1e-1),
'lrf': (1, 0.01, 1.0),
'momentum': (0.3, 0.6, 0.98),
'weight_decay': (1, 0.0, 0.001),
'warmup_epochs': (1, 0.0, 5.0),
'warmup_momentum': (1, 0.0, 0.95),
'warmup_bias_lr': (1, 0.0, 0.2),
'box': (1, 0.02, 0.2),
'cls': (1, 0.2, 4.0),
'cls_pw': (1, 0.5, 2.0),
'obj': (1, 0.2, 4.0),
'obj_pw': (1, 0.5, 2.0),
'iou_t': (0, 0.1, 0.7),
'anchor_t': (1, 2.0, 8.0),
'anchors': (2, 2.0, 10.0),
'fl_gamma': (0, 0.0, 2.0),
'hsv_h': (1, 0.0, 0.1),
'hsv_s': (1, 0.0, 0.9),
'hsv_v': (1, 0.0, 0.9),
'degrees': (1, 0.0, 45.0),
'translate': (1, 0.0, 0.9),
'scale': (1, 0.0, 0.9),
'shear': (1, 0.0, 10.0),
'perspective': (0, 0.0, 0.001),
'flipud': (1, 0.0, 1.0),
'fliplr': (0, 0.0, 1.0),
'mosaic': (1, 0.0, 1.0),
'mixup': (1, 0.0, 1.0),
'copy_paste': (1, 0.0, 1.0)}
with open(opt.hyp, errors='ignore') as f:
hyp = yaml.safe_load(f)
if 'anchors' not in hyp:
hyp['anchors'] = 3
if opt.noautoanchor:
del hyp['anchors'], meta['anchors']
opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir)
evolve_yaml, evolve_csv = save_dir / 'hyp_evolve.yaml', save_dir / 'evolve.csv'
if opt.bucket:
os.system(f'gsutil cp gs://{opt.bucket}/evolve.csv {evolve_csv}')
# 进行指定代数的超参数进化
for _ in range(opt.evolve):
if evolve_csv.exists():
parent = 'single'
x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1)
n = min(5, len(x))
x = x[np.argsort(-fitness(x))][:n]
w = fitness(x) - fitness(x).min() + 1E-6
if parent == 'single' or len(x) == 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()])
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):
# 用于外部调用时，可以通过传入关键字参数来覆盖默认参数
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)
扫一扫
834
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
