本文主要是介绍YOLOV5入门教程day3,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!
一. 导入包和基本配置
import argparse import math import os import random import subprocess import sys import time from copy import deepcopy from datetime import datetime, timedelta from pathlib import Pathtry:import comet_ml # must be imported before torch (if installed) except ImportError:comet_ml = Noneimport numpy as np import torch import torch.distributed as dist import torch.nn as nn import yaml from torch.optim import lr_scheduler from tqdm import tqdm
以下是每个导入包的详细介绍及总结:
1.
argparse
- 介绍:
argparse
是Python标准库中的一个模块,用于处理命令行参数,使得用户可以通过命令行传递参数给脚本。- 用途: 在YOLO模型中,它用于解析用户输入的配置选项,如权重路径、输入源、图像大小等。
2.
math
- 介绍: 这是Python的数学标准库,提供数学运算的基本功能。
- 用途: 用于执行数学计算,例如平方根、三角函数等,可能在模型计算中需要。
3.
os
- 介绍:
os
模块提供与操作系统交互的功能,包括文件和目录操作。- 用途: 在YOLO模型中用于管理文件路径、创建或删除目录等。
4.
random
- 介绍: 该模块用于生成随机数和随机选择。
- 用途: 可能用于数据增强、随机打乱数据集顺序等场景,以提高模型的泛化能力。
5.
subprocess
- 介绍: 用于创建子进程并与其交互,执行外部命令。
- 用途: 在YOLO模型中可能用于调用其他命令行工具,执行文件操作和管理外部进程。
6.
sys
- 介绍:
sys
模块提供与Python解释器相关的功能,如访问命令行参数。- 用途: 用于获取脚本的运行环境或参数,检查模块路径等。
7.
time
- 介绍: 提供时间相关的功能,如获取当前时间、计算时间间隔等。
- 用途: 在YOLO中用于性能监控,例如记录推理和训练的时间。
8.
deepcopy
- 介绍: 从
copy
模块导入,提供深拷贝的功能,允许完全复制对象及其引用的对象。- 用途: 在处理模型配置或多个模型状态时,避免原始对象被意外修改。
9.
datetime
和timedelta
- 介绍: 这两个模块用于处理日期和时间。
- 用途: 在YOLO中可能用于记录操作的时间戳,或计算时间差异,监控训练进度。
10.
pathlib
- 介绍: 提供面向对象的文件系统路径操作方法,取代传统的
os.path
方式。- 用途: 用于处理文件和目录路径,使代码更清晰易读。
11.
comet_ml
- 介绍: 用于实验跟踪和模型监控的第三方库。
- 用途: 如果安装了该库,可以用于收集训练过程中的指标和结果,以帮助分析模型性能和优化。
12.
numpy
- 介绍: 用于数值计算的强大库,提供支持多维数组和矩阵的操作。
- 用途: 在YOLO中广泛使用于数据处理,如数组运算、统计分析等。
13.
torch
- 介绍: PyTorch的核心库,支持张量计算,并提供构建和训练深度学习模型的功能。
- 用途: YOLO模型的核心,负责模型定义、训练和推理。
14.
torch.distributed
- 介绍: PyTorch提供的分布式训练模块。
- 用途: 在YOLO中实现多GPU训练,以加速模型训练过程。
15.
torch.nn
- 介绍: 该模块用于构建深度学习神经网络。
- 用途: 提供层、损失函数和其他构建块,以定义YOLO模型的架构。
16.
torch.optim
- 介绍: 包含多种优化算法,用于更新模型权重。
- 用途: 在YOLO训练过程中用于选择和应用不同的优化策略。
17.
tqdm
- 介绍: 提供快速、可扩展的进度条,用于监控循环或迭代的状态。
- 用途: 在训练和推理过程中显示进度信息,使用户能够直观地了解操作进度。
这段代码提供了YOLO模型在推理或训练中使用的一些必要的导入和设置。具体功能总结如下:
基础库导入:
- 导入标准库(如
argparse
,math
,os
,random
,subprocess
,sys
,time
,datetime
,pathlib
)以支持各种功能,包括参数解析、数学运算、文件和路径管理、随机数生成等。Comet ML导入:
- 尝试导入
comet_ml
以支持实验追踪。如果未安装该库,comet_ml
变量将设置为None
。该库通常用于模型训练过程的实验记录和可视化。科学计算和深度学习库:
- 导入 NumPy,用于高效的数组操作和计算。
- 导入 PyTorch(
torch
)及其相关模块,包括分布式计算 (torch.distributed
)、神经网络模块 (torch.nn
)、学习率调度器 (torch.optim.lr_scheduler
) 和其他深度学习功能。进度条显示:
- 从
tqdm
导入,用于在控制台中提供可视化的进度条,方便监控训练或推理的进度。
1.1获取文件绝对路径
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
代码总结
这段代码的主要功能是设置YOLOv5项目的根目录,并确保该目录被添加到Python的模块搜索路径中。具体功能步骤如下:
获取当前文件的绝对路径:
使用
FILE = Path(__file__).resolve()
Path
对象操作来获取当前执行的脚本的绝对路径,并将其赋值给FILE
变量。确定项目根目录:
通过
ROOT = FILE.parents[0] # YOLOv5 root directory
parents[0]
属性获取当前文件的父目录,通常这是 YOLOv5 项目的根目录,并将其赋值给ROOT
变量。添加根目录到模块搜索路径:
检查
if str(ROOT) not in sys.path:sys.path.append(str(ROOT)) # add ROOT to PATH
ROOT
是否已经在 Python 的模块搜索路径(sys.path
)中。如果不在,则将其添加。这一操作确保后续代码可以正确导入YOLOv5项目中的模块。生成相对路径:
使用
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
os.path.relpath
将ROOT
路径转换为相对于当前工作目录的路径。这为后续操作提供了一种更加简洁和灵活的路径表达方式。
1.2自定义模块导入
import val as validate # for end-of-epoch mAP from models.experimental import attempt_load from models.yolo import Model from utils.autoanchor import check_anchors from utils.autobatch import check_train_batch_size 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, ) from utils.loggers import LOGGERS, Loggers from utils.loggers.comet.comet_utils import check_comet_resume from utils.loss import ComputeLoss from utils.metrics import fitness 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, )
代码中导入包的介绍及总结
这段代码导入了YOLOv5模型所需的多个模块和功能,具体包的介绍如下:
1.
import val as validate
- 介绍: 导入
val
模块并命名为validate
,用于进行验证操作。- 用途: 通常用于在每个训练周期结束后计算平均精度(mAP),以评估模型性能。
2.
from models.experimental import attempt_load
- 介绍: 从
models.experimental
模块导入attempt_load
函数。- 用途: 尝试加载YOLO模型的权重,通常用于初始化模型。
3.
from models.yolo import Model
- 介绍: 从 YOLO 模型定义模块中导入
Model
类。- 用途: 用于创建和操作YOLO深度学习模型。
4.
from utils.autoanchor import check_anchors
- 介绍: 从
utils.autoanchor
模块导入check_anchors
函数。- 用途: 检查并自动设置锚框,以优化目标检测过程。
5.
from utils.autobatch import check_train_batch_size
- 介绍: 从
utils.autobatch
模块导入check_train_batch_size
函数。- 用途: 验证和确定训练时的批量大小,以提高训练效率。
6.
from utils.callbacks import Callbacks
- 介绍: 从
utils.callbacks
模块导入Callbacks
类。- 用途: 定义并管理在训练过程中触发的回调函数,用于监控和调整训练过程。
7.
from utils.dataloaders import create_dataloader
- 介绍: 从
utils.dataloaders
模块导入create_dataloader
函数。- 用途: 创建数据加载器,以便从数据集中有效加载训练和验证数据。
8.
from utils.downloads import attempt_download, is_url
- 介绍: 导入下载相关的工具函数。
- 用途:
attempt_download
用于下载文件,is_url
用于检查字符串是否为URL。9.
from utils.general import ...
- 这里导入了多个常用的工具函数:
LOGGER
: 日志记录器,用于输出训练过程中的信息。check_amp
: 检查是否启用混合精度训练。check_dataset
: 验证数据集的有效性。check_file
: 检查文件路径是否存在。check_img_size
: 校验输入图像大小的有效性。check_requirements
: 检查项目所需的依赖项。colorstr
: 用于格式化和着色输出的字符串。increment_path
: 自动生成唯一的路径名以防止覆盖。init_seeds
: 初始化随机种子,确保训练的可重复性。intersect_dicts
: 用于合并字典,常用于加载模型参数。labels_to_class_weights
: 确定类的权重以平衡数据。print_args
: 打印命令行参数,便于调试。10.
from utils.loggers import LOGGERS, Loggers
- 介绍: 导入日志记录器相关的类和工具。
- 用途: 用于管理记录训练过程中的信息。
11.
from utils.loggers.comet.comet_utils import check_comet_resume
- 介绍: 导入 Comet ML 特定的功能。
- 用途: 检查是否可以恢复之前的实验记录。
12.
from utils.loss import ComputeLoss
- 介绍: 从
utils.loss
模块导入ComputeLoss
类。- 用途: 计算损失函数,用于训练过程中的模型优化。
13.
from utils.metrics import fitness
- 介绍: 从
utils.metrics
模块导入fitness
函数。- 用途: 评估模型的性能指标,如精度和召回率。
14.
from utils.plots import plot_evolve
- 介绍: 从
utils.plots
模块导入plot_evolve
函数。- 用途: 绘制进化过程中的图表,以可视化指定的训练和验证指标。
15.
from utils.torch_utils import ...
- 介绍: 导入 PyTorch 相关工具和类:
EarlyStopping
: 提前停止训练的策略。ModelEMA
: 使用指数移动平均的模型。select_device
: 选择可用的计算设备。smart_DDP
: 智能分布式数据并行工具。smart_optimizer
: 用于设置智能优化器。smart_resume
: 支持从上次训练状态恢复。torch_distributed_zero_first
: 处理分布式训练的零阶段。这段代码通过导入多个模块和工具,构建了YOLOv5模型训练和推理所需的基础。导入的模块涵盖了从模型管理、数据处理、损失计算到日志记录和可视化等方方面面,为后续的训练和评估提供了功能支持。整体结构设计灵活,便于扩展和维护。
1.3分布式训练
LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1)) # https://pytorch.org/docs/stable/elastic/run.html RANK = int(os.getenv("RANK", -1)) WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1)) GIT_INFO = check_git_info()
这段代码的主要功能是处理与分布式训练相关的环境变量,并检查当前代码库的Git信息。具体功能步骤如下:
获取本地进程排名:使用
os.getenv
获取名为LOCAL_RANK
的环境变量,表示在当前节点中的进程编号。如果未设置该变量,则默认值为-1
。这个变量在分布式训练中非常重要,用于区分不同进程的身份。获取全局进程排名:获取名为
RANK
的环境变量,表示当前进程在整个分布式训练中的唯一标识。如果未设置,默认为-1
。获取世界规模:使用
os.getenv
获取名为WORLD_SIZE
的环境变量,表示在分布式训练中总共有多少个进程。如果未设置,默认为1
,表示没有使用分布式训练。检查Git信息:调用
check_git_info()
函数以获取当前代码库的Git信息。这个信息通常用于记录模型训练的版本,确保可复现性,以便于追踪代码的变更和调试。整体而言,这段代码为使用PyTorch进行分布式训练提供了一些基础设置,确保能够正确识别和管理不同进程的身份和规模。这些变量对于分布式模型的训练、资源管理和过程监控至关重要。同时,记录Git信息也有助于在训练中保持版本控制,以便后期的复现和调试。
二. main主函数
if RANK in {-1, 0}:print_args(vars(opt))check_git_status()check_requirements(ROOT / "requirements.txt")
这段代码主要用于在分布式训练环境中做一些初始化检查和设置:
条件判断:
if RANK in {-1, 0}:
- 该条件判断当前进程的排名(
RANK
)是否为 -1 或 0。这通常用于标识主进程或非分布式运行环境,主进程负责执行一些初始化任务。打印参数:
print_args(vars(opt))
- 调用
print_args
函数将所有命令行参数以字典形式打印出来,帮助用户了解当前的训练配置,便于调试。检查Git状态:
check_git_status()
- 调用
check_git_status
函数以检查代码库的Git状态。如果工作树中有未提交的更改,将会给出提示。这确保在训练开始之前,代码是稳定的,并且没有未保存的变更。检查依赖项:
check_requirements(ROOT / "requirements.txt")
- 调用
check_requirements
函数以验证项目所需的所有依赖项是否已经安装。检查的依据是项目根目录下的requirements.txt
文件,确保没有缺失的库,以防在运行时出现错误。这段代码的功能是为YOLOv5的训练过程中的分布式环境做准备。它通过打印参数、检查Git状态以及验证依赖项的完整性,确保在开始训练之前,所有关键条件都符合要求。这有助于提高训练的稳定性和可复现性,确保用户可以监控和控制训练过程。
2.1判断是否断点继续训练
# Resume (from specified or most recent last.pt) if opt.resume and not check_comet_resume(opt) and not opt.evolve:last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run())opt_yaml = last.parent.parent / "opt.yaml" # train options yamlopt_data = opt.data # original datasetif 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 = argparse.Namespace(**d) # replaceopt.cfg, opt.weights, opt.resume = "", str(last), True # reinstateif is_url(opt_data):opt.data = check_file(opt_data) # avoid HUB resume auth timeout 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
代码总结
这段代码在YOLO模型训练过程中,处理从之前训练状态恢复的逻辑,确保用户能够继续未完成的训练或基于之前的训练结果进行改进。具体步骤如下:
恢复训练参数:
if opt.resume and not check_comet_resume(opt) and not opt.evolve:
- 检查用户是否请求恢复训练(
opt.resume
),并确保不使用Comet ML的恢复功能,且不处于进化模式。若满足条件则继续执行恢复过程。选择恢复文件:
last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run())
- 检查
opt.resume
是否为字符串,若是则确保它是一个有效文件,并将其路径赋给last
。如果不是字符串,则获取最新的训练结果。这个路径指向上一次训练的权重文件。找到训练选项 YAML 文件:
opt_yaml = last.parent.parent / "opt.yaml" # train options yaml
- 根据
last
路径,构建对应的训练选项 YAML 文件的路径。加载训练选项:
if opt_yaml.is_file():with open(opt_yaml, errors="ignore") as f:d = yaml.safe_load(f) else:d = torch.load(last, map_location="cpu")["opt"]
- 检查是否存在 YAML 文件。如果存在,则读取并加载训练参数。若不存在,则从
last
加载PyTorch模型,提取训练选项。更新参数:
opt = argparse.Namespace(**d) # replace opt.cfg, opt.weights, opt.resume = "", str(last), True # reinstate
- 将加载的参数更新至
opt
对象。重置cfg
(配置文件)为空,重设weights
为恢复文件,设置resume
为True
。检查数据集路径:
if is_url(opt_data):opt.data = check_file(opt_data) # avoid HUB resume auth timeout
- 如果数据集路径是URL,则检查并确保其为有效路径,以避免在恢复时出现身份验证超时问题。
常规检查:
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
- 如果没有执行恢复,进行一般性检查,以确保所有必要的配置文件(数据集、模型配置、超参数和权重)都存在并有效。
整体而言,这段代码为YOLO模型的训练提供了恢复功能,确保用户能够从之前的模型状态继续训练。通过判断和读取特定的文件,加载和更新训练参数,检查文件有效性,从而保证训练的高效性和连贯性。这种设计增加了训练流程的灵活性,允许开发者在不同时间点恢复工作而不中断进度。
2.2判断是否是分布式训练
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", timeout=timedelta(seconds=10800))
代码总结
这段代码主要负责设置设备(CPU或GPU)和处理与分布式数据并行训练(DDP)相关的初始化。具体功能步骤如下:
选择设备:
device = select_device(opt.device, batch_size=opt.batch_size)
- 调用
select_device
函数,根据用户提供的参数(如要使用的设备和批量大小)选择GPU或CPU,并将选定的设备分配给变量device
。检查本地进程排名是否有效:
if LOCAL_RANK != -1:
- 判断当前进程的本地排名是否有效(即非 -1)。该条件用于确保只在多GPU环境下执行以下代码块,适用于使用分布式训练的场景。
错误检查与断言:
设备兼容性:
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}"
- 确保用户没有启用图像权重(
--image-weights
)和进化参数(--evolve
),因为这些在多GPU DDP训练中可能不兼容。批量大小有效性检查:
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"
- 确保批量大小不为 -1,并且是
WORLD_SIZE
的倍数,以便在多个进程之间分配合理的负载。CUDA设备检查:
assert torch.cuda.device_count() > LOCAL_RANK, "insufficient CUDA devices for DDP command"
- 确保CUDA设备数量充足,以支持当前的DDP命令。
设置CUDA设备:
torch.cuda.set_device(LOCAL_RANK) device = torch.device("cuda", LOCAL_RANK)
- 设置当前CUDA设备为
LOCAL_RANK
,并更新device
变量为所选的CUDA设备,以确保后续计算在正确的GPU上进行。初始化分布式进程组:
dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo", timeout=timedelta(seconds=10800) )
- 调用
torch.distributed.init_process_group
初始化分布式训练的进程组。根据是否支持NCCL后端选择适当的后端(如nccl
用于GPU),并设置超时时间为3小时。这段代码有效地处理了对计算设备的选择和分布式训练的初始化,包括对设备兼容性、批量大小和CUDA设备可用性等多个方面的检查。通过这些步骤,确保YOLO模型能够在多GPU环境下进行有效的训练和推理,从而提高训练效率和资源利用率。
2.3是否进行进化训练
if not opt.evolve:train(opt.hyp, opt, device, callbacks)# Evolve hyperparameters (optional)else:# Hyperparameter evolution metadata (including this hyperparameter True-False, lower_limit, upper_limit)meta = {"lr0": (False, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3)"lrf": (False, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf)"momentum": (False, 0.6, 0.98), # SGD momentum/Adam beta1"weight_decay": (False, 0.0, 0.001), # optimizer weight decay"warmup_epochs": (False, 0.0, 5.0), # warmup epochs (fractions ok)"warmup_momentum": (False, 0.0, 0.95), # warmup initial momentum"warmup_bias_lr": (False, 0.0, 0.2), # warmup initial bias lr"box": (False, 0.02, 0.2), # box loss gain"cls": (False, 0.2, 4.0), # cls loss gain"cls_pw": (False, 0.5, 2.0), # cls BCELoss positive_weight"obj": (False, 0.2, 4.0), # obj loss gain (scale with pixels)"obj_pw": (False, 0.5, 2.0), # obj BCELoss positive_weight"iou_t": (False, 0.1, 0.7), # IoU training threshold"anchor_t": (False, 2.0, 8.0), # anchor-multiple threshold"anchors": (False, 2.0, 10.0), # anchors per output grid (0 to ignore)"fl_gamma": (False, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5)"hsv_h": (True, 0.0, 0.1), # image HSV-Hue augmentation (fraction)"hsv_s": (True, 0.0, 0.9), # image HSV-Saturation augmentation (fraction)"hsv_v": (True, 0.0, 0.9), # image HSV-Value augmentation (fraction)"degrees": (True, 0.0, 45.0), # image rotation (+/- deg)"translate": (True, 0.0, 0.9), # image translation (+/- fraction)"scale": (True, 0.0, 0.9), # image scale (+/- gain)"shear": (True, 0.0, 10.0), # image shear (+/- deg)"perspective": (True, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001"flipud": (True, 0.0, 1.0), # image flip up-down (probability)"fliplr": (True, 0.0, 1.0), # image flip left-right (probability)"mosaic": (True, 0.0, 1.0), # image mixup (probability)"mixup": (True, 0.0, 1.0), # image mixup (probability)"copy_paste": (True, 0.0, 1.0),} # segment copy-paste (probability)# GA configspop_size = 50mutation_rate_min = 0.01mutation_rate_max = 0.5crossover_rate_min = 0.5crossover_rate_max = 1min_elite_size = 2max_elite_size = 5tournament_size_min = 2tournament_size_max = 10with open(opt.hyp, errors="ignore") as f:hyp = yaml.safe_load(f) # load hyps dictif "anchors" not in hyp: # anchors commented in hyp.yamlhyp["anchors"] = 3if opt.noautoanchor:del hyp["anchors"], meta["anchors"]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 indicesevolve_yaml, evolve_csv = save_dir / "hyp_evolve.yaml", save_dir / "evolve.csv"if opt.bucket:# download evolve.csv if existssubprocess.run(["gsutil","cp",f"gs://{opt.bucket}/evolve.csv",str(evolve_csv),])# Delete the items in meta dictionary whose first value is Falsedel_ = [item for item, value_ in meta.items() if value_[0] is False]hyp_GA = hyp.copy() # Make a copy of hyp dictionaryfor item in del_:del meta[item] # Remove the item from meta dictionarydel hyp_GA[item] # Remove the item from hyp_GA dictionary# Set lower_limit and upper_limit arrays to hold the search space boundarieslower_limit = np.array([meta[k][1] for k in hyp_GA.keys()])upper_limit = np.array([meta[k][2] for k in hyp_GA.keys()])# Create gene_ranges list to hold the range of values for each gene in the populationgene_ranges = [(lower_limit[i], upper_limit[i]) for i in range(len(upper_limit))]# Initialize the population with initial_values or random valuesinitial_values = []# If resuming evolution from a previous checkpointif opt.resume_evolve is not None:assert os.path.isfile(ROOT / opt.resume_evolve), "evolve population path is wrong!"with open(ROOT / opt.resume_evolve, errors="ignore") as f:evolve_population = yaml.safe_load(f)for value in evolve_population.values():value = np.array([value[k] for k in hyp_GA.keys()])initial_values.append(list(value))# If not resuming from a previous checkpoint, generate initial values from .yaml files in opt.evolve_populationelse:yaml_files = [f for f in os.listdir(opt.evolve_population) if f.endswith(".yaml")]for file_name in yaml_files:with open(os.path.join(opt.evolve_population, file_name)) as yaml_file:value = yaml.safe_load(yaml_file)value = np.array([value[k] for k in hyp_GA.keys()])initial_values.append(list(value))# Generate random values within the search space for the rest of the populationif initial_values is None:population = [generate_individual(gene_ranges, len(hyp_GA)) for _ in range(pop_size)]elif pop_size > 1:population = [generate_individual(gene_ranges, len(hyp_GA)) for _ in range(pop_size - len(initial_values))]for initial_value in initial_values:population = [initial_value] + population# Run the genetic algorithm for a fixed number of generationslist_keys = list(hyp_GA.keys())for generation in range(opt.evolve):if generation >= 1:save_dict = {}for i in range(len(population)):little_dict = {list_keys[j]: float(population[i][j]) for j in range(len(population[i]))}save_dict[f"gen{str(generation)}number{str(i)}"] = little_dictwith open(save_dir / "evolve_population.yaml", "w") as outfile:yaml.dump(save_dict, outfile, default_flow_style=False)# Adaptive elite sizeelite_size = min_elite_size + int((max_elite_size - min_elite_size) * (generation / opt.evolve))# Evaluate the fitness of each individual in the populationfitness_scores = []for individual in population:for key, value in zip(hyp_GA.keys(), individual):hyp_GA[key] = valuehyp.update(hyp_GA)results = train(hyp.copy(), opt, device, callbacks)callbacks = Callbacks()# Write mutation resultskeys = ("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)fitness_scores.append(results[2])# Select the fittest individuals for reproduction using adaptive tournament selectionselected_indices = []for _ in range(pop_size - elite_size):# Adaptive tournament sizetournament_size = max(max(2, tournament_size_min),int(min(tournament_size_max, pop_size) - (generation / (opt.evolve / 10))),)# Perform tournament selection to choose the best individualtournament_indices = random.sample(range(pop_size), tournament_size)tournament_fitness = [fitness_scores[j] for j in tournament_indices]winner_index = tournament_indices[tournament_fitness.index(max(tournament_fitness))]selected_indices.append(winner_index)# Add the elite individuals to the selected indiceselite_indices = [i for i in range(pop_size) if fitness_scores[i] in sorted(fitness_scores)[-elite_size:]]selected_indices.extend(elite_indices)# Create the next generation through crossover and mutationnext_generation = []for _ in range(pop_size):parent1_index = selected_indices[random.randint(0, pop_size - 1)]parent2_index = selected_indices[random.randint(0, pop_size - 1)]# Adaptive crossover ratecrossover_rate = max(crossover_rate_min, min(crossover_rate_max, crossover_rate_max - (generation / opt.evolve)))if random.uniform(0, 1) < crossover_rate:crossover_point = random.randint(1, len(hyp_GA) - 1)child = population[parent1_index][:crossover_point] + population[parent2_index][crossover_point:]else:child = population[parent1_index]# Adaptive mutation ratemutation_rate = max(mutation_rate_min, min(mutation_rate_max, mutation_rate_max - (generation / opt.evolve)))for j in range(len(hyp_GA)):if random.uniform(0, 1) < mutation_rate:child[j] += random.uniform(-0.1, 0.1)child[j] = min(max(child[j], gene_ranges[j][0]), gene_ranges[j][1])next_generation.append(child)# Replace the old population with the new generationpopulation = next_generation# Print the best solution foundbest_index = fitness_scores.index(max(fitness_scores))best_individual = population[best_index]print("Best solution found:", best_individual)# Plot resultsplot_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}')
代码总结
这段代码的功能主要是实现超参数的进化算法,允许在模型训练过程中自动调整和优化超参数,以提高模型的性能。具体步骤如下:
训练模型:
if not opt.evolve:train(opt.hyp, opt, device, callbacks)
- 如果不需要进行超参数进化,则调用
train
函数,以给定的超参数、选项和设备训练模型。超参数进化:
else:
- 执行超参数进化算法,该过程只在用户指定需要进化时进行。
定义超参数元数据:
meta = {"lr0": (False, 1e-5, 1e-1),"lrf": (False, 0.01, 1.0),... }
- 定义一个字典
meta
,包含需要优化的超参数及其初始状态和范围。这些超参数包括学习率、动量、权重衰减、图像增强参数等。超参数加载和处理:
with open(opt.hyp, errors="ignore") as f:hyp = yaml.safe_load(f)
- 加载用户指定的超参数文件,并确保其中包含锚框参数。如果使用了
noautoanchor
选项,则删除锚框参数。遗传算法配置:
- 设置遗传算法的相关参数,如种群大小、变异率、交叉率等。
pop_size = 50 mutation_rate_min = 0.01 ...
下载历史记录:
if opt.bucket:subprocess.run([...]) # download evolve.csv if exists
- 如果指定了存储桶,则尝试下载之前的进化结果。
准备初始种群:
- 初始化种群,有可能从文件中读取历史进化值,或者生成随机值。
if opt.resume_evolve is not None:...
遗传算法主循环:
for generation in range(opt.evolve):...
- 执行固定次数的世代循环,评估每个个体的适应度,选择适合的个体进行繁殖(交叉和变异),生成下一代的种群。
适应度评估:
results = train(hyp.copy(), opt, device, callbacks)
- 在每个世代中,评估每个超参数组合的适应度,通常通过训练模型并记录性能指标来实现。
选择和繁殖:
- 通过适应度成绩选择最优秀的个体进行交叉繁殖,同时应用随机变异,生成新的个体。
打印最佳解:
best_index = fitness_scores.index(max(fitness_scores))
- 找到最佳解并输出,表示超参数进化的结果。
结果记录和绘图:
plot_evolve(evolve_csv)
- 保存和展示模型训练中超参数的进化过程。
整体而言,这段代码实现了YOLOv5中的超参数进化,运用遗传算法来自动调整模型训练中的超参数,以便提高模型性能。结合训练过程中的适应度评估,代码确保每一代都选择出最优参数组合,从而逐步逼近最佳模型配置。通过这种方式,开发者可以有效地优化模型而不需要手动调节每个超参数,提高了训练的效率和结果的稳健性。
三 设置opt超参数
def parse_opt(known=False):parser = argparse.ArgumentParser()parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path")parser.add_argument("--cfg", type=str, default="", help="model.yaml path")parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path")parser.add_argument("--hyp", type=str, default=ROOT / "data/hyps/hyp.scratch-low.yaml", help="hyperparameters path")parser.add_argument("--epochs", type=int, default=10, 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=640, help="train, val image size (pixels)")parser.add_argument("--rect", action="store_true", help="rectangular training")parser.add_argument("--resume", nargs="?", const=True, default=False, help="resume most recent training")parser.add_argument("--nosave", action="store_true", help="only save final checkpoint")parser.add_argument("--noval", action="store_true", help="only validate final epoch")parser.add_argument("--noautoanchor", action="store_true", help="disable AutoAnchor")parser.add_argument("--noplots", action="store_true", help="save no plot files")parser.add_argument("--evolve", type=int, nargs="?", const=300, help="evolve hyperparameters for x generations")parser.add_argument("--evolve_population", type=str, default=ROOT / "data/hyps", help="location for loading population")parser.add_argument("--resume_evolve", type=str, default=None, help="resume evolve from last generation")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 argumentsparser.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")# NDJSON loggingparser.add_argument("--ndjson-console", action="store_true", help="Log ndjson to console")parser.add_argument("--ndjson-file", action="store_true", help="Log ndjson to file")return parser.parse_known_args()[0] if known else parser.parse_args()
1. 模型及路径相关
--weights
: 指定初始模型权重文件的路径(默认使用yolov5s.pt
)。--cfg
: 模型配置文件的路径,用于定义模型结构。--data
: 数据集配置文件的路径(默认使用coco128.yaml
),定义训练和验证数据集的设置。--hyp
: 超参数配置文件的路径(默认使用hyp.scratch-low.yaml
),用于定义训练时的超参数。2. 训练过程相关
--epochs
: 训练的总轮数(默认 10)。--batch-size
: 总批量大小,对于所有GPU而言(默认 1),特殊值 -1 表示自动计算。--imgsz
,--img
,--img-size
: 图像的训练和验证大小(默认 640 像素)。--rect
: 启用矩形训练,以优化输入图像的填充方式。--resume
: 从最近的训练状态恢复(可指定名或默认恢复最新)。--nosave
: 仅保存最终的检查点,而不是每个epoch的检查点。--noval
: 仅在最终epoch进行验证,避免在训练期间进行每个epoch的验证。3. 数据处理相关
--noautoanchor
: 禁用自动锚框优化。--noplots
: 不保存任何训练过程中的图表。--image-weights
: 使用加权图像选择,以优化训练过程。4. 超参数进化相关
--evolve
: 进行超参数进化的世代数(默认 300)。--evolve_population
: 超参数进化时加载种群的位置。--resume_evolve
: 从上一次进化的状态恢复。5. 硬件和其他设置
--device
: 指定使用的计算设备,如 GPU(0)或 CPU。--multi-scale
: 在训练过程中变更输入图像大小(±50%)。--single-cls
: 将多类别数据视作单类别进行训练。--optimizer
: 选择优化器(SGD、Adam 或 AdamW,默认 SGD)。--sync-bn
: 在分布式训练模式下使用同步批量规范化。6. 日志及保存设置
--project
: 结果保存的主目录(默认使用runs/train
)。--name
: 保存结果的子目录名称。--exist-ok
: 允许在目标目录存在的情况下覆盖,而不自动增加。--save-period
: 每多少个epoch保存一次检查点(当值小于1时不保存)。--seed
: 全局训练种子,用于控制随机性。--local_rank
: 自动分布式数据并行的多GPU参数,通常不应手动修改。7. 其他设置
--bucket
: 指定Google云存储桶,便于上传数据。--cache
: 指定图像缓存方式(内存或磁盘)。--entity
,--upload_dataset
,--bbox_interval
,--artifact_alias
: 用于图像日志、实体标识和数据集相关的上传选项。8. NDJSON 日志记录
--ndjson-console
: 将日志输出为NDJSON格式到控制台。--ndjson-file
: 将日志输出为NDJSON格式到文件。
parse_opt
函数包含了大量参数,覆盖了YOLOv5的训练、验证和测试中的所有核心设置和灵活配置。这些参数使得用户能够根据不同的需求定制训练过程,优化模型性能,提高训练效果。
四 train文件
def train(hyp, opt, device, callbacks):save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = (Path(opt.save_dir),opt.epochs,opt.batch_size,opt.weights,opt.single_cls,opt.evolve,opt.data,opt.cfg,opt.resume,opt.noval,opt.nosave,opt.workers,opt.freeze,)callbacks.run("on_pretrain_routine_start")# Directoriesw = save_dir / "weights" # weights dir(w.parent if evolve else w).mkdir(parents=True, exist_ok=True) # make dirlast, best = w / "last.pt", w / "best.pt"
这段代码的功能是为YOLOv5模型的训练过程进行初始化和目录配置。以下是每个部分的具体分析:
参数解包:
- 在这一行,使用Python的元组解包从
opt
对象中提取多个训练相关的参数,并使用适当的数据类型(如将opt.save_dir
转换为Path
对象)。这些参数包括:
save_dir
: 保存目录的路径。epochs
: 训练的总轮数。batch_size
: 当前的批量大小。weights
: 初始模型权重的路径。single_cls
: 是否将多类别数据当作单类数据来训练。evolve
: 是否进行超参数进化。data
: 数据集配置路径。cfg
: 模型配置路径。resume
: 恢复训练的标志。noval
: 仅在最后一轮进行验证的标志。nosave
: 仅保存最终的检查点,而不是每轮的检查点。workers
: 数据加载时使用的工作线程数量。freeze
: 指定需要冻结的层。调用回调函数:
- 运行回调函数的特定事件,表示训练即将开始。这通常用于记录或监控训练过程。
设置目录:
- 定义权重保存目录
w
,确保在训练过程中能够保存模型的权重。mkdir(parents=True, exist_ok=True)
创建目录,如果父目录不存在也会同时创建。根据是否进行超参数进化选择权重保存路径的父目录。定义权重文件路径:
- 定义两个文件路径:
last.pt
用于保存最新的权重,best.pt
用于保存最佳的权重。这样可以在训练中自动保存模型状态,以便后续检索和使用。这段代码在YOLOv5的训练过程中负责初始化各种参数,并设置必要的目录以存储训练过程中生成的权重文件。通过调用回调函数,代码能够监控训练状态。整体而言,它为后续的训练逻辑提供了一个组织良好的基础,确保模型的权重在训练期间得以妥善管理。
# Hyperparametersif isinstance(hyp, str):with open(hyp, errors="ignore") as f:hyp = yaml.safe_load(f) # load hyps dictLOGGER.info(colorstr("hyperparameters: ") + ", ".join(f"{k}={v}" for k, v in hyp.items()))opt.hyp = hyp.copy() # for saving hyps to checkpoints# Save run settingsif not evolve:yaml_save(save_dir / "hyp.yaml", hyp)yaml_save(save_dir / "opt.yaml", vars(opt))
代码总结
这段代码主要用于处理和管理YOLOv5模型训练中的超参数设置。具体步骤如下:
加载超参数:
- 检查
hyp
是否为字符串类型,如果是,则假设它是一个指向超参数配置文件的路径。打开该文件并使用yaml.safe_load()
函数加载超参数字典,以确保正确读取配置。记录超参数信息:
- 使用日志记录器
LOGGER
输出当前超参数的详细信息。colorstr
函数将文本着色,以提高可读性。所有超参数以key=value
格式串联,并用逗号分隔,便于用户快速了解超参数设置。保存超参数副本:
- 将当前超参数字典的副本保存到
opt.hyp
中,以便在保存模型检查点时能够保存超参数信息。这有助于确保后续的训练或推理可以复现原有的设置。保存运行设置:
- 如果当前不进行超参数进化(
not evolve
),则将超参数和运行配置选项分别保存到两个 YAML 文件中:
hyp.yaml
: 保存超参数设置。opt.yaml
: 保存运行选项设置(如数据集、设备等)。这些文件便于后续的复现与调试。整体而言,这段代码在YOLOv5的训练过程中负责加载、记录和保存超参数设置,确保用户定义的超参数能够被正确管理。通过日志记录超参数信息,用户可以查阅当前训练的配置,保存这些设置为 YAML 文件则为训练的可复现性提供了支持。
# Loggersdata_dict = Noneif RANK in {-1, 0}:include_loggers = list(LOGGERS)if getattr(opt, "ndjson_console", False):include_loggers.append("ndjson_console")if getattr(opt, "ndjson_file", False):include_loggers.append("ndjson_file")loggers = Loggers(save_dir=save_dir,weights=weights,opt=opt,hyp=hyp,logger=LOGGER,include=tuple(include_loggers),)# Register actionsfor k in methods(loggers):callbacks.register_action(k, callback=getattr(loggers, k))# Process custom dataset artifact linkdata_dict = loggers.remote_datasetif resume: # If resuming runs from remote artifactweights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size
代码总结
这段代码主要用于设置和管理训练过程中的日志记录器。具体功能分解如下:
初始化数据字典:
- 初始化一个变量
data_dict
,用于存储将来可能记录的数据。判断是否为主进程:
- 检查当前进程的排名。如果
RANK
是 -1 或 0,表示当前进程是主进程,通常执行日志记录相关的操作。准备日志记录器:
- 创建一个包含所有日志记录器的列表
include_loggers
。- 根据用户选项
ndjson_console
和ndjson_file
,决定是否将控制台和文件日志记录器添加到日志记录器列表中。初始化日志记录器实例:
- 创建
Loggers
类的实例,传入必要的参数,如保存目录、权重、选项、超参数、主日志记录器和包含的日志记录器列表。这些参数用于配置日志记录器的行为。注册回调动作:
- 遍历日志记录器的方法(如记录、输出等),并将这些方法注册为回调动作,使得在训练过程中可以通过回调触发相应的日志记录操作。
处理自定义数据集的链接:
- 从日志记录器提取远程数据集信息并赋值给
data_dict
。如果在恢复之前的训练环境(resume
)时,重新加载权重、轮数、超参数和批量大小,以确保继续训练的配置与之前一致。整体来看,这段代码为YOLOv5训练过程中的日志记录提供了配置和管理机制,确保在主进程中正确初始化各种日志记录器,并通过回调功能将其整合进训练过程。这样做的好处是提高了训练的可观测性,使得用户能够获取和分析训练期间的详细信息,同时支持自定义数据集的追踪与恢复。
# Configplots = not evolve and not opt.noplots # create plotscuda = 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 Nonetrain_path, val_path = data_dict["train"], data_dict["val"]nc = 1 if single_cls else int(data_dict["nc"]) # number of classesnames = {0: "item"} if single_cls and len(data_dict["names"]) != 1 else data_dict["names"] # class namesis_coco = isinstance(val_path, str) and val_path.endswith("coco/val2017.txt") # COCO dataset
代码总结
这段代码的功能是进行训练配置的初始化,并对数据集进行检查和准备。以下是具体步骤的分析:
设置绘图标志:
plots = not evolve and not opt.noplots # create plots
- 检查是否启用绘图。如果当前不是进行超参数进化(
not evolve
)且未设置禁用绘图(not opt.noplots
),则设置plots
为True
,表示将创建训练过程中的可视化图表。CUDA设备检查:
cuda = device.type != "cpu"
- 根据
device
类型判断是否使用CUDA设备(GPU)。如果设备类型不是CPU,则cuda
为True
。初始化随机种子:
init_seeds(opt.seed + 1 + RANK, deterministic=True)
- 调用
init_seeds
函数设置随机种子,以保证训练的可重复性。种子的计算方式为用户指定的种子加上进程排名(RANK
)。deterministic=True
确保生成的随机数在不同运行中一致。nc = 1 if single_cls else int(data_dict["nc"]) # number of classes
检查数据集:
with torch_distributed_zero_first(LOCAL_RANK):data_dict = data_dict or check_dataset(data) # check if None
- 使用
torch_distributed_zero_first
上下文管理器,确保只有主进程进行数据集检查。如果data_dict
为None
,则调用check_dataset(data)
函数检查并验证数据集,返回的数据集信息包括训练和验证的路径、类别数量等。获取训练和验证路径:
train_path, val_path = data_dict["train"], data_dict["val"]
- 从
data_dict
中提取训练集和验证集的路径。确定类别数量:
nc = 1 if single_cls else int(data_dict["nc"]) # number of classes
- 根据
single_cls
标志判断类别数量。如果是单类别训练,则nc
设置为 1;否则,从数据字典中取出类别数量。获取类别名称:
names = {0: "item"} if single_cls and len(data_dict["names"]) != 1 else data_dict["names"] # class names
- 设置类别名称。如果是单类别且类别名称列表的长度不是 1,则将其设置为
{0: "item"}
,否则使用数据字典中定义的类别名称。检测是否为COCO数据集:
is_coco = isinstance(val_path, str) and val_path.endswith("coco/val2017.txt") # COCO dataset
- 判断验证集路径是否为COCO格式的数据集,具体通过检查路径是否为字符串,并以
"coco/val2017.txt"
结尾。整体来说,这段代码为YOLOv5训练过程中的配置初始化提供了支持,包括绘图设置、CUDA设备检测、随机种子初始化和数据集的检查。它确保数据集信息(训练与验证路径、类别数量与名称)在训练开始时已准备好,以便为模型训练提供必要的信息。这种方式有助于提高训练的可重复性、监控过程和数据集的有效使用。
4.2加载模型文件
# Modelcheck_suffix(weights, ".pt") # check weightspretrained = weights.endswith(".pt")if pretrained:with torch_distributed_zero_first(LOCAL_RANK):weights = attempt_download(weights) # download if not found locallyckpt = torch.load(weights, map_location="cpu") # load checkpoint to CPU to avoid CUDA memory leakmodel = Model(cfg or ckpt["model"].yaml, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device) # createexclude = ["anchor"] if (cfg or hyp.get("anchors")) and not resume else [] # exclude keyscsd = ckpt["model"].float().state_dict() # checkpoint state_dict as FP32csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # intersectmodel.load_state_dict(csd, strict=False) # loadLOGGER.info(f"Transferred {len(csd)}/{len(model.state_dict())} items from {weights}") # reportelse:model = Model(cfg, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device) # createamp = check_amp(model) # check AMP
代码总结
这段代码负责检查和加载YOLOv5模型的权重,以便进行训练或推理。具体步骤如下:
检查权重文件后缀:
check_suffix(weights, ".pt") # check weights
- 调用
check_suffix
函数检查权重文件的后缀是否为.pt
,以确保使用了正确的PyTorch模型权重文件。判断是否为预训练模型:
pretrained = weights.endswith(".pt")
- 根据权重文件名判断是否为预训练模型,如果是
.pt
文件,则设置pretrained
为True
。处理预训练模型:
if pretrained:with torch_distributed_zero_first(LOCAL_RANK):weights = attempt_download(weights) # download if not found locally
- 如果是预训练模型,则使用
torch_distributed_zero_first
上下文管理器,确保只有主进程尝试下载权重文件。使用attempt_download
函数尝试下载文件,如果在本地未找到权重文件。加载模型检查点:
ckpt = torch.load(weights, map_location="cpu") # load checkpoint to CPU to avoid CUDA memory leak
- 将下载或指定的权重文件加载为模型检查点,使用
map_location="cpu"
确保先将权重加载到CPU,以避免可能的CUDA内存泄漏。创建模型实例:
model = Model(cfg or ckpt["model"].yaml, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device) # create
- 初始化YOLOv5模型,使用配置文件(
cfg
)或检查点中的模型配置(ckpt["model"].yaml
)。设置输入通道数(ch=3
)和类别数量(nc
),并将模型加载到指定的计算设备(GPU或CPU)。处理排除的键:
exclude = ["anchor"] if (cfg or hyp.get("anchors")) and not resume else [] # exclude keys
- 根据配置和超参数决定在加载模型状态时排除哪些键。这通常是为了避免由于锚框参数不匹配而导致的加载错误。
更新模型状态字典:
csd = ckpt["model"].float().state_dict() # checkpoint state_dict as FP32 csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # intersect model.load_state_dict(csd, strict=False) # load
- 将模型检查点的状态字典转换为FP32格式,然后通过
intersect_dicts
函数提取与当前模型相匹配的参数。这确保仅加载匹配的参数,并且可以忽略未匹配的参数,防止潜在的错误。报告加载的参数数量:
LOGGER.info(f"Transferred {len(csd)}/{len(model.state_dict())} items from {weights}") # report
- 使用
LOGGER
记录从权重文件中成功加载的参数数量,帮助用户理解模型的迁移状态。处理非预训练模型:
else:model = Model(cfg, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device) # create
- 如果不是预训练模型,则直接创建模型实例,使用提供的配置文件。
检查自动混合精度(AMP):
amp = check_amp(model) # check AMP
- 调用
check_amp
函数,检查是否可以使用自动混合精度训练。这有助于提高训练效率并降低模型的内存使用。整体来说,这段代码实现了YOLOv5模型的权重加载与初始化过程,包括处理预训练模型和从配置文件创建新模型。它确保模型能够根据指定参数正确设置,并有效管理模型状态字典的加载,以配送成功的训练结果。同时,报告加载的参数数量和AMP检查可以帮助用户监控模型的配置和性能。
4.3冻结模型
# Freezefreeze = [f"model.{x}." for x in (freeze if len(freeze) > 1 else range(freeze[0]))] # layers to freezefor k, v in model.named_parameters():v.requires_grad = True # train all layers# v.register_hook(lambda x: torch.nan_to_num(x)) # NaN to 0 (commented for erratic training results)if any(x in k for x in freeze):LOGGER.info(f"freezing {k}")v.requires_grad = False# Image sizegs = max(int(model.stride.max()), 32) # grid size (max stride)imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # verify imgsz is gs-multiple# Batch sizeif RANK == -1 and batch_size == -1: # single-GPU only, estimate best batch sizebatch_size = check_train_batch_size(model, imgsz, amp)loggers.on_params_update({"batch_size": batch_size})
代码总结
这段代码负责设置YOLOv5模型的层冻结、图像大小检查和批量大小估算。具体步骤如下:
冻结指定层:
freeze = [f"model.{x}." for x in (freeze if len(freeze) > 1 else range(freeze[0]))] # layers to freeze for k, v in model.named_parameters():v.requires_grad = True # train all layersif any(x in k for x in freeze):LOGGER.info(f"freezing {k}")v.requires_grad = False
- 将要冻结的层名转换为符合PyTorch模型参数命名规范的格式。
- 遍历模型的每个参数(
k
为参数名,v
为参数值)。先将所有参数的requires_grad
属性设置为True
,以使其可训练。- 然后检查当前参数名是否包含在冻结层列表中,如果是,则将该参数的
requires_grad
属性设置为False
,并在日志中记录冻结的参数名。图像大小设置:
gs = max(int(model.stride.max()), 32) # grid size (max stride) imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # verify imgsz is gs-multiple
- 计算模型的最大步幅
gs
,确保其至少为 32(这是YOLO的最小处理步幅)。- 调用
check_img_size
函数检查用户设置的图像大小opt.imgsz
是否为gs
的倍数,并确保它符合模型要求。批量大小设置:
if RANK == -1 and batch_size == -1: # single-GPU only, estimate best batch sizebatch_size = check_train_batch_size(model, imgsz, amp)loggers.on_params_update({"batch_size": batch_size})
- 检查是否处于单GPU模式(
RANK == -1
)且用户未指定批量大小(batch_size == -1
)。- 如果条件满足,则调用
check_train_batch_size
函数估算最佳批量大小,并根据模型和图像大小设置。- 通过
loggers.on_params_update
记录更新后的批量大小,以帮助用户接受相关信息。整体来看,这段代码通过设置层冻结、检查图像大小与自动估算批量大小,有效地为YOLOv5模型的训练准备了必要的配置。冻结层的选择有助于在特定场景下减少训练时间并防止过拟合,而对训练参数(如批量大小)的动态计算又保证了模型训练的灵活性和效率。这种管理方式有助于维持训练过程的可控性与高效性。
# Optimizernbs = 64 # nominal batch sizeaccumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizinghyp["weight_decay"] *= batch_size * accumulate / nbs # scale weight_decayoptimizer = smart_optimizer(model, opt.optimizer, hyp["lr0"], hyp["momentum"], hyp["weight_decay"])
代码总结
这段代码段主要用于配置YOLO模型的优化器,具体步骤如下:
设置名义批量大小:
- 定义名义批量大小
nbs
,这里设置为 64。这是一个基准值,用于后续的优化器参数计算。计算累计梯度数量:
- 根据用户提供的实际批量大小
batch_size
计算梯度累积的倍数。通过计算nbs
和batch_size
的比值,得出在每次优化前需要累积多少次梯度(即实际批次的数量)。确保至少为 1,避免除零错误。调整权重衰减:
- 根据实际的批量大小和累积次数调整超参数
weight_decay
。这一调整有助于在不同批量大小下保持正则化效果的一致性。通过将权重衰减参数与当前配置的比例相乘,确保模型的优化在不同训练条件下有统一的效果。初始化优化器:
- 调用
smart_optimizer
函数初始化所选的优化器,传入模型、优化器类型(如"SGD", "Adam"等)、初始学习率、动量和调整后的权重衰减参数。整体来说,这段代码负责为YOLOv5模型的训练准备优化器的设置,包括调整权重衰减参数和计算梯度累积的需求。这些步骤确保模型在不同批量大小下能够有效训练,并通过智能选择优化器来提高训练性能和效率。这样的设计使得模型训练过程更为灵活,能够适应变化的训练环境与需求。
# Schedulerif opt.cos_lr:lf = one_cycle(1, hyp["lrf"], epochs) # cosine 1->hyp['lrf']else:def lf(x):"""Linear learning rate scheduler function with decay calculated by epoch proportion."""return (1 - x / epochs) * (1.0 - hyp["lrf"]) + hyp["lrf"] # linearscheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)
代码总结
这段代码主要用于配置学习率调度器(scheduler),以便在训练过程中动态调整学习率。具体步骤如下:
选择学习率调度策略:
if opt.cos_lr:lf = one_cycle(1, hyp["lrf"], epochs) # cosine 1->hyp['lrf']
- 检查是否启用余弦学习率调度(
--cos-lr
参数)。如果启用,则定义lf
为余弦学习率调度的函数,使用one_cycle
函数生成从 1 到目标学习率(hyp['lrf']
)的调度。定义线性学习率调度器:
else:def lf(x):"""Linear learning rate scheduler function with decay calculated by epoch proportion."""return (1 - x / epochs) * (1.0 - hyp["lrf"]) + hyp["lrf"] # linear
- 如果不使用余弦调度,则定义一个名为
lf
的线性学习率调度函数。该函数根据训练的当前轮次x
和总轮数epochs
动态计算学习率,采用线性衰减的方式。- 具体逻辑是随着训练轮次的增加,逐渐减小学习率,从初始学习率线性回落到目标学习率
hyp["lrf"]
。初始化学习率调度器:
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)
- 使用 PyTorch 的
LambdaLR
调度器初始化学习率调度器,传入之前定义的optimizer
和调度函数lf
。LambdaLR
将在每个训练轮次调用lf
函数,动态更新优化器的学习率。整体而言,这段代码为YOLOv5训练过程中的学习率调整提供了灵活配置。通过选择余弦或线性调度策略,模型可以适应不同的训练需求,从而提高学习效率并帮助模型更好地收敛。利用学习率调度器,开发者能够有效地控制训练过程中的学习率变化,优化训练结果。
# EMAema = ModelEMA(model) if RANK in {-1, 0} else None# Resumebest_fitness, start_epoch = 0.0, 0if pretrained:if resume:best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume)del ckpt, csd# DP modeif 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)# SyncBatchNormif opt.sync_bn and cuda and RANK != -1:model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)LOGGER.info("Using SyncBatchNorm()")
代码总结
这段代码的主要功能是设置和管理YOLOv5模型的指数移动平均(EMA)、恢复训练、配置数据并行(DP)以及启用同步批量归一化(SyncBatchNorm)。具体步骤如下:
设置EMA:
- 如果当前进程是主进程(
RANK
为 -1 或 0),则创建ModelEMA
实例以跟踪模型的指数移动平均。EMA 是一种用于平滑模型权重的技术,提高模型的稳定性和泛化能力;否则,将ema
设置为None
。恢复训练:
- 初始化
best_fitness
和start_epoch
以便在恢复时使用。- 如果是预训练模型并且设置了恢复选项(
resume
),调用smart_resume
函数获取最佳适应度、开始轮次和重新设置的总轮数。该函数还会处理优化器和EMA的状态。- 在恢复后,删除检查点(
ckpt
)和改进后模型的状态字典(csd
)以释放内存。数据并行模式:
- 如果在使用CUDA并且当前进程是主进程(
RANK == -1
),且检测到可用的CUDA设备数量大于1,则输出警告,提醒用户最佳实践是使用分布式数据并行(DDP)而不是数据并行(DP)。最后,如果选择使用DP,则将模型包装为DataParallel
。同步批量归一化:
- 如果启用了同步批量归一化(
sync_bn
),且在CUDA环境中并且当前进程不是主进程(即正在进行分布式训练),则使用torch.nn.SyncBatchNorm.convert_sync_batchnorm
将模型转换为使用同步批量归一化的版本。然后将模型移到指定的计算设备,并记录日志。总结
这段代码为YOLOv5训练设置了多重要素,包括指数移动平均(EMA)的初始化、训练恢复机制、数据并行最佳实践的警告以及同步批量归一化的配置。通过这些设置,代码确保训练过程中的模型性能最大化,且能够在多GPU环境中高效稳定地运行。在分布式训练中,采取这些措施有助于提高模型的收敛速度和最终性能。
训练前准备总结
在YOLOv5模型的训练前准备中,涉及多个重要步骤,以确保训练过程的顺利进行和模型的有效性。这些步骤可以归纳为以下几个方面:
命令行参数解析:
- 使用
parse_opt
函数解析用户通过命令行输入的参数,包括模型权重、数据集、超参数等设置,以便于灵活管理训练配置。环境和依赖性检查:
- 确保所有必要的依赖项已安装,并检查Git状态,以避免在训练过程中出现未提交的更改。
设备选择与初始化:
- 选择适当的计算设备(CPU或GPU),并确保在分布式训练环境下正确配置各个进程的 ranking。
超参数加载与记录:
- 从 YAML 文件加载超参数,并在日志中记录当前的超参数设置,以便于后续监控和调试。
数据集准备:
- 检查并验证数据集的有效性,从中提取训练和验证数据的路径,以及类别信息。
图像大小与批量设置:
- 验证输入图像大小是否符合模型要求,并根据设定的条件计算或估算最佳的批量大小。
冻结层设置:
- 根据用户设置选择需要冻结的网络层,确保在训练过程中不会更新这些层的权重。
优化器与学习率调度器的配置:
- 初始化优化器和学习率调度器,以便在训练过程中动态调整学习率,优化模型训练。
日志记录器的初始化:
- 设置日志记录功能,以便在训练期间监控程序的运行状态和结果,帮助后续分析和复现。
EMA(指数移动平均)模型设置(如适用):
- 初始化EMA模型,以便在训练期间保持权重的时间平均值,提高模型的稳定性和泛化能力。
4.4训练模型
# Trainloadertrain_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 classassert mlc < nc, f"Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}"
这段代码主要负责创建训练数据加载器,并进行一些数据验证,以确保训练数据的格式和内容符合预期。具体步骤如下:
创建数据加载器:
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, )
- 调用
create_dataloader
函数创建训练数据加载器和数据集,传入各类参数:
train_path
: 训练数据的路径。imgsz
: 输入图像的大小。batch_size // WORLD_SIZE
: 确定每个GPU的批量大小,以支持多GPU训练。gs
: 网格大小(用于调整图像尺寸)。single_cls
: 是否将多类数据视为单类。hyp
: 超参数配置。augment
: 是否启用数据增强(这会在数据加载时应用随机变换)。cache
: 图像缓存设置(控制是否使用RAM或磁盘缓存)。rect
: 是否进行矩形训练。rank
: 当前进程的当地排名。workers
: 数据加载工作线程的数量。image_weights
: 是否采用加权图像选择。quad
: 是否使用四元加载。shuffle
: 是否随机打乱数据集。seed
: 随机种子,用于确保数据加载的一致性。提取标签:
labels = np.concatenate(dataset.labels, 0)
- 从数据集对象中提取标签,将所有样本的标签合并为一个NumPy数组。这有助于后续的统计和分析。
获取最大标签类:
mlc = int(labels[:, 0].max()) # max label class
- 计算数据集中最大标签类(
mlc
),这是通过找到标签数组的第一列的最大值来实现的。验证标签类数量:
assert mlc < nc, f"Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}"
- 使用
assert
语句确保最大标签类mlc
不超过预定义的类别数量nc
。如果最大标签类超过nc
,则抛出错误,以防止在训练时出现类别不匹配的问题。总结
整体而言,这段代码有效地为YOLOv5的训练过程设置了合适的数据加载器,并确保了加载数据的准确性。通过创建训练数据加载器,提取标签并验证类别信息,代码为后续的训练提供了必要的准备。这种设置有助于提高模型在训练过程中的数据处理效率和准确性,确保训练模型时输入的质量。
# Process 0if RANK in {-1, 0}:val_loader = create_dataloader(val_path,imgsz,batch_size // WORLD_SIZE * 2,gs,single_cls,hyp=hyp,cache=None if noval else opt.cache,rect=True,rank=-1,workers=workers * 2,pad=0.5,prefix=colorstr("val: "),)[0]if not resume:if not opt.noautoanchor:check_anchors(dataset, model=model, thr=hyp["anchor_t"], imgsz=imgsz) # run AutoAnchormodel.half().float() # pre-reduce anchor precisioncallbacks.run("on_pretrain_routine_end", labels, names)# DDP modeif cuda and RANK != -1:model = smart_DDP(model)
代码总结
这段代码主要负责处理验证集的数据加载和条件执行在YOLOv5训练中的准备工作,尤其是在分布式训练(DDP)情况下。具体步骤如下:
检查主进程:
if RANK in {-1, 0}:
- 判断当前进程是否为主进程(
RANK
为 -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]
- 调用
create_dataloader
函数创建验证集的数据加载器,传入以下参数:
val_path
: 验证集的路径。imgsz
: 输入图像的大小。batch_size // WORLD_SIZE * 2
: 为验证加载两倍的批量大小,以提高效率。gs
: 网格大小(用于图像调整)。single_cls
: 是否将多类视为单类。hyp
: 超参数配置。cache
: 是否使用缓存,依赖于noval
。rect
: 启用矩形图像填充。rank
: 在DCG中设置为 -1。workers
: 将工作线程数翻倍以加速数据加载。pad
: 设置图像填充比例。prefix
: 用于输出的信息前缀(表示验证)。检查并优化锚框:
if not resume:if not opt.noautoanchor:check_anchors(dataset, model=model, thr=hyp["anchor_t"], imgsz=imgsz) # run AutoAnchor
- 如果不是从上次训练恢复,并且未禁用自动锚框,则运行
check_anchors
函数来验证和调整锚框,以优化目标检测性能。转换模型精度:
model.half().float() # pre-reduce anchor precision
- 将模型的参数转换为半精度(FP16),以降低内存使用并加速计算,然后转换回FP32以处理锚框精度。
运行回调:
callbacks.run("on_pretrain_routine_end", labels, names)
- 调用注册的回调函数,表示训练准备流程的结束,传入标签和类别名称信息。
处理分布式训练模式:
if cuda and RANK != -1:model = smart_DDP(model)
- 如果使用CUDA并且当前进程不是主进程,则使用
smart_DDP
对模型应用分布式数据并行(DDP)设置,以便在多GPU环境中更高效地训练。整体而言,这段代码为YOLOv5的训练过程中的验证集管理提供了必要的支持,通过加载验证数据集、运行锚框检查和配置模型的精度,同时在需要时设置分布式训练模式。这些准备工作确保在训练之前,所有必要的验证步骤都被合理管理,从而提高了训练的准确性和效率。
4.5初始化模型所用参数
# Model attributesnl = de_parallel(model).model[-1].nl # number of detection layers (to scale hyps)hyp["box"] *= 3 / nl # scale to layershyp["cls"] *= nc / 80 * 3 / nl # scale to classes and layershyp["obj"] *= (imgsz / 640) ** 2 * 3 / nl # scale to image size and layershyp["label_smoothing"] = opt.label_smoothingmodel.nc = nc # attach number of classes to modelmodel.hyp = hyp # attach hyperparameters to modelmodel.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weightsmodel.names = names
代码总结
这段代码主要负责处理验证集的数据加载和条件执行在YOLOv5训练中的准备工作,尤其是在分布式训练(DDP)情况下。具体步骤如下:
检查主进程:
if RANK in {-1, 0}:
- 判断当前进程是否为主进程(
RANK
为 -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]
- 调用
create_dataloader
函数创建验证集的数据加载器,传入以下参数:
val_path
: 验证集的路径。imgsz
: 输入图像的大小。batch_size // WORLD_SIZE * 2
: 为验证加载两倍的批量大小,以提高效率。gs
: 网格大小(用于图像调整)。single_cls
: 是否将多类视为单类。hyp
: 超参数配置。cache
: 是否使用缓存,依赖于noval
。rect
: 启用矩形图像填充。rank
: 在DCG中设置为 -1。workers
: 将工作线程数翻倍以加速数据加载。pad
: 设置图像填充比例。prefix
: 用于输出的信息前缀(表示验证)。检查并优化锚框:
if not resume:if not opt.noautoanchor:check_anchors(dataset, model=model, thr=hyp["anchor_t"], imgsz=imgsz) # run AutoAnchor
- 如果不是从上次训练恢复,并且未禁用自动锚框,则运行
check_anchors
函数来验证和调整锚框,以优化目标检测性能。转换模型精度:
model.half().float() # pre-reduce anchor precision
- 将模型的参数转换为半精度(FP16),以降低内存使用并加速计算,然后转换回FP32以处理锚框精度。
运行回调:
callbacks.run("on_pretrain_routine_end", labels, names)
- 调用注册的回调函数,表示训练准备流程的结束,传入标签和类别名称信息。
处理分布式训练模式:
if cuda and RANK != -1:model = smart_DDP(model)
- 如果使用CUDA并且当前进程不是主进程,则使用
smart_DDP
对模型应用分布式数据并行(DDP)设置,以便在多GPU环境中更高效地训练。整体而言,这段代码为YOLOv5的训练过程中的验证集管理提供了必要的支持,通过加载验证数据集、运行锚框检查和配置模型的精度,同时在需要时设置分布式训练模式。这些准备工作确保在训练之前,所有必要的验证步骤都被合理管理,从而提高了训练的准确性和效率。
4.6 训练前准备
# Start trainingt0 = time.time()nb = len(train_loader) # number of batchesnw = max(round(hyp["warmup_epochs"] * nb), 100) # number of warmup iterations, max(3 epochs, 100 iterations)# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of traininglast_opt_step = -1maps = np.zeros(nc) # mAP per classresults = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)scheduler.last_epoch = start_epoch - 1 # do not movescaler = torch.cuda.amp.GradScaler(enabled=amp)stopper, stop = EarlyStopping(patience=opt.patience), Falsecompute_loss = ComputeLoss(model) # init loss classcallbacks.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...')
代码总结
这段代码主要负责设置训练过程的开始阶段,初始化训练相关的变量,并配置训练环境。具体步骤如下:
记录起始时间:
t0 = time.time()
- 记录训练开始时间,以便后续计算训练的持续时间。
获取批次数量
nb = len(train_loader) # number of batches
- 计算训练数据加载器中的批次数(即一轮训练中要处理的批次数量)。
设置热身迭代次数:
nw = max(round(hyp["warmup_epochs"] * nb), 100) # number of warmup iterations, max(3 epochs, 100 iterations)
- 根据超参数配置中的热身周期数和批次数量计算热身迭代次数,确保至少为100次。热身阶段使学习率逐渐增加,帮助模型稳定在初始阶段。
初始化训练状态:
last_opt_step = -1 maps = np.zeros(nc) # mAP per class results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
- 设置最后优化步骤为-1,表示训练尚未开始。
- 创建一个数组
maps
用于记录每个类别的平均精度(mAP)。- 初始化训练结果元组
results
,用于记录精度(P)、召回率(R)、mAP、验证损失等指标。调整调度器的当前轮次:
scheduler.last_epoch = start_epoch - 1 # do not move
- 将学习率调度器的
last_epoch
属性设置为start_epoch - 1
,以确保在开始训练时调度器从正确的状态开始。GradScaler初始化:
scaler = torch.cuda.amp.GradScaler(enabled=amp)
- 初始化自动混合精度(AMP)梯度缩放器,启用与混合精度训练相关的功能,以提高训练效率和减少内存使用。
早停机制:
stopper, stop = EarlyStopping(patience=opt.patience), False
- 创建一个早停实例,设置耐心值(即在多少个epoch无改善后停止训练),并将停止标志初始化为
False
。损失计算初始化:
compute_loss = ComputeLoss(model) # init loss class
- 初始化损失计算类,以便后续计算训练期间的损失值。
运行回调:
callbacks.run("on_train_start")
- 执行开始训练时注册的回调函数,可能用于记录、监控等目的。
日志记录训练信息:
LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n'f"Logging results to {colorstr('bold', save_dir)}\n"f'Starting training for {epochs} epochs...' )
- 使用日志记录器输出训练的图像大小、使用的数据加载线程、结果记录路径和总训练周期数的信息,以便用户了解训练的具体设置和状态。
训练结束判断与模型评估总结
这段描述详细阐述了在YOLOv5训练过程中如何判断是否结束训练以及如何计算最佳模型的标准。具体分析如下:
训练结束判断:
- 首先,判断是否需要终止训练。这通常发生在以下两种情况:
- 如果用户选择在每个训练周期后进行验证。
- 当前训练已达到最后一轮(epoch)。
计算最佳模型的标准:
- 这里的“最好”模型是通过一系列评估指标来判断的,这些指标综合反映了模型的性能。
- Fitness 计算:
- 使用的评估指标为:
P
: 精度R
: 召回率mAP@.5
: 在IoU阈值为0.5时的平均精度mAP@.5-.95
: 在多个IoU阈值(从0.5到0.95,步长为0.05)下的平均精度- 计算公式为:
- fi=0.1⋅mAP@0.5+0.9⋅mAP@0.5:0.95
- 这个加权计算表明,评估标准更倾向于
mAP@0.5:0.95
,强调模型在多个IoU阈值下的表现。该指标的高值意味着模型在物体识别能力上有更强的鲁棒性。mAP@0.5:0.95 的重要性:
- 更高的
mAP@0.5:0.95
表明模型在识别物体时,即使在较难的验证条件下(多种IoU阈值),也能保持较好的性能。这表明该模型能够在更加复杂的场景中有效识别对象。总结
通过这种方式,YOLOv5在训练过程中能够有效地评估模型性能,以科学的方法确定何时结束训练,并确保保留表现最佳的模型。这种灵活的训练结束判断与模型评估机制,有助于提升模型的识别能力和应用有效性。
4.7开始训练
for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------callbacks.run("on_train_epoch_start")model.train()# Update image weights (optional, single-GPU only)if opt.image_weights:cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weightsiw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weightsdataset.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
代码总结
这段代码主要负责YOLOv5模型的每个训练周期(epoch)初始化和图像权重的更新。具体步骤分析如下:
训练循环的开始:
for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------callbacks.run("on_train_epoch_start")model.train()
- 在给定的轮次范围内(从
start_epoch
到epochs
),启动每个训练周期的循环。- 调用回调函数
callbacks.run("on_train_epoch_start")
,执行与训练开始相关的操作,可能是记录或监控。- 将模型设置为训练模式,通过
model.train()
启用特定的训练行为,如启用掉落层(Dropout)和批量规范化(Batch Normalization)。更新图像权重:
if opt.image_weights:cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weightsiw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weightsdataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx
- 如果启用了图像权重选项(
--image-weights
):
- 计算类别权重
cw
,其中通过调整与各类别的平均精度(maps
)相关的权重,确保较低的mAP类得到加权提升。- 调用
labels_to_image_weights
函数计算每个图像的权重iw
,以优化训练过程,使得训练中的样本选择更加有效。- 使用
random.choices
根据计算出的权重随机选择图像索引,以进行加权训练。更新马赛克边界(注释掉的代码部分):
# 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
- 这段注释代码提供了如何在数据集的训练过程中设置马赛克图像的边界信息(可选)。
- 如果启用,则随机计算边界
b
,并根据该边界设置马赛克剪裁边界。这有助于生成多样化的输入图像,以增强模型的鲁棒性,但该部分代码在此未启用。总结
这段代码为YOLOv5模型的每个训练周期设置了必要的初始化步骤,包括图像权重的更新和模型模式的切换。图像权重的动态调整有助于平衡类别之间的训练样本,提高模型对不平衡数据集的适应能力。此外,虽然边界更新的部分被注释掉,但它提供了额外的数据增强方案,以进一步提升模型的表现。整体而言,这些设置是提高训练效果和模型性能的重要组成部分。
mloss = torch.zeros(3, device=device) # mean lossesif RANK != -1:train_loader.sampler.set_epoch(epoch)pbar = enumerate(train_loader)LOGGER.info(("\n" + "%11s" * 7) % ("Epoch", "GPU_mem", "box_loss", "obj_loss", "cls_loss", "Instances", "Size"))if RANK in {-1, 0}:pbar = tqdm(pbar, total=nb, bar_format=TQDM_BAR_FORMAT) # progress baroptimizer.zero_grad()
这段代码是YOLOv5训练脚本中的一部分,主要功能包括初始化平均损失、设置分布式训练中的epoch、初始化进度条以及在每个批次训练前将优化器的梯度清零。
总结:
- 初始化了一个三元素的张量用于存储平均损失。
- 在分布式训练环境中为数据加载器设置当前epoch。
- 使用
tqdm
库创建训练进度条并记录训练过程中的关键信息(如GPU内存使用、各类损失等)。- 每次训练新批次前,将优化器的梯度重置为零。
for i, (imgs, targets, paths, _) in pbar: # batch -------------------------------------------------------------callbacks.run("on_train_batch_start")ni = i + nb * epoch # number integrated batches (since train start)imgs = imgs.to(device, non_blocking=True).float() / 255 # uint8 to float32, 0-255 to 0.0-1.0# Warmupif 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 lr0x["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-scaleif opt.multi_scale:sz = random.randrange(int(imgsz * 0.5), int(imgsz * 1.5) + gs) // gs * gs # sizesf = sz / max(imgs.shape[2:]) # scale factorif 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)
这段代码是YOLOv5训练过程中一个批次(batch)的主要处理逻辑,包含了批次的预处理、学习率和动量的调整、多尺度训练等关键步骤。
总结:
- 批次循环处理:代码遍历每个批次的图像和标签,记录当前批次数。
- 图像预处理:图像被移动到指定设备(如GPU),转换为浮点数,并归一化到0到1之间。
- 学习率和动量的热身(Warmup):在训练的早期阶段,通过插值动态调整学习率和动量,从而逐渐引导模型进入稳定训练阶段。
- 多尺度训练:如果启用了多尺度训练,代码会随机调整输入图像的尺寸,这有助于提高模型的鲁棒性,使其更好地适应不同的输入尺度。
# Forwardwith torch.cuda.amp.autocast(amp):pred = model(imgs) # forwardloss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_sizeif RANK != -1:loss *= WORLD_SIZE # gradient averaged between devices in DDP modeif opt.quad:loss *= 4.0# Backwardscaler.scale(loss).backward()# Optimize - https://pytorch.org/docs/master/notes/amp_examples.htmlif ni - last_opt_step >= accumulate:scaler.unscale_(optimizer) # unscale gradientstorch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0) # clip gradientsscaler.step(optimizer) # optimizer.stepscaler.update()optimizer.zero_grad()if ema:ema.update(model)last_opt_step = ni
这段代码展示了YOLOv5在训练过程中前向传播、反向传播以及优化步骤的具体实现,尤其是应用了混合精度训练(AMP)以提高训练效率。
总结:
- 前向传播:使用
torch.cuda.amp.autocast(amp)
来启用混合精度训练,模型接收经过预处理的图像并生成预测结果。然后通过compute_loss
计算损失,并根据分布式训练(DDP)模式或其他选项(如quad
)调整损失值。- 反向传播:使用
scaler.scale(loss).backward()
对损失进行缩放后计算梯度,以避免数值不稳定性。- 优化步骤:当累积的梯度达到一定步数时,执行优化步骤:
- 首先对梯度进行反缩放和裁剪,以控制梯度的大小。
- 使用缩放后的梯度更新模型参数,并更新缩放器状态。
- 重置优化器的梯度,以便进行下一次迭代。
- 如果使用了EMA(指数滑动平均),会在每次优化后更新模型的EMA权重。
这段代码通过混合精度训练和梯度累积等技术,有效地优化了模型训练的性能和稳定性。
# Logif RANK in {-1, 0}:mloss = (mloss * i + loss_items) / (i + 1) # update mean lossesmem = 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 ------------------------------------------------------------------------------------------------# Schedulerlr = [x["lr"] for x in optimizer.param_groups] # for loggersscheduler.step()
这段代码处理了YOLOv5训练过程中的日志记录和学习率调度,确保在训练期间的信息记录和学习率的调整。
总结:
- 日志记录:在每个批次结束后,代码会更新和记录平均损失(
mloss
),并在进度条中显示当前的训练信息,包括当前的epoch、GPU内存使用情况、平均损失、目标数目和图像尺寸等。如果在分布式训练环境下,只在主进程(RANK
为-1或0)记录日志。- 回调函数执行:在每个批次结束后,执行相关的回调函数(如
on_train_batch_end
),这可能包含自定义的训练中操作。如果在回调中检测到训练停止的信号,会立即结束训练。- 学习率调度:在所有批次结束后,代码会从优化器中提取当前的学习率信息,供日志记录使用,并调用调度器(
scheduler
)更新学习率,为下一轮训练做好准备。这部分代码确保了训练过程中每个批次的信息能够被有效记录,并且在每个epoch结束时动态调整学习率以促进模型更好地收敛。
if RANK in {-1, 0}:# mAPcallbacks.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_stopif not noval or final_epoch: # Calculate mAPresults, 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 mAPfi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95]stop = stopper(epoch=epoch, fitness=fi) # early stop checkif fi > best_fitness:best_fitness = filog_vals = list(mloss) + list(results) + lrcallbacks.run("on_fit_epoch_end", log_vals, epoch, best_fitness, fi)# Save modelif (not nosave) or (final_epoch and not evolve): # if saveckpt = {"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 deletetorch.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 ckptcallbacks.run("on_model_save", last, epoch, final_epoch, best_fitness, fi)# EarlyStoppingif RANK != -1: # if DDP trainingbroadcast_list = [stop if RANK == 0 else None]dist.broadcast_object_list(broadcast_list, 0) # broadcast 'stop' to all ranksif RANK != 0:stop = broadcast_list[0]if stop:break # must break all DDP ranks# end epoch ----------------------------------------------------------------------------------------------------# end training -----------------------------------------------------------------------------------------------------
这段代码处理了YOLOv5训练过程中每个epoch结束时的关键步骤,包括验证、早停、模型保存等内容。以下是总结:
mAP计算:在每个epoch结束时(特别是在最终epoch或启用验证时),代码会运行验证过程,计算模型在验证集上的mAP(平均精度),并更新模型的EMA(指数滑动平均)属性。
最佳mAP更新和早停:通过
fitness
函数计算当前模型的综合性能指标,如果当前指标优于之前的最佳指标,则更新最佳mAP。同时,检查是否满足早停条件,如果达到条件,会触发早停,终止训练。模型保存:在每个epoch结束时,尤其是在最终epoch或满足保存条件时,保存当前模型的状态,包括模型参数、EMA、优化器状态等。保存的模型包括最新的模型和最佳mAP模型。
早停广播:在分布式训练环境下,如果主进程决定早停,会将这一决定广播给其他进程,确保所有进程一致地停止训练。
这段代码确保了每个epoch结束后,模型的性能能够被正确评估,并根据情况决定是否需要早停或保存模型,从而提高训练的效率和稳定性。
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 optimizersif 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.65single_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 plotsif 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
这段代码处理了YOLOv5训练过程结束后的清理和验证步骤,确保在训练完成后对模型进行最后的评估,并释放资源。以下是总结:
训练完成日志:当训练完成时,记录训练所用的时间,方便后续分析。
优化器剥离与模型验证:
- 对保存的模型文件(
last
和best
)进行处理,移除优化器的状态(通过strip_optimizer
),以减小文件大小。- 对于最佳模型(
best
),再次运行验证,尤其是在COCO数据集上,使用更高的IoU阈值进行评估,并生成验证结果和相关的图表。回调函数:在训练结束时,调用
on_train_end
回调函数,传递最终模型、最佳模型、epoch数和最后的验证结果,供进一步的处理或日志记录。清理GPU缓存:调用
torch.cuda.empty_cache()
来释放未使用的显存,防止显存泄漏。返回结果:最终返回验证结果,供后续使用或分析。
这部分代码在训练结束后,对模型进行精简和评估,确保最终保存的模型是经过验证的最佳版本,并清理训练过程中使用的资源。
五,run文件
def run(**kwargs):opt = parse_opt(True)for k, v in kwargs.items():setattr(opt, k, v)main(opt)return opt
这段代码定义了一个名为
run
的函数,用于运行YOLOv5的训练或推理任务,并允许通过关键字参数(kwargs
)动态修改默认配置选项。以下是总结:
配置解析:函数首先调用
parse_opt(True)
,解析默认配置选项并生成一个包含这些配置的对象opt
。动态设置参数:通过遍历
kwargs
中的键值对,将这些值动态地赋给opt
对象的相应属性,从而覆盖默认配置。这使得在调用run
函数时,可以灵活地传递和修改配置选项。调用主函数:调用
main(opt)
,启动主要的训练或推理流程,opt
对象作为配置参数传递给main
函数。返回配置:最后,返回修改后的
opt
对象,以便在调用run
函数后,可以查看或使用最终的配置选项。这段代码提供了一种灵活的方式来运行YOLOv5任务,并且可以在调用时动态调整配置选项,适应不同的需求。
这篇关于YOLOV5入门教程day3的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!