engine.py
yolov6\core\engine.py
目录
engine.py
1.所需的库和模块
2.class Trainer:
1.所需的库和模块
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# 主要负责模型的训练过程。
from ast import Pass
import os
import time
from copy import deepcopy
import os.path as osp
from tqdm import tqdm
import cv2
import numpy as np
import math
import torch
from torch.cuda import amp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.tensorboard import SummaryWriter
import tools.eval as eval
from yolov6.data.data_load import create_dataloader
from yolov6.models.yolo import build_model
from yolov6.models.yolo_lite import build_model as build_lite_model
from yolov6.models.losses.loss import ComputeLoss as ComputeLoss
from yolov6.models.losses.loss_fuseab import ComputeLoss as ComputeLoss_ab
from yolov6.models.losses.loss_distill import ComputeLoss as ComputeLoss_distill
from yolov6.models.losses.loss_distill_ns import ComputeLoss as ComputeLoss_distill_ns
from yolov6.utils.events import LOGGER, NCOLS, load_yaml, write_tblog, write_tbimg
from yolov6.utils.ema import ModelEMA, de_parallel
from yolov6.utils.checkpoint import load_state_dict, save_checkpoint, strip_optimizer
from yolov6.solver.build import build_optimizer, build_lr_scheduler
from yolov6.utils.RepOptimizer import extract_scales, RepVGGOptimizer
from yolov6.utils.nms import xywh2xyxy
from yolov6.utils.general import download_ckpt2.class Trainer:
class Trainer:
# 1.args ：这是一个包含命令行参数的对象，可能包括训练配置、数据路径、模型参数等。
# 2.cfg ：这是一个配置对象，通常从配置文件加载，包含模型结构、训练策略、优化器设置等详细信息。
# 3.device ：这是指定的训练设备，可以是CPU或GPU，用于指示模型和数据应该在哪个设备上进行训练。
def __init__(self, args, cfg, device):
self.args = args
self.cfg = cfg
self.device = device
# 检查 args 中是否有 resume 属性，这个属性通常用于指示是否从之前训练的检查点（checkpoint）恢复训练。
if args.resume:
# torch.load(f, map_location=None, pickle_module=pickle, **pickle_load_args)
# torch.load 是 PyTorch 中用于加载保存的 PyTorch 对象的函数。这个函数可以加载之前使用 torch.save 保存的文件，这些文件可以包含模型参数、优化器状态、张量等。
# 参数说明 ：
# 1. f ：要加载的文件名或文件对象。可以是一个字符串路径，也可以是一个文件对象。
# 2. map_location ：一个字符串或一个函数，指定如何映射存储位置。默认为 None ，意味着使用默认存储位置。如果设置为 'cpu' ，那么所有张量都会映射到CPU上。
# 如果设置为 'cuda:device_id' ，那么张量会被映射到指定的GPU上。也可以是一个函数，该函数接受一个存储位置字符串并返回一个新的存储位置字符串。
# 3. pickle_module ：用于序列化和反序列化的模块。默认为 Python 标准库中的 pickle 。
# 4. **pickle_load_args ：传递给 pickle.load 函数的其他参数。
# 返回值 ：
# 加载的 PyTorch 对象，这可以是张量、模型、优化器状态等。
# 如果 args.resume 为真，使用PyTorch的 torch.load 函数加载检查点文件。 map_location='cpu' 参数确保检查点文件被加载到CPU上，无论它原来保存在哪个设备上。
self.ckpt = torch.load(args.resume, map_location='cpu')
# 获取 args 中的 rank 属性，并赋值给 self.rank 。 rank 通常用于分布式训练中标识每个进程的编号。
self.rank = args.rank
# 获取 args 中的 local_rank 属性，并赋值给 self.local_rank 。 local_rank 标识本地进程的编号，用于GPU上的分布式训练。
self.local_rank = args.local_rank
# 获取 args 中的 world_size 属性，并赋值给 self.world_size 。 world_size 表示分布式训练中总的进程数。
self.world_size = args.world_size
# 设置 self.main_process 为布尔值，表示当前进程是否为主进程。在分布式训练中，通常只有一个主进程负责保存模型和输出日志等任务。
self.main_process = self.rank in [-1, 0]
# 获取 args 中的 save_dir 属性，并赋值给 self.save_dir 。 save_dir 指定了模型检查点和日志文件的保存目录。
self.save_dir = args.save_dir
# get data loader 获取数据加载器。
# def load_yaml(file_path): -> 从 yaml 文件加载数据。函数返回从 YAML 文件中加载的数据，通常是一个字典或列表，具体取决于 YAML 文件的内容。 -> return data_dict
# 这行代码调用了一个名为 load_yaml 的函数，用于加载YAML格式的数据集配置文件。 args.data_path 提供了配置文件的路径。加载后的数据被存储在 self.data_dict 中，这是一个字典，包含了数据集的所有配置信息。
self.data_dict = load_yaml(args.data_path)
# 从 self.data_dict 字典中获取 'nc' 键对应的值，这个值代表数据集中的类别数量（ num_classes ），并将其赋值给类的实例变量 self.num_classes 。
self.num_classes = self.data_dict['nc']
# self.get_data_loader 方法来创建训练和验证数据的加载器。这个方法接受三个参数： args （命令行参数）， cfg （配置信息），和 self.data_dict （数据集配置字典）。
# 方法的返回值是两个数据加载器，分别用于训练和验证，它们被赋值给 self.train_loader 和 self.val_loader 。
# def get_data_loader(args, cfg, data_dict): -> 它用于创建并返回训练和验证数据加载器（ DataLoader ）。 -> return train_loader, val_loader
self.train_loader, self.val_loader = self.get_data_loader(args, cfg, self.data_dict)
# get model and optimizer 获取模型和优化器。
# 这行代码检查是否启用了蒸馏（ distill ）以及模型类型是否为 YOLOv6n 或 YOLOv6s 。如果是，则设置 self.distill_ns 为 True ，表示将使用蒸馏训练小型网络。
self.distill_ns = True if self.args.distill and self.cfg.model.type in ['YOLOv6n','YOLOv6s'] else False
# 调用 self.get_model 方法来创建模型实例。这个方法接受命令行参数、配置、类别数量和设备作为输入，并返回一个模型实例。
# def get_model(self, args, cfg, nc, device): -> 用于根据提供的参数构建和配置YOLOv6模型。返回构建和配置好的模型实例。 -> return model
model = self.get_model(args, cfg, self.num_classes, device)
# 如果启用了蒸馏，执行以下操作。
if self.args.distill:
# 检查是否启用了 fuse_ab 。如果启用了蒸馏并且同时启用了 fuse_ab ，则记录错误并退出程序，因为蒸馏模型不应该启用 fuse_ab 。
if self.args.fuse_ab:
LOGGER.error('ERROR in: Distill models should turn off the fuse_ab.\n') # 错误：蒸馏模型应该关闭fuse_ab。
exit()
# 如果启用了蒸馏，调用 self.get_teacher_model 方法来创建教师模型实例。
# def get_teacher_model(self, args, cfg, nc, device):
# -> 它用于创建和配置一个教师模型，这在知识蒸馏（Knowledge Distillation）的场景中是常见的。返回构建和配置好的教师模型实例。-> return model
self.teacher_model = self.get_teacher_model(args, cfg, self.num_classes, device)
# 如果启用了量化，调用 self.quant_setup 方法来设置模型的量化。
if self.args.quant:
self.quant_setup(model, cfg, device)
# 如果训练模式是 repopt （Re-Optimization），则执行以下操作。
if cfg.training_mode == 'repopt':
# 从预训练模型加载量化缩放因子。
scales = self.load_scale_from_pretrained_models(cfg, device)
# 如果配置中指定了预训练模型，则不重新初始化模型参数。
reinit = False if cfg.model.pretrained is not None else True
# 创建一个 RepVGGOptimizer 优化器实例，用于 repopt 训练模式。
self.optimizer = RepVGGOptimizer(model, scales, args, cfg, reinit=reinit)
# 如果不是 repopt 训练模式，执行以下操作。
else:
# 调用 self.get_optimizer 方法来创建一个优化器实例。
self.optimizer = self.get_optimizer(args, cfg, model)
# 调用 self.get_lr_scheduler 方法来创建学习率调度器和学习率函数。
self.scheduler, self.lf = self.get_lr_scheduler(args, cfg, self.optimizer)
# 如果当前进程是主进程，则创建一个模型指数移动平均（EMA）实例。EMA用于平滑模型参数，提高训练稳定性。
self.ema = ModelEMA(model) if self.main_process else None
# tensorboard
# 这行代码检查当前进程是否为主进程（ self.main_process ）。如果是主进程，则创建一个TensorBoard的 SummaryWriter 实例，用于记录训练过程中的日志信息，这些信息可以用于TensorBoard可视化。
# SummaryWriter 是TensorBoard的一个API，它允许用户将训练过程中的数据（如损失、准确率等）写入日志文件，以便后续使用TensorBoard进行可视化分析。
# self.save_dir 是保存TensorBoard日志文件的目录。
# 如果当前进程不是主进程，则 self.tblogger 被设置为 None ，这意味着非主进程不会记录TensorBoard日志。
self.tblogger = SummaryWriter(self.save_dir) if self.main_process else None
# 这行代码设置训练的起始周期（epoch）为0。在深度学习训练中，一个周期指的是对整个数据集进行一次前向和反向传播的过程。这个变量通常用于记录从哪个周期开始训练，特别是在从检查点（checkpoint）恢复训练时。
self.start_epoch = 0
#resume
# 这段代码处理从检查点（checkpoint）恢复模型、优化器、调度器和指数移动平均（EMA）的状态，并更新训练的起始周期。
# 检查当前实例是否有 ckpt 属性。 ckpt 通常是一个从文件加载的检查点字典，包含了模型参数、优化器状态、训练周期等信息。
if hasattr(self, "ckpt"):
# 从检查点中提取模型状态，并确保所有参数都是32位浮点数（FP32），然后获取模型的状态字典。
resume_state_dict = self.ckpt['model'].float().state_dict() # checkpoint state_dict as FP32 检查点 state_dict 作为 FP32。
# 将检查点中的状态字典加载到模型中， strict=True 确保状态字典中的所有键都必须与模型中的参数名称匹配。
model.load_state_dict(resume_state_dict, strict=True) # load
# 从检查点中提取训练周期，并将其设置为恢复训练的起始周期，即上一周期的下一个周期。
self.start_epoch = self.ckpt['epoch'] + 1
# 从检查点中加载优化器的状态。
self.optimizer.load_state_dict(self.ckpt['optimizer'])
# 从检查点中加载调度器的状态。
self.scheduler.load_state_dict(self.ckpt['scheduler'])
# 检查当前是否为主进程。
if self.main_process:
# 为主进程加载EMA的状态字典。
self.ema.ema.load_state_dict(self.ckpt['ema'].float().state_dict())
# 更新EMA的更新次数。
self.ema.updates = self.ckpt['updates']
# 从数据字典中提取类别数量和类别名称，并更新到模型的属性中。
# def parallel_model(args, model, device): -> 用于根据提供的参数配置模型以在不同的并行计算模式下运行。返回包装后的模型，无论是DP还是DDP模式。 -> return model
self.model = self.parallel_model(args, model, device)
# 从数据字典中提取类别数量和类别名称，并更新到模型的属性中。
self.model.nc, self.model.names = self.data_dict['nc'], self.data_dict['names']
# 将命令行参数中的训练周期数（ epochs ）赋值给 self.max_epoch ，表示训练将要进行的最大周期数。
self.max_epoch = args.epochs
# 计算训练数据加载器（ train_loader ）的长度，即总的批次数，并赋值给 self.max_stepnum ，这通常用于确定训练的总步数。
self.max_stepnum = len(self.train_loader)
# 将命令行参数中的批量大小（ batch_size ）赋值给 self.batch_size 。
self.batch_size = args.batch_size
# 将命令行参数中的图像尺寸（ img_size ）赋值给 self.img_size 。
self.img_size = args.img_size
# 初始化一个空列表，用于存储用于可视化的图像。
self.vis_imgs_list = []
# 将命令行参数中的 write_trainbatch_tb 赋值给 self.write_trainbatch_tb ，表示是否在TensorBoard中记录训练批次的信息。
self.write_trainbatch_tb = args.write_trainbatch_tb
# 设置类名的颜色。
# set color for classnames
# 为每个类别随机生成一个颜色，并存储在 self.color 列表中。这里使用 np.random.choice 从0到255中随机选择三个数字，代表RGB颜色值， self.model.nc 是类别的数量。
self.color = [tuple(np.random.choice(range(256), size=3)) for _ in range(self.model.nc)]
# 初始化损失数量为3，对应于 'Epoch' 、 'iou_loss' 和 'dfl_loss' 。
self.loss_num = 3
# 初始化损失信息列表，包含周期、IoU损失、dfl损失和分类损失。
self.loss_info = ['Epoch', 'iou_loss', 'dfl_loss', 'cls_loss']
# 如果启用了蒸馏（ self.args.distill ）。
if self.args.distill:
# # 蒸馏会额外增加一个损失，因此损失数量加1。
self.loss_num += 1
# 在损失信息列表中添加 'cwd_loss' ，表示蒸馏过程中的类权重蒸馏损失。
self.loss_info += ['cwd_loss']
# 训练过程。
# Training Process
def train(self):
# 开始一个异常处理块，用于捕获训练过程中可能发生的任何异常。
try:
# 在训练循环开始之前执行的操作，可能是一些初始化工作，如准备数据加载器、初始化优化器等。
# def train_before_loop(self): -> 它在训练循环开始之前执行一系列初始化操作。
self.train_before_loop()
# 从 self.start_epoch 开始，遍历到 self.max_epoch 结束， self.epoch 会在每个周期中被更新。
for self.epoch in range(self.start_epoch, self.max_epoch):
# 对每个周期执行的训练操作，这个方法可能包含每个周期的前向传播、损失计算、反向传播和优化器步骤。
self.train_in_loop(self.epoch)
# 在所有周期完成后，执行模型压缩或者简化操作，以去除不必要的参数，减少模型大小。
self.strip_model()
# 如果训练循环或模型保存过程中出现任何异常，捕获该异常并记录错误信息。
except Exception as _:
# 使用日志记录器记录错误信息。
LOGGER.error('ERROR in training loop or eval/save model.')
# 重新抛出捕获的异常，以便外部可以进一步处理。
raise
# 无论是否发生异常，都会执行的代码块。
finally:
# 在训练循环结束后执行的操作，可能是一些清理工作，如释放资源、保存最终模型等。
self.train_after_loop()
# Training loop for each epoch 每个时期的训练循环。
# 它负责在一个训练周期内执行训练步骤，并在周期结束后进行评估和模型保存。
# 1.epoch_num ：当前训练周期的编号。
def train_in_loop(self, epoch_num):
try:
# 在开始训练步骤之前执行的准备工作，可能包括数据加载器的初始化、设备设置等。
# def prepare_for_steps(self): -> 它在每个训练周期开始之前执行一系列准备工作。
self.prepare_for_steps()
# 遍历进度条 self.pbar ，它是一个迭代器，提供了当前步数 self.step 和批次数据 self.batch_data 。
for self.step, self.batch_data in self.pbar:
# 对每个批次执行训练步骤，包括前向传播、损失计算、反向传播和优化器更新。
# def train_in_steps(self, epoch_num, step_num): -> 它负责执行单个训练步骤中的操作。
self.train_in_steps(epoch_num, self.step)
# 打印当前步骤的训练细节，如损失值、精度等。
# def print_details(self): -> 打印每个步骤后的损失。更新进度条的描述信息，显示当前周期和平均损失值。
self.print_details()
# 如果训练步骤中发生任何异常，记录错误信息并重新抛出异常。
except Exception as _:
LOGGER.error('ERROR in training steps.') # 训练步骤中出现问题。
raise
# 开始另一个异常处理块，用于捕获评估和保存模型过程中可能发生的任何异常。
try:
# 在周期结束后执行模型评估，并根据评估结果决定是否保存模型。
self.eval_and_save()
except Exception as _:
# 如果评估和保存模型过程中发生任何异常，记录错误信息并重新抛出异常。
LOGGER.error('ERROR in evaluate and save model.') # 评估和保存模型时出现错误。
raise
# 批量数据的训练循环。
# Training loop for batchdata
# 它负责执行单个训练步骤中的操作。
# 1.epoch_num ：当前训练周期的编号。
# 2.step_num ：当前步骤的编号。
def train_in_steps(self, epoch_num, step_num):
# 对批次数据进行预处理，将数据移动到指定的设备（如GPU），并将其转换为模型训练所需的格式。 self.prepro_data 方法可能包括数据增强、归一化等操作。
# def prepro_data(batch_data, device): -> 它用于将输入的批次数据（ batch_data ）进行预处理，并将其移动到指定的设备（如GPU）。返回预处理后的图像和目标数据。 -> return images, targets
images, targets = self.prepro_data(self.batch_data, self.device)
# plot train_batch and save to tensorboard once an epoch 绘制 train_batch 并每 epoch 保存一次到 tensorboard。
# 检查是否需要将训练批次写入TensorBoard，当前进程是否为主进程，以及当前步骤是否为周期的第一步。
if self.write_trainbatch_tb and self.main_process and self.step == 0:
# 如果条件满足，调用 self.plot_train_batch 方法绘制当前批次的图像和目标，以便可视化。
# def plot_train_batch(self, images, targets, max_size=1920, max_subplots=16): -> 它用于将训练批次中的图像和对应的目标（标签）绘制成一幅马赛克图（mosaic），并保存用于可视化。
self.plot_train_batch(images, targets)
# 使用 write_tbimg 函数将绘制的训练批次图像写入TensorBoard。
# 1.self.tblogger ：是TensorBoard的 SummaryWriter 实例。
# 2.self.vis_train_batch ：是要写入的图像。
# 3.self.step + self.max_stepnum * self.epoch ：计算当前周期和步骤的全局步数。
# 4.type='train' ：指定图像类型为训练。
# def write_tbimg(tblogger, imgs, step, type='train'): -> 将图像数据展示在 TensorBoard 中，用于可视化训练批次和验证过程中的预测结果。
write_tbimg(self.tblogger, self.vis_train_batch, self.step + self.max_stepnum * self.epoch, type='train')
# 模型的前向传播。
# forward
# 使用PyTorch的自动混合精度（AMP）上下文管理器来优化模型的前向传播。如果设备不是CPU（即在GPU上训练），则启用AMP以提高性能和减少内存使用。
with amp.autocast(enabled=self.device != 'cpu'):
# 将图像数据传递给模型进行前向传播，获取模型的预测结果 preds 和特征图 s_featmaps 。
preds, s_featmaps = self.model(images)
# 如果启用了蒸馏（ self.args.distill ）。
if self.args.distill:
# 使用PyTorch的 torch.no_grad() 上下文管理器来禁用梯度计算，这对于教师模型的前向传播是必要的，以避免不必要的计算和内存使用。
with torch.no_grad():
# 将图像数据传递给教师模型进行前向传播，获取教师模型的预测结果 t_preds 和特征图 t_featmaps 。
t_preds, t_featmaps = self.teacher_model(images)
# 从命令行参数中获取温度参数，用于调整蒸馏损失的尺度。
temperature = self.args.temperature
# 计算蒸馏损失，包括学生模型和教师模型的预测结果和特征图，以及目标数据。
total_loss, loss_items = self.compute_loss_distill(preds, t_preds, s_featmaps, t_featmaps, targets, \
epoch_num, self.max_epoch, temperature, step_num)
# 如果启用了AB融合（ self.args.fuse_ab ）
elif self.args.fuse_ab:
# 计算YOLOv6-af部分的损失。
total_loss, loss_items = self.compute_loss((preds[0],preds[3],preds[4]), targets, epoch_num, step_num) # YOLOv6_af
# 计算YOLOv6-ab部分的损失。
total_loss_ab, loss_items_ab = self.compute_loss_ab(preds[:3], targets, epoch_num, step_num) # YOLOv6_ab
# 将两部分的损失相加，得到总损失。
total_loss += total_loss_ab
# 将两部分的损失项合并。
loss_items += loss_items_ab
else:
# 否则，如果既没有启用蒸馏也没有启用AB融合：
total_loss, loss_items = self.compute_loss(preds, targets, epoch_num, step_num) # YOLOv6_af
# 如果不是单GPU训练（ self.rank != -1 ）。
if self.rank != -1:
# 将总损失乘以world_size（即分布式训练中的进程数），这通常用于在分布式训练中平均损失值。
total_loss *= self.world_size
# 负责执行反向传播和更新优化器。
# backward
# 使用自动混合精度（AMP）的梯度缩放器 self.scaler 来缩放总损失 total_loss ，然后调用 backward() 方法执行反向传播。这一步计算损失相对于模型参数的梯度。
self.scaler.scale(total_loss).backward()
# 将计算得到的损失项 loss_items 保存到实例变量 self.loss_items 中。这些损失项可以用于后续的日志记录、可视化或其他分析。
self.loss_items = loss_items
# 调用 update_optimizer 方法来更新优化器的状态。这通常包括调用优化器的 step() 方法来更新模型参数，并在需要时执行梯度裁剪或其他自定义的优化步骤。
# def update_optimizer(self): -> 负责在训练过程中更新优化器的状态，包括学习率的调整、执行优化步骤和梯度的清零。
self.update_optimizer()
# 负责决定是否进行模型评估，并在适当的时候保存模型。
def eval_and_save(self):
# 计算从当前周期到最大周期之前还剩下多少个训练周期。这里 -1 是因为 self.epoch 是从0开始的。
remaining_epochs = self.max_epoch - 1 - self.epoch # self.epoch is start from 0
# 根据剩余周期数和配置参数 heavy_eval_range 确定评估间隔。如果在 heavy_eval_range 范围内，则使用配置的评估间隔 eval_interval ，否则使用一个较小的间隔（这里是3）。
eval_interval = self.args.eval_interval if remaining_epochs >= self.args.heavy_eval_range else 3
# 判断当前周期是否应该进行评估。如果只剩下最后一个周期，或者不是只评估最终周期且当前周期是评估间隔的倍数，则需要进行评估。
is_val_epoch = (remaining_epochs == 0) or ((not self.args.eval_final_only) and ((self.epoch + 1) % eval_interval == 0))
# 只有主进程执行评估和保存模型的操作，以避免在分布式训练中重复执行这些操作。
if self.main_process:
# 更新指数移动平均（EMA）模型的属性，确保EMA模型的这些属性与主模型保持一致。
self.ema.update_attr(self.model, include=['nc', 'names', 'stride']) # update attributes for ema model
# 如果当前周期是评估周期（ is_val_epoch 为 True ）。
if is_val_epoch:
# 调用 eval_model 方法对模型进行评估。
self.eval_model()
# 从评估结果中获取当前的平均精度（AP）值，并更新实例变量 self.ap 。
self.ap = self.evaluate_results[1]
# 比较当前的AP值和之前保存的最佳AP值，如果当前AP值更高，则更新最佳AP值。
self.best_ap = max(self.ap, self.best_ap)
# 保存检查点。
# save ckpt
# ckpt = {...} ：创建一个字典 ckpt ，用于保存当前训练状态。
ckpt = {
# 保存模型的状态字典， deepcopy 用于创建模型参数的深拷贝， de_parallel 用于去除模型的并行包装（如 DataParallel 或 DistributedDataParallel ）， .half() 将模型参数转换为半精度（FP16）。
'model': deepcopy(de_parallel(self.model)).half(),
# 保存EMA模型的状态字典，并转换为半精度。
'ema': deepcopy(self.ema.ema).half(),
# 保存EMA的更新次数。
'updates': self.ema.updates,
# 保存优化器的状态字典。
'optimizer': self.optimizer.state_dict(),
# 保存学习率调度器的状态字典。
'scheduler': self.scheduler.state_dict(),
# 保存当前的周期数。
'epoch': self.epoch,
# 保存评估结果。
'results': self.evaluate_results,
}
# 使用 osp.join （ os.path.join 的简写）函数将保存根目录 self.save_dir 与子目录 'weights' 连接，形成检查点文件的保存目录。
save_ckpt_dir = osp.join(self.save_dir, 'weights')
# 调用 save_checkpoint 函数保存检查点。
# ckpt ：要保存的检查点字典，包含模型参数、优化器状态等。
# (is_val_epoch) and (self.ap == self.best_ap) ：一个布尔值，指示当前周期是否是评估周期且当前的平均精度（AP）是否等于最佳AP，如果是，则只保存这个检查点。
# save_ckpt_dir ：检查点的保存目录。
# model_name='last_ckpt' ：检查点的文件名前缀。
# def save_checkpoint(ckpt, is_best, save_dir, model_name=""): -> 将检查点保存到磁盘。使用 shutil.copyfile 函数将当前检查点文件复制为 best_ckpt.pt ，这样可以保留最佳模型的副本。 -> shutil.copyfile(filename, best_filename)
save_checkpoint(ckpt, (is_val_epoch) and (self.ap == self.best_ap), save_ckpt_dir, model_name='last_ckpt')
# 检查当前周期是否在训练的最后 self.args.save_ckpt_on_last_n_epoch 个周期内。
if self.epoch >= self.max_epoch - self.args.save_ckpt_on_last_n_epoch:
# 如果是，则调用 save_checkpoint 函数保存检查点，使用当前周期数作为文件名前缀。
save_checkpoint(ckpt, False, save_ckpt_dir, model_name=f'{self.epoch}_ckpt')
# 默认保存最佳 ap ckpt 以停止强 aug 时期
#default save best ap ckpt in stop strong aug epochs
# 检查当前周期是否在训练的最后 self.args.stop_aug_last_n_epoch 个周期内。 self.max_epoch 是总的训练周期数， self.args.stop_aug_last_n_epoch 是在训练结束前需要停止增强的周期数。
if self.epoch >= self.max_epoch - self.args.stop_aug_last_n_epoch:
# 如果当前周期的AP值 self.ap 大于之前记录的最佳AP值 self.best_stop_strong_aug_ap ，则执行以下操作。
if self.best_stop_strong_aug_ap < self.ap:
# 更新 self.best_stop_strong_aug_ap 为当前周期的AP值和之前记录的最佳AP值中的较大者。
self.best_stop_strong_aug_ap = max(self.ap, self.best_stop_strong_aug_ap)
# 调用 save_checkpoint 函数保存当前周期的模型检查点。
# ckpt ：要保存的检查点字典，包含模型参数、优化器状态等。
# False ：这个参数通常用于指示是否仅在达到更好的评估结果时才保存检查点。在这里，我们忽略这个条件，始终保存检查点。
# save_ckpt_dir ：检查点的保存目录。
# model_name='best_stop_aug_ckpt' ：检查点的文件名前缀，标识这是一个在停止增强后达到最佳AP值的检查点。
save_checkpoint(ckpt, False, save_ckpt_dir, model_name='best_stop_aug_ckpt')
# 删除 ckpt 变量以释放内存。这通常在检查点已经保存到磁盘后执行，以避免不必要的内存占用。
del ckpt
# 记录学习率。
# log for learning rate
# 从优化器的参数组中提取当前的学习率，并将其存储在列表 lr 中。
lr = [x['lr'] for x in self.optimizer.param_groups]
# 将学习率列表 lr 追加到评估结果列表 self.evaluate_results 中。这样，评估结果就包含了模型性能指标和学习率信息。
self.evaluate_results = list(self.evaluate_results) + lr
# tensorboard 日志。
# log for tensorboard
# 调用 write_tblog 函数将当前周期 self.epoch 、评估结果 self.evaluate_results 和平均损失 self.mean_loss 写入TensorBoard。这有助于在训练过程中可视化性能指标和损失变化。
write_tblog(self.tblogger, self.epoch, self.evaluate_results, self.mean_loss)
# 将验证预测保存到 tensorboard。
# save validation predictions to tensorboard
# 调用 write_tbimg 函数将验证阶段的预测图像 self.vis_imgs_list 保存到TensorBoard。这些图像通常展示了模型在验证集上的预测结果，有助于直观地评估模型性能。
write_tbimg(self.tblogger, self.vis_imgs_list, self.epoch, type='val')
# 它负责执行模型的评估过程。
def eval_model(self):
# 检查配置对象 self.cfg 是否包含 eval_params 属性。如果没有，说明评估参数未被特别指定，将使用默认的评估流程。
if not hasattr(self.cfg, "eval_params"):
# 调用 eval.run 函数执行评估。
# 1.elf.data_dict ：包含数据集信息的字典。
# 2.batch_size=self.batch_size // self.world_size * 2 ：设置评估时的批量大小，通常是训练时批量大小的两倍，但考虑到分布式训练中的 world_size 。
# 3.img_size=self.img_size ：设置图像的尺寸。
# 4.model=self.ema.ema if self.args.calib is False else self.model ：选择使用EMA模型或原始模型进行评估。如果不在校准（ calib ）阶段，则使用EMA模型；否则使用原始模型。
# 5.conf_thres=0.03 ：设置置信度阈值，用于过滤低置信度的检测结果。
# 6.dataloader=self.val_loader ：使用验证数据加载器。
# 7.save_dir=self.save_dir ：设置保存评估结果的目录。
# 8.task='train' ：设置任务类型，这里为 'train' ，可能用于区分训练和测试阶段的评估。
# 函数返回评估结果 results ，可视化输出 vis_outputs 和可视化路径 vis_paths 。
results, vis_outputs, vis_paths = eval.run(self.data_dict,
batch_size=self.batch_size // self.world_size * 2,
img_size=self.img_size,
model=self.ema.ema if self.args.calib is False else self.model,
conf_thres=0.03,
dataloader=self.val_loader,
save_dir=self.save_dir,
task='train')
else:
# 它用于从配置字典 cfg_dict 中获取指定键 value_str 对应的值。如果该键存在，则返回其值；如果不存在，则返回默认值 default_value 。
# 1.cfg_dict ：配置信息的字典。
# 2.value_str ：要获取值的键的字符串表示。
# 3.default_value ：如果键不存在或键对应的值为 None ，则返回的默认值。
def get_cfg_value(cfg_dict, value_str, default_value):
# 检查 value_str 是否作为键存在于 cfg_dict 中。
if value_str in cfg_dict:
# 如果键存在，进一步检查该键对应的值的类型：
if isinstance(cfg_dict[value_str], list):
# 如果值是列表类型，进一步检查列表中的第一个元素是否为 None 。
# 如果列表的第一个元素不是 None ，则返回该元素；如果是 None ，则返回 default_value 。
return cfg_dict[value_str][0] if cfg_dict[value_str][0] is not None else default_value
# 否则，如果值不是列表类型。
else:
# 如果值不是 None ，则直接返回该值；如果是 None ，则返回 default_value 。
return cfg_dict[value_str] if cfg_dict[value_str] is not None else default_value
else:
# 如果键不存在于字典中，则返回 default_value 。
return default_value
# 使用 get_cfg_value 函数从评估参数中获取图像尺寸 img_size ，如果未指定，则使用实例变量 self.img_size 作为默认值。
eval_img_size = get_cfg_value(self.cfg.eval_params, "img_size", self.img_size)
# 调用 eval.run 函数执行模型评估，传递以下参数 。
# 1.self.data_dict ：包含数据集信息的字典。
# 2.batch_size ：评估时的批量大小，从评估参数中获取，如果未指定，则使用 self.batch_size // self.world_size * 2 作为默认值。
# 3.img_size ：评估时的图像尺寸，使用 eval_img_size 。
# 4.model ：选择使用EMA模型或原始模型进行评估。
# 5.conf_thres ：置信度阈值，从评估参数中获取，如果未指定，则使用0.03作为默认值。
# 6.dataloader ：验证数据加载器。
# 7.save_dir ：保存评估结果的目录。
# 8.task ：任务类型，这里为 'train' 。
# 9.shrink_size ：图像缩小尺寸，从评估参数中获取，如果未指定，则使用 eval_img_size 作为默认值。
# 10.infer_on_rect ：是否在矩形区域进行推理，从评估参数中获取，如果未指定，则使用 False 作为默认值。
# 11.verbose ：是否输出详细信息，从评估参数中获取，如果未指定，则使用 False 作为默认值。
# 12.do_coco_metric ：是否计算COCO评估指标，从评估参数中获取，如果未指定，则使用 True 作为默认值。
# 13.do_pr_metric ：是否计算PR曲线评估指标，从评估参数中获取，如果未指定，则使用 False 作为默认值。
# 14.plot_curve ：是否绘制评估曲线，从评估参数中获取，如果未指定，则使用 False 作为默认值。
# 15.plot_confusion_matrix ：是否绘制混淆矩阵，从评估参数中获取，如果未指定，则使用 False 作为默认值。
results, vis_outputs, vis_paths = eval.run(self.data_dict,
batch_size=get_cfg_value(self.cfg.eval_params, "batch_size", self.batch_size // self.world_size * 2),
img_size=eval_img_size,
model=self.ema.ema if self.args.calib is False else self.model,
conf_thres=get_cfg_value(self.cfg.eval_params, "conf_thres", 0.03),
dataloader=self.val_loader,
save_dir=self.save_dir,
task='train',
shrink_size=get_cfg_value(self.cfg.eval_params, "shrink_size", eval_img_size),
infer_on_rect=get_cfg_value(self.cfg.eval_params, "infer_on_rect", False),
verbose=get_cfg_value(self.cfg.eval_params, "verbose", False),
do_coco_metric=get_cfg_value(self.cfg.eval_params, "do_coco_metric", True),
do_pr_metric=get_cfg_value(self.cfg.eval_params, "do_pr_metric", False),
plot_curve=get_cfg_value(self.cfg.eval_params, "plot_curve", False),
plot_confusion_matrix=get_cfg_value(self.cfg.eval_params, "plot_confusion_matrix", False),
)
# 使用日志记录器记录当前周期和评估结果，包括 mAP@0.5 和 mAP@0.50:0.95。
LOGGER.info(f"Epoch: {self.epoch} | mAP@0.5: {results[0]} | mAP@0.50:0.95: {results[1]}")
# 更新实例变量 self.evaluate_results ，存储前两个评估结果。
self.evaluate_results = results[:2]
# plot validation predictions
# 调用 plot_val_pred 方法绘制验证预测结果。
# def plot_val_pred(self, vis_outputs, vis_paths, vis_conf=0.3, vis_max_box_num=5): -> 它用于绘制验证阶段的预测结果，并将绘制好的图像保存到一个列表中。
self.plot_val_pred(vis_outputs, vis_paths)
# 它在训练循环开始之前执行一系列初始化操作。
def train_before_loop(self):
# 训练开始...。
LOGGER.info('Training start...')
# 记录训练开始的时间戳。
self.start_time = time.time()
# 根据配置中的预热周期数和最大步数计算预热步数。如果启用了量化，则预热步数设置为0。
self.warmup_stepnum = max(round(self.cfg.solver.warmup_epochs * self.max_stepnum), 1000) if self.args.quant is False else 0
# 设置学习率调度器的初始周期，通常用于从检查点恢复时保持学习率的正确状态。
self.scheduler.last_epoch = self.start_epoch - 1
# 初始化最后一个优化步骤的计数器。
self.last_opt_step = -1
# 初始化自动混合精度（AMP）的梯度缩放器，用于FP16训练。
self.scaler = amp.GradScaler(enabled=self.device != 'cpu')
# 初始化最佳平均精度（AP）和当前AP值。
self.best_ap, self.ap = 0.0, 0.0
# 初始化最佳停止增强平均精度。
self.best_stop_strong_aug_ap = 0.0
# 初始化评估结果，通常包含AP50和AP50_95。
self.evaluate_results = (0, 0) # AP50, AP50_95
# resume results 恢复结果。
# 如果存在检查点（ if hasattr(self, "ckpt") ），则从检查点中恢复评估结果和最佳AP值。
# 这行代码检查当前类的实例（ self ）是否有一个名为 ckpt 的属性。 ckpt 通常是一个字典，包含了之前训练过程中保存的状态信息，比如模型参数、优化器状态、评估结果等。
if hasattr(self, "ckpt"):
# 如果存在检查点，从检查点中获取 results 键对应的值，并将其赋值给 self.evaluate_results 。这里的 results 可能是一个包含多个评估指标的元组或列表。
self.evaluate_results = self.ckpt['results']
# 从评估结果中取出第二个元素（索引为1），并将其赋值给 self.best_ap 。这个值通常代表最佳平均精度（AP）。
self.best_ap = self.evaluate_results[1]
# 同样地，将第二个元素的值也赋值给 self.best_stop_strong_aug_ap 。这个值可能代表在停止训练时的最佳增强平均精度。
self.best_stop_strong_aug_ap = self.evaluate_results[1]
# 根据配置和数据字典初始化损失计算模块。
self.compute_loss = ComputeLoss(num_classes=self.data_dict['nc'],
ori_img_size=self.img_size,
warmup_epoch=self.cfg.model.head.atss_warmup_epoch,
use_dfl=self.cfg.model.head.use_dfl,
reg_max=self.cfg.model.head.reg_max,
iou_type=self.cfg.model.head.iou_type,
fpn_strides=self.cfg.model.head.strides)
if self.args.fuse_ab:
# 如果启用了AB融合，则初始化AB融合损失计算模块。
self.compute_loss_ab = ComputeLoss_ab(num_classes=self.data_dict['nc'],
ori_img_size=self.img_size,
warmup_epoch=0,
use_dfl=False,
reg_max=0,
iou_type=self.cfg.model.head.iou_type,
fpn_strides=self.cfg.model.head.strides)
if self.args.distill :
# 如果启用了蒸馏，则根据模型类型初始化蒸馏损失计算模块。
if self.cfg.model.type in ['YOLOv6n','YOLOv6s']:
Loss_distill_func = ComputeLoss_distill_ns
else:
Loss_distill_func = ComputeLoss_distill
self.compute_loss_distill = Loss_distill_func(num_classes=self.data_dict['nc'],
ori_img_size=self.img_size,
fpn_strides=self.cfg.model.head.strides,
warmup_epoch=self.cfg.model.head.atss_warmup_epoch,
use_dfl=self.cfg.model.head.use_dfl,
reg_max=self.cfg.model.head.reg_max,
iou_type=self.cfg.model.head.iou_type,
distill_weight = self.cfg.model.head.distill_weight,
distill_feat = self.args.distill_feat,
)
# 它在每个训练周期开始之前执行一系列准备工作。
def prepare_for_steps(self):
# 如果当前周期 self.epoch 大于起始周期 self.start_epoch ，则调用 self.scheduler.step() 更新学习率。
if self.epoch > self.start_epoch:
self.scheduler.step()
# 如果当前周期等于起始周期且存在检查点（ self.ckpt ），则为优化器的每个参数组设置学习率，这是从检查点恢复时必要的步骤。
elif hasattr(self, "ckpt") and self.epoch == self.start_epoch: # resume first epoch, load lr 恢复第一个 epoch，加载 lr。
for k, param in enumerate(self.optimizer.param_groups):
param['lr'] = self.scheduler.get_lr()[k]
#stop strong aug like mosaic and mixup from last n epoch by recreate dataloader 通过重新创建数据加载器来停止上个 n 个时期的马赛克和混合等强增强。
# 如果当前周期是最后一个周期减去配置中指定的在最后几个周期停止增强的周期数 self.args.stop_aug_last_n_epoch ，则将数据增强的概率设置为0，并重新创建数据加载器以应用这些更改。
if self.epoch == self.max_epoch - self.args.stop_aug_last_n_epoch:
self.cfg.data_aug.mosaic = 0.0
self.cfg.data_aug.mixup = 0.0
# def get_data_loader(args, cfg, data_dict): -> 它用于创建并返回训练和验证数据加载器（ DataLoader ）。这个方法不依赖于类的实例，因此被定义为静态方法。
# -> 这个方法返回训练和验证数据加载器，但在这段代码中，验证数据加载器只在主进程中创建，因此可能返回 None 。 -> return train_loader, val_loader
self.train_loader, self.val_loader = self.get_data_loader(self.args, self.cfg, self.data_dict)
# 将模型设置为训练模式，这对于某些特定层（如Dropout和BatchNorm）的行为是必要的。
self.model.train()
# 如果不是单GPU训练（ self.rank != -1 ），则为训练数据加载器的采样器设置当前周期，这对于确保数据在多个GPU间正确分配是必要的。
if self.rank != -1:
self.train_loader.sampler.set_epoch(self.epoch)
# 初始化一个张量来存储平均损失值。
self.mean_loss = torch.zeros(self.loss_num, device=self.device)
# 清空（重置）优化器的梯度，这是每次训练步骤之前必须的操作。
self.optimizer.zero_grad()
# 使用日志记录器记录损失信息的标题。
LOGGER.info(('\n' + '%10s' * (self.loss_num + 1)) % (*self.loss_info,))
# 初始化一个进度条，用于遍历训练数据加载器。
self.pbar = enumerate(self.train_loader)
if self.main_process:
# 如果是主进程（ self.main_process ），则使用 tqdm 库来创建一个可视化的进度条，并设置总步数、列宽和进度条格式。
self.pbar = tqdm(self.pbar, total=self.max_stepnum, ncols=NCOLS, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}')
# 打印每个步骤后的损失。
# Print loss after each steps
def print_details(self):
# 这个条件检查当前进程是否为主进程。在分布式训练环境中，只有主进程负责打印日志信息，以避免重复的日志输出。
if self.main_process:
# 更新平均损失值。
# 这里使用了一个在线更新平均值的公式，其中 self.mean_loss 是之前的平均损失值， self.step 是当前步骤数， self.loss_items 是当前步骤的损失值。
# 这个公式可以有效地计算新的平均损失，而不需要存储所有历史损失值。
self.mean_loss = (self.mean_loss * self.step + self.loss_items) / (self.step + 1)
# 更新进度条的描述信息，显示当前周期和平均损失值。
# %10s 和 %10.4g 是格式化字符串，分别用于格式化文本和浮点数， %10s 表示左对齐的字符串， %10.4g 表示保留4位小数的浮点数。
# self.loss_num 是损失值的数量，用于确定需要多少个 %10.4g 格式化占位符。
# f'{self.epoch}/{self.max_epoch - 1}' 表示当前周期和总周期数（不包括当前周期）。
# *(self.mean_loss) 是一个解包操作，将 self.mean_loss 数组中的每个损失值作为参数传递给格式化字符串。
self.pbar.set_description(('%10s' + '%10.4g' * self.loss_num) % (f'{self.epoch}/{self.max_epoch - 1}', \
*(self.mean_loss)))
def strip_model(self):
if self.main_process:
LOGGER.info(f'\nTraining completed in {(time.time() - self.start_time) / 3600:.3f} hours.')
save_ckpt_dir = osp.join(self.save_dir, 'weights')
strip_optimizer(save_ckpt_dir, self.epoch) # strip optimizers for saved pt model
# Empty cache if training finished
def train_after_loop(self):
if self.device != 'cpu':
torch.cuda.empty_cache()
# 负责在训练过程中更新优化器的状态，包括学习率的调整、执行优化步骤和梯度的清零。
def update_optimizer(self):
# 计算全局步骤数，这是当前周期数乘以每个周期的步数加上当前步数。
curr_step = self.step + self.max_stepnum * self.epoch
# 设置梯度积累的步数，这是基于批量大小来确定的。积累步数用于在小批量训练时模拟大批量训练的效果。
self.accumulate = max(1, round(64 / self.batch_size))
# 如果当前步骤在预热步数内（ curr_step <= self.warmup_stepnum ）
if curr_step <= self.warmup_stepnum:
# 使用 NumPy 的 np.interp 函数在预热期间线性插值计算积累步数。这个值从1开始，逐渐增加到基于批量大小计算的目标积累步数（ 64 / self.batch_size ），并向上取整到最接近的整数。
self.accumulate = max(1, np.interp(curr_step, [0, self.warmup_stepnum], [1, 64 / self.batch_size]).round())
# 遍历优化器的每个参数组。
for k, param in enumerate(self.optimizer.param_groups):
# 如果当前是第三个参数组（索引为2），则使用配置中的预热偏置学习率 warmup_bias_lr ；否则，使用0.0。
warmup_bias_lr = self.cfg.solver.warmup_bias_lr if k == 2 else 0.0
# 使用 np.interp 函数在预热期间线性插值计算学习率。这个值从 warmup_bias_lr 开始，逐渐增加到初始学习率乘以当前周期的学习率因子 self.lf(self.epoch) 。
param['lr'] = np.interp(curr_step, [0, self.warmup_stepnum], [warmup_bias_lr, param['initial_lr'] * self.lf(self.epoch)])
# 如果参数组中包含动量参数。
if 'momentum' in param:
# 使用 np.interp 函数在预热期间线性插值计算动量值。这个值从配置中的预热动量 warmup_momentum 开始，逐渐增加到正常动量 momentum 。
param['momentum'] = np.interp(curr_step, [0, self.warmup_stepnum], [self.cfg.solver.warmup_momentum, self.cfg.solver.momentum])
# 检查自上次优化器更新以来是否已经累积了足够的步骤。 curr_step 是当前的全局步骤数， self.last_opt_step 是上次优化器更新时的步骤数， self.accumulate 是梯度积累的步数。
if curr_step - self.last_opt_step >= self.accumulate:
# 如果条件满足，使用自动混合精度（AMP）的梯度缩放器 self.scaler 来执行优化器的 step 方法，这会应用累积的梯度来更新模型的参数。
self.scaler.step(self.optimizer)
# 更新梯度缩放器的状态，准备下一次梯度缩放。
self.scaler.update()
# 清空优化器的梯度，为下一次梯度积累做准备。
self.optimizer.zero_grad()
# 如果启用了EMA。
if self.ema:
# 更新模型参数的EMA。EMA是一种技术，用于平滑模型参数，通常可以提高模型的泛化能力。
self.ema.update(self.model)
# 更新最后优化步骤的计数器，记录当前步骤数。
self.last_opt_step = curr_step
# 它用于创建并返回训练和验证数据加载器（ DataLoader ）。这个方法不依赖于类的实例，因此被定义为静态方法。
@staticmethod
# 1.args ：包含命令行参数的对象，可能包含图像尺寸、批量大小、工作进程数等。
# 2.cfg ：配置信息，可能包含数据增强设置等。
# 3.data_dict ：包含数据集信息的字典，如训练路径、验证路径、类别数量等。
def get_data_loader(args, cfg, data_dict):
# 从 data_dict 中提取训练和验证数据的路径。
train_path, val_path = data_dict['train'], data_dict['val']
# check data 检查数据。
# 提取类别数量并转换为整数。
nc = int(data_dict['nc'])
# 提取类别名称列表。
class_names = data_dict['names']
# 确保类别名称的数量与定义的类别数量相匹配。
assert len(class_names) == nc, f'the length of class names does not match the number of classes defined' # 类名的长度与定义的类数不匹配。
# 计算网格尺寸，这是基于模型头部的步长和32的较大值。
grid_size = max(int(max(cfg.model.head.strides)), 32)
# create train dataloader 创建训练数据加载器。
# create_dataloader 函数用于创建数据加载器（Dataloader），它负责在训练或验证过程中提供图像和标签数据。
# 1. path ： 图像数据的加载路径，通常是训练或验证数据集的根目录。
# 2. img_size ： 图像的输入尺寸，即图像在送入模型之前被缩放的尺寸。
# 3. batch_size ： 每个批次中的图像数量。
# 4. stride ： 模型下采样的步长，用于确定特征图的尺寸。
# 5. hyp （默认值为 None） ： 超参数字典，包含了训练过程中的各种参数，如学习率、迭代次数等。
# 6. augment （默认值为 False） ： 是否启用数据增强。
# 7. check_images （默认值为 False） ： 是否检查图像文件的完整性。
# 8. check_labels （默认值为 False） ： 是否检查标签文件的完整性和格式。
# 9. pad （默认值为 0.0） ： 在进行矩形训练时，用于调整图像尺寸的填充比例。
# 10. rect （默认值为 False） ： 是否使用矩形训练，即图像的宽高比可以不同。
# 11. rank （默认值为 -1） ： 在分布式训练中，指定当前进程的编号。如果是 -1，则不使用分布式训练。
# 12. workers （默认值为 8） ： 加载数据时使用的子进程数量。
# 13. shuffle （默认值为 False） ： 是否在每个epoch开始时打乱数据。
# 14. data_dict （默认值为 None） ： 包含数据集信息的字典，如类别名称、类别数量等。
# 15. task （默认值为 "Train"） ： 指定数据集用于训练（"train"）还是验证（"val"）。
# 调用 create_dataloader 函数创建训练数据加载器。
# 返回值是一个包含数据加载器和其他可能信息的元组，这里只取第一个元素（即数据加载器）。
train_loader = create_dataloader(train_path, args.img_size, args.batch_size // args.world_size, grid_size,
hyp=dict(cfg.data_aug), augment=True, rect=False, rank=args.local_rank,
workers=args.workers, shuffle=True, check_images=args.check_images,
check_labels=args.check_labels, data_dict=data_dict, task='train')[0]
# create val dataloader 创建验证数据加载器
# 初始化验证数据加载器为 None 。
val_loader = None
# 检查当前进程是否为主进程（在分布式训练中，只有主进程创建验证数据加载器）。
if args.rank in [-1, 0]:
# 如果当前进程为主进程，调用 create_dataloader 函数创建验证数据加载器。参数与训练数据加载器类似，但批量大小可能不同，且不进行数据增强。
val_loader = create_dataloader(val_path, args.img_size, args.batch_size // args.world_size * 2, grid_size,
hyp=dict(cfg.data_aug), rect=True, rank=-1, pad=0.5,
workers=args.workers, check_images=args.check_images,
check_labels=args.check_labels, data_dict=data_dict, task='val')[0]
# 这个方法返回训练和验证数据加载器，但在这段代码中，验证数据加载器只在主进程中创建，因此可能返回 None 。
return train_loader, val_loader
# 它用于将输入的批次数据（ batch_data ）进行预处理，并将其移动到指定的设备（如GPU）。
@staticmethod
# 1.batch_data ：包含图像和目标（标签）的批次数据。通常， batch_data[0] 是图像张量， batch_data[1] 是对应的目标张量。
# 2.device ：数据需要被移动到的设备，例如 torch.device("cuda:0") 或 torch.device("cpu") 。
def prepro_data(batch_data, device):
# 将图像数据移动到指定的设备上，并转换为浮点类型。这里使用 non_blocking=True 参数来异步执行数据传输，这可以提高数据加载的效率。然后，将图像数据除以255，将其值从[0, 255]范围归一化到[0, 1]范围。
images = batch_data[0].to(device, non_blocking=True).float() / 255
# 将目标数据（标签）移动到指定的设备上。
targets = batch_data[1].to(device)
# 返回预处理后的图像和目标数据。
return images, targets
# 用于根据提供的参数构建和配置YOLOv6模型。
# 1.self ：类的实例。
# 2.args ：包含命令行参数的对象。
# 3.cfg ：配置信息，可能包含模型类型、预训练权重等。
# 4.nc ：类别数量（ num_classes ）。
# 5.device ：模型运行的设备，比如CPU或GPU。
def get_model(self, args, cfg, nc, device):
# 如果配置中的模型类型包含 'YOLOv6-lite' ，则执行以下操作：
if 'YOLOv6-lite' in cfg.model.type:
# 断言不启用 fuse_ab 和 distill 模式，因为 YOLOv6-lite 模型不支持这些模式。
assert not self.args.fuse_ab, 'ERROR in: YOLOv6-lite models not support fuse_ab mode.'
assert not self.args.distill, 'ERROR in: YOLOv6-lite models not support distill mode.'
# 调用 build_lite_model 函数来构建轻量级模型。
model = build_lite_model(cfg, nc, device)
else:
# 如果不是 YOLOv6-lite 模型，则调用 build_model 函数来构建常规模型，并传递 fuse_ab 和 distill_ns 参数。
model = build_model(cfg, nc, device, fuse_ab=self.args.fuse_ab, distill_ns=self.distill_ns)
# 检查 cfg.model.pretrained 是否设置了预训练权重路径。
weights = cfg.model.pretrained
# 如果设置了预训练权重路径：
if weights: # finetune if pretrained model is set 如果设置了预训练模型则进行微调。
# 检查权重文件是否存在，如果不存在则下载权重文件。
if not os.path.exists(weights):
# def download_ckpt(path): -> 下载预训练模型的检查点。
download_ckpt(weights)
# 从 {weights} 加载 state_dict 以进行微调......
LOGGER.info(f'Loading state_dict from {weights} for fine-tuning...')
# def load_state_dict(weights, model, map_location=None): -> # 从检查点文件加载权重，仅为那些层的名称和形状匹配的分配权重。 -> return model
model = load_state_dict(weights, model, map_location=device)
# 使用 LOGGER.info 记录模型信息。
LOGGER.info('Model: {}'.format(model))
# 返回构建和配置好的模型实例。
return model
# 它用于创建和配置一个教师模型，这在知识蒸馏（Knowledge Distillation）的场景中是常见的。
# 1.self ：类的实例。
# 2.args ：包含命令行参数的对象。
# 3.cfg ：配置信息，可能包含模型类型、预训练权重等。
# 4.nc ：类别数量（ num_classes ）。
# 5.device ：模型运行的设备，比如CPU或GPU。
def get_teacher_model(self, args, cfg, nc, device):
# 这行代码检查配置中模型头部的层数是否为3，如果是，则设置 teacher_fuse_ab 为 True ，否则为 False 。
teacher_fuse_ab = False if cfg.model.head.num_layers != 3 else True
# 调用 build_model 函数来构建教师模型，并传递融合参数 teacher_fuse_ab 。
model = build_model(cfg, nc, device, fuse_ab=teacher_fuse_ab)
# 从命令行参数中获取教师模型的预训练权重路径。
weights = args.teacher_model_path
# 如果 weights 非空（即指定了预训练权重路径）。
if weights: # finetune if pretrained model is set 如果设置了预训练模型，则进行微调。
# 使用 LOGGER.info 记录正在从指定路径加载状态字典（ state_dict ）作为教师模型的信息。
LOGGER.info(f'Loading state_dict from {weights} for teacher') # 从 {weights} 为教师加载 state_dict。
# 调用 load_state_dict 函数加载预训练权重到教师模型中，并指定映射位置为 device 。
model = load_state_dict(weights, model, map_location=device)
# 使用 LOGGER.info 记录教师模型的信息。
LOGGER.info('Model: {}'.format(model))
# Do not update running means and running vars 不更新运行方式和运行变量。
# 遍历模型中的所有模块（ module ）。
# 如果模块是 torch.nn.BatchNorm2d 类型（即二维批量归一化层），则将 track_running_stats 设置为 False 。这样做可以防止在推理或蒸馏过程中更新批量归一化的运行均值和方差。
for module in model.modules():
if isinstance(module, torch.nn.BatchNorm2d):
# module.track_running_stats
# 在PyTorch中， module.track_running_stats 是 torch.nn.modules.batchnorm._BatchNorm 类（包括 torch.nn.BatchNorm1d , torch.nn.BatchNorm2d , 和 torch.nn.BatchNorm3d ）的一个属性。
# 这个属性控制着批量归一化（Batch Normalization）层是否更新其运行均值（running mean）和运行方差（running variance）。
# 具体来说 ：
# track_running_stats=True （默认值） ：当设置为 True 时，批量归一化层会在训练过程中更新其运行均值和运行方差。这些运行统计数据是基于在该层上传递的所有数据批次计算得出的，并用于推理（即模型评估）时的归一化。这是训练模式下的典型行为。
# track_running_stats=False ：当设置为 False 时，批量归一化层不会更新其运行均值和运行方差。这意味着在推理时，该层将只使用在训练过程中计算并保存的均值和方差，而不会根据新的输入数据进一步调整这些统计数据。
# 这在某些情况下很有用，比如在知识蒸馏或者模型微调时，我们希望保持批量归一化层的统计数据不变，以确保教师模型和学生模型的行为一致。
# 在推理或模型部署时，通常将 track_running_stats 设置为 False ，因为此时不需要进一步更新统计数据，而且可以减少计算量。
# 此外，在某些特定的模型微调或蒸馏场景中，保持批量归一化层的统计数据不变可以确保模型的行为与训练时一致，这对于模型性能的一致性是有益的。
module.trac