trainer.py
ultralytics\engine\trainer.py
目录
1.所需的库和模块
# Ultralytics YOLO 🚀, AGPL-3.0 license
"""
Train a model on a dataset.
Usage:
$ yolo mode=train model=yolov8n.pt data=coco8.yaml imgsz=640 epochs=100 batch=16
"""
import gc
import math
import os
import subprocess
import time
import warnings
from copy import deepcopy
from datetime import datetime, timedelta
from pathlib import Path
import numpy as np
import torch
from torch import distributed as dist
from torch import nn, optim
from ultralytics.cfg import get_cfg, get_save_dir
from ultralytics.data.utils import check_cls_dataset, check_det_dataset
from ultralytics.nn.tasks import attempt_load_one_weight, attempt_load_weights
from ultralytics.utils import (
DEFAULT_CFG,
LOCAL_RANK,
LOGGER,
RANK,
TQDM,
__version__,
callbacks,
clean_url,
colorstr,
emojis,
yaml_save,
)
from ultralytics.utils.autobatch import check_train_batch_size
from ultralytics.utils.checks import check_amp, check_file, check_imgsz, check_model_file_from_stem, print_args
from ultralytics.utils.dist import ddp_cleanup, generate_ddp_command
from ultralytics.utils.files import get_latest_run
from ultralytics.utils.torch_utils import (
TORCH_2_4,
EarlyStopping,
ModelEMA,
autocast,
convert_optimizer_state_dict_to_fp16,
init_seeds,
one_cycle,
select_device,
strip_optimizer,
torch_distributed_zero_first,
)
2.class BaseTrainer:
# 这段代码是一个Python类的实现,名为 BaseTrainer ,它看起来是训练深度学习模型的基础类。
class BaseTrainer:
# 用于创建训练器的基类。
# 属性:
# args (SimpleNamespace) :训练器的配置。
# validator (BaseValidator) :验证器实例。
# model (nn.Module) :模型实例。
# callbacks (defaultdict) :回调字典。
# save_dir (Path) :保存结果的目录。
# wdir (Path) :保存权重的目录。
# last (Path) :最后一个检查点的路径。
# best (Path) :最佳检查点的路径。
# save_period (int) :每 x 个 epoch 保存一次检查点(如果 < 1,则禁用)。
# batch_size (int) :训练的批次大小。
# epochs (int) :要训练的 epoch 数。
# start_epoch (int) :训练的起始 epoch。
# device (torch.device) :用于训练的设备。
# amp (bool) :启用 AMP(自动混合精度)的标志。
# scaler (amp.GradScaler) :AMP 的梯度缩放器。
# data (str) :数据路径。
# trainset (torch.utils.data.Dataset) :训练数据集。
# testset (torch.utils.data.Dataset) :测试数据集。
# ema (nn.Module) :模型的 EMA(指数移动平均线)。
# resume (bool) :从检查点恢复训练。
# lf (nn.Module) :损失函数。
# scheduler (torch.optim.lr_scheduler._LRScheduler) :学习率调度程序。
# best_fitness (float) :达到的最佳适应度值。
# fitness (float) :当前适应度值。
# loss (float) :当前损失值。
# tloss (float) :总损失值。
# loss_names (list) :损失名称列表。
# csv (Path) :结果 CSV 文件的路径。
"""
A base class for creating trainers.
Attributes:
args (SimpleNamespace): Configuration for the trainer.
validator (BaseValidator): Validator instance.
model (nn.Module): Model instance.
callbacks (defaultdict): Dictionary of callbacks.
save_dir (Path): Directory to save results.
wdir (Path): Directory to save weights.
last (Path): Path to the last checkpoint.
best (Path): Path to the best checkpoint.
save_period (int): Save checkpoint every x epochs (disabled if < 1).
batch_size (int): Batch size for training.
epochs (int): Number of epochs to train for.
start_epoch (int): Starting epoch for training.
device (torch.device): Device to use for training.
amp (bool): Flag to enable AMP (Automatic Mixed Precision).
scaler (amp.GradScaler): Gradient scaler for AMP.
data (str): Path to data.
trainset (torch.utils.data.Dataset): Training dataset.
testset (torch.utils.data.Dataset): Testing dataset.
ema (nn.Module): EMA (Exponential Moving Average) of the model.
resume (bool): Resume training from a checkpoint.
lf (nn.Module): Loss function.
scheduler (torch.optim.lr_scheduler._LRScheduler): Learning rate scheduler.
best_fitness (float): The best fitness value achieved.
fitness (float): Current fitness value.
loss (float): Current loss value.
tloss (float): Total loss value.
loss_names (list): List of loss names.
csv (Path): Path to results CSV file.
"""
# __init__ 方法。这是类的构造函数,用于初始化类的实例。它接受三个参数。
# cfg :配置。是一个参数,用于传递训练过程中的配置信息。这个参数有一个默认值 DEFAULT_CFG ,这意味着如果调用构造函数时没有提供 cfg 参数,它将使用 DEFAULT_CFG 作为默认配置。 DEFAULT_CFG 是一个预定义的配置对象或字典,包含了训练过程中所需的默认设置。
# overrides :是一个可选参数,用于覆盖 cfg 中提供的默认配置。它允许用户在不修改 DEFAULT_CFG 的情况下,提供特定的配置值来定制训练过程。这个参数默认为 None ,意味着如果没有提供覆盖配置,将完全使用 cfg 中的设置。
# _callbacks :用于传递一个回调函数列表,这些回调函数可以在训练过程中的特定点被调用,以执行自定义的操作,如记录训练进度、执行评估等。这个参数默认为 None ,意味着如果没有提供回调函数,将不会使用任何额外的回调。
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
# 初始化 BaseTrainer 类。
"""
Initializes the BaseTrainer class.
Args:
cfg (str, optional): Path to a configuration file. Defaults to DEFAULT_CFG.
overrides (dict, optional): Configuration overrides. Defaults to None.
"""
# 这段代码是 BaseTrainer 类的构造函数( __init__ 方法)的一部分,它负责在创建类的实例时进行一些基础的初始化工作。
# 这行代码调用 get_cfg 函数来获取配置信息。 cfg 参数是一个默认配置, overrides 是一个可选参数,用于覆盖默认配置中的某些值。
# def get_cfg(cfg: Union[str, Path, Dict, SimpleNamespace] = DEFAULT_CFG_DICT, overrides: Dict = None):
# -> 从一个配置源(可以是字符串、路径、字典或 SimpleNamespace 对象)获取配置,并允许通过 overrides 字典来覆盖默认配置。返回一个 IterableSimpleNamespace 对象,它是一个可迭代的命名空间对象,其属性由 cfg 字典中的键值对初始化。
# -> return IterableSimpleNamespace(**cfg)
self.args = get_cfg(cfg, overrides)
# 这行代码调用 check_resume 方法,这个方法可能用于检查是否存在先前的训练状态,如果存在,则可能从该状态恢复训练。 overrides 参数用于提供额外的信息,以便在恢复时进行必要的调整。
self.check_resume(overrides)
# 这行代码调用 select_device 函数来确定训练将使用的设备(CPU或GPU)。 self.args.device 指定了设备类型, self.args.batch 指定了批量大小,这影响设备选择。
# def select_device(device="", batch=0, newline=False, verbose=True):
# -> 用于根据用户提供的设备字符串选择并返回一个合适的 PyTorch 设备对象。返回设备对象。返回一个 PyTorch 设备对象,可以是 GPU、MPS 或 CPU。
# -> return torch.device(arg)
self.device = select_device(self.args.device, self.args.batch)
# 初始化 validator 属性为 None 。用于验证模型性能的对象,稍后会在训练过程中被赋值。
self.validator = None
# 初始化 metrics 属性为 None 。用于跟踪训练过程中的各种指标的对象,稍后会被赋值。
self.metrics = None
# 初始化 plots 字典为空。用于存储训练过程中的绘图数据。
self.plots = {}
# 调用 init_seeds 函数来设置随机种子,以确保实验的可重复性。 self.args.seed 是基础随机种子, RANK 是当前进程的排名(在分布式训练中使用), deterministic 是一个布尔值,指示是否需要确保所有随机操作都是确定性的。
# def init_seeds(seed=0, deterministic=False): -> 初始化随机数生成器(RNG)的种子,以确保实验的可重复性。
init_seeds(self.args.seed + 1 + RANK, deterministic=self.args.deterministic)
# 这段代码的主要目的是为训练过程设置初始状态,包括配置、设备选择、随机种子等。这些设置对于确保训练的一致性和可重复性至关重要。
# 这段代码是 BaseTrainer 类构造函数中的一部分,它负责设置和初始化与文件和目录相关的属性。
# Dirs
# 调用 get_save_dir 函数,传入 self.args (包含配置参数的对象)来确定保存目录的路径。这个目录用于存放训练过程中的各种文件,如日志、权重等。
# def get_save_dir(args, name=None): -> 获取或创建一个用于保存特定数据(如模型权重、日志等)的目录路径。返回保存目录的路径。返回一个 Path 对象,表示保存目录的路径。 -> return Path(save_dir)
self.save_dir = get_save_dir(self.args)
# 更新 self.args 中的 name 属性,将其设置为 self.save_dir 的名称。这通常用于日志记录和其他需要标识符的地方。
self.args.name = self.save_dir.name # update name for loggers
# 设置权重目录 self.wdir 为 self.save_dir 下的一个子目录,名称为 weights 。
self.wdir = self.save_dir / "weights" # weights dir
# 这个条件检查 RANK 的值。在分布式训练中, RANK 用于标识当前进程的编号。 RANK 为-1通常表示非分布式环境,而 RANK 为0表示分布式环境中的主进程。
if RANK in {-1, 0}:
# 如果 RANK 为-1或0, 会创建权重目录 self.wdir 。 parents=True 参数表示如果父目录不存在也会创建, exist_ok=True 表示如果目录已存在不会抛出异常。
self.wdir.mkdir(parents=True, exist_ok=True) # make dir
# 将 self.save_dir 的路径转换为字符串,并更新 self.args 中的 save_dir 属性。
self.args.save_dir = str(self.save_dir)
# 将 self.args 中的配置参数保存到 self.save_dir 下的 args.yaml 文件中。 yaml_save 函数是一个自定义函数,用于将Python对象保存为YAML格式的文件。
# def yaml_save(file="data.yaml", data=None, header=""): -> 将 Python 数据结构保存为 YAML 格式的文件。
yaml_save(self.save_dir / "args.yaml", vars(self.args)) # save run args
# 设置最后和最佳检查点文件的路径。 last.pt 用于存储最近的模型权重,而 best.pt 用于存储验证集上表现最佳的模型权重。
self.last, self.best = self.wdir / "last.pt", self.wdir / "best.pt" # checkpoint paths
# 设置模型保存周期,即每隔多少个epoch保存一次模型。
self.save_period = self.args.save_period
# 设置批量大小,从 self.args 中获取。
self.batch_size = self.args.batch
# 设置训练周期总数,从 self.args 中获取。
self.epochs = self.args.epochs
# 初始化起始周期为0,表示训练从第一个周期开始。
self.start_epoch = 0
# 检查 RANK 是否为-1,如果是,则执行下面的代码。
if RANK == -1:
# vars(object)
# vars() 函数在 Python 中用于获取对象的属性字典。这个字典包含了对象的大部分属性,但不包括方法和其他一些特殊的属性。对于用户自定义的对象, vars() 返回的字典包含了对象的 __dict__ 属性,这是一个包含对象所有属性的字典。
# 参数说明 :
# object :要获取属性字典的对象。
# 返回值 :
# 返回指定对象的属性字典。
# 注意事项 :
# vars() 对于内置类型(如 int 、 float 、 list 等)返回的是一个包含魔术方法和特殊属性的字典,这些属性通常是不可访问的。
# 对于自定义对象, vars() 返回的是对象的 __dict__ 属性,如果对象没有定义 __dict__ ,则可能返回一个空字典或者抛出 TypeError 。
# 在 Python 3 中, vars() 也可以用于获取内置函数的全局变量字典。
# vars() 函数是一个内置函数,通常用于调试和访问对象的内部状态,但在处理复杂对象时应该谨慎使用,因为直接修改对象的属性可能会导致不可预测的行为。
# 如果 RANK 为-1,调用 print_args 函数打印出所有配置参数。 vars(self.args) 将 self.args 对象转换为字典,以便打印。
# def print_args(args: Optional[dict] = None, show_file=True, show_func=False): -> 打印函数的参数,这在调试和记录函数调用时非常有用。
print_args(vars(self.args))
# 这段代码的主要目的是设置训练过程中文件和目录的路径,并保存配置参数,以便在训练过程中使用和记录。
# 这段代码继续描述了 BaseTrainer 类构造函数中的一些初始化步骤,涉及设备选择、模型和数据集的加载,以及优化器工具的初始化。
# Device
# 设备选择。
# 检查 self.device.type 是否为CPU或MPS(苹果的Metal Performance Shaders,用于在iOS和macOS上进行GPU加速)。
if self.device.type in {"cpu", "mps"}:
# 如果是,将 self.args.workers 设置为0。这意味着在CPU或MPS上训练时,不会使用额外的工作线程来加载数据,因为数据加载的时间相对于模型推理来说并不显著,减少工作线程可以加快训练速度。
self.args.workers = 0 # faster CPU training as time dominated by inference, not dataloading
# Model and Dataset
# 模型加载。调用 check_model_file_from_stem 函数来检查并设置模型文件的路径。这个函数可能接受模型名称的茎(不带后缀的部分),并添加适当的后缀(如 .pt )来形成完整的模型文件路径。
# def check_model_file_from_stem(model="yolov8n"): -> 根据模型的名称(不带扩展名)检查并返回一个有效的模型文件名。 -> return Path(model).with_suffix(".pt") / return model
self.model = check_model_file_from_stem(self.args.model) # add suffix, i.e. yolov8n -> yolov8n.pt
# 数据集加载。
# 使用 torch_distributed_zero_first 上下文管理器来确保在分布式训练环境中,只有主进程( LOCAL_RANK 为0的进程)执行数据集的自动下载。这样可以避免多个进程同时下载相同的数据集。
# def torch_distributed_zero_first(local_rank: int): -> torch_distributed_zero_first 上下文管理器(context manager),用于在分布式训练中确保所有进程等待主进程(通常是 rank 0)完成特定任务后再继续执行。
with torch_distributed_zero_first(LOCAL_RANK): # avoid auto-downloading dataset multiple times
# self.get_dataset 方法被调用来获取训练集和测试集。
self.trainset, self.testset = self.get_dataset()
# 指数移动平均(EMA)初始化。 初始化指数移动平均(EMA)为 None 。EMA是一种用于跟踪模型参数的平滑版本的技术,通常用于模型训练中以提高性能。
self.ema = None
# Optimization utils init
# 优化器工具初始化。
# 分别初始化学习率函数( lf )和学习率调度器( scheduler )为 None 。学习率函数用于调整学习率,而学习率调度器用于在训练过程中根据预定义的策略调整学习率。
self.lf = None
self.scheduler = None
# 这段代码的目的是为训练过程设置必要的环境,包括确定使用的设备、加载模型和数据集,以及准备优化器相关的工具。这些步骤是训练任何机器学习模型之前的标准准备工作。
# 这段代码继续描述了 BaseTrainer 类构造函数中的初始化步骤,这次关注的是周期级别的度量(epoch level metrics)、模型中心(HUB)会话、以及回调函数的设置。
# Epoch level metrics
# 周期级别的度量。
# best_fitness 用于存储最佳适应度值。
self.best_fitness = None
# fitness 用于存储当前周期的适应度值。
self.fitness = None
# loss 和 tloss 分别用于存储当前周期的损失和总损失。
self.loss = None
self.tloss = None
# loss_names 是一个列表,用于存储损失名称,这里初始化为包含"Loss"的列表。
self.loss_names = ["Loss"]
# CSV文件路径。设置CSV文件的路径,该文件将用于保存训练结果。 self.save_dir 是之前设置的保存目录,这里在该目录下创建或指定一个名为"results.csv"的文件。
self.csv = self.save_dir / "results.csv"
# 绘图索引。初始化了一个绘图索引列表,用于指定在训练过程中哪些度量指标需要被绘制成图表。这里的索引0、1、2可能对应于 loss_names 中的损失名称。
self.plot_idx = [0, 1, 2]
# HUB
# 模型中心(HUB)会话。初始化了一个模型中心(HUB)会话对象,这里暂时设置为 None 。模型中心用于共享和下载预训练模型、权重等。
self.hub_session = None
# Callbacks
# 回调函数。设置了回调函数列表。如果提供了 _callbacks 参数,则使用提供的回调函数;如果没有提供,则调用 callbacks.get_default_callbacks() 获取默认的回调函数列表。
# def get_default_callbacks():
# -> 返回默认的回调函数列表。这行代码返回一个 defaultdict 对象,它使用 list 作为默认工厂函数,这意味着如果访问的键不存在,将会创建一个空列表作为该键的值。
# -> return defaultdict(list, deepcopy(default_callbacks))
self.callbacks = _callbacks or callbacks.get_default_callbacks()
# 添加集成回调。
# 这行代码检查 RANK 值,如果是-1(非分布式环境)或0(分布式环境中的主进程),则调用 callbacks.add_integration_callbacks(self) 方法来添加集成回调。这些回调包括日志记录、模型保存、评估等集成功能。
if RANK in {-1, 0}:
# def add_integration_callbacks(instance): -> 向一个实例添加集成的回调函数。这些回调函数通常用于训练过程中的不同阶段,例如记录训练日志、跟踪实验进度等。
callbacks.add_integration_callbacks(self)
# 这段代码的目的是为训练过程设置度量指标、结果记录、回调函数等,这些都是监控和管理训练过程的重要组成部分。通过这些设置,可以确保训练过程中的关键信息被适当记录和处理。
# BaseTrainer 类提供了一个框架,用于设置和初始化训练过程中的各种组件,包括数据加载、模型初始化、优化器设置、保存和恢复训练状态等。
# 这段代码定义了一个名为 add_callback 的方法,它是 BaseTrainer 类的一个成员方法。这个方法的目的是将一个回调函数添加到特定事件的回调列表中。
# 这是一个方法定义, self 参数代表类的实例本身。
# 1.event : str 是一个参数,它期望一个字符串类型的值,表示要添加回调的事件名称。
# 2.callback :是另一个参数,它期望一个函数或可调用对象,这个对象将在指定的事件发生时被调用。
def add_callback(self, event: str, callback):
# 附加给定的回调。
"""Appends the given callback."""
# 执行实际的回调添加操作。 self.callbacks 是一个字典,它的键是事件名称,值是对应事件的回调函数列表。
# event 是指定的事件名称, callback 是要添加的回调函数。
# append 方法用于将 callback 添加到与 event 关联的列表中。
self.callbacks[event].append(callback)
# 这个方法允许用户在训练过程中的特定事件(如每个epoch结束、模型保存等)添加自定义的回调函数。这些回调函数可以执行各种任务,如打印日志、调整学习率、保存模型等。通过这种方式,用户可以在不修改 BaseTrainer 类代码的情况下,灵活地扩展训练过程中的行为。
# 这段代码定义了一个名为 set_callback 的方法,它是 BaseTrainer 类的一个成员方法。这个方法用于设置特定事件的回调函数,覆盖任何先前为该事件设置的回调。
# 这是一个方法定义, self 参数代表类的实例本身。
# 1.event : str 是一个参数,它期望一个字符串类型的值,表示要设置回调的事件名称。
# 2.callback :是另一个参数,它期望一个函数或可调用对象,这个对象将在指定的事件发生时被调用。
def set_callback(self, event: str, callback):
# 使用给定的回调覆盖现有的回调。
"""Overrides the existing callbacks with the given callback."""
# 执行实际的回调设置操作。 self.callbacks 是一个字典,它的键是事件名称,值是对应事件的回调函数列表。
# event 是指定的事件名称, callback 是要设置的回调函数。
# 这行代码将与 event 关联的回调列表设置为只包含 callback 的新列表,从而覆盖任何之前为该事件设置的回调。
self.callbacks[event] = [callback]
# 这个方法允许用户为训练过程中的特定事件(如每个epoch结束、模型保存等)设置一个新的回调函数,替换掉任何之前为该事件设置的回调函数。这使得用户可以精确控制特定事件的行为,确保只有特定的回调函数被执行。通过这种方式,用户可以在不修改 BaseTrainer 类代码的情况下,灵活地定制训练过程中的行为。
# 这段代码定义了一个名为 run_callbacks 的方法,它是 BaseTrainer 类的一个成员方法。这个方法用于触发与特定事件关联的所有回调函数。
# 这是一个方法定义, self 参数代表类的实例本身。
# event : str 是一个参数,它期望一个字符串类型的值,表示要触发回调的事件名称。
def run_callbacks(self, event: str):
# 运行与特定事件相关的所有现有回调。
"""Run all existing callbacks associated with a particular event."""
# 从 self.callbacks 字典中获取与 event 事件关联的回调函数列表。 self.callbacks.get(event, []) 表示如果 event 不存在于字典中,则返回一个空列表,这样可以避免在没有注册回调函数的情况下出现错误。
for callback in self.callbacks.get(event, []):
# 对于获取到的回调列表中的每个回调函数,这行代码调用该函数并传递 self 作为参数。这意味着每个回调函数都会接收到 BaseTrainer 实例的引用,允许回调函数访问和修改训练器的状态或行为。
callback(self)
# 这个方法允许 BaseTrainer 类在特定的事件点执行所有注册的回调函数。这可以用于执行一系列自定义操作,例如在每个训练周期结束后记录日志、评估模型性能、保存模型状态等。通过这种方式,用户可以在训练过程中的不同阶段插入自定义逻辑,而无需修改 BaseTrainer 类的核心代码。
# 这段代码定义了一个名为 train 的方法,它是 BaseTrainer 类的一个成员方法。这个方法负责根据设备的配置启动训练过程。
# 这是 train 方法的定义,它没有参数,只接受 self ,代表类的实例本身。
def train(self):
# 允许多 GPU 系统上的 device=''、device=None 默认为 device=0。
"""Allow device='', device=None on Multi-GPU systems to default to device=0."""
# 检查 self.args.device 是否是一个非空字符串,如果是,它可能包含一个或多个GPU设备的编号。
if isinstance(self.args.device, str) and len(self.args.device): # i.e. device='0' or device='0,1,2,3'
# 如果 self.args.device 是一个字符串,通过逗号分隔设备编号来计算GPU的数量,并赋值给 world_size 。
world_size = len(self.args.device.split(","))
# 检查 self.args.device 是否是一个元组或列表,这通常用于指定多个GPU设备。
elif isinstance(self.args.device, (tuple, list)): # i.e. device=[0, 1, 2, 3] (multi-GPU from CLI is list)
# 如果 self.args.device 是一个元组或列表,计算其中的元素数量,并赋值给 world_size 。
world_size = len(self.args.device)
# 检查 self.args.device 是否是"cpu"或"mps",如果是,设置 world_size 为0,表示不使用GPU。
elif self.args.device in {"cpu", "mps"}: # i.e. device='cpu' or 'mps'
world_size = 0
# 检查是否有可用的CUDA设备(即GPU),如果有,但 self.args.device 是 None 或空字符串,设置 world_size 为1,表示默认使用设备0。
elif torch.cuda.is_available(): # i.e. device=None or device='' or device=number
world_size = 1 # default to device 0
# 如果以上条件都不满足,设置 world_size 为0,表示不使用GPU。
else: # i.e. device=None or device=''
world_size = 0
# 这段代码是 BaseTrainer 类中的 train 方法的一部分,它处理多GPU训练(使用分布式数据并行,DDP)的情况。
# Run subprocess if DDP training, else train normally
# 检查是否应该进行多GPU训练。 world_size 表示可用的GPU数量,如果大于1,并且环境变量中没有 LOCAL_RANK (这通常在分布式训练中设置,用于标识进程),则表示我们正在处理一个多GPU训练场景,但不是作为分布式训练的一部分(即不是由外部分布式训练框架启动的子进程)。
if world_size > 1 and "LOCAL_RANK" not in os.environ:
# Argument checks
# 检查 self.args.rect 是否为 True 。 rect 是一个配置参数,用于某种特定的训练模式。
if self.args.rect:
# 如果 rect=True ,记录一个警告,说明这种模式与多GPU训练不兼容,并通知用户 rect 将被设置为 False 。
LOGGER.warning("WARNING ⚠️ 'rect=True' is incompatible with Multi-GPU training, setting 'rect=False'") # 警告⚠️'rect=True' 与多 GPU 训练不兼容,请设置'rect=False'。
# 将 rect 参数设置为 False ,以确保多GPU训练不会因为这个参数而受到影响。
self.args.rect = False
# 检查 self.args.batch 参数是否小于1.0。 batch 代表批量大小,如果是小于1的值,表示使用了某种自动批量大小调整策略(AutoBatch)。
if self.args.batch < 1.0:
# 如果 batch 小于1,记录一个警告,说明AutoBatch与多GPU训练不兼容,并通知用户批量大小将被设置为默认值16。
LOGGER.warning(
"WARNING ⚠️ 'batch<1' for AutoBatch is incompatible with Multi-GPU training, setting "
"default 'batch=16'" # 警告⚠️ AutoBatch 的“batch<1”与多 GPU 训练不兼容,设置默认“batch=16”。
)
# 将批量大小设置为16,这是一个常见的默认批量大小,适用于大多数多GPU训练场景。
self.args.batch = 16
# 这段代码的主要目的是确保在多GPU训练环境中,某些特定的配置参数(如 rect 和 batch )被设置为兼容的值,以避免潜在的冲突或错误。通过这种方式, BaseTrainer 类可以更加健壮地处理不同的训练配置和环境。
# 这段代码继续处理 BaseTrainer 类中的 train 方法,特别是在多GPU训练场景下,使用分布式数据并行(DDP)的情况。
# Command
# 调用 generate_ddp_command 函数,传入 world_size (GPU数量)和 self (当前 BaseTrainer 实例)。这个函数负责生成用于启动分布式训练的命令行参数和临时文件路径,返回命令列表 cmd 和文件路径 file 。
cmd, file = generate_ddp_command(world_size, self)
# 开始一个 try 块,用于捕获在执行DDP命令时可能发生的任何异常。
try:
# 使用 LOGGER 记录器打印一条信息级别的日志,显示DDP训练的调试命令。 colorstr 函数用于给日志消息添加颜色,以便在终端中突出显示。
LOGGER.info(f'{colorstr("DDP:")} debug command {" ".join(cmd)}') # {colorstr("DDP:")} 调试命令 {" ".join(cmd)}。
# 使用 subprocess.run 执行生成的DDP命令。 cmd 是一个命令列表, check=True 参数表示如果命令执行失败(即返回非零退出状态),将抛出一个异常。
subprocess.run(cmd, check=True)
# 如果在执行DDP命令时发生任何异常, except 块将捕获这个异常,并将其重新抛出。这样可以让调用者知道训练过程中发生了错误。
except Exception as e:
raise e
# finally 块确保无论是否发生异常,都会执行 ddp_cleanup 函数。
finally:
# 调用 ddp_cleanup 函数进行清理工作,传入当前 BaseTrainer 实例和临时文件路径。这个函数负责删除临时文件、释放资源等。
ddp_cleanup(self, str(file))
# 如果不满足多GPU训练的条件(即 world_size 不大于1或环境变量中存在 LOCAL_RANK ),则执行 else 块中的代码。
else:
# 调用 _do_train 方法进行正常的单GPU或CPU训练。这个方法将处理训练循环,包括数据加载、模型前向传播、损失计算、反向传播和参数更新等。
self._do_train(world_size)
# 这段代码的主要目的是根据是否进行多GPU训练来决定执行DDP命令还是单GPU/CPU训练。在多GPU训练的情况下,它负责生成和执行DDP命令,并在训练结束后进行清理。在单GPU/CPU训练的情况下,它直接调用 _do_train 方法来执行训练过程。
# 这个方法的目的是根据不同的设备配置来启动训练过程,无论是单GPU、多GPU还是CPU。通过这种方式, BaseTrainer 类可以灵活地适应不同的训练环境。
# 这段代码定义了 BaseTrainer 类中的一个私有方法 _setup_scheduler ,它用于初始化训练过程中的学习率调度器。
# 这是 _setup_scheduler 方法的定义,它是一个实例方法,只接受 self 参数。
def _setup_scheduler(self):
# 初始化训练学习率调度程序。
"""Initialize training learning rate scheduler."""
# 检查 self.args.cos_lr 是否为 True ,这是一个配置参数,用于决定是否使用余弦退火学习率调度器。
if self.args.cos_lr:
# 如果使用余弦退火学习率调度器,创建一个 one_cycle 函数,它是一个周期性的学习率调整函数,参数包括初始学习率、最终学习率( self.args.lrf )和总训练周期数( self.epochs )。
# def one_cycle(y1=0.0, y2=1.0, steps=100):
# -> 生成并返回一个 lambda 函数,该 lambda 函数实现了一个从 y1 到 y2 的正弦波形(sinusoidal ramp)变化。返回一个 lambda 函数,该函数接受一个参数 x , x 表示当前的步骤。
# -> return lambda x: max((1 - math.cos(x * math.pi / steps)) / 2, 0) * (y2 - y1) + y1
self.lf = one_cycle(1, self.args.lrf, self.epochs) # cosine 1->hyp['lrf']
# 如果不使用余弦退火学习率调度器,执行 else 块中的代码。
else:
# 定义了一个匿名函数(lambda函数),用于线性学习率调度。这个函数根据当前的训练周期 x 计算学习率,其中 x 的范围是从0到 self.epochs 。
# 这个 lambda 函数实现了一个从初始学习率线性递减到最终学习率的学习率调度策略。这种策略在训练深度学习模型时很常见,可以帮助模型在训练初期快速收敛,然后随着训练的进行逐渐减小学习率,以细化模型的权重调整。
self.lf = lambda x: max(1 - x / self.epochs, 0) * (1.0 - self.args.lrf) + self.args.lrf # linear
# 创建一个 LambdaLR 学习率调度器实例,传入优化器( self.optimizer )和学习率lambda函数( self.lf )。 LambdaLR 调度器允许使用自定义的lambda函数来动态调整学习率。
# torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch=-1, verbose=False)
# torch.optim.lr_scheduler.LambdaLR 是 PyTorch 中的一个学习率调度器,它允许你根据一个给定的函数来调整学习率。这个函数可以自定义,以适应不同的训练需求和策略。
# 参数 :
# optimizer:被包装的优化器。
# lr_lambda:一个函数,它接受一个整数参数 epoch ,并返回一个乘法因子。这个因子将用于调整学习率。也可以是一个函数列表,其中每个函数对应于优化器中的一个参数组。
# last_epoch:整数,表示最后一个epoch的索引。默认为 -1 ,意味着从初始学习率开始。
# verbose:布尔值,如果为 True ,则在每次更新时打印一条消息到标准输出。默认为 False 。
# 工作原理 :
# LambdaLR 调度器将每个参数组的学习率设置为初始学习率乘以给定函数的值。这个函数可以是任何形式,例如,可以是一个简单的线性衰减、指数衰减,或者更复杂的自定义函数。
self.scheduler = optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=self.lf)
# 这个方法的主要目的是根据配置参数初始化学习率调度器,学习率调度器在训练过程中根据预定义的策略调整学习率,这可以帮助模型更好地收敛。余弦退火调度器是一种常见的学习率调整策略,它在训练周期中模拟余弦波形来逐渐减小学习率。如果不使用余弦退火,这里使用的是一个简单的线性衰减策略。
# 这段代码定义了 BaseTrainer 类中的一个私有方法 _setup_ddp ,它用于初始化分布式数据并行(Distributed Data Parallel,简称DDP)环境的参数。
# 这是 _setup_ddp 方法的定义,它接受两个参数 : self 和 world_size 。
# 1.world_size :表示参与分布式训练的进程总数。
def _setup_ddp(self, world_size):
# 初始化并设置用于训练的 DistributedDataParallel 参数。
"""Initializes and sets the DistributedDataParallel parameters for training."""
# 设置当前进程的CUDA设备为 RANK 指定的GPU。 RANK 是一个全局变量,用于标识当前进程在分布式训练中的编号。
torch.cuda.set_device(RANK)
# 将 self.device 设置为指定编号 RANK 的CUDA设备。
self.device = torch.device("cuda", RANK)
# LOGGER.info(f'DDP info: RANK {RANK}, WORLD_SIZE {world_size}, DEVICE {self.device}')
# 设置环境变量 TORCH_NCCL_BLOCKING_WAIT 为 1 ,这个环境变量用于控制NCCL(NVIDIA Collective Communications Library)的行为。设置为 1 表示NCCL操作将等待所有进程完成,否则可能会超时。
os.environ["TORCH_NCCL_BLOCKING_WAIT"] = "1" # set to enforce timeout
# torch.distributed.init_process_group(backend, init_method='env://', timeout=None, world_size=None, rank=None, store=None, group_name='', pg_options=None)
# torch.distributed.init_process_group 是 PyTorch 分布式训练中用于初始化进程组的函数。这个函数设置了多个进程间通信和同步的环境。
# 参数解释 :
# backend :指定用于进程间通信的后端。常用的后端包括 'nccl' (NVIDIA Collective Communication Library,适用于NVIDIA GPU环境), 'gloo' (适用于CPU或不支持NCCL的环境),和 'mpi' (需要MPI环境支持)。
# init_method :指定初始化进程组的方法。默认为 'env://' ,表示从环境变量中读取配置信息。其他选项包括 'file:///path/to/file' (通过文件系统), 'tcp://hostname:port' (通过TCP连接)等。
# timeout :设置进程组操作的超时时间。如果未指定,则使用默认值。
# world_size :指定参与分布式训练的总进程数。
# rank :指定当前进程的唯一标识符(排名),在 0 到 world_size-1 之间。
# store :一个用于存储和检索进程组信息的对象,用于初始化方法,但在大多数情况下不需要直接使用。
# group_name :为进程组指定一个名称,主要用于调试和日志记录。
# pg_options :一个可选的配置对象,用于设置进程组的特定选项。
# 功能描述 :
# init_process_group 函数会创建一个新的进程组,或者让当前进程加入到一个已存在的进程组中。这个进程组允许多个进程(可能分布在不同的机器上)进行通信和同步。在分布式训练中,所有进程都必须调用这个函数来初始化通信环境。
# 使用场景 :
# 在开始分布式训练之前,每个进程都需要调用 init_process_group 来设置通信环境。这个函数通常只在主进程(通常是 rank=0 的进程)中显式调用,其他进程在通过 torch.distributed.launch 工具启动时会自动调用。
# 注意事项 :
# 所有参与分布式训练的进程必须使用相同的 backend 和 init_method 参数。
# world_size 和 rank 必须在所有进程中保持一致。
# 在使用 init_process_group 之后,需要确保所有进程都正确地加入了进程组,才能进行后续的分布式操作。
# torch.distributed.is_nccl_available()
# torch.distributed.is_nccl_available() 是 PyTorch 分布式训练库中的一个函数,用于检查当前环境是否支持 NCCL(NVIDIA Collective Communications Library)后端。NCCL 是一个专门设计用于 NVIDIA GPU 的高性能通信库,能够提供高效的多GPU和多节点之间的通信。
# 返回值 :
# 返回一个布尔值,如果 NCCL 后端可用则返回 True ,否则返回 False 。
# 使用场景 :
# 这个函数通常在分布式训练设置中使用,以确定是否可以使用 NCCL 作为通信后端。如果 is_nccl_available() 返回 True ,则可以在调用 init_process_group 时指定 backend='nccl' 来初始化进程组。
# 注意事项 :
# NCCL 只能在 NVIDIA GPU 上使用,并且需要相应的驱动和 NCCL 库支持。
# 在某些系统上,可能需要设置环境变量 LD_LIBRARY_PATH 来指向 NCCL 库的安装路径,以便 PyTorch 能够找到并使用 NCCL。
# is_nccl_available() 的返回值也受到 PyTorch 版本和系统配置的影响。
# 初始化进程组,这是进行分布式训练的必要步骤。
dist.init_process_group(
# backend 参数指定了用于分布式通信的后端,如果可用,优先选择 nccl ,否则使用 gloo 。
backend="nccl" if dist.is_nccl_available() else "gloo",
# datetime.timedelta(days=0, seconds=0, microseconds=0, milliseconds=0, minutes=0, hours=0, weeks=0)
# 在Python中, timedelta 是 datetime 模块中的一个类,用于表示两个日期或时间之间的差值。
# 参数 :
# days : 整数,表示时间间隔的天数。
# seconds : 整数,表示时间间隔的秒数。
# microseconds : 整数,表示时间间隔的微秒数。
# milliseconds : 整数,表示时间间隔的毫秒数。
# minutes : 整数,表示时间间隔的分钟数。
# hours : 整数,表示时间间隔的小时数。
# weeks : 整数,表示时间间隔的周数。
# 功能描述 :
# timedelta 对象表示一个时间段,可以用于日期和时间的计算。它通常用于添加或减去某个时间段,或者在两个日期时间对象之间计算差值。
# timedelta 对象在 Python 的 datetime 模块中用于表示两个日期或时间之间的时间差。以下是一些常用的 timedelta 对象方法:
# + (加法) :
# 将两个 timedelta 对象相加,或者将 timedelta 对象加到一个 date 或 datetime 对象上。
# - (减法) :
# 从一个 timedelta 对象中减去另一个 timedelta 对象,或者从一个 date 或 datetime 对象中减去一个 timedelta 对象。
# total_seconds() :
# 返回 timedelta 对象表示的总秒数。
# days 、 seconds 、 microseconds :
# 这些属性分别返回 timedelta 对象的天数、秒数和微秒数。
# max 和 min :
# 这两个类方法返回 timedelta 对象可能的最大值和最小值。
# resolution :
# 这是一个类属性,表示 timedelta 对象能够表示的最小时间单位,通常是微秒。
# timedelta 对象通常与 datetime 对象一起使用,用于日期和时间的计算,例如计算两个日期之间的差异、添加或减去特定的时间段等。这些方法使得 timedelta 成为处理日期和时间计算的强大工具。
# timeout 参数设置了进程组操作的超时时间,这里设置为3小时。
timeout=timedelta(seconds=10800), # 3 hours
# rank 参数是当前进程的编号。world_size 是总进程数。
rank=RANK,
world_size=world_size,
)
# 这个方法的主要目的是为DDP训练设置必要的环境和参数,包括设置CUDA设备、初始化进程组等。这些步骤是进行分布式训练的基础,确保每个进程都能正确地与其他进程通信和同步。
# 这段代码定义了 BaseTrainer 类中的一个私有方法 _setup_train ,它负责在正确的进程上构建数据加载器和优化器,为训练做准备。
# 这是 _setup_train 方法的定义,它接受两个参数 : self 和 world_size 。
# 1.world_size :表示参与分布式训练的进程总数。
def _setup_train(self, world_size):
# 在正确的排序过程上构建数据加载器和优化器。
"""Builds dataloaders and optimizer on correct rank process."""
# Model
# 运行预训练开始的回调函数。
self.run_callbacks("on_pretrain_routine_start")
# 设置模型并返回检查点(checkpoint)信息。
ckpt = self.setup_model()
# 将模型移动到指定的设备(GPU或CPU)。
self.model = self.model.to(self.device)
# 设置模型属性。
self.set_model_attributes()
# 这段代码处理模型中某些层的冻结逻辑,即设置这些层的参数在训练过程中不被更新。
# Freeze layers
# 定义了一个变量 freeze_list ,它根据 self.args.freeze 的类型动态决定其值。
freeze_list = (
# 如果 self.args.freeze 是一个列表,则直接使用这个列表。
self.args.freeze
if isinstance(self.args.freeze, list)
# 如果是一个整数,则创建一个从0到该整数(不包括整数)的 range 对象
else range(self.args.freeze)
if isinstance(self.args.freeze, int)
# 否则,使用空列表。
else []
)
# 定义了一个列表 always_freeze_names ,包含了总是需要被冻结的层的名称后缀。
always_freeze_names = [".dfl"] # always freeze these layers
# 创建一个新的列表 freeze_layer_names ,它包含了所有需要被冻结的层的名称。这个列表由两部分组成。一部分是通过 freeze_list 生成的,另一部分是 always_freeze_names 中的层。
freeze_layer_names = [f"model.{x}." for x in freeze_list] + always_freeze_names
# 遍历模型的所有参数, k 是参数的名称, v 是参数的值。
for k, v in self.model.named_parameters():
# v.register_hook(lambda x: torch.nan_to_num(x)) # NaN to 0 (commented for erratic training results)
# 对于每个参数,检查其名称是否包含在 freeze_layer_names 列表中。
if any(x in k for x in freeze_layer_names):
# 如果参数名称包含在冻结列表中,记录一条信息,表示该层将被冻结。
LOGGER.info(f"Freezing layer '{k}'") # 冻结层‘{k}’。
# 设置参数 v 的 requires_grad 属性为 False ,这意味着在训练过程中,这个参数的梯度不会被计算,从而不会被更新。
v.requires_grad = False
# 如果参数 v 的 requires_grad 属性已经是 False ,并且参数类型是浮点数(只有浮点数参数可以有梯度),则执行以下操作。
elif not v.requires_grad and v.dtype.is_floating_point: # only floating point Tensor can require gradients
# 记录一条警告信息,表示某个本应被冻结的层的 requires_grad 属性被设置为 True 。这可能是由于用户自定义了冻结层的行为。
LOGGER.info(
f"WARNING ⚠️ setting 'requires_grad=True' for frozen layer '{k}'. " # 警告⚠️为冻结层“{k}”设置“requires_grad=True”。
"See ultralytics.engine.trainer for customization of frozen layers." # 请参阅 ultralytics.engine.trainer 以了解冻结层的定制。
)
# 将参数 v 的 requires_grad 属性设置为 True ,以确保它可以在训练中更新。
v.requires_grad = True
# 这段代码的目的是确保在训练过程中,某些特定的层不会被更新,这通常用于微调预训练模型时冻结某些层,或者在某些特定的训练策略中固定某些层的权重。通过设置 requires_grad 属性,可以控制哪些参数在反向传播时会计算梯度。
# 这段代码处理自动混合精度(Automatic Mixed Precision, AMP)的设置,这是一种用于加速训练并减少内存使用的技術,特别是在支持CUDA的GPU上。
# Check AMP
# 将命令行参数 self.args.amp (一个布尔值,指示是否启用AMP)转换为PyTorch张量,并将其移动到指定的设备(GPU或CPU)上。
self.amp = torch.tensor(self.args.amp).to(self.device) # True or False
# 检查是否启用了AMP,并且当前进程是单GPU环境( RANK == -1 )或分布式训练的主进程( RANK == 0 )。
if self.amp and RANK in {-1, 0}: # Single-GPU and DDP
# 在检查AMP支持性之前,备份当前的默认回调函数,因为 check_amp 函数可能会重置它们。
callbacks_backup = callbacks.default_callbacks.copy() # backup callbacks as check_amp() resets them
# 调用 check_amp 函数检查模型是否支持AMP,并将结果转换为张量,然后移动到指定的设备上。
# def check_amp(model): -> 用于检查在当前系统上是否支持 Automatic Mixed Precision (AMP) 并验证其准确性。如果没有异常发生,函数返回 True ,表示 AMP 检查通过。 -> return True
self.amp = torch.tensor(check_amp(self.model), device=self.device)
# 恢复之前备份的默认回调函数。
callbacks.default_callbacks = callbacks_backup # restore callbacks
# 检查是否处于分布式训练环境中( RANK 大于-1且 world_size 大于1)。
if RANK > -1 and world_size > 1: # DDP
# torch.distributed.broadcast(tensor, src, group=None, async_op=False)
# torch.distributed.broadcast 是 PyTorch 分布式通信包中的一个函数,用于将一个张量从指定的源进程(src)广播到所有其他进程。
# 参数 :
# tensor :要广播的张量。这个张量将从源进程复制到所有其他进程。
# src :源进程的排名(ID),它将张量广播给其他进程。
# group :(可选)指定通信组。如果不指定,则使用默认的全局通信组。通信组允许你创建不同的进程子集进行通信。
# async_op :(可选)布尔值,指定是否异步执行广播操作。如果设置为 True ,则返回一个 DistTensor 对象,该对象可以用于查询操作的完成状态。
# 功能描述 :
# torch.distributed.broadcast 函数在分布式环境中将一个张量从源进程广播到所有其他进程。在广播过程中,源进程发送张量,其他进程接收张量,并将其写入自己的内存中。这样,所有进程都能够获得相同的张量副本,从而可以在分布式训练或其他任务中进行协同计算。
# 使用场景 :
# 在分布式训练或其他分布式任务中,通常需要在所有进程上同时执行 torch.distributed.broadcast 函数,以确保所有进程都能够接收到相同的张量副本。这在初始化模型参数、同步全局变量等场景中非常有用。
# 从rank 0广播AMP标志到所有其他进程,确保所有进程都知道是否启用AMP。
dist.broadcast(self.amp, src=0) # broadcast the tensor from rank 0 to all other ranks (returns None)
# 将AMP标志转换为布尔值。
self.amp = bool(self.amp) # as boolean
# torch.cuda.amp.GradScaler(init_scale=2**16, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)
# torch.cuda.amp.GradScaler 是 PyTorch 中用于自动混合精度(AMP)训练的一个工具,它帮助动态缩放梯度以避免在混合精度训练时可能发生的数值溢出或下溢问题。
# 参数 :
# init_scale :初始缩放因子,用于缩放梯度。默认值为 2**16 ,即65536。这个值用于开始训练时的梯度缩放,可以根据实际情况调整。
# growth_factor :缩放因子增长因子。当梯度没有溢出时,缩放因子会按照这个因子增长。默认值为 2.0 。
# backoff_factor :缩放因子回退因子。当梯度溢出时,缩放因子会按照这个因子减少。默认值为 0.5 。
# growth_interval :增长间隔。在连续多少次训练迭代中没有遇到梯度溢出后,缩放因子会增长。默认值为 2000 。
# enabled :布尔值,指示是否启用 GradScaler 。如果设置为 False ,则 GradScaler 的调用将成为无操作。默认值为 True 。
# 功能描述 :
# GradScaler 的主要作用是 :
# 动态调整缩放因子,以避免在反向传播中因梯度值太大而导致的溢出。
# 如果在反向传播中检测到溢出,则跳过优化步骤并降低缩放因子。
# 自动管理精度,根据训练过程的动态变化调整缩放因子,确保模型的数值稳定性。
# 使用场景 :
# 在混合精度训练中, GradScaler 与 autocast 上下文管理器一起使用,以实现自动混合精度训练。 GradScaler 负责缩放损失和梯度,而 autocast 负责将部分计算切换到半精度。
# 根据AMP标志创建一个 GradScaler 实例,用于在AMP训练中调整梯度。
self.scaler = (
torch.amp.GradScaler("cuda", enabled=self.amp) if TORCH_2_4 else torch.cuda.amp.GradScaler(enabled=self.amp)
)
# 再次检查是否处于分布式训练环境中。
if world_size > 1:
# 如果是分布式训练环境,将模型包装为 DistributedDataParallel 实例,以便在多个GPU上并行训练。 device_ids=[RANK] 指定了当前进程应该使用的GPU编号, find_unused_parameters=True 允许在某些参数未在前向传播中使用时仍然更新它们。
self.model = nn.parallel.DistributedDataParallel(self.model, device_ids=[RANK], find_unused_parameters=True)
# 这段代码的目的是确保在训练过程中正确地设置和使用AMP,以及在分布式训练环境中正确地配置模型。通过使用AMP,可以提高训练效率并减少内存消耗,而 DistributedDataParallel 则允许模型在多个GPU上并行训练,进一步提高训练速度。
# 这段代码处理图像大小的检查和批量大小的调整,以确保训练的顺利进行。
# Check imgsz
# 计算网格大小( grid size ),即模型的最大步长(stride)。如果模型有 stride 属性,它取模型 stride 中的最大值并转换为整数,否则默认为32。然后,它取这个值和32之间的最大值,确保网格大小至少为32。
gs = max(int(self.model.stride.max() if hasattr(self.model, "stride") else 32), 32) # grid size (max stride)
# 调用 check_imgsz 函数来检查和调整图像大小。 self.args.imgsz 是原始图像大小, stride=gs 指定了步长, floor=gs 表示图像大小应该向下取整到最接近的步长的倍数, max_dim=1 表示最大维度为1。
# def check_imgsz(imgsz, stride=32, min_dim=1, max_dim=2, floor=0): -> 用于检查和调整图像尺寸( imgsz ),以确保它们符合特定的要求。返回调整后的图像尺寸。 -> return sz
self.args.imgsz = check_imgsz(self.args.imgsz, stride=gs, floor=gs, max_dim=1)
# 将计算出的网格大小赋值给 self.stride ,用于后续的多尺度训练。
self.stride = gs # for multiscale training
# Batch size
# 检查批量大小是否小于1,并且当前进程是单GPU环境( RANK == -1 )。
if self.batch_size < 1 and RANK == -1: # single-GPU only, estimate best batch size
# 如果满足上述条件,调用 check_train_batch_size 函数来估计最佳的批量大小。这个函数考虑了模型、图像大小、是否使用AMP以及当前的批量大小。然后,它更新 self.args.batch 和 self.batch_size 为计算出的最佳批量大小。
# def check_train_batch_size(model, imgsz=640, amp=True, batch=-1):
# -> 用于确定在给定图像尺寸和内存使用率下,模型训练时的最佳批量大小。返回值。函数返回 autobatch 函数计算出的最佳批量大小。
# -> return autobatch(deepcopy(model).train(), imgsz, fraction=batch if 0.0 < batch < 1.0 else 0.6)
self.args.batch = self.batch_size = check_train_batch_size(
model=self.model,
imgsz=self.args.imgsz,
amp=self.amp,
batch=self.batch_size,
)
# 这段代码的目的是确保在训练开始之前,图像大小和批量大小都被适当地设置和调整。图像大小的检查和调整有助于确保模型可以处理不同尺寸的输入,而批量大小的估计则有助于在单GPU上找到最佳的批量大小,以平衡训练速度和内存使用。
# 这段代码负责设置数据加载器(dataloaders),这对于训练和验证模型是必要的。
# Dataloaders
# 计算每个GPU上的批量大小。如果 world_size (参与训练的GPU数量)大于1,则将总批量大小 self.batch_size 除以 world_size ,否则保持 self.batch_size 不变。这是为了在分布式训练中平均分配数据到每个GPU。
batch_size = self.batch_size // max(world_size, 1)
# 创建训练数据加载器。 self.get_dataloader 是一个方法,它接受训练数据集 self.trainset 、批量大小 batch_size 、当前GPU的排名 LOCAL_RANK 和模式 "train" 。
self.train_loader = self.get_dataloader(self.trainset, batch_size=batch_size, rank=LOCAL_RANK, mode="train")
# 检查当前进程是否为主进程( RANK 为-1表示单GPU环境, RANK 为0表示分布式训练的主进程)。
if RANK in {-1, 0}:
# Note: When training DOTA dataset, double batch size could get OOM on images with >2000 objects.
# 创建验证数据加载器。如果任务是 "obb" (指 Oriented Bounding Box,方向边界框),则使用与训练相同的批量大小;否则,批量大小加倍。这是因为对于某些数据集,如DOTA,当图像中的对象超过2000个时,较大的批量大小可能会导致内存不足(OOM)。
self.test_loader = self.get_dataloader(
self.testset, batch_size=batch_size if self.args.task == "obb" else batch_size * 2, rank=-1, mode="val"
)
# 获取验证器,这是一个用于评估模型性能的对象。
self.validator = self.get_validator()
# 获取验证器的度量指标键,并将其与前缀为 "val" 的标签损失项合并。
metric_keys = self.validator.metrics.keys + self.label_loss_items(prefix="val")
# 初始化一个字典 self.metrics ,它将每个度量指标键映射到初始值0。
self.metrics = dict(zip(metric_keys, [0] * len(metric_keys)))
# 创建模型的指数移动平均(EMA)实例,这是一种技术,用于平滑模型参数,通常用于提高模型的泛化能力。
# class ModelEMA:
# -> 实现了模型指数移动平均(Exponential Moving Average, EMA)的功能。EMA是一种技术,用于在训练过程中平滑模型参数,通常可以提高模型的稳定性和性能。
# ->def __init__(self, model, decay=0.9999, tau=2000, updates=0):
self.ema = ModelEMA(self.model)
# 检查是否需要绘制训练过程中的图表。
if self.args.plots:
# 如果需要,调用方法绘制训练标签。
self.plot_training_labels()
# 这段代码的目的是为模型训练和验证准备必要的数据加载器,并初始化相关的度量指标和EMA。这些步骤是训练和评估模型性能的关键部分。通过调整批量大小和创建数据加载器,可以确保模型在训练和验证时能够高效地处理数据。
# 这段代码负责设置优化器,这是训练机器学习模型时用于调整模型参数的关键组件。
# Optimizer
# 计算累积梯度的次数。 self.args.nbs 是指定的累积批次大小, self.batch_size 是实际的批次大小。这个计算结果表示在执行一次优化步骤之前需要累积多少个批次的梯度。结果至少为1,意味着即使 self.args.nbs 小于 self.batch_size ,也会执行至少一次梯度累积。
self.accumulate = max(round(self.args.nbs / self.batch_size), 1) # accumulate loss before optimizing
# 计算权重衰减( weight decay ),这是一种正则化技术,用于防止过拟合。这里,权重衰减被缩放以考虑实际的批次大小和累积梯度的次数。
weight_decay = self.args.weight_decay * self.batch_size * self.accumulate / self.args.nbs # scale weight_decay
# math.ceil(x)
# math.ceil 是 Python 标准库 math 模块中的一个函数,用于返回大于或等于给定数字的最小整数。这个函数通常用于将浮点数向上舍入到最接近的整数。
# 参数 :
# x :要向上舍入的数字,可以是整数或浮点数。
# 返回值 :
# 返回大于或等于 x 的最小整数。
# math.ceil 函数在处理需要整数结果的数值计算时非常有用,例如在确定数组大小、分配内存或进行其他需要整数运算的场景中。
# 计算总的迭代次数,即整个训练过程中的批次数。 len(self.train_loader.dataset) 是训练数据集中的样本总数, max(self.batch_size, self.args.nbs) 确保使用的是较大的批次大小。结果乘以 self.epochs ,得到整个训练周期的迭代次数。
iterations = math.ceil(len(self.train_loader.dataset) / max(self.batch_size, self.args.nbs)) * self.epochs
# 创建优化器实例。 self.build_optimizer 是一个方法,它接受以下参数。这个方法将返回一个优化器对象,用于在训练过程中更新模型的参数。
self.optimizer = self.build_optimizer(
# model :要优化的模型。
model=self.model,
# name :优化器的名称,根据 self.args.optimizer 确定。
name=self.args.optimizer,
# lr :学习率,根据 self.args.lr0 确定。
lr=self.args.lr0,
# momentum :动量,根据 self.args.momentum 确定。
momentum=self.args.momentum,
# decay :权重衰减,使用上面计算的 weight_decay 。
decay=weight_decay,
# iterations :总迭代次数,使用上面计算的 iterations 。
iterations=iterations,
)
# 这段代码的目的是配置优化器,包括设置累积梯度的次数、调整权重衰减和计算总迭代次数。这些设置对于模型的训练动态和最终性能至关重要。通过合理配置优化器,可以确保模型以有效的方式学习,并且有助于模型收敛到较好的性能。
# 这段代码负责设置学习率调度器、早停机制,并在训练前进行一些准备工作。
# Scheduler
# 调用 _setup_scheduler 方法来初始化学习率调度器。这个方法根据之前的设置(例如,是否使用余弦退火学习率调度器)来配置学习率如何随着训练的进行而变化。
self._setup_scheduler()
# 创建一个早停(EarlyStopping)对象 self.stopper ,并设置其耐心值(patience),即在多少个epoch没有改进时停止训练。 self.args.patience 是从配置参数中获取的耐心值。 self.stop 是一个布尔值,用于跟踪训练是否应该停止。
# class EarlyStopping:
# -> 用于实现早停(early stopping)策略,这是一种在训练机器学习模型时用于避免过拟合的技术。
# -> def __init__(self, patience=50):
self.stopper, self.stop = EarlyStopping(patience=self.args.patience), False
# 如果提供了检查点(checkpoint) ckpt ,调用 resume_training 方法从检查点恢复训练。这包括加载模型权重、优化器状态和训练epoch的索引。
self.resume_training(ckpt)
# 设置学习率调度器的 last_epoch 属性,以便从 self.start_epoch - 1 开始。这样做是为了确保在第一个epoch开始时,学习率调度器不会更新学习率,这通常在第一个epoch之后发生。
self.scheduler.last_epoch = self.start_epoch - 1 # do not move
# 运行预训练例程结束时的回调函数。这些回调可能包括日志记录、模型检查点的保存或其他自定义行为。
self.run_callbacks("on_pretrain_routine_end")
# 这段代码的目的是完成训练前的最后准备工作,包括配置学习率调度器、设置早停机制、从检查点恢复(如果提供)以及运行预训练结束的回调。这些步骤有助于确保训练过程的顺利进行,并提供在训练过程中可能需要的灵活性和自定义功能。
# 这个方法的主要目的是为训练过程准备必要的组件,包括模型、数据加载器、优化器、学习率调度器等,并根据配置参数进行相应的设置和调整。
# 这段代码是 BaseTrainer 类中的 _do_train 方法,它负责执行模型的训练过程。
# 这是 _do_train 方法的定义,它接受两个参数 : self 和 world_size (默认为1)。
# 1.world_size :表示参与分布式训练的进程总数。
def _do_train(self, world_size=1):
# 训练完成,如果参数指定,则进行评估和绘图。
"""Train completed, evaluate and plot if specified by arguments."""
# 这段代码是 BaseTrainer 类中的 _do_train 方法的一部分,它负责初始化训练过程。
# 这个条件判断当前是否是多GPU训练环境(即 world_size 大于1)。
if world_size > 1:
# 如果是,则调用 _setup_ddp 方法来设置分布式数据并行(DDP)环境。
self._setup_ddp(world_size)
# 调用 _setup_train 方法来完成训练前的准备工作,包括构建数据加载器、优化器、学习率调度器等。
self._setup_train(world_size)
# 获取训练数据加载器中的批次总数,并将其存储在变量 nb 中。
nb = len(self.train_loader) # number of batches
# 计算预热(warmup)迭代次数。如果 self.args.warmup_epochs 大于0,则计算出预热周期数乘以批次总数,否则设置为-1。预热迭代次数至少为100,或者根据 self.args.warmup_epochs 的值动态计算。
nw = max(round(self.args.warmup_epochs * nb), 100) if self.args.warmup_epochs > 0 else -1 # warmup iterations
# 初始化一个变量 last_opt_step ,用于跟踪上一次执行优化器步骤的迭代次数。
last_opt_step = -1
# 初始化一个变量 epoch_time ,用于存储每个epoch的时间。
self.epoch_time = None
# 记录当前epoch开始的时间。
self.epoch_time_start = time.time()
# 记录整个训练过程开始的时间。
self.train_time_start = time.time()
# 运行训练开始的回调函数。
self.run_callbacks("on_train_start")
# 使用 LOGGER 记录训练的基本信息,包括图像大小、数据加载器工作进程数、结果日志保存目录以及训练周期或时间。
LOGGER.info(
f'Image sizes {self.args.imgsz} train, {self.args.imgsz} val\n' # 图像大小 {self.args.imgsz} train, {self.args.imgsz} val。
f'Using {self.train_loader.num_workers * (world_size or 1)} dataloader workers\n' # 使用 {self.train_loader.num_workers * (world_size 或 1)} 个数据加载器工作程序。
f"Logging results to {colorstr('bold', self.save_dir)}\n" # 将结果记录到 {colorstr('bold', self.save_dir)}。
f'Starting training for ' + (f"{self.args.time} hours..." if self.args.time else f"{self.epochs} epochs...") # 开始训练 ' + (f"{self.args.time} 小时..." if self.args.time else f"{self.epochs} epochs...
)
# 这段代码的主要目的是设置训练环境,包括分布式训练的配置、预热迭代次数的计算、时间记录的初始化以及训练开始的日志记录。这些步骤为后续的训练过程提供了必要的初始化和配置。
# 这段代码继续描述了 BaseTrainer 类中的 _do_train 方法,特别关注于训练周期的开始、优化器的梯度清零以及学习率调度器的步骤更新。
# 这个条件判断是否设置了关闭镶嵌(mosaic)数据增强的参数。如果 self.args.close_mosaic 为真,则执行以下操作。
if self.args.close_mosaic:
# 计算在训练周期结束前多少个批次开始关闭镶嵌数据增强。 base_idx 表示在总的训练周期中,从哪个批次开始不再使用mosaic数据增强。
base_idx = (self.epochs - self.args.close_mosaic) * nb
# 将计算出的 base_idx 及其后两个批次索引添加到 plot_idx 列表中。 plot_idx 用于指定在训练过程中哪些批次的图像需要被绘制或记录。
self.plot_idx.extend([base_idx, base_idx + 1, base_idx + 2])
# 设置当前的训练周期 epoch 为起始周期 self.start_epoch 。
epoch = self.start_epoch
# torch.optim.Optimizer.zero_grad(set_to_none=True)
# 在PyTorch中, zero_grad() 函数是 torch.optim.Optimizer 类中的一个方法,用于将所有优化器中的参数梯度清零。这是一个非常重要的步骤,因为它确保了在每次迭代或优化步骤之前,旧的梯度不会影响新的梯度计算。
# 参数 :
# set_to_none :一个布尔值,指定是否将梯度设置为 None 。默认值为 True 。如果设置为 True ,则梯度会被设置为 None ,这有助于减少内存消耗。如果设置为 False ,则梯度会被设置为零。
# 功能描述 :
# zero_grad() 方法遍历优化器中的所有参数,并将其梯度清零。这通常在每次迭代的开始或在反向传播之前进行,以确保旧的梯度值不会累积到新的梯度上。
# 使用场景 :
# 在训练神经网络时,每次进行前向传播和反向传播之前,都需要清零梯度,以避免梯度累积。这可以通过直接调用优化器实例的 zero_grad() 方法来实现。
# 清空优化器中的梯度。这是为了确保在训练开始时,任何从先前训练周期恢复的梯度都被清零,从而保证训练的稳定性。
self.optimizer.zero_grad() # zero any resumed gradients to ensure stability on train start
# 开始一个无限循环,用于执行训练周期,直到满足某个条件(如达到最大训练周期数或满足早停条件)。
while True:
# 在每个训练周期开始时,更新当前的训练周期 epoch 。
self.epoch = epoch
# 运行每个训练周期开始的回调函数。
self.run_callbacks("on_train_epoch_start")
# 使用 warnings 模块的上下文管理器来捕获和处理警告。
with warnings.catch_warnings():
# 设置警告过滤器,忽略特定类型的警告。在这里,忽略的是关于 lr_scheduler.step() 在 optimizer.step() 之前被调用的警告。
warnings.simplefilter("ignore") # suppress 'Detected lr_scheduler.step() before optimizer.step()'
# torch.optim.Optimizer.step()
# 在 PyTorch 中, step() 函数是优化器(Optimizer)类的一个方法,用于执行单步参数更新。这个方法会根据优化器内部的参数和学习率调度器来更新模型的参数。
# 参数说明 :无参数。
# 返回值 :无返回值。
# step() 方法通常在每次迭代后调用,它会对所有之前通过 zero_grad() 方法清零的参数进行梯度更新。在调用 step() 之前,需要确保已经计算了模型的梯度(通过反向传播)并且清零了之前的梯度。
# 需要注意的是, step() 方法不应该在没有梯度的情况下被调用,因为这可能会导致错误或未定义的行为。
# 此外,如果你使用了学习率调度器(LR Scheduler),通常在 step() 调用之后立即更新学习率 :
# scheduler.step() # 在每个 epoch 结束时更新学习率。
# 这里的 scheduler.step() 会根据学习率调度器的策略更新学习率,为下一个 epoch 的训练做好准备。
# 更新学习率调度器。这个步骤通常在每个周期开始或结束时调用,以根据预设的策略调整学习率。
self.scheduler.step()
# 这段代码的主要目的是准备训练环境,包括关闭数据增强的设置、梯度清零以及学习率调度器的更新,并进入主训练循环。这些步骤是训练过程中的标准初始化和准备工作,确保训练可以顺利进行。
# 这段代码是 BaseTrainer 类中的 _do_train 方法的一部分,它负责在每个训练周期开始时设置模型为训练模式,并可能更新数据加载器的属性。
# 将模型设置为训练模式。这是必要的,因为模型在训练和评估时的行为可能不同(例如,Dropout和Batch Normalization层在训练和评估时的行为不同)。
self.model.train()
# 这个条件判断是否处于分布式训练环境中。 RANK 是一个全局变量,用于标识当前进程在分布式训练中的编号。如果 RANK 不等于-1(即不是单GPU训练),则执行以下操作。
if RANK != -1:
# 如果数据加载器的采样器(sampler)支持设置epoch,那么在每个训练周期开始时更新采样器的epoch。这对于某些类型的采样器(如分布式采样器)是必要的,因为它们可能需要根据当前的训练周期来改变采样行为。
self.train_loader.sampler.set_epoch(epoch)
# 创建一个枚举器 pbar ,用于遍历训练数据加载器 self.train_loader 中的批次。这允许在训练循环中同时获取批次索引和数据。
pbar = enumerate(self.train_loader)
# Update dataloader attributes (optional)
# 这个条件判断是否到达了关闭镶嵌(mosaic)数据增强的周期。 self.args.close_mosaic 是一个参数,指定在训练的最后几个周期关闭mosaic数据增强。
if epoch == (self.epochs - self.args.close_mosaic):
# 如果到达了关闭mosaic数据增强的周期,调用 _close_dataloader_mosaic 方法来执行关闭mosaic数据增强所需的操作。
self._close_dataloader_mosaic()
# 重置训练数据加载器,这可能包括重置采样器状态或重新洗牌数据集,以确保在新的训练周期中数据的顺序和采样行为符合预期。
self.train_loader.reset()
# 这段代码的主要目的是在每个训练周期开始时确保模型处于正确的模式,并根据训练配置更新数据加载器的状态。这对于维持训练过程的一致性和有效性至关重要。
# 这段代码是 BaseTrainer 类中的 _do_train 方法的一部分,它负责在训练过程中处理日志记录、进度条显示、批次处理、预热期的学习率和动量调整。
# 这个条件判断当前进程是否是主进程(在单GPU或分布式训练的主节点中, RANK 为-1或0)。
if RANK in {-1, 0}:
# 如果是主进程,使用 LOGGER 记录当前的训练进度字符串。 progress_string 方法可能返回一个包含当前epoch、损失、学习率等信息的字符串。
LOGGER.info(self.progress_string())
# 创建一个 TQDM 进度条,用于显示训练过程中的批次处理进度。 enumerate(self.train_loader) 枚举训练数据加载器中的批次, total=nb 指定进度条的总批次数。
pbar = TQDM(enumerate(self.train_loader), total=nb)
# 初始化总损失 self.tloss 为 None ,这将在后续的批次处理中更新为累计损失。
self.tloss = None
# 遍历训练数据加载器中的每个批次。
for i, batch in pbar:
# 在处理每个批次之前运行回调函数。
self.run_callbacks("on_train_batch_start")
# Warmup
# 计算当前迭代的编号,包括当前批次索引 i 和当前周期的批次总数 nb * epoch 。
ni = i + nb * epoch
# 如果当前迭代编号小于或等于预热期的迭代次数 nw ,则执行以下操作。
if ni <= nw:
# 设置线性插值的x轴范围,从0到 nw 。
xi = [0, nw] # x interp
# 计算累积梯度的次数,使用 np.interp 在1和 self.args.nbs / self.batch_size 之间进行线性插值。
self.accumulate = max(1, int(np.interp(ni, xi, [1, self.args.nbs / self.batch_size]).round()))
# 遍历优化器中的参数组。
for j, x in enumerate(self.optimizer.param_groups):
# Bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
# np.interp(x, xp, fp, left=None, right=None, period=None)
# np.interp 是 NumPy 库中的一个函数,用于一维线性插值。给定一组数据点 x 和相应的值 xp ,以及一个新的查询点 x , np.interp 函数会找到 xp 中 x 值所在的区间,并使用线性插值来估计 x 对应的值。
# 参数说明 :
# x :查询点,即你想要插值的点。
# xp :数据点,一个一维数组,包含数据点的横坐标。
# fp : xp 对应的值,一个一维数组,包含数据点的纵坐标。
# left :可选参数,如果 x 中的值小于 xp 中的最小值,则使用这个值作为插值结果。
# right :可选参数,如果 x 中的值大于 xp 中的最大值,则使用这个值作为插值结果。
# period :可选参数,表示周期性,如果指定, xp 将被视为周期性的。
# 返回值 :
# 插值结果,一个与 x 形状相同的数组。
# 对于每个参数组,根据当前迭代编号 ni 和预热期的范围 xi ,使用线性插值计算学习率。对于第一个参数组(通常是偏置),学习率从 self.args.warmup_bias_lr 开始;对于其他参数组,学习率从0.0开始。
x["lr"] = np.interp(
ni, xi, [self.args.warmup_bias_lr if j == 0 else 0.0, x["initial_lr"] * self.lf(epoch)]
)
# 如果参数组中包含动量项。
if "momentum" in x:
# 使用线性插值计算动量值,从 self.args.warmup_momentum 开始,逐渐增加到 self.args.momentum 。
x["momentum"] = np.interp(ni, xi, [self.args.warmup_momentum, self.args.momentum])
# 这段代码的主要目的是在训练过程中处理批次级别的操作,包括日志记录、进度条更新、预热期的学习率和动量调整。这些步骤有助于确保训练过程的稳定性和效率,特别是在预热期,通过逐渐调整学习率和动量可以帮助模型更平滑地开始训练。
# 这段代码是 BaseTrainer 类中的 _do_train 方法的一部分,它负责执行模型的前向传播和反向传播。
# Forward
# 前向传播。
# 使用PyTorch的 autocast 上下文管理器,它会自动将模型的前向传播转换为混合精度(FP16和FP32)计算。 self.amp 是 GradScaler 实例,用于自动混合精度训练。
with autocast(self.amp):
# 对批次数据进行预处理。 self.preprocess_batch 方法包括数据增强、归一化等操作,以准备输入模型的数据。
batch = self.preprocess_batch(batch)
# 将预处理后的数据 batch 输入模型,计算损失。模型返回 损失值 和 损失项的详细列表 。
self.loss, self.loss_items = self.model(batch)
# 如果不是单GPU训练(即在分布式训练环境中),则将损失值乘以 world_size 。这是为了在反向传播时考虑到所有GPU上的数据,确保梯度的正确缩放。
if RANK != -1:
self.loss *= world_size
# 更新累计损失 self.tloss 。如果 self.tloss 不是 None ,则累加当前批次的损失并计算平均值;如果是 None ,则直接设置 self.tloss 为当前批次的损失项。
self.tloss = (
(self.tloss * i + self.loss_items) / (i + 1) if self.tloss is not None else self.loss_items
)
# Backward
# 反向传播。使用 GradScaler 的 scale 方法对损失进行缩放,以防止在混合精度训练中梯度下溢。然后调用 backward() 方法进行反向传播,计算梯度。
self.scaler.scale(self.loss).backward()
# 这段代码的主要目的是在每个训练批次中执行模型的前向传播和反向传播。前向传播计算损失,而反向传播计算参数的梯度。使用 autocast 和 GradScaler 可以提高训练的效率和稳定性,特别是在使用混合精度训练时。
# 这段代码是 BaseTrainer 类中的 _do_train 方法的一部分,它负责在训练过程中确定何时执行优化器步骤(即更新模型参数)以及根据时间条件停止训练。
# Optimize - https://pytorch.org/docs/master/notes/amp_examples.html
# 这个条件判断自从上次优化器步骤以来是否已经累积了足够的梯度。 ni 是当前迭代编号, last_opt_step 是上次优化器步骤的迭代编号, self.accumulate 是累积梯度的次数。如果条件满足,则执行优化器步骤。
if ni - last_opt_step >= self.accumulate:
# 调用 optimizer_step 方法来执行优化器步骤,更新模型参数。这个方法通常包括调用 GradScaler 的 step 方法来处理缩放后的梯度,并更新优化器的状态。
self.optimizer_step()
# 更新 last_opt_step 为当前迭代编号 ni ,标记最后一次优化器步骤的位置。
last_opt_step = ni
# Timed stopping
# 检查是否设置了训练时间限制。 self.args.time 是一个参数,表示训练的最大时间(以小时为单位)。
if self.args.time:
# 计算从训练开始到现在的时间是否超过了设定的时间限制。如果是,设置 self.stop 为 True ,表示训练应该停止。
self.stop = (time.time() - self.train_time_start) > (self.args.time * 3600)
# 这个条件判断是否处于分布式训练环境中。 RANK 是一个全局变量,用于标识当前进程在分布式训练中的编号。如果 RANK 不等于-1(即不是单GPU训练),则执行以下操作。
if RANK != -1: # if DDP training
# 创建一个列表,用于存储停止条件。如果当前进程是主进程( RANK 为0),则将 self.stop 的值放入列表;否则,放入 None 。
broadcast_list = [self.stop if RANK == 0 else None]
# torch.distributed.broadcast_object_list(object_list, src=0, group=None, device=None)
# torch.distributed.broadcast_object_list 是 PyTorch 分布式通信包中的一个函数,用于将一个对象列表从一个进程(源进程)广播到所有其他进程。
# 参数 :
# object_list :要广播的输入对象列表。每个对象都必须是可pickle的(即可以通过Python的 pickle 模块进行序列化和反序列化)。
# src :源进程的排名,默认为0。只有 src 进程上的对象会被广播,但每个进程都必须提供大小相等的列表。
# group :(可选)要处理的进程组。如果没有指定,则使用默认进程组。默认为 None 。
# device :(可选)如果指定了 device ,则将对象序列化并转换为张量,在广播之前将其移动到 device 上。默认为 None 。
# 返回值 :该函数没有返回值。如果进程是组的一部分, object_list 将包含来自 src 进程广播的对象。
# 功能描述 :
# torch.distributed.broadcast_object_list 将 object_list 中的可pickle对象广播给整个进程组。这与 broadcast 函数类似,但可以传入Python对象。注意, object_list 中的所有对象都必须是可pickle的才能进行广播。
# 注意事项 :
# 对于基于NCCL的进程组,对象的内部张量表示必须在通信发生之前移动到GPU设备上。在这种情况下,使用的设备由 torch.cuda.current_device() 给出,用户有责任通过 torch.cuda.set_device() 确保将其设置为每个进程都有单独的GPU。
# 请注意,此API与 all_gather() 略有不同,因为它不提供 async_op 句柄,因此将是阻塞调用。
# broadcast_object_list() 隐式使用 pickle 模块,已知这是不安全的。可以构造恶意的pickle数据,该数据将在unpickling期间执行任意代码。仅使用您信任的数据调用此函数。
# 使用 dist.broadcast_object_list 方法将停止条件从主进程( RANK 为0)广播到所有其他进程。
dist.broadcast_object_list(broadcast_list, 0) # broadcast 'stop' to all ranks
# 更新 self.stop 为广播后的值,确保所有进程的停止条件一致。
self.stop = broadcast_list[0]
# 如果 self.stop 为 True ,表示训练时间已超过设定限制。
if self.stop: # training time exceeded
# 如果训练时间已超过设定限制,则跳出训练循环,结束训练。
break
# 这段代码的主要目的是在训练过程中累积足够的梯度后更新模型参数,并根据时间条件判断是否停止训练。在分布式训练环境中,它还确保所有进程的停止条件一致。
# 在训练过程中累积足够的梯度后更新模型参数的做法,通常称为梯度累积(Gradient Accumulation),这是一种用于训练深度学习模型的技术。以下是梯度累积的几个主要原因和好处 :
# 更大的有效批量大小 :
# 在某些情况下,由于硬件限制(如GPU内存不足),无法一次性处理非常大的批量大小。梯度累积允许模型在较小的批量上进行训练,但通过累积梯度来模拟更大批量的效果,从而可以在不增加硬件负担的情况下训练更深或更宽的模型。
# 内存和计算效率 :
# 对于大型模型或复杂的数据集,一次性处理大批量数据可能会导致内存溢出。梯度累积可以在不牺牲模型性能的情况下,通过小批量训练减少内存消耗。
# 训练稳定性 :
# 梯度累积可以帮助稳定训练过程,因为它减少了每次参数更新时的梯度方差。这对于使用较小批量大小训练时特别有用,因为小批量可能会导致梯度估计的高方差。
# 模拟更大批量的效果 :
# 研究表明,使用较大的批量大小可以提高模型的泛化能力。通过梯度累积,可以在不直接使用大批量的情况下,模拟大批量训练的效果。
# 灵活的训练策略 :
# 梯度累积提供了一种灵活的训练策略,允许研究人员和工程师根据具体的硬件限制和模型需求调整训练过程。
# 减少训练时间 :
# 对于需要大量计算资源的任务,梯度累积可以减少训练时间,因为它允许使用更大的有效批量大小,从而减少达到一定迭代次数所需的训练周期数。
# 适应不同的硬件配置 :
# 不同的硬件配置可能对批量大小有不同的限制。梯度累积使得训练过程可以适应不同的硬件配置,而不需要对模型架构或训练策略进行重大更改。
# 总之,梯度累积是一种实用的技术,可以在不牺牲模型性能的情况下,适应不同的训练需求和硬件限制。通过累积梯度,可以在较小的批量上进行训练,同时模拟较大批量的效果,从而提高训练效率和模型性能。
# 在训练过程中实现梯度积累通常涉及以下步骤 :
# 设置累积次数 :确定梯度积累的次数,即在执行一次优化器步骤( optimizer.step() )之前要累积多少个批次的梯度。
# 禁用每个批次的优化器更新 :在每个批次的末尾,而不是立即执行 optimizer.step() ,需要禁用自动梯度更新。这可以通过在PyTorch中使用 torch.no_grad() 上下文管理器或简单地不调用 optimizer.step() 来实现。
# 累积梯度 :在每个批次的末尾,执行 loss.backward() 来计算梯度,但不立即更新参数。这会导致梯度在模型的参数上累积。
# 检查累积次数 :在每个批次结束后,检查是否已经累积了足够的梯度。这通常是通过跟踪累积的批次数来完成的。
# 执行优化器步骤 :一旦累积了足够数量的梯度,执行一次 optimizer.step() 来更新模型参数,然后重置梯度。
# 可选的梯度裁剪 :在执行优化器步骤之前,可以使用梯度裁剪( torch.nn.utils.clip_grad_norm_ 或 torch.nn.utils.clip_grad_value_ )来防止梯度爆炸。
# 以下是使用PyTorch实现梯度积累的示例代码:
# from torch.cuda.amp import autocast, GradScaler
# # 假设 model, optimizer, train_loader, epochs, accumulate_steps 已经定义
# scaler = GradScaler() # 如果使用自动混合精度,则需要
# accumulate_steps = 4 # 累积梯度的批次数
# for epoch in range(epochs):
# for batch_idx, (inputs, targets) in enumerate(train_loader):
# with autocast():
# outputs = model(inputs)
# loss = criterion(outputs, targets)
# scaler.scale(loss).backward() # 累积梯度
# scaler.unscale_(optimizer) # 可选的,如果使用GradScaler
# if (batch_idx + 1) % accumulate_steps == 0 or batch_idx == len(train_loader) - 1:
# scaler.step(optimizer) # 更新参数
# scaler.update() # 更新缩放器
# optimizer.zero_grad() # 重置梯度
# 在这个示例中, accumulate_steps 定义了在执行优化器步骤之前要累积多少个批次的梯度。 scaler.step(optimizer) 在累积足够的梯度后更新模型参数, scaler.update() 更新GradScaler的状态, optimizer.zero_grad() 重置梯度,为下一次梯度累积做准备。
# 梯度积累是一种灵活的技术,可以根据具体的训练需求和硬件限制进行调整。
# 这段代码是 BaseTrainer 类中的 _do_train 方法的一部分,它负责在训练过程中进行日志记录、进度更新、回调执行以及模型参数的指数移动平均(EMA)更新。
# Log
# 日志记录和进度更新。
# 这个条件判断当前进程是否是主进程(在单GPU或分布式训练的主节点中, RANK 为-1或0)。
if RANK in {-1, 0}:
# 计算总损失 self.tloss 的维度长度,如果 self.tloss 是一个多维张量,则取第一个维度的长度;如果是标量,则长度为1。
loss_length = self.tloss.shape[0] if len(self.tloss.shape) else 1
# 更新进度条描述,显示 当前epoch 、 GPU内存使用情况 、 损失值 、 批次大小 和 图像大小 。 % 操作符用于格式化字符串。
pbar.set_description(
("%11s" * 2 + "%11.4g" * (2 + loss_length))
% (
f"{epoch + 1}/{self.epochs}",
f"{self._get_memory():.3g}G", # (GB) GPU memory util
*(self.tloss if loss_length > 1 else torch.unsqueeze(self.tloss, 0)), # losses
batch["cls"].shape[0], # batch size, i.e. 8
batch["img"].shape[-1], # imgsz, i.e 640
)
)
# 在每个批次结束时运行回调函数。
self.run_callbacks("on_batch_end")
# 如果设置了绘图参数 self.args.plots 为真,并且当前迭代编号 ni 在绘图索引列表 self.plot_idx 中,则执行绘图操作。
if self.args.plots and ni in self.plot_idx:
# 绘制训练样本,用于可视化训练过程中的图像或特征图。
self.plot_training_samples(batch, ni)
# 在每个训练批次结束时运行回调函数。
self.run_callbacks("on_train_batch_end")
# 学习率记录和EMA更新。
# 记录每个参数组的 学习率 ,以便后续的日志记录或分析。
self.lr = {f"lr/pg{ir}": x["lr"] for ir, x in enumerate(self.optimizer.param_groups)} # for loggers
# 在每个训练周期结束时运行回调函数。
self.run_callbacks("on_train_epoch_end")
# 再次检查当前进程是否是主进程。
if RANK in {-1, 0}:
# 判断是否是最后一个训练周期。
final_epoch = epoch + 1 >= self.epochs
# 更新模型的指数移动平均(EMA)。EMA是一种技术,用于计算模型参数的移动平均值,这通常有助于提高模型的稳定性和性能。
# def update_attr(self, model, include=(), exclude=("process_group", "reducer")): -> ModelEMA 类中的 update_attr 方法,它用于将模型的某些属性复制到 EMA 模型中,并保存一个去除优化器的精简模型。
self.ema.update_attr(self.model, include=["yaml", "nc", "args", "names", "stride", "class_weights"])
# 这段代码的主要目的是在训练过程中提供详细的日志记录、进度更新和模型参数的EMA更新。这些步骤有助于监控训练过程、调试问题以及最终模型的性能评估。通过这种方式, _do_train 方法提供了一个完整的训练循环,可以根据不同的训练需求进行定制和扩展。
# 这段代码是 BaseTrainer 类中的 _do_train 方法的一部分,它负责在训练过程中执行验证(validation)和根据验证结果以及训练时间限制来决定是否停止训练。
# Validation
# 验证执行。
# 这个条件判断是否需要执行验证。这可能发生在以下几种情况 :
# self.args.val 为真,表示根据配置需要定期执行验证。
# final_epoch 为真,表示当前是最后一个训练周期。
# self.stopper.possible_stop 为真,表示早停机制可能需要执行验证来决定是否停止训练。
# self.stop 为真,表示训练可能需要停止。
if self.args.val or final_epoch or self.stopper.possible_stop or self.stop:
# 调用 validate 方法执行验证,并获取验证指标 metrics 和适应度 fitness 。适应度通常是验证集上的性能指标,用于评估模型的泛化能力。
self.metrics, self.fitness = self.validate()
# 指标和模型保存。
# 保存训练和验证指标。这包括 损失项 、 验证指标 和 学习率 。
self.save_metrics(metrics={**self.label_loss_items(self.tloss), **self.metrics, **self.lr})
# 早停和时间停止。
# 调用早停机制的 stopper 方法,传入当前周期和适应度,以决定是否需要停止训练。如果早停机制建议停止或当前是最后一个周期,则设置 self.stop 为真。
self.stop |= self.stopper(epoch + 1, self.fitness) or final_epoch
# 检查是否设置了训练时间限制。
if self.args.time:
# 如果设置了训练时间限制,计算从训练开始到现在的时间是否超过了这个限制。如果超过了,设置 self.stop 为真。
self.stop |= (time.time() - self.train_time_start) > (self.args.time * 3600)
# 这段代码的主要目的是在训练过程中定期执行验证,根据验证结果和训练时间限制来决定是否停止训练。通过这种方式,可以确保模型在达到预定的性能目标或时间限制后停止训练,从而避免不必要的计算资源浪费,并确保模型不会过拟合。
# 这段代码是 BaseTrainer 类中的 _do_train 方法的一部分,它负责在训练过程中保存模型。
# Save model
# 这个条件判断是否需要保存模型。这可能发生在以下几种情况 : self.args.save 为真,表示根据配置需要定期保存模型。 final_epoch 为真,表示当前是最后一个训练周期,需要保存最终的模型状态。
if self.args.save or final_epoch:
# 调用 save_model 方法来保存模型的当前状态。这个方法可能会将模型的参数、优化器状态和其他训练相关的状态保存到文件中,通常是一个 .pt 或 .pth 文件。
self.save_model()
# 运行模型保存后的回调函数。这些回调可能包括日志记录、验证模型的完整性、清理临时文件等操作。
self.run_callbacks("on_model_save")
# 这段代码的主要目的是确保在训练过程中或训练结束时保存模型的状态,以便后续可以恢复训练或进行模型评估。通过这种方式,可以避免因意外中断或其他原因导致的训练进度丢失,并确保模型的最佳状态被保留下来。
# 这段代码是 BaseTrainer 类中的 _do_train 方法的一部分,它负责管理学习率调度器、计算训练周期时间、调整训练周期数以及清理内存。
# Scheduler
# 学习率调度器和时间管理。
# 获取当前时间(以秒为单位)。
t = time.time()
# 计算当前训练周期(epoch)的持续时间,并将其存储在 self.epoch_time 中。
self.epoch_time = t - self.epoch_time_start
# 更新 self.epoch_time_start 为当前时间,为下一个周期的时间计算做准备。
self.epoch_time_start = t
# 检查是否设置了训练时间限制。
if self.args.time:
# 计算从训练开始到当前周期的平均周期时间。
mean_epoch_time = (t - self.train_time_start) / (epoch - self.start_epoch + 1)
# 如果设置了训练时间限制,重新计算 基于当前平均周期时间的训练周期数 ,并将其设置为 self.epochs 和 self.args.epochs 。这确保了训练不会超过预定的时间限制。
self.epochs = self.args.epochs = math.ceil(self.args.time * 3600 / mean_epoch_time)
# 调用 _setup_scheduler 方法重新设置学习率调度器,这涉及到根据新的周期数调整学习率计划。
self._setup_scheduler()
# 更新学习率调度器的 last_epoch 属性,以确保学习率调整与当前周期保持一致。
self.scheduler.last_epoch = self.epoch # do not move
# 如果当前周期数达到或超过了调整后的训练周期数,设置 self.stop 为真,以停止训练。
self.stop |= epoch >= self.epochs # stop if exceeded epochs
# 回调和内存管理。
# 在每个训练周期结束时运行回调函数。
self.run_callbacks("on_fit_epoch_end")
# 清理内存,可能包括释放不再需要的变量或张量,以减少内存消耗。
self._clear_memory()
# 这段代码的主要目的是在训练过程中管理学习率调度器、监控训练时间,并在必要时调整训练周期数。通过这种方式,可以确保训练过程既高效又符合预定的时间限制。同时,通过清理内存,可以避免潜在的内存泄漏问题,确保训练过程的稳定性。
# 这段代码是 BaseTrainer 类中的 _do_train 方法的一部分,它负责处理早停(Early Stopping)逻辑,特别是在分布式数据并行(DDP)训练环境中。
# Early Stopping
# 早停逻辑。
# 这个条件判断当前是否处于分布式训练环境中。 RANK 是一个全局变量,用于标识当前进程在分布式训练中的编号。如果 RANK 不等于-1(即不是单GPU训练),则执行以下操作。
if RANK != -1: # if DDP training
# 创建一个列表 broadcast_list ,用于存储早停信号。如果当前进程是主进程( RANK 为0),则将 self.stop 的值放入列表;否则,放入 None 。
broadcast_list = [self.stop if RANK == 0 else None]
# 使用 dist.broadcast_object_list 方法将早停信号从主进程( RANK 为0)广播到所有其他进程。这样,所有进程都会接收到是否需要停止训练的信号。
dist.broadcast_object_list(broadcast_list, 0) # broadcast 'stop' to all ranks
# 更新 self.stop 为广播后的值,确保所有进程的早停条件一致。
self.stop = broadcast_list[0]
# 如果 self.stop 为真,表示训练应该停止。
if self.stop:
# 如果训练需要停止,跳出训练循环,结束训练。在分布式训练环境中,这将确保所有进程都会停止训练。
break # must break all DDP ranks
# 周期递增。如果训练没有停止,增加训练周期数,以便开始下一个训练周期。
epoch += 1
# 这段代码的主要目的是在分布式训练环境中同步所有进程的早停决策。通过广播早停信号,可以确保在任何一个进程决定停止训练时,所有进程都会停止,从而保持训练过程的一致性。这种方法特别重要,因为在分布式训练中,每个进程可能独立地满足早停条件,需要同步这些决策以避免不一致的状态。
# 这段代码是 BaseTrainer 类中的 _do_train 方法的最后一部分,它负责在训练结束后执行最终的验证、日志记录、绘图、回调执行以及内存清理。
# 最终验证和日志记录。
# 这个条件判断当前进程是否是主进程(在单GPU或分布式训练的主节点中, RANK 为-1或0)。
if RANK in {-1, 0}:
# Do final val with best.pt
# 使用 LOGGER 记录训练过程中的重要信息,包括已完成的训练周期数和总训练时间(以小时为单位)。
LOGGER.info(
f"\n{epoch - self.start_epoch + 1} epochs completed in "
f"{(time.time() - self.train_time_start) / 3600:.3f} hours." # {epoch - self.start_epoch + 1} 个时期在 {(time.time() - self.train_time_start) / 3600:.3f} 小时内完成。
)
# 调用 final_eval 方法执行最终的验证。这通常是在所有训练周期完成后,使用最佳模型权重(例如,保存在 best.pt 文件中)对验证集进行评估。
self.final_eval()
# 绘图和回调执行。
# 检查是否设置了绘图参数 self.args.plots 为真。
if self.args.plots:
# 如果需要绘图,调用 plot_metrics 方法来绘制训练和验证过程中的指标,如损失和准确率等。
self.plot_metrics()
# 运行训练结束的回调函数。这些回调可能包括保存模型、日志记录、评估模型性能等操作。
self.run_callbacks("on_train_end")
# 内存清理和回调执行。
# 清理内存,释放不再需要的资源,如删除临时变量或释放GPU内存。
self._clear_memory()
# 运行“teardown”回调函数,这些函数在训练完全结束时执行,可能包括清理资源、保存最终结果等操作。
self.run_callbacks("teardown")
# 这段代码的主要目的是在训练结束后执行一系列收尾工作,包括最终验证、日志记录、绘图、回调执行和内存清理。这些步骤有助于确保训练过程的完整性,以及资源的有效管理和释放。通过这种方式, _do_train
# 方法提供了一个完整的训练循环,可以根据不同的训练需求进行定制和扩展。
# 这个方法的主要目的是执行模型的训练过程,包括前向传播、损失计算、反向传播和优化器更新。它还处理分布式训练、学习率调度、预热期和停止条件。通过这种方式, _do_train 方法提供了一个完整的训练循环,可以根据不同的训练需求进行定制和扩展。
# 这段代码定义了一个名为 _get_memory 的方法,它用于获取加速器(如GPU或MPS)的内存利用率,单位为GB(吉字节)。
# 方法定义。
def _get_memory(self):
# 获取加速器内存利用率(以 GB 为单位)。
"""Get accelerator memory utilization in GB."""
# 这个条件判断当前设备是否是MPS(Metal Performance Shaders),这是苹果用于在iOS和macOS上进行GPU加速的技术。
if self.device.type == "mps":
# torch.mps.driver_allocated_memory()
# torch.mps.driver_allocated_memory() 是 PyTorch 在支持 Apple Silicon(如 M1 或 M2 芯片)的 Mac 上进行深度学习训练时,用于获取 MPS(Metal Performance Shaders)后端驱动程序分配的内存量的一个函数。这个函数返回的值是当前 MPS 驱动程序为 PyTorch 分配的内存总量,单位是字节。
# 返回值 :
# 返回 MPS 驱动程序为 PyTorch 分配的内存总量,单位是字节。
# 使用场景 :
# 这个函数通常用于监控和调试在 Apple Silicon 芯片上使用 PyTorch 进行训练时的内存使用情况。它可以帮助开发者了解模型训练对设备内存的占用情况,从而优化模型或调整训练策略。
# 需要注意的是, torch.mps.driver_allocated_memory() 函数只在 PyTorch 的夜间构建版本(nightly build)中可用,且只能在支持 MPS 的系统上使用。如果你的系统或 PyTorch 版本不支持 MPS,这个函数将不可用。
# 如果设备是MPS,调用 torch.mps.driver_allocated_memory() 函数获取MPS驱动程序分配的内存量。
memory = torch.mps.driver_allocated_memory()
# 这个条件判断当前设备是否是CPU。
elif self.device.type == "cpu":
# 如果设备是CPU,将内存量设置为0,因为CPU内存不通过这种方式测量。
memory = 0
# 如果设备既不是MPS也不是CPU,那么它可能是CUDA支持的GPU。
else:
# torch.cuda.memory_reserved(device=None)
# torch.cuda.memory_reserved() 是 PyTorch 提供的一个函数,用于查询当前 GPU 设备上由 PyTorch 保留的显存总量。这个函数返回的值是 PyTorch 预先分配但尚未实际使用的显存量,单位是字节。
# 参数 :
# device :(可选)指定要查询的设备。如果为 None ,则默认查询当前设备。
# 返回值 :
# 返回由 PyTorch 保留的显存量,单位是字节。
# 使用场景 :
# 这个函数通常用于监控和调试 GPU 显存的使用情况。它可以帮助开发者了解 PyTorch 在 GPU 上预留的显存量,这包括 PyTorch 为未来操作预先分配的显存缓冲区。
# 需要注意的是, torch.cuda.memory_reserved() 返回的显存量可能大于实际分配给张量和模型参数的显存量,因为 PyTorch 会预留一部分显存作为缓冲区,以便在需要时快速分配。这部分预留的显存不会直接被占用,因此不计入 torch.cuda.memory_allocated() 的返回值。
# 对于CUDA设备,调用 torch.cuda.memory_reserved() 函数获取CUDA保留的内存量。
memory = torch.cuda.memory_reserved()
# 将获取的内存量(以字节为单位)转换为GB(吉字节),并返回这个值。
# 返回值。方法返回加速器的内存利用率,单位为GB。
return memory / 1e9
# 这个方法的主要目的是监控和报告加速器的内存使用情况,这对于调试和优化训练过程中的内存使用非常有用。通过这种方式,用户可以了解模型训练对硬件资源的占用情况,并据此做出相应的调整。
# 这段代码定义了一个名为 _clear_memory 的方法,它用于在不同的平台上清理加速器内存。
# 这是 _clear_memory 方法的定义,它是一个实例方法,只接受 self 参数。
def _clear_memory(self):
# 清除不同平台上的加速器内存。
"""Clear accelerator memory on different platforms."""
# gc.collect(generation=2)
# gc.collect() 是 Python 标准库中的 gc (垃圾收集器)模块的一个函数,用于强制执行一次垃圾收集。Python 使用自动内存管理,其中包含一个周期性的垃圾收集器来回收不再使用的内存。 gc.collect() 函数可以用来手动触发这一过程。
# 参数 :
# generation :(可选)整数,指定要收集的垃圾代。Python 的垃圾收集器将对象分为三代 :
# 第0代 :新创建的对象。
# 第1代 :已经存活一段时间的对象。
# 第2代 :长期存活的对象。
# 默认值为2,即收集所有代的对象。
# 返回值 :
# 返回一个元组,包含三个整数,分别表示 :
# 收集到的对象数量。
# 回收的对象数量。
# 未回收的对象数量。
# 使用场景 :
# 当你确定程序中不再需要某些大型对象时,可以调用 gc.collect() 来尝试释放内存。
# 在内存紧张或需要优化内存使用的情况下,可以手动触发垃圾收集。
# 在调试内存相关问题时,可以用来检查是否有内存泄漏。
# 需要注意的是, gc.collect() 并不总是能立即释放所有内存,因为 Python 的垃圾收集器可能不会立即处理所有代的对象,特别是对于长时间存活的对象。此外,即使调用了 gc.collect() ,操作系统也可能不会立即释放内存,因为操作系统有自己的内存管理策略。
# 调用 Python 的垃圾收集器 gc 模块的 collect 方法,强制进行一次垃圾收集。这有助于释放不再使用的 Python 对象所占用的内存。
gc.collect()
# 这个条件判断当前设备是否是 MPS(Metal Performance Shaders),这是苹果用于在 iOS 和 macOS 上进行 GPU 加速的技术。
if self.device.type == "mps":
# torch.mps.empty_cache()
# torch.mps.empty_cache() 是 PyTorch 在 Apple Silicon(如 M1 或 M2 芯片)Mac 上使用 MPS(Metal Performance Shaders)后端时,用于释放 MPS 缓存的函数。这个函数可以释放 MPS 分配的未占用缓存内存,使其可以在其他 GPU 应用程序中使用。
# 参数 :无参数。
# 返回值 :无返回值。
# 使用场景 :
# 这个函数通常用于以下情况 :
# 当你希望手动释放 MPS 分配的未使用内存时,以便这些内存可以被系统或其他应用程序使用。
# 在内存紧张的情况下,或者在训练过程中需要频繁地分配和释放大量内存时,可以帮助减少内存碎片和溢出的风险。
# 需要注意的是, torch.mps.empty_cache() 函数只在 PyTorch 的 MPS 后端中可用,且只能在支持 MPS 的系统上使用。如果你的系统或 PyTorch 版本不支持 MPS,这个函数将不可用。此外,释放缓存并不总是能够立即减少 nvidia-smi 或 gpustat 显示的显存占用,因为这些工具可能显示的是进程的总显存占用,包括 PyTorch 预留的显存缓冲区。
# 如果设备是 MPS,调用 torch.mps.empty_cache() 函数来释放 MPS 后端的缓存内存。这有助于释放 MPS 分配的未使用内存,使其可供其他应用使用。
torch.mps.empty_cache()
# 这个条件判断当前设备是否是 CPU。
elif self.device.type == "cpu":
# 如果设备是 CPU,则不执行任何操作,直接返回。这是因为 CPU 内存管理是由操作系统负责的,不需要手动释放。
return
# 如果设备既不是 MPS 也不是 CPU,那么它可能是 CUDA 支持的 GPU。
else:
# 对于 CUDA 设备,调用 torch.cuda.empty_cache() 函数来释放 PyTorch 缓存的未使用 GPU 内存。这有助于释放 PyTorch 分配的未使用内存,使其可供其他 CUDA 程序使用。
torch.cuda.empty_cache()
# 这个方法的主要目的是在不同的硬件平台上清理加速器内存,以确保不再使用的内存被释放,从而提高内存使用效率。这对于管理有限的 GPU 资源尤为重要,可以帮助避免内存泄漏和不必要的内存占用。
# 这段代码定义了一个名为 read_results_csv 的方法,它用于从一个 CSV 文件中读取结果,并将这些结果存储在一个字典中。
# 这是 read_results_csv 方法的定义,它是一个实例方法,只接受 self 参数。
def read_results_csv(self):
# 使用 pandas 将 results.csv 读入字典。
"""Read results.csv into a dict using pandas."""
# 这行代码导入了 pandas 库,并将其别名设置为 pd 。注释表明这个导入是为了更快地导入 ultralytics ,这可能是一个特定的库或框架。
import pandas as pd # scope for faster 'import ultralytics'
# 这行代码执行了以下操作 :
# pandas.read_csv(file_name, **kwargs)
# pd.read_csv() 是 Python 中 pandas 库的一个函数,用于读取 CSV(逗号分隔值)文件,并将其转换为 pandas 的 DataFrame 对象。DataFrame 是 pandas 中用于存储和操作结构化数据的主要数据结构。
# 参数 :
# file_name :文件路径或文件对象,指向要读取的 CSV 文件。
# 可选参数(kwargs) :
# header :指定作为列名的第一行,默认为0(第一行)。
# sep :指定字段的分隔符,默认为逗号( , )。
# delimiter :与 sep 类似,但在某些情况下可以提供更多灵活性。
# names :列名列表,用于没有列名的 CSV 文件。
# index_col :指定作为行索引的列。
# usecols :需要读取的列。
# squeeze :如果为 True ,并且文件只有一列,则返回 Series 而不是 DataFrame。
# dtype :指定列的数据类型。
# skiprows :跳过文件开始的行数或需要跳过的行的列表。
# nrows :需要读取的行数。
# na_values :将这些值视为缺失值。
# keep_default_na :是否保留默认的缺失值(例如 'nan' 、 'null' 等)。
# skipfooter :从文件末尾跳过的行数。
# verbose :是否打印详细的进度信息。
# parse_dates :是否解析日期。
# infer_datetime_format :是否推断日期时间格式。
# keep_date_col :是否在解析日期后保留原始列。
# dayfirst :日期格式中天是否在前。
# cache_dates :是否缓存日期解析。
# iterator :返回一个迭代器而不是 DataFrame。
# chunksize :返回一个迭代器时,每次迭代的行数。
# compression :文件压缩类型。
# quotechar :用于引用字段的字符。
# escapechar :用于转义的字符。
# encoding :文件编码。
# decimal :小数点字符。
# lineterminator :行终止符。
# quotechar :引用字符。
# 返回值 :
# 返回一个 DataFrame 对象,其中包含了 CSV 文件的内容。
# 使用场景 :
# 读取 CSV 文件并进行数据分析。
# 数据清洗和预处理。
# 数据科学和机器学习中的数据处理。
# DataFrame.to_dict(orient='dict', into=dict, index=False, columns=False)
# to_dict() 是 pandas 库中 DataFrame 对象的一个方法,它将 DataFrame 转换为字典格式。这个方法提供了多种方式来指定转换的方向和层次结构。
# 参数 :
# orient :指定返回字典的格式,可以是以下值之一 :
# 'dict' :默认值,返回一个字典,其键是列名,值是每个列的数据,以 Series 的形式存在。
# 'list' :返回一个字典,其键是列名,值是每个列的数据,以列表的形式存在。
# 'series' :返回一个字典,其键是索引值,值是每个索引对应的行数据,以 Series 的形式存在。
# 'split' :返回一个字典,包含 'index' 、 'columns' 和 'data' 三个键,分别对应索引、列名和数据。
# 'records' :返回一个字典列表,每个字典代表一行数据。
# 'index' :返回一个字典,其键是索引值,值是以 Series 形式存在的列数据。
# into :指定返回字典的类型,可以是 dict 或 OrderedDict (如果需要保持顺序)。
# index :布尔值,指定是否将索引作为字典的键。
# columns :布尔值,指定是否将列名作为字典的键。
# 返回值 :
# 返回一个字典,其结构取决于 orient 参数的设置。
# 使用场景 :
# 当需要将 DataFrame 数据转换为字典格式以便进行进一步处理时。
# 在数据转换和序列化时,例如将 DataFrame 数据转换为 JSON。
# pd.read_csv(self.csv) :使用 pandas 的 read_csv 函数读取 self.csv 指定的 CSV 文件。
# .to_dict(orient="list") :将 pandas DataFrame 转换为字典,其中 orient="list" 参数指定了字典的值应该是列表形式。
# .items() :将字典转换为键值对元组的迭代器。
# {k.strip(): v for k, v in ...} :使用字典推导式创建一个新的字典,其中键 k 被 strip() 方法处理,移除了任何前后空白字符。值 v 保持不变。
return {k.strip(): v for k, v in pd.read_csv(self.csv).to_dict(orient="list").items()}
# 这个方法的主要目的是将 CSV 文件中的数据读取为一个更易于处理的字典格式,其中每个键对应 CSV 文件中的一列,每个值是一个列表,包含该列的所有数据。这种格式对于数据分析和结果处理非常有用。
# 这段代码定义了一个名为 save_model 的方法,它用于保存模型训练的检查点(checkpoints)以及附加的元数据。
# 这是 save_model 方法的定义,它是一个实例方法,只接受 self 参数。
def save_model(self):
# 使用附加元数据保存模型训练检查点。
"""Save model training checkpoints with additional metadata."""
# 导入 Python 的 io 模块,它提供了处理二进制数据流的接口。
import io
# Serialize ckpt to a byte buffer once (faster than repeated torch.save() calls)
# io.BytesIO(initial_bytes=None)
# io.BytesIO 是 Python 标准库 io 模块中的一个类,它提供了一个类文件对象,用于在内存中处理二进制数据流,而不需要实际的文件。这个类实现了一个缓冲区接口,可以用于存储和读取二进制数据。
# 参数 :
# initial_bytes :(可选)一个二进制序列,用于初始化 BytesIO 对象的内容。如果不提供, BytesIO 对象将为空。
# 方法 :
# read() :从缓冲区中读取数据。
# write(b) :将二进制数据写入缓冲区。
# seek(offset, whence=0) :移动文件指针到指定位置。
# tell() :返回当前文件指针的位置。
# truncate(size=None) :截断缓冲区,使其大小为指定大小。
# getvalue() :返回缓冲区中的当前内容。
# close() :关闭缓冲区。
# 返回值 :
# BytesIO 对象实例。
# 使用场景 :
# 当需要一个临时的二进制文件对象时,例如在内存中处理数据流,而不是在磁盘上。
# 在读写二进制数据时,需要类似于文件的操作,但不希望或不需要实际的文件系统交互。
# 创建一个 BytesIO 对象,它是一个在内存中的二进制流,用于存储序列化的检查点数据。
buffer = io.BytesIO()
# 使用 PyTorch 的 torch.save 函数将检查点数据序列化到 buffer 中。这里的数据包括当前周期、最佳适应度、EMA(指数移动平均)模型、优化器状态、训练参数、训练指标、训练结果、日期、版本、许可证和文档链接等。
# 这段代码使用 PyTorch 的 torch.save 函数将模型训练的检查点(checkpoint)和附加元数据保存到一个内存中的二进制缓冲区 buffer 。
# torch.save(...) :这是 PyTorch 的保存函数,用于序列化 Python 对象。
# {...} : 一个字典,包含了要保存的各种信息。
torch.save(
{
# 保存当前的训练周期。
"epoch": self.epoch,
# 保存当前最佳适应度值。
"best_fitness": self.best_fitness,
# 这里没有保存模型权重,可能是因为使用指数移动平均(EMA)来恢复和最终检查点。
"model": None, # resume and final checkpoints derive from EMA
# 保存 EMA 模型的权重,并且转换为半精度(FP16)。
"ema": deepcopy(self.ema.ema).half(),
# 保存 EMA 更新的次数。
"updates": self.ema.updates,
# 保存优化器的状态,并且转换为半精度(FP16)。 convert_optimizer_state_dict_to_fp16 是一个自定义函数,用于将优化器状态转换为 FP16。
# def convert_optimizer_state_dict_to_fp16(state_dict): -> 用于将优化器的状态字典( state_dict )中的浮点数张量转换为半精度浮点数(FP16)。返回状态字典。返回更新后的状态字典。 -> return state_dict
"optimizer": convert_optimizer_state_dict_to_fp16(deepcopy(self.optimizer.state_dict())),
# 保存训练参数,以字典形式保存。
"train_args": vars(self.args), # save as dict
# 保存训练指标,包括额外的适应度值。
"train_metrics": {**self.metrics, **{"fitness": self.fitness}},
# 保存从 CSV 文件中读取的训练结果。
"train_results": self.read_results_csv(),
# isoformat(sep='T', timespec='auto')
# isoformat() 函数是 Python datetime 模块中 datetime 类的一个方法,它用于将 datetime 对象格式化为 ISO 8601 标准的日期和时间字符串。
# 参数 :
# sep :一个单字符分隔符,用于日期和时间部分之间的分隔。默认值为 'T' 。
# timespec :指定要包含的时间附加组件的数量。默认值为 'auto' 。如果微秒为 0,则与 'seconds' 相同,否则为微秒。
# 返回值 :
# 返回类型为字符串,表示日期和时间的 ISO 8601 格式。
# 说明 :
# 如果 datetime 对象的微秒部分不为 0, isoformat() 方法会包含微秒部分,格式为 YYYY-MM-DDTHH:MM:SS.ffffff 。
# 如果 datetime 对象的微秒部分为 0, isoformat() 方法将返回没有微秒部分的字符串,格式为 YYYY-MM-DDTHH:MM:SS 。
# 如果 datetime 对象是一个“感知”对象(即包含时区信息),则还会包含 UTC 偏移量,格式为 YYYY-MM-DDTHH:MM:SS.ffffff+HH:MM 或 YYYY-MM-DDTHH:MM:SS+HH:MM 。
# isoformat() 方法提供了一种无歧义的方式来表示日期和时间,符合国际标准 ISO 8601,这在数据交换和存储中非常有用。
# 保存保存检查点的日期和时间。
"date": datetime.now().isoformat(),
# 保存当前软件的版本号。
"version": __version__,
# 保存许可证信息。
"license": "AGPL-3.0 (https://ultralytics.com/license)",
# 保存文档链接。
"docs": "https://docs.ultralytics.com",
},
# 指定保存的目标为 buffer ,这是一个 io.BytesIO 对象,用于在内存中存储序列化数据。
buffer,
)
# 从 buffer 中获取序列化后的检查点内容。
serialized_ckpt = buffer.getvalue() # get the serialized content to save
# Save checkpoints
# Path.write_bytes(data, *, encoding=None, errors=None)
# Path.write_bytes() 是 Python pathlib 模块中 Path 类的一个方法,用于将字节数据写入到指定的文件路径。这个方法提供了一个方便的方式来直接将字节流写入文件,而不需要打开文件和关闭文件的显式调用。
# 参数 :
# data :要写入文件的字节数据。
# encoding :(可选)如果提供字符串数据而不是字节数据,指定字符串的编码方式。
# errors :(可选)指定如何处理编码错误。
# 返回值 :无返回值,方法直接将数据写入文件。
# 使用场景 :
# 当你需要将字节数据持久化到文件系统时。
# 当你需要一个简单的方式来写入二进制文件时。
# Path.write_bytes() 方法是 pathlib 模块提供的一个现代文件操作方法,它简化了文件的写入过程,特别是在处理二进制数据时。这个方法在 Python 3.5 及以上版本中可用。
# 将序列化后的检查点数据保存到 last.pt 文件中。
self.last.write_bytes(serialized_ckpt) # save last.pt
# 如果当前的适应度等于最佳适应度,则执行以下操作。
if self.best_fitness == self.fitness:
# 将序列化后的检查点数据保存到 best.pt 文件中。
self.best.write_bytes(serialized_ckpt) # save best.pt
# 如果设置了保存周期( save_period )并且当前周期是保存周期的整数倍,则执行以下操作。
if (self.save_period > 0) and (self.epoch % self.save_period == 0):
# 将序列化后的检查点数据保存到以当前周期命名的文件中,例如 epoch3.pt 。
(self.wdir / f"epoch{self.epoch}.pt").write_bytes(serialized_ckpt) # save epoch, i.e. 'epoch3.pt'
# 这个方法的主要目的是在训练过程中定期保存模型的状态,包括模型参数、优化器状态和其他训练相关的元数据。这些检查点可以在训练过程中或训练后用于恢复训练、分析结果或进行模型评估。通过这种方式,可以确保训练过程中的关键信息被适当记录和保存。
# 这段代码定义了一个名为 get_dataset 的方法,它用于根据指定的任务和数据源加载和验证数据集。
# 这是 get_dataset 方法的定义,它是一个实例方法,只接受 self 参数。
def get_dataset(self):
# 如果数据字典中存在,则获取 train、val 路径。
# 如果无法识别数据格式,则返回 None。
"""
Get train, val path from data dict if it exists.
Returns None if data format is not recognized.
"""
# 开始一个 try 块,用于捕获在加载数据集过程中可能发生的任何异常。
try:
# 检查当前任务是否为分类("classify")。
if self.args.task == "classify":
# 如果是分类任务,调用 check_cls_dataset 函数来验证和加载分类数据集。
# def check_cls_dataset(dataset, split=""):
# -> 检查分类数据集的完整性,并在必要时下载数据集。函数返回一个字典,包含 训练集 、 验证集 和 测试集 的 路径 、 类别数量 和 类别名称 。
# -> return {"train": train_set, "val": val_set, "test": test_set, "nc": nc, "names": names}
data = check_cls_dataset(self.args.data)
# 检查数据文件的扩展名是否为 .yaml 或 .yml ,或者任务是否为检测("detect")、分割("segment")、姿态估计("pose")或方向边界框("obb")。
elif self.args.data.split(".")[-1] in {"yaml", "yml"} or self.args.task in {
"detect",
"segment",
"pose",
"obb",
}:
# 如果条件满足,调用 check_det_dataset 函数来验证和加载检测数据集。
# def check_det_dataset(dataset, autodownload=True):
# -> 检查和处理数据集(dataset),确保数据集的YAML配置文件符合要求,并在需要时下载数据集。函数返回处理后的 data 字典,其中包含了 数据集的 配置信息 和 路径 。
# -> return data # dictionary
data = check_det_dataset(self.args.data)
# 检查加载的数据中是否包含 yaml_file 键。
if "yaml_file" in data:
# 如果包含,更新 self.args.data 为 yaml_file 的值,这通常用于验证数据集的 URL。
self.args.data = data["yaml_file"] # for validating 'yolo train data=url.zip' usage
# 如果在尝试加载数据集时发生异常,捕获这个异常。
except Exception as e:
# 抛出一个 RuntimeError ,包含错误信息和数据集的 URL。 emojis 函数可能用于在错误消息中添加表情符号。 clean_url 函数可能用于清理或格式化 URL。
# def clean_url(url):
# -> 清除 URL 中的认证信息(如用户名和密码)以及其他可能的查询参数,只保留 URL 的基本部分。返回值。函数返回一个清理后的 URL 字符串,不包含认证信息和查询参数。
# -> return urllib.parse.unquote(url).split("?")[0] # '%2F' to '/', split https://url.com/file.txt?auth
raise RuntimeError(emojis(f"Dataset '{clean_url(self.args.data)}' error ❌ {e}")) from e
# 将加载的数据集赋值给 self.data 属性。
self.data = data
# 返回训练数据集和验证数据集(如果存在)。如果验证数据集不存在,则尝试返回测试数据集。
return data["train"], data.get("val") or data.get("test")
# 这个方法的主要目的是根据任务类型和数据源加载相应的数据集,并处理可能出现的错误。它确保了数据集的正确性和可用性,为后续的训练和验证过程提供了必要的数据。
# 这段代码定义了一个名为 setup_model 的方法,它用于加载、创建或下载与任务相关的模型。
# 这是 setup_model 方法的定义,它是一个实例方法,只接受 self 参数。
def setup_model(self):
# 为任何任务加载/创建/下载模型。
"""Load/create/download model for any task."""
# 检查 self.model 是否已经是一个 PyTorch 的 nn.Module 实例。如果是,意味着模型已经被加载,不需要进一步设置。
if isinstance(self.model, torch.nn.Module): # if model is loaded beforehand. No setup needed
# 如果模型已经加载,直接返回。
return
# 初始化配置 cfg 为 self.model ,权重 weights 为 None 。
cfg, weights = self.model, None
# 初始化检查点 ckpt 为 None 。
ckpt = None
# 检查 self.model 的字符串表示是否以 .pt 结尾,这通常意味着它是一个模型权重文件的路径。
if str(self.model).endswith(".pt"):
# 如果是 .pt 文件,调用 attempt_load_one_weight 函数尝试加载权重,并返回权重和检查点。
# def attempt_load_one_weight(weight, device=None, inplace=True, fuse=False):
# -> 用于加载单个模型权重,并进行一些兼容性更新和配置。返回模型和检查点。函数返回 加载的模型 和 检查点对象 。
# -> return model, ckpt
weights, ckpt = attempt_load_one_weight(self.model)
# 如果成功加载权重,从权重中获取配置信息,并更新 cfg 。
cfg = weights.yaml
# 检查 self.args.pretrained 是否为字符串或 Path 实例,这表示预训练模型的路径。
elif isinstance(self.args.pretrained, (str, Path)):
# 如果提供了预训练模型的路径,调用 attempt_load_one_weight 函数尝试加载预训练权重。
# def attempt_load_one_weight(weight, device=None, inplace=True, fuse=False):
# -> 用于加载单个模型权重,并进行一些兼容性更新和配置。返回模型和检查点。函数返回 加载的模型 和 检查点对象 。
# -> return model, ckpt
weights, _ = attempt_load_one_weight(self.args.pretrained)
# 调用 get_model 方法创建模型实例,传入配置 cfg 、权重 weights 和一个布尔值 verbose ,指示是否打印详细信息。
self.model = self.get_model(cfg=cfg, weights=weights, verbose=RANK == -1) # calls Model(cfg, weights)
# 返回检查点 ckpt ,它可能包含有关加载的权重的额外信息。
return ckpt
# 这个方法的主要目的是确保模型被正确加载或创建,无论是从预训练权重、配置文件还是从头开始。它处理了不同情况下模型的加载逻辑,并提供了灵活性以适应不同的训练需求。
# 这段代码定义了一个名为 optimizer_step 的方法,它用于执行单步训练优化器的操作,包括梯度缩放、梯度裁剪、优化器步骤执行以及指数移动平均(EMA)更新。
# 这是 optimizer_step 方法的定义,它是一个实例方法,只接受 self 参数。
def optimizer_step(self):
# 使用梯度剪辑和 EMA 更新执行训练优化器的单步。
"""Perform a single step of the training optimizer with gradient clipping and EMA update."""
# GradScaler.unscale_(optimizer)
# torch.cuda.amp.GradScaler 类中的 unscale_() 方法用于在反向传播和 optimizer.step() 之间撤销梯度的缩放。
# 参数 :
# optimizer :需要撤销梯度缩放的优化器实例。
# 功能描述 :
# unscale_() 方法将优化器中的梯度张量除以缩放因子,以撤销之前由 scale() 方法应用的缩放。这是在自动混合精度(AMP)训练中使用的,用于防止由于梯度值过小而在半精度浮点数下发生的下溢。
# 使用场景 :
# 在执行 optimizer.step() 之前,如果你需要修改或检查梯度,可以显式调用 unscale_() 方法。
# 如果在调用 scaler.step(optimizer) 之前没有显式调用 unscale_() ,则 step() 方法内部会自动执行梯度的撤销缩放。
# 注意事项 :
# unscale_() 方法不会产生 CPU-GPU 同步,这意味着它不会等待所有设备上的梯度撤销缩放完成。
# 每个优化器每次 step() 调用只能调用一次 unscale_() ,并且只能在该优化器分配的参数的所有梯度都已累积之后调用。在每个 step() 之间为给定优化器调用 unscale_() 两次会触发运行时错误。
# unscale_() 可能会撤销缩放稀疏梯度,从而替换 .grad 属性。
# 使用 GradScaler 的 unscale_ 方法来撤销梯度的缩放。这是在自动混合精度(AMP)训练中使用的,用于处理梯度的动态缩放。
self.scaler.unscale_(self.optimizer) # unscale gradients
# torch.nn.utils.clip_grad_norm_(parameters, max_norm, norm_type=2)
# torch.nn.utils.clip_grad_norm_() 是 PyTorch 中的一个函数,用于裁剪(限制)模型参数梯度的范数,以防止梯度爆炸问题。这个函数会缩放梯度,使得它们的最大范数不超过指定的阈值。
# 参数 :
# parameters :PyTorch 参数(或参数组)的迭代器,通常是模型的参数。
# max_norm :梯度的最大范数。如果计算出的范数超过这个值,梯度将被缩放。
# norm_type :(可选)用于计算梯度范数的类型,默认为2,表示L2范数。也可以使用其他值,如1(L1范数)或 np.inf (无穷范数,即最大值)。
# 返回值 :无返回值。该函数直接修改输入的参数梯度。
# 使用场景 :
# 在训练神经网络时,特别是在使用RNN或其他容易产生梯度爆炸的结构时,使用梯度裁剪来保持训练的稳定性。
# 使用 PyTorch 的 clip_grad_norm_ 函数来裁剪模型参数的梯度,以防止梯度爆炸。 max_norm=10.0 指定了梯度的最大范数。
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=10.0) # clip gradients
# GradScaler.step(optimizer)
# torch.cuda.amp.GradScaler 类中的 step() 方法用于执行优化器步骤,同时考虑了梯度缩放。
# 参数 :
# optimizer :需要更新的优化器实例。
# 功能描述 :
# step() 方法执行以下操作 :
# 撤销梯度缩放 :首先, step() 方法会撤销之前通过 scale() 方法应用的梯度缩放。这是必要的,因为缩放后的梯度可能非常大,直接用于参数更新会导致过大的步长。
# 检查梯度值 :然后,它会检查撤销缩放后的梯度是否包含无穷大(Inf)或非数值(NaN)。
# 更新参数 :如果梯度是有效的(即不包含 Inf 或 NaN), step() 方法会调用优化器的 step() 方法来更新模型参数。
# 跳过更新 :如果梯度包含 Inf 或 NaN,则不会执行参数更新。
# 使用场景 :在自动混合精度(AMP)训练中, step() 方法通常在执行反向传播后调用,用于更新模型参数。
# 注意事项 :
# step() 方法应该在 scale() 方法和 update() 方法之间调用。
# 如果在调用 step() 之前需要修改或检查梯度,应先调用 unscale_() 方法来撤销梯度缩放。
# 使用 GradScaler 的 step 方法来执行优化器步骤,更新模型参数。这是在AMP训练中使用的,以确保梯度缩放后的正确参数更新。
self.scaler.step(self.optimizer)
# GradScaler.update()
# torch.cuda.amp.GradScaler 类中的 update() 方法用于在每次优化器步骤之后更新缩放器的状态,特别是动态调整用于梯度缩放的缩放因子。
# 参数 :无参数。
# 功能描述 :
# update() 方法执行以下操作 :
# 动态调整缩放因子 :根据最近的训练步骤中的数值稳定性情况(即是否出现梯度为 NaN 或 Inf)动态调整缩放因子。如果连续多次训练步骤中梯度都是有限值(即没有 NaN 或 Inf),缩放因子会逐步增大,以减少梯度下溢。如果检测到梯度为 NaN 或 Inf,则缩放因子会减小,以防止梯度溢出。
# 更新状态 : update() 方法还会更新 GradScaler 的内部状态,以准备下一次训练步骤。
# 使用场景 :
# 在自动混合精度(AMP)训练中,每次调用 scaler.step(optimizer) 之后,都需要调用 update() 方法来更新缩放器的状态。
# 注意事项 :
# update() 方法必须在 scaler.step(optimizer) 之后调用,以确保正确地根据最近的训练步骤调整缩放因子。
# update() 方法不应该在出现梯度为 NaN 或 Inf 的情况下调用,因为这种情况下通常需要跳过参数更新,并且不更新缩放因子。
# 更新 GradScaler ,准备下一次反向传播。
self.scaler.update()
# torch.optim.Optimizer.zero_grad(set_to_none=True)
# 在PyTorch中, zero_grad() 函数是 torch.optim.Optimizer 类中的一个方法,用于将所有优化器中的参数梯度清零。这是一个非常重要的步骤,因为它确保了在每次迭代或优化步骤之前,旧的梯度不会影响新的梯度计算。
# 参数 :
# set_to_none :一个布尔值,指定是否将梯度设置为 None 。默认值为 True 。如果设置为 True ,则梯度会被设置为 None ,这有助于减少内存消耗。如果设置为 False ,则梯度会被设置为零。
# 功能描述 :
# zero_grad() 方法遍历优化器中的所有参数,并将其梯度清零。这通常在每次迭代的开始或在反向传播之前进行,以确保旧的梯度值不会累积到新的梯度上。
# 使用场景 :
# 在训练神经网络时,每次进行前向传播和反向传播之前,都需要清零梯度,以避免梯度累积。这可以通过直接调用优化器实例的 zero_grad() 方法来实现。
# 清空优化器的梯度,为下一次迭代准备。
self.optimizer.zero_grad()
# 检查是否存在EMA(指数移动平均)实例。
if self.ema:
# 如果存在EMA,调用其 update 方法来更新EMA副本的参数。
# def update(self, model): -> 用于更新模型的指数移动平均(EMA)参数。这个方法的核心作用是将模型的参数以一定的衰减率 d 融合到 EMA 参数中,使得 EMA 参数成为模型参数的指数移动平均。
self.ema.update(self.model)
# 这个方法的主要目的是在训练过程中执行优化器的单步更新,同时确保梯度缩放和裁剪的正确处理,以及EMA的更新。这些步骤对于保持训练稳定性和提高模型性能至关重要。通过这种方式, optimizer_step 方法提供了一个完整的优化器步骤,可以根据不同的训练需求进行定制和扩展。
# 这段代码定义了一个名为 preprocess_batch 的方法,它用于在训练或推理之前对模型的输入和目标(ground truths)进行预处理。
# 这是 preprocess_batch 方法的定义,它是一个实例方法,接受两个参数: self 和 batch 。
# 1.batch :参数通常是一个包含输入数据和目标的批次。
def preprocess_batch(self, batch):
# 允许根据任务类型自定义预处理模型输入和基本事实。
"""Allows custom preprocessing model inputs and ground truths depending on task type."""
# 这行代码简单地返回传入的 batch 。这意味着在这个实现中,没有对批次进行任何预处理,直接返回原始数据。
return batch
# 在实际应用中,这个方法可能会被重写或扩展,以包含特定的预处理步骤,例如 :
# 归一化输入数据。
# 调整图像大小或裁剪。
# 应用数据增强技术,如随机旋转、翻转或颜色变换。
# 编码目标标签,例如,将类别名称转换为整数索引或one-hot编码。
# 将数据转换为模型所需的格式,例如,将图像从 NumPy 数组转换为 PyTorch 张量。
# 这段代码定义了一个名为 validate 的方法,它用于在模型验证阶段评估模型性能,并更新最佳适应度(fitness)记录。
# 这是 validate 方法的定义,它是一个实例方法,只接受 self 参数。
def validate(self):
# 使用 self.validator 对测试集运行验证。
# 返回的字典应包含“fitness”键。
"""
Runs validation on test set using self.validator.
The returned dict is expected to contain "fitness" key.
"""
# 调用 self.validator ,它是一个验证器对象,传入 self (模型实例),并返回一个包含各种验证指标的字典。
metrics = self.validator(self)
# pop()
# pop() 是 Python 中列表(list)和字典(dict)数据结构的一个方法,用于移除并返回容器中的一个元素。
# 对于字典 :
# 在字典中, pop() 方法用于删除并返回字典中指定键(key)的值。如果键存在,则删除该键值对并返回其值;如果键不存在,除非指定了第二个参数,否则会抛出 KeyError 。
# 字典的 pop() 方法定义:
# value = dict.pop(key, default=None)
# key :要删除的键。
# default :如果键不存在时返回的默认值,默认为 None 。
# pop() 方法是处理列表和字典时常用的方法,它允许在需要时动态地修改容器的内容。
# 从 metrics 字典中弹出 "fitness" 键对应的值,并将其赋值给 fitness 变量。如果 "fitness" 键不存在,则使用模型的损失值( self.loss )作为适应度度量。
# 这里假设 self.loss 是一个 PyTorch 张量, .detach() 方法用于分离张量以避免梯度计算, .cpu() 将张量移动到 CPU, .numpy() 将张量转换为 NumPy 数组。负号是因为在某些情况下,损失值越低越好,所以用其负值作为适应度。
fitness = metrics.pop("fitness", -self.loss.detach().cpu().numpy()) # use loss as fitness measure if not found
# 检查 self.best_fitness 是否未定义或当前适应度 fitness 是否优于之前记录的最佳适应度 self.best_fitness 。
if not self.best_fitness or self.best_fitness < fitness:
# 如果当前适应度更优,则更新 self.best_fitness 。
self.best_fitness = fitness
# 返回包含验证指标的字典 metrics 和适应度 fitness 。
return metrics, fitness
# 这个方法的主要目的是在模型验证阶段提供一种机制来评估模型性能,并根据验证结果更新最佳适应度记录。这有助于在训练过程中监控模型的泛化能力,并在模型性能不再提高时及时停止训练(早停)。
# 这段代码定义了一个名为 get_model 的方法,它是一个类实例方法,用于获取模型实例。
# 这是 get_model 方法的定义,它是一个实例方法,包含四个参数。
# 参数定义 :
# 1.self :类的实例引用,用于访问类的属性和方法。
# 2.cfg :一个可选参数,用于传递模型的配置信息,默认值为 None 。这可以是一个配置文件路径、配置对象或其他形式的配置数据。
# 3.weight :一个可选参数,用于传递预训练模型的权重,默认值为 None 。这可以是一个权重文件路径或其他形式的权重数据。
# 4.verbose :一个可选参数,用于控制是否输出额外的信息,默认值为 True 。如果设置为 True ,则在执行过程中可能会打印出更多的日志信息。
def get_model(self, cfg=None, weights=None, verbose=True):
# 获取模型并引发 NotImplementedError 以加载 cfg 文件。
"""Get model and raise NotImplementedError for loading cfg files."""
# 这行代码引发了一个 NotImplementedError 异常,表示当前的任务训练器类(包含这个方法的类)不支持加载配置文件的功能。这意味着这个方法在当前类中没有实现具体的功能,需要在子类中重写这个方法以提供具体的模型加载逻辑。
raise NotImplementedError("This task trainer doesn't support loading cfg files") # 此任务训练器不支持加载 cfg 文件。
# 这个方法的目的是提供一个框架,用于在子类中实现获取模型实例的功能,包括加载配置和预训练权重。然而,在当前类中,这个方法并没有实现具体的逻辑,而是直接抛出异常,提示用户需要在子类中实现该功能。
# 这段代码定义了一个名为 get_validator 的方法,它用于返回一个验证器实例,但在当前的类中并没有实现具体的功能。
# 这是 get_validator 方法的定义,它是一个实例方法,只接受 self 参数。
def get_validator(self):
# 当调用 get_validator 函数时返回 NotImplementedError。
"""Returns a NotImplementedError when the get_validator function is called."""
# 这行代码引发了一个 NotImplementedError 异常,表示当前的任务训练器类(包含这个方法的类)不支持加载配置文件的功能。这意味着这个方法在当前类中没有实现具体的功能,需要在子类中重写这个方法以提供具体的模型加载逻辑。
raise NotImplementedError("get_validator function not implemented in trainer") # get_validator 函数未在训练器中实现。
# 这个方法的目的是提供一个接口,用于在子类中实现获取验证器的功能。验证器通常用于评估模型在验证集上的性能,例如计算准确率、召回率等指标。由于当前类中这个方法只是抛出异常,所以它不会返回任何值。在实际应用中,需要在子类中重写这个方法,以提供具体的验证器实例。
# 这段代码定义了一个名为 get_dataloader 的方法,它用于创建和返回一个数据加载器(dataloader),但在当前的类中并没有实现具体的功能。
# 这是 get_dataloader 方法的定义,它是一个实例方法,包含五个参数。
# 参数定义 :
# 1.self :类的实例引用,用于访问类的属性和方法。
# 2.dataset_path :数据集的路径,可能是一个文件路径或目录路径,指向存储数据的位置。
# 3.batch_size :一个可选参数,指定每个批次的样本数量,默认值为16。
# 4.rank :一个可选参数,通常用于分布式训练中标识进程的编号,默认值为0。
# 5.mode :一个可选参数,指定数据加载器的模式,例如训练("train")、验证("val")或测试("test"),默认值为"train"。
def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode="train"):
# 返回从 torch.data.Dataloader 派生的数据加载器。
"""Returns dataloader derived from torch.data.Dataloader."""
# 这行代码引发了一个 NotImplementedError 异常。
raise NotImplementedError("get_dataloader function not implemented in trainer") # get_dataloader 函数未在训练器中实现。
# 这个方法的目的是提供一个接口,用于在子类中实现创建数据加载器的功能。数据加载器是机器学习中用于批量加载数据的常用工具,它可以帮助开发者高效地从磁盘读取数据,并在训练过程中提供给模型。
# 由于当前类中这个方法只是抛出异常,所以它不会返回任何值。在实际应用中,需要在子类中重写这个方法,以提供具体的数据加载逻辑,例如使用 PyTorch 的 torch.utils.data.DataLoader 类来创建数据加载器。
# 这段代码定义了一个名为 build_dataset 的方法,它用于构建数据集,但在当前的类中并没有实现具体的功能。
# 这是 build_dataset 方法的定义,它是一个实例方法,包含四个参数。
# 参数定义 :
# 1.self :类的实例引用,用于访问类的属性和方法。
# 2.img_path :图像数据的路径,可能是一个文件路径或目录路径,指向存储图像的位置。
# 3.mode :一个可选参数,指定构建数据集的模式,例如训练("train")、验证("val")或测试("test"),默认值为"train"。
# 4.batch :一个可选参数,可能用于指定批次大小或相关批次处理的信息,默认值为 None 。
def build_dataset(self, img_path, mode="train", batch=None):
# 建立数据集。
"""Build dataset."""
# 这行代码引发了一个 NotImplementedError 异常。
raise NotImplementedError("build_dataset function not implemented in trainer") # build_dataset 函数未在训练器中实现。
# 这个方法的目的是提供一个接口,用于在子类中实现构建数据集的功能。在机器学习和深度学习中,构建数据集是一个重要的步骤,它涉及从原始数据中加载图像、标签和其他相关信息,并将它们组织成适合模型训练和评估的格式。
# 由于当前类中这个方法只是抛出异常,所以它不会返回任何值。在实际应用中,需要在子类中重写这个方法,以提供具体的数据集构建逻辑,例如使用 PyTorch 的 torch.utils.data.Dataset 类来创建自定义数据集。
# 这段代码定义了一个名为 label_loss_items 的方法,它用于标记损失项,可能是为了日志记录或可视化。
# 这是 label_loss_items 方法的定义,它是一个实例方法,包含三个参数。
# 1.self :类的实例引用,用于访问类的属性和方法。
# 2.loss_items :一个可选参数,用于传递损失项的数据,默认值为 None 。
# 3.prefix :一个可选参数,用于给损失项添加前缀,默认值为 "train" 。
def label_loss_items(self, loss_items=None, prefix="train"):
# 返回带有标记训练损失项张量的损失字典。
# 注意:
# 这对于分类来说不是必需的,但对于分割和检测来说是必需的。
"""
Returns a loss dict with labelled training loss items tensor.
Note:
This is not needed for classification but necessary for segmentation & detection
"""
# 这行代码是一个条件表达式,用于返回损失项的标签。
# 如果 loss_items 不是 None ,则返回一个字典,其中键为 "loss" ,值为 loss_items 。
# 如果 loss_items 是 None ,则返回一个包含单个字符串 "loss" 的列表。
return {"loss": loss_items} if loss_items is not None else ["loss"]
# 这个方法的主要作用是提供一种方式来标记损失项,这在训练过程中记录和监控模型性能时非常有用。通过这种方式,可以轻松地识别和访问不同的损失项。
# 这段代码定义了一个名为 set_model_attributes 的方法,它用于在训练开始前设置或更新模型的参数。
# 这是 set_model_attributes 方法的定义,它是一个实例方法,只接受 self 参数。
def set_model_attributes(self):
# 在训练之前设置或更新模型参数。
"""To set or update model parameters before training."""
# 将模型的 names 属性设置为 self.data 字典中的 "names" 键对应的值。这里假设 self.data 是一个包含数据集信息的字典,而 "names" 键存储了类别名称或其他相关信息。
self.model.names = self.data["names"]
# 这个方法的主要作用是确保模型在训练前具有正确的属性,例如类别名称,这对于模型的输出和预测后处理非常重要。通过设置这些属性,模型可以正确地解释其输出,并在训练和评估过程中使用正确的标签。
# 这段代码定义了一个名为 build_targets 的方法,它用于构建用于训练 YOLO 模型的目标张量。
# 这是 build_targets 方法的定义,它是一个实例方法,包含三个参数。
# 1.self :类的实例引用,用于访问类的属性和方法。
# 2.preds :预测结果,可能是模型的输出,包含了预测的边界框、置信度和类别概率。
# 3.targets :真实标签,包含了真实边界框、置信度和类别标签。
def build_targets(self, preds, targets):
# 构建用于训练 YOLO 模型的目标张量。
"""Builds target tensors for training YOLO model."""
# 这行代码是一个空操作,表示这个方法目前没有实现具体的功能。
pass
# 在 YOLO 模型的训练过程中,需要将真实标签转换为模型可以理解的目标张量,这通常包括边界框的坐标、对象的置信度和类别标签。这些目标张量将用于计算损失函数,指导模型的学习过程。
# 由于这个方法目前只是定义了一个框架并没有实现具体的逻辑,因此在实际应用中,需要在子类中重写这个方法,以提供具体的构建目标张量的逻辑。这可能涉及到对预测结果和真实标签进行处理,以匹配 YOLO 模型的损失函数所需的格式。
# 这段代码定义了一个名为 progress_string 的方法,它用于生成和返回一个描述训练进度的字符串。
# 这是 progress_string 方法的定义,它是一个实例方法,只接受 self 参数。
def progress_string(self):
# 返回描述训练进度的字符串。
"""Returns a string describing training progress."""
# 这行代码返回一个空字符串。这表明在当前的实现中, progress_string 方法并没有提供任何具体的进度信息。
return ""
# 在实际应用中,这个方法通常会被重写以包含实际的训练进度信息。例如,它可能返回当前的epoch数、已处理的批次数、损失值、准确率或其他重要的训练指标。这样的信息可以帮助用户直观地了解训练过程的进展情况。
# TODO: may need to put these following functions into callback TODO:可能需要将以下函数放入回调中。
# 这段代码定义了一个名为 plot_training_samples 的方法,它用于在 YOLO 模型训练期间绘制训练样本。
# 这是 plot_training_samples 方法的定义,它是一个实例方法,包含三个参数。
# 1.self :类的实例引用,用于访问类的属性和方法。
# 2.batch :一个批次的数据,可能包含图像和对应的标签。
# 3.ni :当前的迭代编号,可能是批次索引或全局迭代次数。
def plot_training_samples(self, batch, ni):
# 返回描述训练进度的字符串。
"""Plots training samples during YOLO training."""
# 这行代码是一个空操作,表示这个方法目前没有实现具体的功能。
pass
# 在实际应用中, plot_training_samples 方法可能会被用来可视化训练过程中的图像和预测结果,以帮助开发者理解模型的行为和性能。这通常涉及到使用绘图库(如 Matplotlib)来显示图像和绘制边界框、类别标签等。
# 这段代码定义了一个名为 plot_training_labels 的方法,它用于在 YOLO 模型训练期间绘制训练标签。
# 这是 plot_training_labels 方法的定义,它是一个实例方法,只接受 self 参数。
def plot_training_labels(self):
# 绘制 YOLO 模型的训练标签。
"""Plots training labels for YOLO model."""
# 这行代码是一个空操作,表示这个方法目前没有实现具体的功能。
pass
# 在实际应用中, plot_training_labels 方法可能会被用来可视化训练过程中的标签分布,以帮助开发者理解数据集的组成和模型的学习进度。这通常涉及到使用绘图库(如 Matplotlib)来显示标签的统计信息,例如类别的概率、每个类别的样本数量等。
# 这段代码定义了一个名为 save_metrics 的方法,它用于将训练指标(metrics)保存到 CSV 文件中。
# 这是 save_metrics 方法的定义,它是一个实例方法,包含两个参数。
# 1.self :类的实例引用,用于访问类的属性和方法。
# 2.metrics :一个字典,包含了要保存的训练指标。
def save_metrics(self, metrics):
# 将训练指标保存到 CSV 文件。
"""Saves training metrics to a CSV file."""
# 获取 metrics 字典的 键 (指标名称)和 值 (指标数值),并将它们转换为列表。
keys, vals = list(metrics.keys()), list(metrics.values())
# 计算 CSV 文件中的列数,包括 epoch 在内。
n = len(metrics) + 1 # number of cols
# str.rstrip([chars])
# rstrip 是 Python 中字符串( str )对象的一个方法,用于移除字符串末尾的特定字符。如果没有指定字符,则默认移除空白字符(包括空格、制表符、换行符等)。
# chars :一个字符串,包含需要从末尾移除的字符集合。如果省略此参数,则默认移除空白字符。返回值这个方法返回一个新的字符串,原字符串末尾的指定字符被移除。
# 注意事项 :
# rstrip 方法不会修改原始字符串,而是返回一个新的字符串。
# 如果需要移除的字符在字符串中没有出现,则返回原始字符串的一个副本。
# 如果需要从字符串的开头移除字符,可以使用 lstrip 方法;如果需要从两端移除字符,可以使用 strip 方法。
# rstrip 是处理字符串时常用的方法之一,特别是在处理文件路径、用户输入或任何可能包含不需要的尾随字符的情况。
# 如果 CSV 文件已存在,则 s 为空字符串;如果文件不存在,则创建一个包含列标题的字符串。 %23s 指定每个列标题占用的宽度为 23 个字符, tuple(["epoch"] + keys) 将 epoch 和指标名称组合成一个元组。
s = "" if self.csv.exists() else (("%23s," * n % tuple(["epoch"] + keys)).rstrip(",") + "\n") # header
# 使用 with 语句和 open 函数以追加模式("a")打开 CSV 文件。 self.csv 是一个 Path 对象,指向 CSV 文件的路径。
with open(self.csv, "a") as f:
# 将列标题和指标数值写入 CSV 文件。 %23.5g 指定每个指标数值占用的宽度为 23 个字符,保留 5 位小数。 tuple([self.epoch + 1] + vals) 将 epoch 和指标数值组合成一个元组。 rstrip(",") 用于移除字符串末尾的逗号。
f.write(s + ("%23.5g," * n % tuple([self.epoch + 1] + vals)).rstrip(",") + "\n")
# 这个方法的主要作用是将训练过程中的指标保存到 CSV 文件中,便于后续的分析和可视化。通过这种方式,可以轻松地跟踪训练进度和性能变化。
# 这段代码定义了一个名为 plot_metrics 的方法,它用于以图形化的方式展示和显示训练过程中的指标。
# 这是 plot_metrics 方法的定义,它是一个实例方法,只接受 self 参数。
def plot_metrics(self):
# 以视觉方式绘制和显示指标。
"""Plot and display metrics visually."""
# 这行代码是一个空操作,表示这个方法目前没有实现具体的功能。
pass
# 在实际应用中, plot_metrics 方法可能会被用来绘制训练和验证过程中的各种指标,例如损失、准确率、召回率等,以便开发者可以直观地观察模型性能的变化。这通常涉及到使用绘图库(如 Matplotlib、Seaborn 或 Plotly)来创建图表。
# 这段代码定义了一个名为 on_plot 的方法,它用于注册绘图操作,以便在回调函数中使用。
# 这是 on_plot 方法的定义,它是一个实例方法,包含三个参数。
# 1.self :类的实例引用,用于访问类的属性和方法。
# 2.name :绘图的名称或路径。
# 3.data :一个可选参数,用于传递绘图所需的数据,默认值为 None 。
def on_plot(self, name, data=None):
# 注册图表(例如在回调中使用)。
"""Registers plots (e.g. to be consumed in callbacks)."""
# 将 name 参数转换为 Path 对象, Path 是 pathlib 模块中的一个类,用于表示文件系统路径。
path = Path(name)
# 将绘图信息存储在 self.plots 字典中。
# self.plots 是一个字典,其键是绘图的路径,值是另一个字典,包含绘图数据和时间戳。
# "data": data :存储传递给 on_plot 方法的 data 参数。
# "timestamp": time.time() :存储当前时间的时间戳, time.time() 返回当前时间(自 Unix 纪元以来的秒数)。
self.plots[path] = {"data": data, "timestamp": time.time()}
# 这个方法的主要作用是将绘图数据和相关信息注册到类的 self.plots 字典中,以便在训练过程中的回调函数中使用。这可以用于在训练过程中动态地跟踪和可视化各种指标。
# 这段代码定义了一个名为 final_eval 的方法,它用于在 YOLO 模型训练结束后执行最终的评估和验证。
# 这是 final_eval 方法的定义,它是一个实例方法,只接受 self 参数。
def final_eval(self):
# 对目标检测 YOLO 模型进行最终评估和验证。
"""Performs final evaluation and validation for object detection YOLO model."""
# 初始化一个空字典 ckpt ,用于存储检查点信息。
ckpt = {}
# 遍历 self.last 和 self.best ,这两个属性可能分别指向最后和最佳的模型检查点文件。
for f in self.last, self.best:
# 检查检查点文件是否存在。
if f.exists():
# 如果当前检查点是最后的检查点,执行以下操作。
if f is self.last:
# 调用 strip_optimizer 函数加载最后的检查点,并移除优化器状态,只保留模型权重。
# def strip_optimizer(f: Union[str, Path] = "best.pt", s: str = "", updates: dict = None) -> dict:
# -> 从 PyTorch 模型的检查点文件中移除优化器状态,并保存一个精简的模型文件。返回结果。函数返回合并后的检查点字典 combined ,这样调用者可以进一步处理或使用这个字典。
# -> return combined
ckpt = strip_optimizer(f)
# 如果当前检查点是最佳的检查点,执行以下操作。
elif f is self.best:
# 设置键 k 为 "train_results" ,用于更新 best.pt 中的训练结果。
k = "train_results" # update best.pt train_metrics from last.pt
# 调用 strip_optimizer 函数加载最佳的检查点,并移除优化器状态。如果 ckpt 中存在 k 键,则更新 best.pt 中的训练结果。
strip_optimizer(f, updates={k: ckpt[k]} if k in ckpt else None)
# 使用 LOGGER 记录器打印信息,表示开始验证 f 检查点。
LOGGER.info(f"\nValidating {f}...") # 正在验证 {f}...
# 设置验证器的 plots 参数为当前训练器的 plots 参数。
self.validator.args.plots = self.args.plots
# 调用验证器的 validator 方法,传入最佳的检查点文件 f ,获取验证指标。
self.metrics = self.validator(model=f)
# 从验证指标中移除 "fitness" 键。
self.metrics.pop("fitness", None)
# 运行在训练周期结束时的回调函数。
self.run_callbacks("on_fit_epoch_end")
# 这个方法的主要作用是在训练结束后,使用最后和最佳的检查点对模型进行最终的评估和验证,获取验证指标,并运行相关的回调函数。这些步骤有助于评估模型的最终性能,并为后续的模型部署和应用提供依据。
# 这段代码定义了一个名为 check_resume 的方法,它用于检查是否存在恢复训练的检查点,并根据检查点更新参数。
# 这是 check_resume 方法的定义,它是一个实例方法,包含两个参数。
# 1.self :类的实例引用,用于访问类的属性和方法。
# 2.overrides :一个字典或对象,包含要覆盖的参数。
def check_resume(self, overrides):
# 检查恢复检查点是否存在并相应地更新参数。
"""Check if resume checkpoint exists and update arguments accordingly."""
# 从类的属性中获取 resume 标志,它指示是否需要从检查点恢复训练。
resume = self.args.resume
# 检查 resume 标志是否为真,如果是,则执行恢复操作。
if resume:
# 开始一个 try 块,用于捕获在恢复过程中可能发生的任何异常。
try:
# 检查 resume 参数是否是一个字符串或 Path 对象,并且对应的文件是否存在。
exists = isinstance(resume, (str, Path)) and Path(resume).exists()
# 调用 check_file 函数检查文件,如果文件存在,则获取最新的运行检查点。
# def check_file(file, suffix="", download=True, download_dir=".", hard=True):
# -> 如果文件路径为空,或者文件已经存在于本地,或者文件路径以 grpc:// 开头(表示 gRPC Triton 图像),则直接返回文件路径。
# -> 如果文件路径以 http:// 、 https:// 、 rtsp:// 、 rtmp:// 或 tcp:// 开头,并且 download 参数为 True ,则下载文件。
# -> 如果文件不需要下载,或者下载失败,则在文件系统中搜索文件。
# -> return file / return str(file) / return files[0] if len(files) else [] # return file
# def get_latest_run(search_dir="."): -> 在指定目录及其子目录中查找最新的名为 last.pt 的文件。 -> return max(last_list, key=os.path.getctime) if last_list else ""
last = Path(check_file(resume) if exists else get_latest_run())
# Check that resume data YAML exists, otherwise strip to force re-download of dataset
# 尝试加载检查点的权重,并获取检查点的参数。
# def attempt_load_weights(weights, device=None, inplace=True, fuse=False): -> 加载一个模型集合(ensemble)的权重,或者单个模型的权重。返回模型集合。函数返回创建的模型集合对象。 -> return ensemble
ckpt_args = attempt_load_weights(last).args
# 检查检查点参数中的 data 路径是否存在。
if not Path(ckpt_args["data"]).exists():
# 如果检查点的 data 路径不存在,则更新为当前的 data 路径。
ckpt_args["data"] = self.args.data
# 设置 resume 标志为真,表示将从检查点恢复训练。
resume = True
# 使用检查点的参数更新当前的参数。
# def get_cfg(cfg: Union[str, Path, Dict, SimpleNamespace] = DEFAULT_CFG_DICT, overrides: Dict = None):
# -> 从一个配置源(可以是字符串、路径、字典或 SimpleNamespace 对象)获取配置,并允许通过 overrides 字典来覆盖默认配置。返回一个 IterableSimpleNamespace 对象,它是一个可迭代的命名空间对象,其属性由 cfg 字典中的键值对初始化。
# -> return IterableSimpleNamespace(**cfg)
self.args = get_cfg(ckpt_args)
# 更新 model 和 resume 参数,指向最新的检查点。
self.args.model = self.args.resume = str(last) # reinstate model
# 遍历 imgsz 、 batch 和 device 参数。
for k in "imgsz", "batch", "device": # allow arg updates to reduce memory or update device on resume
# 检查这些参数是否在 overrides 中。
if k in overrides:
# 如果在 overrides 中,则更新当前参数。
setattr(self.args, k, overrides[k])
# 如果在恢复过程中发生异常,捕获异常。
except Exception as e:
# 抛出 FileNotFoundError 异常,提示检查点未找到,并提供恢复命令的示例。
raise FileNotFoundError(
"Resume checkpoint not found. Please pass a valid checkpoint to resume from, "
"i.e. 'yolo train resume model=path/to/last.pt'" # 未找到恢复检查点。请传递有效的检查点以进行恢复,即‘yolo train resume model=path/to/last.pt’。
) from e
# 更新类的 resume 属性。
self.resume = resume
# 这个方法的主要作用是在训练开始前检查是否存在恢复检查点,并根据检查点更新参数。这允许从之前的训练状态恢复训练,而不是从头开始。通过这种方式,可以在训练过程中断后继续训练,或者在不同的设备上恢复训练。
# 这段代码定义了一个名为 resume_training 的方法,它用于从给定的检查点恢复 YOLO 模型的训练。
# 这是 resume_training 方法的定义,它是一个实例方法,包含两个参数。
# 1.self :类的实例引用,用于访问类的属性和方法。
# 2.ckpt :检查点字典,包含了训练状态的信息。
def resume_training(self, ckpt):
# 从给定的时期和最佳适应度恢复 YOLO 训练。
"""Resume YOLO training from given epoch and best fitness."""
# 检查 ckpt 是否为 None 或者 self.resume 标志是否为 False 。如果是,则直接返回,不执行恢复操作。
if ckpt is None or not self.resume:
return
# 初始化 best_fitness 变量,用于存储最佳适应度值。
best_fitness = 0.0
# 从检查点中获取训练周期 epoch ,并将其转换为下一次训练的起始周期。
start_epoch = ckpt.get("epoch", -1) + 1
# 检查检查点中是否包含优化器状态。
if ckpt.get("optimizer", None) is not None:
# torch.nn.Module.load_state_dict(state_dict, strict=True)
# load_state_dict() 是 PyTorch 中 torch.nn.Module 类(即所有神经网络模型的基类)的一个方法,用于加载模型的参数。这个方法将传入的状态字典(state dictionary)中的参数加载到模型中,使得模型的权重和偏差与状态字典中的相匹配。
# 参数 :
# state_dict :一个包含模型参数的字典对象。通常由 torch.save() 保存的模型参数或通过 model.state_dict() 获取。
# strict :(可选)一个布尔值,默认为 True 。如果为 True ,则要求状态字典中的每个键都必须与模型中的参数匹配。如果为 False ,则忽略不匹配的键。
# 返回值 :无返回值。该方法直接修改模型的参数。
# 使用场景 :
# 当你从文件中加载模型权重或者在训练过程中恢复模型状态时。
# load_state_dict() 方法是 PyTorch 中管理和迁移模型权重的重要工具,特别是在模型保存、加载和迁移学习场景中。
# torch.optim.Optimizer.load_state_dict(state_dict)
# torch.optim.Optimizer.load_state_dict() 是 PyTorch 中优化器类的实例方法,用于加载优化器的状态字典。这个方法允许你将优化器恢复到之前保存的状态,这在训练过程中断后恢复训练时非常有用。
# 参数 :
# state_dict :一个字典,包含了优化器的状态信息。这个字典通常是通过调用 torch.optim.Optimizer.state_dict() 方法获得的。
# 返回值 :无返回值。该方法直接修改优化器的状态。
# 使用场景 :当你保存了优化器的状态并希望在以后的训练中恢复这个状态时。
# 注意事项 :
# 确保保存的状态字典与优化器兼容。如果你更改了模型的结构或者优化器的参数,直接加载旧的状态字典可能会导致错误。
# 在加载优化器状态之前,通常不需要手动清空优化器的当前状态,因为 load_state_dict() 方法会自动处理状态的更新。
# 如果包含优化器状态,则加载优化器状态。
# self.optimizer 是当前模型训练器中的优化器实例,它负责在训练过程中更新模型的参数。
# load_state_dict() 是 PyTorch 的 torch.optim.Optimizer 类的一个方法,用于从状态字典(state dictionary)中加载优化器的状态。
# ckpt 是一个字典,包含了检查点中的所有信息,包括模型权重、优化器状态等。
# ckpt["optimizer"] 指的是检查点字典中键为 "optimizer" 的值,它包含了优化器的状态字典。
self.optimizer.load_state_dict(ckpt["optimizer"]) # optimizer
# 更新最佳适应度值。
best_fitness = ckpt["best_fitness"]
# 检查是否存在 EMA(指数移动平均)模型,并且检查点中是否包含 EMA 状态。
if self.ema and ckpt.get("ema"):
# ema.load_state_dict(state_dict)
# ema.load_state_dict(state_dict) 用于恢复指数移动平均(Exponential Moving Average, EMA)模型的状态。在使用 torch_ema 库时, ExponentialMovingAverage 类(通常简写为 EMA)用来维护模型参数的指数移动平均。 load_state_dict() 方法在这里的作用是加载之前保存的 EMA 状态字典,恢复 EMA 模型的参数。
# 参数 :
# state_dict :一个包含 EMA 模型状态的字典对象。这个字典通常是通过调用 ema.state_dict() 方法获得的,包含了 EMA 参数的当前状态。
# 返回值 :无返回值。该方法直接修改 EMA 模型的状态。
# 使用场景 :
# 当你保存了 EMA 模型的状态并希望在以后的训练或验证中恢复这个状态时。
# 注意事项 :
# 确保保存的状态字典与 EMA 模型兼容。如果你更改了模型的结构或者 EMA 的参数,直接加载旧的状态字典可能会导致错误。
# 在加载 EMA 状态之前,通常不需要手动清空 EMA 模型的当前状态,因为 load_state_dict() 方法会自动处理状态的更新。
# 如果包含 EMA 状态,则加载 EMA 状态。
# .state_dict() 是 PyTorch 中的一个方法,用于获取模型或优化器等对象的状态字典。
self.ema.ema.load_state_dict(ckpt["ema"].float().state_dict()) # EMA
# 更新 EMA 更新次数。
self.ema.updates = ckpt["updates"]
# 确保 start_epoch 大于 0,否则抛出异常,表示训练已经完成,无需恢复。
assert start_epoch > 0, (
f"{self.args.model} training to {self.epochs} epochs is finished, nothing to resume.\n" # {self.args.model} 训练 {self.epochs} 个时期已完成,无需恢复。
f"Start a new training without resuming, i.e. 'yolo train model={self.args.model}'" # 开始新的训练而不恢复,即'yolo train model={self.args.model}'。
)
# 记录恢复训练的信息。
LOGGER.info(f"Resuming training {self.args.model} from epoch {start_epoch + 1} to {self.epochs} total epochs") # 恢复训练 {self.args.model},从第 {start_epoch + 1} 个周期到第 {self.epochs} 个周期。
# 这段代码是恢复训练过程中的一部分,用于处理模型微调(fine-tuning)的情况。
# 这个条件判断当前设置的训练周期数( self.epochs )是否小于从检查点恢复的训练周期数( start_epoch )。这可能发生在检查点已经包含了一定数量的训练周期,而用户希望在此基础上继续训练更多周期。
if self.epochs < start_epoch:
# 如果条件为真,使用 LOGGER 记录器打印一条信息,说明模型已经训练了多少周期,并且接下来将进行多少周期的微调。
# 这里的 self.model 表示模型的名称或描述, ckpt['epoch'] 表示检查点中已经训练的周期数, self.epochs 表示用户希望额外训练的周期数。
LOGGER.info(
f"{self.model} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {self.epochs} more epochs." # {self.model} 已训练了 {ckpt['epoch']} 个时期。 还将进行 {self.epochs} 个时期的微调。
)
# 更新 self.epochs 的值,将其增加检查点中的周期数。这样做是为了确保模型总共训练的周期数是用户期望的周期数加上检查点中已经完成的周期数。
self.epochs += ckpt["epoch"] # finetune additional epochs
# 这段代码的目的是确保在恢复训练时,如果检查点已经包含了一定数量的训练周期,那么模型将继续训练直到达到用户指定的总周期数。这允许用户在已有的检查点基础上进行微调,而不是从头开始训练,从而节省时间和计算资源。
# 更新最佳适应度。
self.best_fitness = best_fitness
# 更新起始周期。
self.start_epoch = start_epoch
# 检查是否需要关闭 Mosaic 数据增强。
if start_epoch > (self.epochs - self.args.close_mosaic):
# 如果需要,则关闭 Mosaic 数据增强。
self._close_dataloader_mosaic()
# 这个方法的主要作用是恢复 YOLO 模型的训练状态,包括优化器状态、EMA 状态、最佳适应度和起始周期。这允许从中断点继续训练,而不是从头开始。通过这种方式,可以有效地利用之前的训练成果,并节省训练资源。
# 这段代码定义了一个名为 _close_dataloader_mosaic 的方法,它用于在训练过程中关闭镶嵌(mosaic)数据增强。
# 这是 _close_dataloader_mosaic 方法的定义,它是一个实例方法,只接受 self 参数。
def _close_dataloader_mosaic(self):
# 更新数据加载器以停止使用马赛克增强。
"""Update dataloaders to stop using mosaic augmentation."""
# 检查 self.train_loader.dataset 是否有一个名为 mosaic 的属性。这个属性通常用于控制是否应用镶嵌数据增强。
if hasattr(self.train_loader.dataset, "mosaic"):
# 如果 mosaic 属性存在,将其设置为 False ,以禁用镶嵌数据增强。
self.train_loader.dataset.mosaic = False
# 检查 self.train_loader.dataset 是否有一个名为 close_mosaic 的方法。这个方法通常用于执行关闭镶嵌数据增强的相关操作。
if hasattr(self.train_loader.dataset, "close_mosaic"):
# 如果 close_mosaic 方法存在,使用 LOGGER 记录器打印一条信息,表示正在关闭数据加载器的镶嵌数据增强。
LOGGER.info("Closing dataloader mosaic") # 关闭数据加载器马赛克。
# 调用 close_mosaic 方法,并传入 self.args 作为参数。 self.args 可能包含关闭镶嵌数据增强所需的配置参数。
self.train_loader.dataset.close_mosaic(hyp=self.args)
# 这个方法的主要作用是在训练过程中的某个点停止使用镶嵌数据增强,这通常是基于训练策略或性能考虑。例如,可能在训练的后期阶段关闭镶嵌数据增强,以避免过度拟合或为了评估模型在更简单数据上的性能。通过这种方式,可以灵活地控制数据增强策略,以适应不同的训练需求。
# 这段代码定义了一个名为 build_optimizer 的方法,它用于构建和配置模型的优化器。
# 这是 build_optimizer 方法的定义,它是一个实例方法,包含六个参数。
# 1.model :要训练的模型。
# 2.name :优化器的名称,默认为 "auto" ,表示自动选择优化器。
# 3.lr :学习率,默认为 0.001 。
# 4.momentum :动量,默认为 0.9 。
# 5.decay :权重衰减,默认为 1e-5 。
# 6.iterations :迭代次数,用于确定优化器的选择和学习率的调整,默认为 1e5 。
def build_optimizer(self, model, name="auto", lr=0.001, momentum=0.9, decay=1e-5, iterations=1e5):
# 根据指定的优化器名称、学习率、动量、权重衰减和迭代次数,为给定模型构建优化器。
"""
Constructs an optimizer for the given model, based on the specified optimizer name, learning rate, momentum,
weight decay, and number of iterations.
Args:
model (torch.nn.Module): The model for which to build an optimizer.
name (str, optional): The name of the optimizer to use. If 'auto', the optimizer is selected
based on the number of iterations. Default: 'auto'.
lr (float, optional): The learning rate for the optimizer. Default: 0.001.
momentum (float, optional): The momentum factor for the optimizer. Default: 0.9.
decay (float, optional): The weight decay for the optimizer. Default: 1e-5.
iterations (float, optional): The number of iterations, which determines the optimizer if
name is 'auto'. Default: 1e5.
Returns:
(torch.optim.Optimizer): The constructed optimizer.
"""
# 这段代码是 build_optimizer 方法的一部分,它负责根据条件自动选择优化器及其参数。
g = [], [], [] # optimizer parameter groups
# 从 PyTorch 的 nn 模块中获取所有包含 "Norm" 的类,通常是归一化层,如 BatchNorm2d 。这些层的权重通常不应用权重衰减。
bn = tuple(v for k, v in nn.__dict__.items() if "Norm" in k) # normalization layers, i.e. BatchNorm2d()
# 如果优化器名称设置为 "auto" ,则自动选择优化器和相关参数。
if name == "auto":
# 使用 LOGGER 记录器打印信息,说明正在自动选择优化器,并忽略用户通过 self.args.lr0 和 self.args.momentum 设置的值。
LOGGER.info(
f"{colorstr('optimizer:')} 'optimizer=auto' found, " # {colorstr('optimizer:')} 发现'optimizer=auto',忽略'lr0={self.args.lr0}'和'momentum={self.args.momentum}'并自动确定最佳'optimizer'、'lr0'和'momentum'...
f"ignoring 'lr0={self.args.lr0}' and 'momentum={self.args.momentum}' and "
f"determining best 'optimizer', 'lr0' and 'momentum' automatically... "
)
# 尝试从模型中获取类别数 nc ,默认值为 10 。
nc = getattr(model, "nc", 10) # number of classes
# 根据类别数 nc 计算一个合适的学习率 lr_fit ,这里使用了一种经验公式。
lr_fit = round(0.002 * 5 / (4 + nc), 6) # lr0 fit equation to 6 decimal places
# 根据迭代次数 iterations 决定使用哪种优化器。如果迭代次数超过 10000 ,则使用 SGD 优化器,学习率为 0.01 ,动量为 0.9 。否则,使用 AdamW 优化器,学习率为 lr_fit ,动量为 0.9 。
name, lr, momentum = ("SGD", 0.01, 0.9) if iterations > 10000 else ("AdamW", lr_fit, 0.9)
# 设置 Adam 优化器的偏置学习率为 0.0 ,这是为了避免在预热阶段偏置的学习率过高。
self.args.warmup_bias_lr = 0.0 # no higher than 0.01 for Adam
# 这段代码的目的是提供一个智能的选择机制,根据模型的类别数和迭代次数自动选择优化器和调整学习率,以及为不同类型的层设置不同的参数组。这样的设计使得优化器的选择更加灵活,能够根据具体的训练任务自动调整,以期获得更好的训练效果。
# 这段代码是在构建优化器时用于分组模型参数的。它将模型的参数分为三组:带有权重衰减的权重、不带有权重衰减的权重(通常是归一化层的权重),以及偏置。
# torch.nn.Module.named_modules()
# named_modules() 是 PyTorch 中 torch.nn.Module 类的一个方法,用于遍历模型的所有模块(包括子模块),并返回它们的名称和模块本身。这个方法类似于 named_parameters() ,但是返回的是模块而不是参数。
# 参数 :无参数。
# 返回值 :
# 返回一个生成器(generator),包含模块的名称和模块对象。
# 使用场景 :
# 当你需要访问模型的每个模块,例如进行自定义处理、修改或记录模块信息时。
# named_modules() 方法在模型结构分析、调试和修改时非常有用,因为它允许你逐个访问模型的每个部分。
# 遍历模型中的所有模块及其名称。 named_modules() 方法返回一个生成器,包含模块的名称和模块本身。
for module_name, module in model.named_modules():
# 对于每个模块,遍历其参数及其名称。 named_parameters(recurse=False) 方法返回当前模块的参数名称和参数张量, recurse=False 表示不递归遍历子模块。
for param_name, param in module.named_parameters(recurse=False):
# 构造参数的完整名称。如果 module_name 不为空,则将模块名称和参数名称用点 . 连接起来;如果 module_name 为空(即参数属于模型的根级别),则只使用参数名称。
fullname = f"{module_name}.{param_name}" if module_name else param_name
# 检查参数的完整名称中是否包含 "bias" 。如果是,这意味着参数是偏置。
if "bias" in fullname: # bias (no decay)
# 如果参数是偏置,则将其添加到第三组 g[2] 中。这组参数在优化器中通常不应用权重衰减。
g[2].append(param)
# 检查当前模块是否是归一化层(如 BatchNorm2d)。 bn 是一个包含所有归一化层类的元组。
elif isinstance(module, bn): # weight (no decay)
# 如果参数属于归一化层,则将其添加到第二组 g[1] 中。这组参数在优化器中也不应用权重衰减。
g[1].append(param)
# 如果参数既不是偏置也不是归一化层的权重,则将其添加到第一组 g[0] 中。这组参数在优化器中应用权重衰减。
else: # weight (with decay)
# 将参数添加到第一组 g[0] 中。
g[0].append(param)
# 通过这种方式,模型的参数被分为三组,每组在优化器中有不同的处理方式。这种分组通常用于在训练过程中对不同类型的参数应用不同的优化策略,例如对偏置和归一化层的权重不使用权重衰减,而对其他权重使用权重衰减。这样的策略有助于提高模型训练的稳定性和性能。
# 这段代码是构建优化器的一部分,它根据指定的优化器名称来创建相应的优化器实例。
# 这个条件判断优化器名称是否为 Adam 或其变体之一。
if name in {"Adam", "Adamax", "AdamW", "NAdam", "RAdam"}:
# 如果条件为真,使用 getattr 函数从 torch.optim 模块中获取指定名称的优化器类。如果该名称不存在,则默认使用 optim.Adam 。
# 创建优化器实例,传入参数组 g[2] (通常是包含偏置的参数组),学习率 lr ,动量和衰减率 betas (对于 Adam 优化器),以及权重衰减 weight_decay (这里设置为 0.0,表示不应用权重衰减)。
optimizer = getattr(optim, name, optim.Adam)(g[2], lr=lr, betas=(momentum, 0.999), weight_decay=0.0)
# 如果优化器名称为 "RMSProp"。
elif name == "RMSProp":
# 创建 RMSProp 优化器实例,传入参数组 g[2] ,学习率 lr ,和动量 momentum 。
optimizer = optim.RMSprop(g[2], lr=lr, momentum=momentum)
# 如果优化器名称为 "SGD"。
elif name == "SGD":
# 创建 SGD 优化器实例,传入参数组 g[2] ,学习率 lr ,动量 momentum ,以及 nesterov=True 以启用 Nesterov 动量。
optimizer = optim.SGD(g[2], lr=lr, momentum=momentum, nesterov=True)
# 如果优化器名称不在支持的列表中。
else:
# 抛出 NotImplementedError 异常,提示未找到指定的优化器,并提供支持的优化器列表和请求支持的链接。
raise NotImplementedError(
f"Optimizer '{name}' not found in list of available optimizers " # 可用优化器 [Adam、AdamW、NAdam、RAdam、RMSProp、SGD、auto] 列表中未找到优化器“{name}”。
f"[Adam, AdamW, NAdam, RAdam, RMSProp, SGD, auto]."
"To request support for addition optimizers please visit https://github.com/ultralytics/ultralytics." # 如需请求对其他优化器的支持,请访问 https://github.com/ultralytics/ultralytics。
)
# 这段代码的目的是根据不同的优化器名称创建相应的优化器实例,并为每种优化器提供适当的参数。这样的设计使得优化器的选择和配置更加灵活,可以根据具体的训练任务和需求选择合适的优化器。
# 这段代码是构建优化器过程的最后一部分,它将不同的参数组添加到优化器中,并记录优化器的配置。
# 调用向优化器添加一个新的参数组。 g[0] 包含了需要应用权重衰减的参数, decay 是权重衰减的值。
optimizer.add_param_group({"params": g[0], "weight_decay": decay}) # add g0 with weight_decay
# 调用向优化器添加另一个参数组, g[1] 包含了不需要权重衰减的参数(通常是归一化层的参数),这里权重衰减被设置为 0.0 。
optimizer.add_param_group({"params": g[1], "weight_decay": 0.0}) # add g1 (BatchNorm2d weights)
# 使用 LOGGER 记录器打印优化器的配置信息。这包括 优化器的类型 、 学习率 lr 、 动量 momentum ,以及各参数组的数量和权重衰减设置。
# {len(g[1])} 表示 g[1] 参数组中的参数数量,这些参数不应用权重衰减。
# {len(g[0])} 表示 g[0] 参数组中的参数数量,这些参数应用了 decay 指定的权重衰减。
# {len(g[2])} 表示 g[2] 参数组中的参数数量,这些参数(通常是偏置)不应用权重衰减。
LOGGER.info(
f"{colorstr('optimizer:')} {type(optimizer).__name__}(lr={lr}, momentum={momentum}) with parameter groups "
f'{len(g[1])} weight(decay=0.0), {len(g[0])} weight(decay={decay}), {len(g[2])} bias(decay=0.0)'
)
# 返回构建并配置好的优化器实例。
return optimizer
# 这段代码的目的是确保优化器能够根据不同参数的特性(如是否需要权重衰减)来正确地更新参数。通过将参数分为不同的组,并为每组设置不同的优化策略,可以更精细地控制训练过程,从而可能提高模型的性能。此外,记录优化器的配置信息有助于调试和复现实验结果。
# 这个方法的主要作用是根据模型和给定的参数构建一个优化器实例,包括选择合适的优化器类型、学习率、动量和权重衰减。通过这种方式,可以灵活地配置模型的训练过程。