YOLOv11-ultralytics-8.3.67部分代码阅读笔记-train.py
train.py
ultralytics\models\yolo\detect\train.py
目录
train.py
1.所需的库和模块
2.class DetectionTrainer(BaseTrainer):
1.所需的库和模块
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
import math
import random
from copy import copy
import numpy as np
import torch.nn as nn
from ultralytics.data import build_dataloader, build_yolo_dataset
from ultralytics.engine.trainer import BaseTrainer
from ultralytics.models import yolo
from ultralytics.nn.tasks import DetectionModel
from ultralytics.utils import LOGGER, RANK
from ultralytics.utils.plotting import plot_images, plot_labels, plot_results
from ultralytics.utils.torch_utils import de_parallel, torch_distributed_zero_first
2.class DetectionTrainer(BaseTrainer):
# 这段代码定义了一个名为 DetectionTrainer 的类,它继承自 BaseTrainer ,用于基于检测模型的训练任务。
# 定义了一个名为 DetectionTrainer 的类,继承自 BaseTrainer 。这意味着它将继承 BaseTrainer 的所有属性和方法,并可以在此基础上进行扩展和定制。
class DetectionTrainer(BaseTrainer):
# 扩展 BaseTrainer 类的类,用于基于检测模型进行训练。
"""
A class extending the BaseTrainer class for training based on a detection model.
Example:
```python
from ultralytics.models.yolo.detect import DetectionTrainer
args = dict(model="yolo11n.pt", data="coco8.yaml", epochs=3)
trainer = DetectionTrainer(overrides=args)
trainer.train()
```
"""
# 这段代码定义了一个名为 build_dataset 的方法,它是 DetectionTrainer 类的一部分,用于构建用于训练或验证的 YOLO 数据集。
# 定义了 build_dataset 方法,它接收以下参数 :
# 1.img_path :图像文件夹的路径,包含用于训练或验证的图像。
# 2.mode :数据集的模式,默认为 "train" ,也可以是 "val" ,用于区分训练集和验证集。
# 3.batch :可选参数,指定批量大小,主要用于矩形训练( rect )。
def build_dataset(self, img_path, mode="train", batch=None):
# 构建 YOLO 数据集。
"""
Build YOLO Dataset.
Args:
img_path (str): Path to the folder containing images.
mode (str): `train` mode or `val` mode, users are able to customize different augmentations for each mode.
batch (int, optional): Size of batches, this is for `rect`. Defaults to None.
"""
# 计算了 模型的网格大小 ( gs ),它是模型最大步长( stride )的整数值。
# de_parallel(self.model) 是一个工具函数,用于处理模型是否并行化的情况,确保可以正确访问模型的属性。
# 如果模型未定义( self.model 为 None ),则网格大小默认为 0。
# 最后,取模型步长和 32 的较大值,确保网格大小至少为 32。这是因为 YOLO 模型通常需要将图像尺寸调整为 32 的倍数。
gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
# 调用了一个外部函数 build_yolo_dataset ,用于实际构建 YOLO 数据集 。
# self.args :包含训练配置的参数字典。
# img_path :图像文件夹路径。
# batch :批量大小,用于矩形训练。
# self.data :数据集配置,通常包含类别信息、数据路径等。
# mode :数据集模式( train 或 val )。
# rect :布尔值,指示是否启用矩形训练。在验证模式下( mode == "val" ),通常启用矩形训练以提高效率。
# stride :网格大小( gs ),用于确保图像尺寸与模型的步长兼容。
return build_yolo_dataset(self.args, img_path, batch, self.data, mode=mode, rect=mode == "val", stride=gs)
# 这段代码的核心功能是构建 YOLO 数据集,适用于训练或验证阶段。它通过以下步骤实现。计算模型的网格大小( gs ),确保图像尺寸与模型的步长兼容。调用 build_yolo_dataset 函数,传入必要的参数(如图像路径、批量大小、数据集配置、模式等),并根据模式决定是否启用矩形训练。返回构建好的数据集对象,供后续的数据加载器( DataLoader )使用。通过这种方式, build_dataset 方法为 YOLO 模型的训练和验证提供了灵活且高效的数据准备流程。
# 这段代码定义了 get_dataloader 方法,它是 DetectionTrainer 类的一部分,用于构建并返回一个 PyTorch 数据加载器( DataLoader ),以便在训练或验证阶段加载和处理数据。
# 定义了 get_dataloader 方法,接收以下参数 :
# 1.dataset_path :数据集的路径,包含图像和标注文件。
# 2.batch_size :每个批次的样本数量,默认为 16。
# 3.rank :分布式训练中当前进程的排名,默认为 0(单机训练时为 0)。
# 4.mode :数据加载器的模式,可以是 "train" 或 "val" ,默认为 "train" 。
def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode="train"):
# 构造并返回数据加载器。
"""Construct and return dataloader."""
# 检查 mode 参数是否为 "train" 或 "val" 。如果不是,则抛出异常,确保数据加载器的模式是有效的。
assert mode in {"train", "val"}, f"Mode must be 'train' or 'val', not {mode}." # 模式必须是“train”或“val”,而不是{mode}。
# 使用 torch_distributed_zero_first 上下文管理器,确保在分布式训练(DDP)中,数据集的缓存文件(如 .cache 文件)只被初始化一次。 rank 参数用于指定当前进程的排名。
# def torch_distributed_zero_first(local_rank: int): -> 用于在分布式训练中确保所有进程等待本地主进程(rank 0)完成某个任务后再继续执行。这种机制常用于在分布式训练中同步操作,例如在主进程完成数据加载或模型初始化后再让其他进程继续执行。
with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP
# 调用 self.build_dataset 方法,根据传入的 路径 、 模式 和 批量大小 构建 数据集对象 。 build_dataset 方法会根据模式(训练或验证)返回一个相应的数据集实例。
dataset = self.build_dataset(dataset_path, mode, batch_size)
# 根据模式设置数据加载器是否需要随机打乱数据。
# 如果是训练模式( mode == "train" ),则设置 shuffle=True ,以便在每个 epoch 内随机打乱数据。
# 如果是验证模式,则不需要随机打乱数据,因此 shuffle=False 。
shuffle = mode == "train"
# 检查数据集是否启用了矩形训练( rect=True )。矩形训练是一种优化技术,用于在验证阶段减少填充,提高效率。
if getattr(dataset, "rect", False) and shuffle:
# 如果启用了矩形训练并且 shuffle=True ,则发出警告,因为矩形训练与数据加载器的随机性不兼容。随后将 shuffle 设置为 False 。
LOGGER.warning("WARNING ⚠️ 'rect=True' is incompatible with DataLoader shuffle, setting shuffle=False") # 警告⚠️'rect=True' 与 DataLoader shuffle 不兼容,请设置 shuffle=False” 。
shuffle = False
# 根据模式设置数据加载器的工作进程数。 在训练模式下,使用 self.args.workers 中定义的工作进程数。 在验证模式下,为了加快数据加载速度,将工作进程数设置为训练模式的两倍。
workers = self.args.workers if mode == "train" else self.args.workers * 2
# 调用 build_dataloader 函数,根据构 建好的数据集 、 批量大小 、 工作进程数 、 随机性设置 以及 当前进程的排名 ,构建并返回一个 PyTorch 数据加载器( DataLoader )。
return build_dataloader(dataset, batch_size, workers, shuffle, rank) # return dataloader
# 这段代码的核心功能是构建一个适用于训练或验证阶段的 PyTorch 数据加载器。它通过以下步骤实现。验证模式参数的有效性。在分布式训练中,确保数据集缓存只被初始化一次。根据路径、模式和批量大小构建数据集。根据模式设置数据加载器的随机性,并检查矩形训练与随机性的兼容性。设置工作进程数,优化数据加载效率。调用 build_dataloader 函数,构建并返回数据加载器。通过这种方式, get_dataloader 方法为 YOLO 模型的训练和验证提供了高效、灵活且兼容分布式训练的数据加载机制。
# 这段代码定义了 preprocess_batch 方法,它是 DetectionTrainer 类的一部分,用于对批量图像数据进行预处理,以便它们能够被模型正确地接收和处理。
# 定义了 preprocess_batch 方法,接收一个参数。
# 1.batch :一个包含图像和标注信息的字典,通常由数据加载器( DataLoader )提供。
def preprocess_batch(self, batch):
# 通过缩放和转换为浮点数对一批图像进行预处理。
"""Preprocesses a batch of images by scaling and converting to float."""
# 将 批量图像数据 移动到指定的设备( self.device ),例如 GPU 或 CPU。 使用 non_blocking=True 参数,允许在非阻塞模式下进行数据传输,从而提高效率。 将图像数据转换为浮点类型,并归一化到 [0, 1] 范围(通过除以 255)。
batch["img"] = batch["img"].to(self.device, non_blocking=True).float() / 255
# 检查是否启用了多尺度训练。多尺度训练是一种技术,通过在训练过程中随机改变图像的尺寸,使模型能够更好地适应不同尺度的目标。
if self.args.multi_scale:
# 从输入的批量数据 batch 中提取 图像数据 ,并将其赋值给变量 imgs 。
imgs = batch["img"]
# 随机生成一个 目标尺寸 sz ,范围在原始图像尺寸的 50% 到 150% 之间。
# 确保目标尺寸是步长( self.stride )的倍数,因为 YOLO 模型要求输入图像的尺寸必须是步长的整数倍。
# random.randrange 用于生成随机整数范围, self.args.imgsz 是预定义的图像尺寸。
sz = (
random.randrange(int(self.args.imgsz * 0.5), int(self.args.imgsz * 1.5 + self.stride))
// self.stride
* self.stride
) # size
# 计算 缩放因子 sf ,即目标尺寸与图像最大边长的比例。 imgs.shape[2:] 表示图像的高和宽。
sf = sz / max(imgs.shape[2:]) # scale factor
# 如果缩放因子不为 1,则需要对图像进行缩放。
if sf != 1:
# 根据缩放因子计算 图像的新尺寸 ns 。 确保新尺寸是步长( self.stride )的倍数,通过 math.ceil 向上取整并调整为步长的整数倍。
ns = [
math.ceil(x * sf / self.stride) * self.stride for x in imgs.shape[2:]
] # new shape (stretched to gs-multiple)
# 使用 PyTorch 的 nn.functional.interpolate 函数对图像进行 双线性插值 ,调整图像的尺寸为新尺寸 ns 。
imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False)
# 将 缩放后的图像数据 更新到 批量字典 中。
batch["img"] = imgs
# 返回预处理后的批量数据,供后续的模型训练或验证使用。
return batch
# 这段代码的核心功能是对批量图像数据进行预处理,主要包括以下步骤。图像归一化和设备迁移:将图像数据归一化到 [0, 1] 范围,并将其移动到指定设备(如 GPU)。多尺度训练:如果启用了多尺度训练,随机生成目标尺寸,并根据缩放因子调整图像尺寸。双线性插值:使用双线性插值对图像进行缩放,确保图像尺寸符合模型的要求。通过这种方式, preprocess_batch 方法为 YOLO 模型的训练提供了灵活的图像预处理机制,同时支持多尺度训练以增强模型的泛化能力。
# 这段代码定义了 set_model_attributes 方法,它是 DetectionTrainer 类的一部分,用于将一些关键属性从训练配置和数据集信息附加到模型上。
# 定义了 set_model_attributes 方法,该方法不接收额外参数,仅操作类实例自身的属性。
def set_model_attributes(self):
# Nl = de_parallel(self.model).model[-1].nl # 检测层的数量(以 hyps 为单位)。
"""Nl = de_parallel(self.model).model[-1].nl # number of detection layers (to scale hyps)."""
# 这些注释代码展示了如何根据模型的层数( nl )、类别数量( self.data["nc"] )和图像尺寸( self.args.imgsz )动态调整超参数(如框损失权重 box 和类别损失权重 cls )。这些调整通常用于优化模型的训练过程,使其适应不同的数据集和模型架构。然而,这些代码在当前版本中被注释掉了,可能是因为这些调整在某些情况下不是必需的,或者需要根据具体需求手动启用。
# self.args.box *= 3 / nl # scale to layers
# self.args.cls *= self.data["nc"] / 80 * 3 / nl # scale to classes and layers
# self.args.cls *= (self.args.imgsz / 640) ** 2 * 3 / nl # scale to image size and layers
# 将类别数量附加到模型。
# self.data["nc"] 是数据集的 类别数量 。 将类别数量附加到模型的 nc 属性上,确保模型知道它需要处理多少个类别。
self.model.nc = self.data["nc"] # attach number of classes to model
# 将类别名称附加到模型。
# self.data["names"] 是一个 包含类别名称的列表 。 将类别名称附加到模型的 names 属性上,使得模型在推理时能够输出可读的类别名称,而不仅仅是类别索引。
self.model.names = self.data["names"] # attach class names to model
# 将超参数附加到模型。
# self.args 是一个 包含训练超参数的字典或对象 。 将超参数附加到模型的 args 属性上,使得模型能够访问训练时使用的超参数配置,例如学习率、图像尺寸等。
self.model.args = self.args # attach hyperparameters to model
# 这是一条待办事项(TODO),表示未来可能会实现的功能。它计划将类别权重( class_weights )附加到模型上,这些权重通常用于处理类别不平衡问题。 labels_to_class_weights 是一个假设的函数,用于根据数据集的标签计算类别权重。权重会被乘以类别数量( nc ),并移动到设备(如 GPU)上。
# TODO: self.model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc
# 这段代码的核心功能是将一些关键属性从训练配置和数据集信息附加到模型上,以便模型能够正确地进行训练和推理。这些属性包括。类别数量( nc ):模型需要知道数据集的类别数量,以便正确地构建输出层。类别名称( names ):模型需要知道类别名称,以便在推理时输出可读的类别标签。超参数( args ):模型需要访问训练时使用的超参数配置,以便在推理或后续训练中保持一致性。通过这种方式, set_model_attributes 方法确保模型能够正确地与数据集和训练配置对接,从而提高模型的灵活性和可扩展性。
# 这段代码定义了 get_model 方法,它是 DetectionTrainer 类的一部分,用于实例化并返回一个 YOLO 检测模型。该方法允许用户根据配置文件、预训练权重和日志输出需求来定制模型。
# 定义了 get_model 方法,接收以下参数 :
# 1.cfg :模型的配置文件路径或配置字典,用于定义模型的结构和参数。
# 2.weights :预训练权重的路径。如果提供,模型将加载这些权重。
# 3.verbose :布尔值,表示是否打印详细的初始化信息。默认为 True 。
def get_model(self, cfg=None, weights=None, verbose=True):
# 返回一个 YOLO 检测模型。
"""Return a YOLO detection model."""
# 实例化一个 DetectionModel 对象。
# cfg :模型的配置文件或字典,定义了模型的架构和超参数。
# nc :数据集的类别数量,从 self.data["nc"] 获取。
# verbose :控制是否打印模型初始化的详细信息。这里使用了 verbose and RANK == -1 ,表示只有在非分布式训练( RANK == -1 )时才打印详细信息,以避免在分布式训练中重复打印。
# class DetectionModel(BaseModel):
# -> 用于实现基于 YOLO 的目标检测模型,支持模型初始化、权重初始化、增强预测等功能,并根据模型类型选择合适的损失函数。
# -> def __init__(self, cfg="yolo11n.yaml", ch=3, nc=None, verbose=True): # model, input channels, number of classes
model = DetectionModel(cfg, nc=self.data["nc"], verbose=verbose and RANK == -1)
# 检查是否提供了预训练权重路径( weights )。
if weights:
# 如果提供了权重路径,则调用模型的 load 方法加载这些权重。这一步通常用于微调预训练模型或恢复训练。
model.load(weights)
# 返回实例化并加载权重(如果有的话)的检测模型对象。
return model
# 这段代码的核心功能是根据用户提供的配置和权重实例化一个 YOLO 检测模型。它通过以下步骤实现。实例化模型:根据配置文件和数据集的类别数量创建一个 DetectionModel 对象。加载权重:如果提供了权重路径,则加载预训练权重,以便进行微调或恢复训练。返回模型:返回配置好的检测模型对象,供后续训练或推理使用。通过这种方式, get_model 方法为 YOLO 模型的训练和推理提供了灵活的初始化机制,支持从头开始训练、微调预训练模型或恢复训练。
# 这段代码定义了 get_validator 方法,它是 DetectionTrainer 类的一部分,用于创建并返回一个用于验证 YOLO 模型性能的验证器( DetectionValidator )。验证器通常用于在验证集上评估模型的性能,并记录相关指标。
# 定义了 get_validator 方法,该方法不接收额外参数,仅操作类实例自身的属性。
def get_validator(self):
# 返回用于 YOLO 模型验证的 DetectionValidator。
"""Returns a DetectionValidator for YOLO model validation."""
# 定义了一个元组 self.loss_names ,包含三种损失的名称。
# box_loss : 边界框回归损失 ,用于衡量预测框与真实框的匹配程度。
# cls_loss : 分类损失 ,用于衡量预测类别与真实类别的匹配程度。
# dfl_loss : 分布焦点损失 (Distribution Focal Loss),通常用于改进分类损失的计算。
# 这些损失名称可能用于后续的性能评估和日志记录。
self.loss_names = "box_loss", "cls_loss", "dfl_loss"
# 实例化了一个 DetectionValidator 对象,并将其返回。
# self.test_loader :验证集的数据加载器( DataLoader ),用于在验证阶段加载数据。
# save_dir :验证结果的保存路径,通常是一个文件夹,用于存储日志、图表和验证指标。
# args :训练配置的副本(通过 copy(self.args) )。这确保了验证器可以访问与训练相同的超参数配置。
# _callbacks :回调函数列表,这些回调函数在验证过程中被调用,用于执行额外的任务,例如日志记录或性能评估。
return yolo.detect.DetectionValidator(
self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
)
# 这段代码的核心功能是创建一个用于验证 YOLO 模型性能的 DetectionValidator 对象。它通过以下步骤实现。定义损失名称:明确在验证过程中需要关注的损失类型。实例化验证器:根据验证集的数据加载器、保存路径、训练配置和回调函数实例化 DetectionValidator 。返回验证器:返回配置好的验证器对象,供后续的验证流程使用。通过这种方式, get_validator 方法为 YOLO 模型的验证阶段提供了标准化的工具,确保模型的性能可以在统一的框架下被评估和记录。
# 这段代码定义了 label_loss_items 方法,它是 DetectionTrainer 类的一部分,用于将损失项( loss_items )与对应的损失名称( self.loss_names )进行匹配,并返回一个格式化的字典。如果没有提供损失项,则返回损失名称的键列表。这个方法主要用于在训练或验证过程中记录和展示损失信息。
# 定义了 label_loss_items 方法,接收以下参数 :
# 1.loss_items :一个可选的列表,包含当前批次或阶段的损失值(通常是张量)。
# 2.prefix :一个字符串,默认为 "train" ,用于标识损失值的来源(例如 "train" 或 "val" )。
def label_loss_items(self, loss_items=None, prefix="train"):
# 返回带有标记训练损失项张量的损失字典。
# 分类不需要,但分割和检测需要
"""
Returns a loss dict with labelled training loss items tensor.
Not needed for classification but necessary for segmentation & detection
"""
# 通过列表推导式生成 损失名称 的键( keys ),每个键由 prefix 和 self.loss_names 中的损失名称拼接而成。
# 例如,如果 prefix="train" 且 self.loss_names=["box_loss", "cls_loss", "dfl_loss"] ,则生成的键为 :
# ["train/box_loss", "train/cls_loss", "train/dfl_loss"]
keys = [f"{prefix}/{x}" for x in self.loss_names]
# 检查是否提供了损失项( loss_items )。如果提供了损失项,则继续处理;否则,返回损失名称的键列表。
if loss_items is not None:
# 遍历 loss_items 列表,将每个损失值(通常是张量)转换为浮点数,并保留 5 位小数。 这一步确保损失值以统一的格式存储和展示。
loss_items = [round(float(x), 5) for x in loss_items] # convert tensors to 5 decimal place floats
# 使用 zip(keys, loss_items) 将 损失名称 ( keys )与 损失值 ( loss_items )进行匹配。
# 将匹配后的结果转换为字典并返回。例如 :
# {"train/box_loss": 0.12345, "train/cls_loss": 0.67890, "train/dfl_loss": 0.23456}
return dict(zip(keys, loss_items))
# 如果没有提供损失项( loss_items=None ),则直接返回损失名称的键列表( keys )。这通常用于初始化或定义日志记录的键。
else:
return keys
# 这段代码的核心功能是将损失项与对应的损失名称进行匹配,并返回一个格式化的字典。如果没有提供损失项,则返回损失名称的键列表。这个方法的主要用途包括。格式化损失信息:将损失值转换为浮点数并保留 5 位小数,便于记录和展示。动态生成键:根据 prefix 参数动态生成损失名称的键,支持训练和验证阶段的区分。灵活的返回值:根据是否提供损失项,返回字典或键列表,适应不同的使用场景。通过这种方式, label_loss_items 方法为 YOLO 模型的训练和验证过程提供了标准化的损失记录机制,便于跟踪和分析模型的性能。
# 这段代码定义了 progress_string 方法,它是 DetectionTrainer 类的一部分,用于生成一个格式化的字符串,用于在训练过程中显示训练进度。这个字符串包含了多个关键指标的占位符,用于在每个 epoch 或每个 batch 的日志中输出当前的训练状态。
# 定义了 progress_string 方法,该方法不接收额外参数,仅操作类实例自身的属性。
def progress_string(self):
# 返回带有纪元、GPU 内存、损失、实例和大小的训练进度格式化字符串。
"""Returns a formatted string of training progress with epoch, GPU memory, loss, instances and size."""
# "\n" + "%11s" * (4 + len(self.loss_names)) :
# 这是一个格式化字符串,用于定义输出的格式。
# "\n" :在字符串开头添加一个换行符,确保输出内容在新的一行开始。
# "%11s" * (4 + len(self.loss_names)) :根据损失项的数量( self.loss_names )动态生成占位符。每个占位符宽度为 11 个字符,用于对齐输出。
# 4 + len(self.loss_names) :固定部分(Epoch、GPU_mem、Instances、Size)加上动态损失项的数量。
# % ("Epoch", "GPU_mem", *self.loss_names, "Instances", "Size") :
# 使用 % 操作符将占位符替换为具体的列标题。
# "Epoch" :表示当前训练的轮数。
# "GPU_mem" :表示 GPU 内存使用情况。
# *self.loss_names :展开损失名称列表(例如 "box_loss" , "cls_loss" , "dfl_loss" ),这些是动态部分。
# "Instances" :表示当前批次中的实例数量(例如目标框的数量)。
# "Size" :表示当前批次图像的尺寸。
return ("\n" + "%11s" * (4 + len(self.loss_names))) % (
"Epoch",
"GPU_mem",
*self.loss_names,
"Instances",
"Size",
)
# 输出示例 :
# 假设 self.loss_names = ["box_loss", "cls_loss", "dfl_loss"] ,则生成的格式化字符串为 :
# \n%11s%11s%11s%11s%11s%11s%11s
# 替换后的输出为 :
# \nEpoch GPU_mem box_loss cls_loss dfl_loss Instances Size
# 这段代码的核心功能是生成一个格式化的字符串,用于在训练过程中输出训练进度。它通过以下方式实现。动态生成占位符:根据损失项的数量动态生成格式化字符串。固定和动态列标题:包含固定的列标题(Epoch、GPU_mem、Instances、Size)和动态的损失名称。对齐输出:每个列标题宽度为 11 个字符,确保输出内容对齐。通过这种方式, progress_string 方法为 YOLO 模型的训练过程提供了一个清晰、标准化的日志输出格式,便于跟踪训练进度和性能指标。
# 这段代码定义了 plot_training_samples 方法,它是 DetectionTrainer 类的一部分,用于可视化训练过程中的样本图像。该方法通过调用 plot_images 函数,将带有标注(边界框和类别)的图像绘制出来,并保存到指定路径。这有助于开发者直观地检查模型的训练效果和数据标注的准确性。
# 定义了 plot_training_samples 方法,接收以下参数 :
# 1.batch :一个包含训练数据的字典,通常由数据加载器提供,包含图像、类别、边界框等信息。
# 2.ni :当前批次的索引,用于在保存图像时标识不同的批次。
def plot_training_samples(self, batch, ni):
# 绘制训练样本及其注释。
"""Plots training samples with their annotations."""
# 调用 plot_images 函数。
plot_images(
# 提取批量中的图像数据。 batch["img"] 是一个张量,包含批量图像的像素值。
images=batch["img"],
# 提取批量索引,用于标识每个图像在批次中的位置。
batch_idx=batch["batch_idx"],
# 提取类别标签。 batch["cls"] 是一个张量,包含每个目标框的类别索引。 使用 .squeeze(-1) 去掉多余的维度(如果存在),确保类别标签的形状与边界框匹配。
cls=batch["cls"].squeeze(-1),
# 提取边界框坐标。 batch["bboxes"] 是一个张量,包含每个目标框的坐标信息(通常是 [x1, y1, x2, y2] 格式)。
bboxes=batch["bboxes"],
# 提取图像文件路径。 batch["im_file"] 是一个列表,包含每个图像的原始文件路径。
paths=batch["im_file"],
# 指定保存可视化图像的文件名。图像将被保存到 self.save_dir 目录下,文件名格式为 "train_batch{ni}.jpg" ,其中 {ni} 是当前批次的索引。
fname=self.save_dir / f"train_batch{ni}.jpg",
# 提供一个回调函数( self.on_plot ),在绘制图像时被调用。这个回调函数可以用于执行额外的绘图操作或自定义行为。
on_plot=self.on_plot,
)
# 这段代码的核心功能是将训练过程中的样本图像及其标注(类别和边界框)可视化,并保存到指定路径。它通过以下步骤实现。提取批量数据:从 batch 字典中提取图像、类别、边界框和文件路径。调用 plot_images 函数:将提取的数据传递给 plot_images 函数,生成可视化图像。保存图像:将可视化图像保存到指定路径,文件名包含批次索引,便于区分不同的批次。通过这种方式, plot_training_samples 方法为训练过程提供了直观的可视化工具,帮助开发者检查数据标注的准确性、模型的预测效果以及训练过程中的异常情况。
# 这段代码定义了 plot_metrics 方法,它是 DetectionTrainer 类的一部分,用于从训练过程中记录的 CSV 文件中提取指标,并将这些指标可视化为图表(如 results.png )。这个方法通常用于展示训练过程中的性能变化,例如损失值、准确率、mAP 等。
# 定义了 plot_metrics 方法,该方法不接收额外参数,仅操作类实例自身的属性。
def plot_metrics(self):
# 从 CSV 文件绘制指标。
"""Plots metrics from a CSV file."""
# 调用 plot_results 函数。
# file=self.csv :指定包含训练指标的 CSV 文件路径。 self.csv 是类实例的一个属性,通常 保存了训练过程中记录的 损失值 、 准确率 、 mAP 等指标 。 这个文件通常在训练过程中动态生成,每一行对应一个 epoch 的指标。
# on_plot=self.on_plot :提供一个回调函数( self.on_plot ),在绘图过程中被调用。这个回调函数可以用于执行额外的绘图操作或自定义行为,例如调整图表的样式、添加注释等。
# # save results.png :这是一个注释,说明调用 plot_results 函数后会生成一个名为 results.png 的图表文件。这个文件通常保存在训练日志目录中,用于展示训练过程中的性能变化。
plot_results(file=self.csv, on_plot=self.on_plot) # save results.png
# 这段代码的核心功能是从训练过程中记录的 CSV 文件中提取指标,并将这些指标可视化为图表。它通过以下步骤实现。调用 plot_results 函数:从指定的 CSV 文件( self.csv )中读取训练指标。生成可视化图表:将指标绘制为图表(如 results.png ),展示训练过程中的性能变化。支持自定义回调:通过 on_plot 回调函数,用户可以自定义图表的样式或添加额外的绘图操作。通过这种方式, plot_metrics 方法为训练过程提供了一个直观的可视化工具,帮助开发者分析模型的性能变化,优化训练过程,并快速发现潜在问题。
# 这段代码定义了 plot_training_labels 方法,它是 DetectionTrainer 类的一部分,用于可视化训练数据集中目标框(bounding boxes)和类别标签(class labels)的分布情况。此方法可以帮助开发者了解数据集的标注情况,例如目标框的尺寸分布、类别分布等,从而更好地调整模型训练策略。
# 定义了 plot_training_labels 方法,该方法不接收额外参数,仅操作类实例自身的属性。
def plot_training_labels(self):
# 创建 YOLO 模型的标记训练图。
"""Create a labeled training plot of the YOLO model."""
# 提取并合并目标框数据。
# self.train_loader.dataset.labels :从训练数据加载器( DataLoader )的数据集中提取 标注信息 。 labels 是一个列表,每个元素包含一个批次的标注信息。
# lb["bboxes"] :从每个标注信息中提取 目标框坐标 (通常是 [x1, y1, x2, y2] 或 [x_center, y_center, width, height] 格式)。
# np.concatenate(..., 0) :将所有批次的目标框坐标合并为一个 NumPy 数组,按批次顺序拼接。
boxes = np.concatenate([lb["bboxes"] for lb in self.train_loader.dataset.labels], 0)
# 提取并合并类别标签。
# lb["cls"] :从每个标注信息中提取 类别标签 (通常是类别索引)。
# np.concatenate(..., 0) :将所有批次的类别标签合并为一个 NumPy 数组,按批次顺序拼接。
cls = np.concatenate([lb["cls"] for lb in self.train_loader.dataset.labels], 0)
# 调用 plot_labels 函数。
# boxes :合并后的目标框坐标数组。
# cls.squeeze() :合并后的类别标签数组。使用 .squeeze() 去掉多余的维度(如果存在),确保类别标签的形状为一维数组。
# names=self.data["names"] :类别名称列表,用于在可视化中显示类别标签的名称。
# save_dir=self.save_dir :指定保存可视化结果的目录。
# on_plot=self.on_plot :回调函数,用于在绘图过程中执行自定义操作(例如调整图表样式或添加注释)。
plot_labels(boxes, cls.squeeze(), names=self.data["names"], save_dir=self.save_dir, on_plot=self.on_plot)
# 这段代码的核心功能是可视化训练数据集中目标框和类别标签的分布情况。它通过以下步骤实现。提取标注信息:从训练数据加载器中提取目标框坐标和类别标签。将所有批次的标注信息合并为两个 NumPy 数组: boxes 和 cls 。调用 plot_labels 函数:使用 plot_labels 函数将目标框和类别标签的分布情况绘制为图表。提供类别名称列表( self.data["names"] ),以便在图表中显示可读的类别标签。指定保存目录( self.save_dir )和回调函数( self.on_plot ),用于自定义图表的样式或行为。通过这种方式, plot_training_labels 方法为开发者提供了一个直观的工具,用于分析训练数据集的标注情况,从而更好地调整模型训练策略。例如,开发者可以通过可视化结果发现类别不平衡问题或目标框尺寸分布的异常情况。
# 这段代码定义了 auto_batch 方法,它是 DetectionTrainer 类的一部分,用于自动计算适合当前模型和硬件的最优批量大小。这个方法的核心思想是根据训练数据集中目标框的数量动态调整批量大小,以确保训练过程中的内存占用在合理范围内。
# 定义了 auto_batch 方法,该方法不接收额外参数,仅操作类实例自身的属性。
def auto_batch(self):
# 通过计算模型的内存占用来获取批次大小。
"""Get batch size by calculating memory occupation of model."""
# 调用 self.build_dataset 方法,构建一个 训练数据集 。
# self.trainset :训练数据集的路径或配置信息。
# mode="train" :指定数据集模式为训练模式。
# batch=16 :指定初始批量大小为 16。这个值通常用于初始化数据集,后续会根据实际计算结果调整。
train_dataset = self.build_dataset(self.trainset, mode="train", batch=16)
# 4 for mosaic augmentation
# 计算最大目标框数量。
# train_dataset.labels :从构建好的训练数据集中提取标注信息。 labels 是一个列表,每个元素包含一个图像的标注信息(例如目标框和类别标签)。
# len(label["cls"]) :计算每个图像中的目标框数量(通过类别标签的数量间接获取)。
# max(...) :找到数据集中单个图像中目标框的最大数量。
# * 4 :将最大目标框数量乘以 4。这是因为使用了 Mosaic 数据增强技术,Mosaic 会将 4 张图像拼接在一起,因此目标框的数量会增加 4 倍。
max_num_obj = max(len(label["cls"]) for label in train_dataset.labels) * 4
# 调用父类( BaseTrainer )的 auto_batch 方法,传入计算得到的 最大目标框数量 ( max_num_obj )。
# 父类的 auto_batch 方法会根据硬件资源(如 GPU 内存)和目标框数量动态计算最优的批量大小。
# 返回计算得到的 最优批量大小 。
return super().auto_batch(max_num_obj)
# 这段代码的核心功能是自动计算适合当前模型和硬件的最优批量大小。它通过以下步骤实现。构建训练数据集:使用初始批量大小构建训练数据集。计算最大目标框数量:统计训练数据集中单个图像的最大目标框数量,并考虑 Mosaic 数据增强的影响。调用父类方法:将最大目标框数量传递给父类的 auto_batch 方法,动态计算最优批量大小。通过这种方式, auto_batch 方法能够根据训练数据集的特点和硬件资源自动调整批量大小,确保训练过程中的内存占用在合理范围内,同时提高训练效率。
# DetectionTrainer 类是一个用于目标检测模型训练的高级工具类,继承自 BaseTrainer 并针对 YOLO 检测模型进行了扩展。它提供了从数据集构建、模型初始化、训练过程可视化到性能评估的全流程支持。通过灵活的配置和丰富的功能, DetectionTrainer 能够高效地处理训练数据、动态调整训练参数,并实时监控训练进度,同时支持多尺度训练、自动批量大小调整等高级特性。此外,它还通过可视化工具帮助开发者直观地分析模型性能和数据分布,从而优化训练策略,提升模型的泛化能力和准确性。