yolov8逐步分解(1)--默认参数&超参配置文件加载
yolov8逐步分解(2)_DetectionTrainer类初始化过程
接上2篇文章，继续讲解yolov8训练过程中的模型加载过程。
使用默认参数完成训练器trainer的初始化后，执行训练函数train()开始YOLOV8的训练。
1. train()方法实现代码如下所示：
def train(self):
"""Allow device='', device=None on Multi-GPU systems to default to device=0."""
#判断设置使用设备总数
if isinstance(self.args.device, int) or self.args.device: # i.e.device=0 or device=[0,1,2,3]
world_size = torch.cuda.device_count() #计算当前可用设备数
elif torch.cuda.is_available(): # i.e. device=None or device=''#判断cuda是否可用
world_size = 1 # default to device 0
else: # i.e. device='cpu' or 'mps'
world_size = 0
# Run subprocess if DDP training, else train normally 分布式训练
#分布式训练
if world_size > 1 and 'LOCAL_RANK' not in os.environ:
# Argument checks
if self.args.rect:
LOGGER.warning("WARNING ⚠️ 'rect=True' is incompatible with Multi-GPU training, setting rect=False")
self.args.rect = False
# Command
cmd, file = generate_ddp_command(world_size, self)
try:
LOGGER.info(f'DDP command: {cmd}')
subprocess.run(cmd, check=True)
except Exception as e:
raise e
finally:
ddp_cleanup(self, str(file))
else:
self._do_train(world_size)上述代码中，主要实现一个功能：判断本次训练使用的机器的个数即world_size的值。world_size 可用于配置分布式训练系统，确保所有设备都能够正确地参与到训练过程中。详细解释请参考博客yolo中RANK、LOACL_RANK以及WORLD_SIZE的介绍-CSDN博客。
若world_size > 1,则进行分布式训练，本次训练使用的机器为单机单卡，所以world_size = 1，不支持分布式训练。直接进入单机训练函数self._do_train(world_size)。
2. _do_train()函数实现代码（部分）截图如下：
下面将逐步讲解_setup_train()函数的功能.
3. _setup_train()代码实现如下
def _setup_train(self, world_size):
"""
Builds dataloaders and optimizer on correct rank process. #构建数据加载器和优化器
"""
# Model
self.run_callbacks('on_pretrain_routine_start')
ckpt = self.setup_model()#加载模型
self.model = self.model.to(self.device)
self.set_model_attributes()
# Check AMP
self.amp = torch.tensor(self.args.amp).to(self.device) # True or False
if self.amp and RANK in (-1, 0): # Single-GPU and DDP
callbacks_backup = callbacks.default_callbacks.copy() # backup callbacks as check_amp() resets them,
self.amp = torch.tensor(check_amp(self.model), device=self.device)
callbacks.default_callbacks = callbacks_backup # restore callbacks,
if RANK > -1 and world_size > 1:
dist.broadcast(self.amp, src=0) # broadcast the tensor from rank 0 to all other ranks (returns None)
self.amp = bool(self.amp) # as boolean
self.scaler = amp.GradScaler(enabled=self.amp) if world_size > 1:
self.model = DDP(self.model, device_ids=[RANK])
# Check imgsz
gs = max(int(self.model.stride.max() if hasattr(self.model, 'stride') else 32), 32) # grid size (max stride)
self.args.imgsz = check_imgsz(self.args.imgsz, stride=gs, floor=gs, max_dim=1)
# Batch size
if self.batch_size == -1: #表示批量大小需要自动估计
if RANK == -1: # single-GPU only, estimate best batch size
self.args.batch = self.batch_size = check_train_batch_size(self.model, self.args.imgsz, self.amp)#估计最佳批量大小
else:
SyntaxError('batch=-1 to use AutoBatch is only available in Single-GPU training. '
'Please pass a valid batch size value for Multi-GPU DDP training, i.e. batch=16')
# Dataloaders
batch_size = self.batch_size // max(world_size, 1)
self.train_loader = self.get_dataloader(self.trainset, batch_size=batch_size, rank=RANK, mode='train')#获取训练集
if RANK in (-1, 0):
self.test_loader = self.get_dataloader(self.testset, batch_size=batch_size * 2, rank=-1, mode='val') #获取测试集
self.validator = self.get_validator() #创建验证器（validator），用于评估模型在验证数据集上的性能。
metric_keys = self.validator.metrics.keys + self.label_loss_items(prefix='val')
self.metrics = dict(zip(metric_keys, [0] * len(metric_keys))) # TODO: init metrics for plot_results()?
self.ema = ModelEMA(self.model)
if self.args.plots and not self.args.v5loader: #如果 self.args.plots 为真且 self.args.v5loader 为假
self.plot_training_labels() #绘制训练标签的图表
# Optimizer
self.accumulate = max(round(self.args.nbs / self.batch_size), 1) # accumulate loss before optimizing
weight_decay = self.args.weight_decay * self.batch_size * self.accumulate / self.args.nbs # scale weight_decay
iterations = math.ceil(len(self.train_loader.dataset) / max(self.batch_size, self.args.nbs)) * self.epochs
self.optimizer = self.build_optimizer(model=self.model,
name=self.args.optimizer,
lr=self.args.lr0,
momentum=self.args.momentum,
decay=weight_decay,
iterations=iterations)
# Scheduler
if self.args.cos_lr:
self.lf = one_cycle(1, self.args.lrf, self.epochs) # cosine 1->hyp['lrf']
else:
self.lf = lambda x: (1 - x / self.epochs) * (1.0 - self.args.lrf) + self.args.lrf # linear
self.scheduler = optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=self.lf)
self.stopper, self.stop = EarlyStopping(patience=self.args.patience), False
self.resume_training(ckpt) #恢复训练过程。ckpt 是一个检查点文件，用于加载之前保存的模型和训练状态。
self.scheduler.last_epoch = self.start_epoch - 1 # do not move
self.run_callbacks('on_pretrain_routine_end') #运行预训练过程结束时的回调函数。4.1 模型加载配置
# Model
self.run_callbacks('on_pretrain_routine_start')
ckpt = self.setup_model()#加载模型
self.model = self.model.to(self.device)
self.set_model_attributes()这段代码是在训练过程中对模型进行加载和设置。
首先，通过 self.run_callbacks('on_pretrain_routine_start') 调用了一个名为 'on_pretrain_routine_start' 的回调函数。这可能是在训练过程中的某个特定时间点执行的回调。
然后，通过 self.setup_model() 加载模型，并将返回的模型断点（checkpoint）保存在变量 ckpt 中。
接下来，通过 self.model = self.model.to(self.device) 将模型移动到指定的设备（self.device）上。
最后，调用 self.set_model_attributes() 来设置模型的属性。这可能是根据特定需求对模型进行自定义设置的函数。
4.1.1 模型加载函数setup_model()详解
def setup_model(self):
"""
load/create/download model for any task.
"""
#判断是否是module模块,如果是直接退出
if isinstance(self.model, torch.nn.Module): # if model is loaded beforehand. No setup needed
return
model, weights = self.model, None
ckpt = None
if str(model).endswith('.pt'): #判断是否是pt格式
weights, ckpt = attempt_load_one_weight(model) #加载模型
cfg = ckpt['model'].yaml
else:
cfg = model
self.model = self.get_model(cfg=cfg, weights=weights, verbose=RANK == -1) # calls Model(cfg, weights)
return ckptsetup_model()用于加载、创建或下载任务的模型。 以下是该方法的逐条详细说明：
它首先检查 self.model 属性是否已经是 torch.nn.Module 的实例，若是，表明模型已经预先加载。 在这种情况下，该方法只是返回而不进行任何进一步的设置。
如果模型尚未加载，它将继续检查模型文件的格式。 如果模型文件为.pt格式，则调用attempt_load_one_weight函数加载模型权重和检查点。 加载权重和检查点后，它从检查点中提取模型配置。
如果模型文件不是 .pt 格式，则假定模型变量直接保存模型配置。
然后，它使用提取或提供的模型配置 (cfg) 和加载的权重 (weights) 调用 get_model 方法来创建模型的实例。
最后，它将创建的模型实例分配给 self.model 属性并返回加载的检查点（ckpt）。
4.1.1.1 模型加载函数attempt_load_one_weight()详解
def attempt_load_one_weight(weight, device=None, inplace=True, fuse=False):
"""Loads a single model weights."""
ckpt, weight = torch_safe_load(weight)
args = {**DEFAULT_CFG_DICT, **(ckpt.get('train_args', {}))}
model = (ckpt.get('ema') or ckpt['model']).to(device).float()
model.args = {k: v for k, v in args.items() if k in DEFAULT_CFG_KEYS}
model.pt_path = weight # attach *.pt file path to model, 值为'yolov8n.pt'
model.task = guess_model_task(model) # 'detect'
if not hasattr(model, 'stride'): #如果没有该属性,则添加一个
model.stride = torch.tensor([32.])
model = model.fuse().eval() if fuse and hasattr(model, 'fuse') else model.eval() # model.eval(): 将模型设置为评估模式
# Module compatibility updates
for m in model.modules():
t = type(m)
if t in (nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU, Detect, Segment):
m.inplace = inplace
elif t is nn.Upsample and not hasattr(m, 'recompute_scale_factor'):
m.recompute_scale_factor = None # torch 1.11.0 compatibility
# Return model and ckpt
return model, ckpt上述代码定义了一个名为 attempt_load_one_weight 的函数，功能如下：
函数参数 weight 表示模型权重的路径或文件，device 表示设备（默认为 None），inplace 表示是否原地操作（默认为 True），fuse 表示是否进行融合操作（默认为 False）。
a. 函数首先调用 torch_safe_load 函数来加载权重，并将返回的结果赋值给 ckpt 和 weight 变量。torch_safe_load 函数内部功能为：判断yolov8*.pt文件是否存在,不存在从github下载,最后通过torch.load加载模型。其中weight类型为字符串，值为’yolov8.pt’ , ckpt 类型为字典，内部为模型结构以及 训练时保存的参数，内容如下：
b. 根据加载的权重，将默认配置字典 DEFAULT_CFG_DICT 和 ckpt 字典中的 'train_args' 键对应的值进行合并，得到 args 字典。键值相同，以后面字典的键值为最终结果。 ckpt字典的内容（后半段数据）如下：
c. 将 ckpt 字典中的 'ema' 键对应的值或者 'model' 键对应的值赋给 model 变量，并将其转移到指定的设备上（默认为 device 参数指定的设备），然后将其转换为浮点型。
d. 将 args 字典中的键值对赋值给 model.args 属性，但仅保留键在 DEFAULT_CFG_KEYS 列表中的键值对。
e. 将 weight 的值赋给 model.pt_path 属性，将模型文件的路径附加到模型中。值为model.pt_path = ‘yolov8.pt’
f .推测模型的任务类型，并将结果赋给 model.task 属性。本次该值为‘detect’.
g. 如果模型中没有名为 stride 的属性，则添加一个名为 stride 的属性，其值为 torch.tensor([32.])。
h. 如果 fuse 为 True 并且模型具有 fuse 属性，则对模型进行融合操作，并将模型设置为评估模式。 此处设置为eval模式，是因为下面要对网络的一些值进行更换，必须进入该模式。后面开始训练时会重新将model设置回训练模式
i. 对于模型中的每个模块，根据模块的类型进行相应的操作：
如果模块的类型是 nn.Hardswish、nn.LeakyReLU、nn.ReLU、nn.ReLU6、nn.SiLU、Detect 或 Segment，则设置 m.inplace 属性为 inplace 参数的值。
如果模块的类型是 nn.Upsample 并且模块不存在 recompute_scale_factor 属性，则将其设置为 None，以保持与 Torch 1.11.0 的兼容性。
j. 返回经过以上处理的模型实例和加载的检查点，即 model 和 ckpt。
4.1.1.2 提取模型配置
cfg = ckpt['model'].yamlcfg类型为字典，内部存放为模型配置参数,即yolov8.yaml中的内容，使用的为YOLOV8n,内容如下图所示：
工程中yolov8.yaml文件位置及内容如下：
4.1.1.3 获取目标检测模型
self.model = self.get_model(cfg=cfg, weights=weights, verbose=RANK == -1) # calls Model(cfg, weights)get_model实现代码如下：
def get_model(self, cfg=None, weights=None, verbose=True):
"""Return a YOLO detection model."""
model = DetectionModel(cfg, nc=self.data['nc'], verbose=verbose and RANK == -1)
if weights:
model.load(weights)
return model该方法用于返回一个 YOLO 检测模型。
方法的参数包括：
self：表示类的实例对象，即调用该方法的对象本身。
cfg：表示模型的配置参数（默认为 None）。
weights：表示模型的权重文件路径（默认为 None）。
verbose：表示是否显示详细信息（默认为 True）。
a. 创建一个 DetectionModel 实例，并传入 cfg 参数和 nc 参数（self.data['nc']）作为构造函数的参数。verbose 参数的值为 verbose and RANK == -1，其中 RANK 是全局变量，用于判断是否为主进程（多进程环境中的一个标识）。
b. 如果 weights 参数不为 None，则调用 model.load(weights) 方法加载权重文件。
c. 返回创建的 model 对象。
至此模型加载完成。
下一章节将详细介绍yolov8检测模型DetectionModel()的实例化过程，内部包含模型的解析构造过程。
