YOLOV5_Train_py代码阅读_01文章链接
https://blog.csdn.net/smilejfy/article/details/140353391?csdn_share_tail=%7B%22type%22%3A%22blog%22%2C%22rType%22%3A%22article%22%2C%22rId%22%3A%22140353391%22%2C%22source%22%3A%22smilejfy%22%7D
3 train()函数
3.1 函数的参数入口 ，logger初始化，configure
def train(hyp, opt, device, tb_writer=None):
- hyp表示超参数
- opt表示配置的命令参数
- device表示的是设备
设置日志文件logger
logger.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
save_dir, epochs, batch_size, total_batch_size, weights, rank = \
Path(opt.save_dir), opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank
文件路径配置
# Directories
wdir = save_dir / 'weights' # 权重保存路径
wdir.mkdir(parents=True, exist_ok=True) # make dir 创建一个目录，parents=True表示父目录不存在时，创建父目录
last = wdir / 'last.pt' # 最近一次权重保存路径
best = wdir / 'best.pt' # 最好一次权重保存路径
results_file = save_dir / 'results.txt' # 结果保存路径
# Save run settings
with open(save_dir / 'hyp.yaml', 'w') as f:
yaml.dump(hyp, f, sort_keys=False) # 保存hyp参数
with open(save_dir / 'opt.yaml', 'w') as f:
yaml.dump(vars(opt), f, sort_keys=False) # 保存opt参数
这段代码设置了权重文件、结果文件、超参数文件、命令参数文件的保存。
3.2 configure
# Configure
plots = not opt.evolve # create plots
cuda = device.type != 'cpu' # 判断是否使用GPU设备，cuda是否安装
init_seeds(2 + rank) # 初始化随机种子
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.SafeLoader) # data dict加载数据
is_coco = opt.data.endswith('coco.yaml') # 解析yaml文件，判断是否为coco数据集
3.3 Logging
# Logging- Doing this before checking the dataset. Might update data_dict
loggers = {'wandb': None} # loggers dict
if rank in [-1, 0]:
opt.hyp = hyp # add hyperparameters
'''
这段代码的作用是检查给定的 weights 文件是否是一个 .pt 文件并且是一个存在的文件。
如果是，它会加载这个文件并尝试获取其中的 wandb_id 键对应的值。如果不是，则返回 None。
具体解释如下：
1.weights.endswith('.pt'):检查 weights 字符串是否以 .pt 结尾，即判断它是否是一个 PyTorch 模型文件。
2.os.path.isfile(weights):检查 weights 是否是一个存在的文件。
3.torch.load(weights):如果前两个条件都满足，则加载这个文件。
4.get('wandb_id'):从加载的文件中获取 wandb_id 键对应的值。如果该键不存在，则返回 None。
5.if ... else None:如果前两个条件不满足，则直接返回 None。
总结起来，这段代码的目的是从指定的 PyTorch 模型文件中获取 wandb_id,
如果文件不存在或不是 .pt 文件，则返回 None。
'''
run_id = torch.load(weights).get('wandb_id') if weights.endswith('.pt') and os.path.isfile(weights) else None
wandb_logger = WandbLogger(opt, Path(opt.save_dir).stem, run_id, data_dict)
loggers['wandb'] = wandb_logger.wandb
data_dict = wandb_logger.data_dict
if wandb_logger.wandb:
weights, epochs, hyp = opt.weights, opt.epochs, opt.hyp # WandbLogger might update weights, epochs if resuming
'''
这段代码的目的是根据 opt.single_cls 的值来决定类别数量 nc,
如果使用单一类别模式，则 nc 为 1, 否则从 data_dict 中获取类别数量
'''
nc = 1 if opt.single_cls else int(data_dict['nc']) # number of classes
names = ['item'] if opt.single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data) # check
主要涉及日志记录和数据字典的更新。以下是对这段代码的详细解释：
-
日志记录初始化：
loggers = {'wandb': None} # loggers dict if rank in [-1, 0]: opt.hyp = hyp # add hyperparameters这段代码初始化了一个日志记录器字典，并检查当前进程是否为主进程（
rank in [-1, 0]）。如果是主进程，则将超参数添加到选项中。 -
检查和加载权重文件：
run_id = torch.load(weights).get('wandb_id') if weights.endswith('.pt') and os.path.isfile(weights) else None这段代码检查给定的
weights文件是否是一个.pt文件并且是一个存在的文件。如果是，它会加载这个文件并尝试获取其中的wandb_id键对应的值。如果不是，则返回None。 -
初始化 WandbLogger：
wandb_logger = WandbLogger(opt, Path(opt.save_dir).stem, run_id, data_dict) loggers['wandb'] = wandb_logger.wandb data_dict = wandb_logger.data_dict这段代码初始化了一个
WandbLogger对象，并将其添加到日志记录器字典中。同时，更新了data_dict为WandbLogger对象中的数据字典。 -
更新权重、epochs 和超参数：
if wandb_logger.wandb: weights, epochs, hyp = opt.weights, opt.epochs, opt.hyp # WandbLogger might update weights, epochs if resuming如果
WandbLogger对象存在，则更新权重、epochs 和超参数。这通常发生在恢复训练时。 -
确定类别数量和类别名称：
nc = 1 if opt.single_cls else int(data_dict['nc']) # number of classes names = ['item'] if opt.single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data) # check这段代码根据
opt.single_cls的值来决定类别数量nc。如果使用单一类别模式，则nc为 1，否则从data_dict中获取类别数量。同时，确定类别名称names，并检查类别名称数量是否与类别数量一致。
总结来说，这段代码主要负责初始化日志记录器、检查和加载权重文件、更新数据字典和确定类别数量及类别名称。
3.4 Model
# Model
'''
这部分代码首先是判断有没有传入预训练权重文件，分两类：
1.如果传入权重文件,直接model.load_state_dict加载模型
2.如果没有传入权重文件，创建模型实例，从头开始训练
'''
pretrained = weights.endswith('.pt') # 判断是否为pt文件，True表示传入了预训练权重文件
if pretrained: # True是存在预训练权重文件
with torch_distributed_zero_first(rank):
attempt_download(weights) # 加载权重文件，如果不存在则下载
'''
这段代码的作用是加载一个 PyTorch 模型检查点(checkpoint),并将其加载到指定的设备上。具体解释如下:
1. torch.load(weights, map_location=device): 这是PyTorch提供的load函数,用于加载保存的模型检查点。
weights:这是一个字符串,表示模型检查点文件的路径。
map_location=device:这是一个参数,用于指定加载模型检查点时使用的设备。
device 通常是一个字符串，例如 'cpu' 或 'cuda:0'，表示将模型加载到 CPU 或 GPU 上。
总结起来,这段代码的目的是将指定的模型检查点文件加载到指定的设备上,并返回模型实例。
并将其赋值给变量ckpt。
'''
ckpt = torch.load(weights, map_location=device) # load checkpoint
# 创建模型
model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
'''
这段代码的作用是根据条件决定是否排除某些键(keys), 具体来说，是否排除 'anchor' 键。
以下是对这段代码的详细解释：
1、opt.cfg:这是一个配置文件路径或内容，表示模型的配置信息。
2、hyp.get('anchors')：从 hyp 字典中获取 'anchors' 键对应的值,hyp 通常包含超参数。
3、opt.resume: 这是一个布尔值, 表示是否从之前的检查点恢复训练。
代码逻辑如下：
1、如果 opt.cfg 存在或者 hyp 字典中有 'anchors' 键对应的值，并且 opt.resume 为 False，否则 exclude 列表中包含 'anchor' 键。
2、否则，exclude 列表为空。
总结起来，这段代码的目的是在特定条件下（即有配置文件或超参数中有锚点信息，并且不是从检查点恢复训练），
排除 'anchor' 键，否则不排除任何键
'''
exclude = ['anchor'] if (opt.cfg or hyp.get('anchors')) and not opt.resume else [] # exclude keys
# 将加载的模型检查点中的模型参数转换为浮点型，并获取其状态字典，以便后续操作（如模型的加载和更新）
state_dict = ckpt['model'].float().state_dict() # to FP32
# 获取加载的模型检查点状态字典和当前模型状态字典的交集，并排除指定的键，
state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude) # intersect
# 加载状态字典
model.load_state_dict(state_dict, strict=False) # load
logger.info('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights)) # report
else:
model = Model(opt.cfg,
ch=3,
nc=nc,
anchors=hyp.get('anchors')).to(device) # create 进入Model类，通过yaml文件创建yolo模型
主要功能是根据是否传入了预训练权重文件来创建模型。以下是对这段代码的详细解释：
-
判断是否为预训练模型：
pretrained = weights.endswith('.pt') # 判断是否为pt文件，True表示传入了预训练权重文件这行代码检查
weights是否以.pt结尾，如果是，则表示传入了预训练权重文件。 -
加载预训练权重文件：
if pretrained: # True是存在预训练权重文件 with torch_distributed_zero_first(rank): attempt_download(weights) # 加载权重文件，如果不存在则下载如果
pretrained为True，则进入这个条件块。torch_distributed_zero_first(rank)是一个上下文管理器，用于确保在分布式训练中，只有主节点（rank为0）会执行某些操作。attempt_download(weights)尝试下载权重文件，如果文件不存在。 -
加载检查点：
ckpt = torch.load(weights, map_location=device) # load checkpoint这行代码加载权重文件，并将其映射到指定的设备（如CPU或GPU）。
-
创建模型：
model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create这行代码创建模型。如果
opt.cfg存在，则使用opt.cfg作为配置文件；否则使用检查点中的模型配置文件。ch=3表示输入通道数为3（RGB图像），nc=nc表示类别数量，anchors=hyp.get('anchors')表示锚点。 -
排除某些键：
exclude = ['anchor'] if (opt.cfg or hyp.get('anchors')) and not opt.resume else [] # exclude keys这行代码确定是否需要排除某些键（如
'anchor'）。如果opt.cfg存在或hyp.get('anchors')存在，并且不是恢复训练，则排除'anchor'。 -
处理状态字典：
state_dict = ckpt['model'].float().state_dict() # to FP32 state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude) # intersect model.load_state_dict(state_dict, strict=False) # load这几行代码将检查点中的模型状态字典转换为浮点型（FP32），然后与当前模型的状态字典进行交集操作，排除指定的键。最后，将处理后的状态字典加载到模型中。
-
报告加载情况：
logger.info('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights)) # report这行代码记录日志，报告从预训练权重文件中加载了多少项到当前模型中。
-
从头创建模型：
else: model = Model(opt.cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create 进入Model类，通过yaml文件创建yolo模型如果
pretrained为False，则从头创建模型，使用opt.cfg作为配置文件。
总结来说，这段代码根据是否传入了预训练权重文件来决定是加载预训练模型还是从头创建模型，并处理相关的配置和状态字典。
3.5 训练数据和val数据
with torch_distributed_zero_first(rank):
check_dataset(data_dict) # check
train_path = data_dict['train'] # 训练的path
test_path = data_dict['val'] # 验证的path
主要功能是检查数据集配置并获取训练和测试数据的路径。以下是对这段代码的详细解释：
-
上下文管理器
torch_distributed_zero_first(rank)：with torch_distributed_zero_first(rank):这是一个上下文管理器，用于确保在分布式训练环境中，只有主节点（rank为0）会执行某些操作。这样可以避免多个节点同时执行相同的操作，从而提高效率并避免冲突。
-
检查数据集配置：
check_dataset(data_dict) # check这行代码调用
check_dataset函数来检查data_dict中的数据集配置是否正确。data_dict通常包含数据集的路径、类别信息等。 -
获取训练数据路径：
train_path = data_dict['train'] # 训练的path这行代码从
data_dict中获取训练数据的路径，并将其赋值给train_path。 -
获取测试数据路径：
test_path = data_dict['val'] # 测试的path这行代码从
data_dict中获取测试数据的路径，并将其赋值给test_path。通常，val表示验证集，但在某些情况下，验证集也可以用作测试集。
总结来说，这段代码确保在分布式训练环境中只有主节点检查数据集配置，并从数据集配置中提取训练和测试数据的路径。
3.6 Freeze
# Freeze
'''
这部分代码是设置冻结层的。
简单来说，就是冻结指定的层的权重，让他们在反向传播的时候不更新权重大小。
作者这里列出来这部分代码的目的其实并不是鼓励使用冻结指定层，因为作者认为这样效果其实并不是很好。
这段代码的目的是根据 freeze 列表中的参数名称，冻结模型中相应的层，使其在训练过程中不进行梯度更新
'''
freeze = [] # parameter names to freeze (full or partial),冻结层的名称
for k, v in model.named_parameters():
v.requires_grad = True # 将所有参数的requires_grad设置为True，表示所有层都需要进行梯度更新，即训练
if any(x in k for x in freeze): # 如果参数名称中包含freeze列表中的参数名称，则冻结该层
print('freezing %s' % k) # 打印冻结层的名称
v.requires_grad = False # 关闭梯度更新
3.7 优化器
# Optimizer
'''
nbs为模拟的batch_size;
就比如默认的话上面设置的opt.batch_size为16,这个nbs就为64,
也就是模型梯度累积了64/16=4(accumulate)次之后再更新一次模型,变相的扩大了batch_size。
这段代码的作用是根据总批量大小和累积次数来调整权重衰减(weight decay)的值。以下是对这段代码的详细解释：
1、nbs = 64:这是一个名义上的批量大小(nominal batch size)，表示在理想情况下每个批次的样本数量。
2、accumulate = max(round(nbs / total_batch_size), 1):计算累积次数。total_batch_size 是实际的总批量大小。
这个公式表示将名义批量大小除以实际总批量大小,然后取整,确保累积次数至少为1。
3、hyp['weight_decay'] *= total_batch_size * accumulate / nbs:根据总批量大小和累积次数来调整权重衰减的值。
具体来说，将原始的权重衰减值乘以总批量大小、累积次数，再除以名义批量大小。
4、logger.info(f"Scaled weight_decay = {hyp['weight_decay']}")：打印调整后的权重衰减值。
总结起来，这段代码的目的是根据实际的总批量大小和累积次数来调整权重衰减的值，
以确保在不同批量大小和累积策略下，权重衰减的效果保持一致。
'''
nbs = 64 # nominal batch size
accumulate = max(round(nbs / total_batch_size), 1) # accumulate loss before optimizing
hyp['weight_decay'] *= total_batch_size * accumulate / nbs # scale weight_decay
logger.info(f"Scaled weight_decay = {hyp['weight_decay']}")
'''
model.modules()迭代遍历模型的所有子层,而model.named_modules()不但返回模型的所有子层，还会返回这些层的名字。
还有一个 model.parameters()，它的作用是迭代地返回模型的所有参数。
然后,用hasattr函数来判断遍历的每个层对象是否拥有相对应的属性,将所有参数分成三:weight、bn, bias。
'''
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_modules():
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
pg2.append(v.bias) # biases 添加偏置，bias
if isinstance(v, nn.BatchNorm2d):
pg0.append(v.weight) # no decay batchNormal2d层
elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
pg1.append(v.weight) # apply decay 添加权重
# 设置优化器Adam或SGD优化方法
if opt.adam:
optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum
else:
optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']}) # add pg1 with weight_decay
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0)))
del pg0, pg1, pg2
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
# https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
# 选择学习率策略，这里作者使用的是cosine学习率策略，学习率随着训练的进行逐渐减小，
# 而学习率的初始值是hyp['lr0']，而学习率的最小值是hyp['lrf']，学习率的最大值是hyp['lr0']。
if opt.linear_lr:
lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear
else:
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# plot_lr_scheduler(optimizer, scheduler, epochs)
# EMA
# 在单机单卡模式或分布式训练的主节点上创建一个指数移动平均模型(Exponetial Moving Average Model)，
# 用于在训练过程中动态更新模型参数，而在其他情况下不创建。
# 这里作者使用的是ModelEMA类，它是PyTorch提供的指数移动平均模型类。
ema = ModelEMA(model) if rank in [-1, 0] else None
这段代码主要用于设置深度学习模型的优化器和学习率调度器，并且包含了模型参数的分组策略。以下是对代码的详细解释：
- 优化器设置：
nbs是名义上的批量大小（nominal batch size）。accumulate是累积损失的次数，用于模拟更大的批量大小。hyp['weight_decay']根据实际批量大小和累积次数进行缩放。- 使用
logger.info记录缩放后的权重衰减值。
- 参数分组：
pg0,pg1,pg2是优化器的参数组，分别用于不同的参数类型。- 遍历模型的所有模块，将偏置（biases）添加到
pg2，将批归一化层的权重添加到pg0，将其他层的权重添加到pg1。 - 根据
opt.adam选择使用 Adam 或 SGD 优化器，并设置学习率和动量。 - 将
pg1和pg2分别添加到优化器中，pg1带有权重衰减，pg2不带权重衰减。 - 使用
logger.info记录优化器组的数量。
- 学习率调度器：
- 根据
opt.linear_lr选择线性学习率调度器或余弦学习率调度器。 - 使用
LambdaLR设置学习率调度器。
- 根据
- EMA（指数移动平均）：
- 如果
rank为-1或0，则使用ModelEMA对模型进行指数移动平均。
- 如果
总结来说，这段代码主要完成了以下任务：
- 设置优化器和学习率调度器。
- 对模型参数进行分组，以便对不同类型的参数应用不同的优化策略。
- 使用指数移动平均（EMA）来平滑模型参数
3.8 Resume
# Resume
start_epoch, best_fitness = 0, 0.0 # start epoch, best fitness
if pretrained: # 如果存在预训练模型，则从预训练模型开始训练
# Optimizer
if ckpt['optimizer'] is not None:
optimizer.load_state_dict(ckpt['optimizer'])
best_fitness = ckpt['best_fitness']
# EMA
if ema and ckpt.get('ema'):
ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
ema.updates = ckpt['updates']
# Results
if ckpt.get('training_results') is not None:
results_file.write_text(ckpt['training_results']) # write results.txt
# Epochs
start_epoch = ckpt['epoch'] + 1
if opt.resume:
assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (weights, epochs)
if epochs < start_epoch:
logger.info('%s has been trained for %g epochs. Fine-tuning for %g additional epochs.' %
(weights, ckpt['epoch'], epochs))
epochs += ckpt['epoch'] # finetune additional epochs
del ckpt, state_dict
# Image sizes
# 获取网络的最大步长，预测头数量和训练及测试图片长宽大小
gs = max(int(model.stride.max()), 32) # grid size (max stride)
nl = model.model[-1].nl # number of detection layers (used for scaling hyp['obj'])
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size] # verify imgsz are gs-multiples
3.9 DP mode
if cuda and rank == -1 and torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
DP mode解释
DP mode
30个数据点，但每个GPU上只能容纳10个数据点，一共3块GPU
首先，我们定义主GPU，然后，我们执行以下步骤：
1、将10个数据点和模型副本从主GPU移动到其他2个GPU
2、在每个GPU上进行前向传递并将输出传递给主GPU
3、在主 GPU 上计算总损失，然后将损失发送回每个 GPU 以计算参数的梯度
4、将梯度发送回Master GPU（这些是所有训练示例的梯度平均值），将它们相加得到整批30个的平均梯度
5、更新主 GPU 上的参数并将这些更新发送到其他 2 个 GPU 以进行下一次迭代
这个过程中存在一些问题和低效率
1、数据-从主 GPU 传递，然后在其他 GPU 之间分配。此外，主 GPU 的利用率高于其他 GPU，
因为总损失的计算和参数更新发生在主 GPU 上
2、我们需要在每次迭代时同步其他 GPU 上的模型，这会减慢训练速度
这段代码的作用是检查是否使用 CUDA 以及是否有多张 GPU，并在满足条件的情况下将模型转换为数据并行模式。
以下是对这段代码的详细解释：
1 if cuda and rank == -1 and torch.cuda.device_count() > 1:：检查以下三个条件是否同时满足：
1、cuda：表示是否使用 CUDA 进行 GPU 加速。
2、rank == -1：表示当前处于单机单卡模式。
3、torch.cuda.device_count() > 1：表示当前系统中有多张 GPU。
2 model = torch.nn.DataParallel(model)：如果上述条件都满足，则将模型转换为数据并行模式。
torch.nn.DataParallel 是一个类，用于在多个 GPU 上并行处理数据。
总结起来，这段代码的目的是在单机多卡模式下，将模型转换为数据并行模式，以便在多个GPU上并行处理数据。
3.10 SyncBatchNorm
'''
这段代码的作用是检查是否需要使用同步批归一化(SyncBatchNorm),
并在满足条件的情况下将模型中的批归一化层转换为同步批归一化层。以下是对这段代码的详细解释：
1 if opt.sync_bn and cuda and rank != -1:：检查以下三个条件是否同时满足：
opt.sync_bn:表示是否启用同步批归一化。
cuda:表示是否使用 CUDA 进行 GPU 加速。
rank != -1:表示当前处于分布式训练模式，而不是单机单卡模式。
2 model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)：如果上述条件都满足，
则将模型中的批归一化层转换为同步批归一化层。torch.nn.SyncBatchNorm.convert_sync_batchnorm 是一个函数，
用于将模型中的批归一化层转换为同步批归一化层。
3 logger.info('Using SyncBatchNorm()')：打印日志信息，表示正在使用同步批归一化。
总结起来，这段代码的目的是在分布式训练模式下，将模型中的批归一化层转换为同步批归一化层，
以便在多个 GPU 之间同步批归一化的统计信息。
'''
# SyncBatchNorm
if opt.sync_bn and cuda and rank != -1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
logger.info('Using SyncBatchNorm()')
3.11 Trainloader
# Trainloader
'''
opt.image_weights 是一个选项(option)，通常用于图像分类或目标检测任务中，表示是否使用图像权重来进行训练。
图像权重是一种数据增强技术，通过给不同图像赋予不同的权重，使得模型在训练过程中更加关注某些特定的图像或类别。
以下是一些可能的使用场景和解释：
类别不平衡问题：在目标检测任务中，某些类别的样本可能非常稀少，
导致模型在训练过程中难以学习到这些类别的特征。
通过使用图像权重，可以增加这些稀有类别样本的权重，使得模型更加关注这些样本。
难例挖掘：在训练过程中，模型可能会遇到一些难以分类的图像（即难例）。
通过给这些难例赋予更高的权重，可以使得模型在训练过程中更加关注这些难例，
从而提高模型的泛化能力。
数据增强：图像权重也可以作为一种数据增强手段，通过随机调整图像的权重，
使得模型在训练过程中能够看到更多样化的数据，从而提高模型的鲁棒性。
'''
dataloader, dataset = create_dataloader(train_path, # 训练数据的路径
imgsz, # 训练图片的大小
batch_size, # 训练的batch_size
gs, # 网络的最大步长
opt, # 训练的配置参数
hyp=hyp, # 超参数
augment=True, # 是否进行数据增强
cache=opt.cache_images, # 是否缓存图片
rect=opt.rect, # 是否使用矩形图像
rank=rank, # 节点编号，分布式训练
world_size=opt.world_size, # 节点数量，分布式训练
workers=opt.workers, # 工作进程数量
image_weights=opt.image_weights, # 类权重系数开关，如果为True那么在每次训练的时候都会遍历数据集，然后计算各类数据的权重，然后开始训练
quad=opt.quad, # 是否使用四边形图像
prefix=colorstr('train: ')) # 训练日志前缀
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class 标签最大值
nb = len(dataloader) # number of batches
# 这里判断mlc是否小于nc，如果小于nc，则报错，可能是标签类别超过了nc，导致模型训练出错。
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)
3.12 Process 0
# Process 0
if rank in [-1, 0]: # 如果处理单卡模式或多卡模式
testloader = create_dataloader(test_path, imgsz_test, batch_size * 2, gs, opt, # testloader
hyp=hyp, cache=opt.cache_images and not opt.notest, rect=True, rank=-1,
world_size=opt.world_size, workers=opt.workers,
pad=0.5, prefix=colorstr('val: '))[0]
if not opt.resume:
labels = np.concatenate(dataset.labels, 0)
c = torch.tensor(labels[:, 0]) # classes
# cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency
# model._initialize_biases(cf.to(device))
if plots:
plot_labels(labels, names, save_dir, loggers)
if tb_writer:
tb_writer.add_histogram('classes', c, 0)
# Anchors
if not opt.noautoanchor:
check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
model.half().float() # pre-reduce anchor precision
主要用于处理单卡或多卡模式下的测试数据加载和一些初始化操作。以下是对代码的详细解释：
- 测试数据加载器：
- 如果
rank为-1或0，表示当前处于单卡或多卡模式。 - 使用
create_dataloader函数创建测试数据加载器testloader。 test_path是测试数据的路径。imgsz_test是测试图像的大小。batch_size * 2是测试时的批量大小，通常是训练批量的两倍。gs是网格大小。opt是命令行参数对象。hyp是超参数。cache表示是否缓存图像。rect表示是否使用矩形推理。rank为-1表示单卡模式。world_size是分布式训练中的总进程数。workers是数据加载器的工作线程数。pad是图像填充比例。prefix是日志前缀。
- 如果
- 标签处理和可视化：
- 如果
opt.resume为False，表示不从检查点恢复训练，需要进行一些初始化操作。 - 将数据集的标签拼接成一个数组
labels。 - 提取标签中的类别信息
c。 - 如果需要绘制标签，调用
plot_labels函数进行绘制，并将结果保存到save_dir。 - 如果启用了 TensorBoard，将类别信息
c添加到 TensorBoard 的直方图中。
- 如果
- 锚点检查：
- 如果
opt.noautoanchor为False，表示需要检查锚点。 - 调用
check_anchors函数检查锚点，确保锚点适合当前数据集。 - 将模型转换为半精度（
half()），然后转换回单精度（float()），以减少锚点的精度损失。
- 如果
总结来说，这段代码主要完成了以下任务：
- 创建测试数据加载器。
- 处理标签数据并进行可视化。
- 检查和调整锚点以适应数据集。
3.13 DDP mode
'''
# DDP mode
引入分布式数据并行是为了改善数据并行算法的低效率。我们仍然采用与之前相同的设置—每批30个数据点，使用 3个GPU。
差异如下：
1、它没有主 GPU
2、因为我们不再拥有主 GPU，所以我们直接从磁盘/RAM以非重叠方式并行加载每个GPU上的数据—DistributedSampler为我们完成这项工作。
在底层，它使用本地等级 (GPU id) 在 GPU 之间分配数据 - 给定 30 个数据点，第一个 GPU 将使用点 [0, 3, 6, ... , 27]，
第二个 GPU [1, 4, 7, .., 28] 和第三个 GPU [2, 5, 8, .. , 29]
3、前向传递、损失计算和后向传递在每个 GPU 上独立执行，异步减少梯度计算平均值，然后在所有 GPU 上进行更新
'''
if cuda and rank != -1:
model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank,
# nn.MultiheadAttention incompatibility with DDP https://github.com/pytorch/pytorch/issues/26698
find_unused_parameters=any(isinstance(layer, nn.MultiheadAttention) for layer in model.modules()))
3.14 超参数–model parameters
'''
这段代码的作用是调整超参数 hyp['box'] 的值，以适应模型的层数。以下是对这段代码的详细解释：
1 hyp['box']：这是一个超参数，通常用于控制目标检测任务中的边界框损失(box loss)的权重。
2 nl:表示模型的层数(number of layers)。
3 3. / nl:计算一个缩放因子，其中 3. 是一个常数,nl 是模型的层数。这个缩放因子用于调整 hyp['box'] 的值。
4 hyp['box'] *= 3. / nl:将 hyp['box'] 乘以缩放因子 3. / nl,以调整其值。
总结起来，这段代码的目的是根据模型的层数来调整边界框损失的权重，以确保在不同层数的模型中，
边界框损失的权重保持在一个合适的范围内。
'''
# Model parameters
hyp['box'] *= 3. / nl # scale to layers nl表示模型的层数
hyp['cls'] *= nc / 80. * 3. / nl # scale to classes and layers
hyp['obj'] *= (imgsz / 640) ** 2 * 3. / nl # scale to image size and layers
hyp['label_smoothing'] = opt.label_smoothing
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou)
'''
这段代码的作用是为模型附加类别权重，以便在训练过程中对不同类别的样本进行加权处理。具体解释如下：
1 labels_to_class_weights(dataset.labels, nc)：这个函数根据数据集的标签（dataset.labels）
和类别数量（nc）计算每个类别的权重。类别权重的计算通常是基于每个类别的样本数量，
样本数量较少的类别会获得更高的权重，以平衡不同类别之间的样本分布。
2 .to(device)：将计算得到的类别权重转移到指定的设备（例如GPU）上，以便在模型训练时进行加速。
3 * nc：将计算得到的类别权重乘以类别数量（nc），这可能是为了调整权重的比例，使其更适合模型的训练。
4 model.class_weights = ...：将计算并调整后的类别权重附加到模型上，
以便在训练过程中使用这些权重进行损失函数的计算。
总结来说，这段代码的目的是为了在模型训练时考虑不同类别的样本数量，
通过附加类别权重来平衡不同类别之间的样本分布，从而提高模型的泛化能力和训练效果。
'''
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights
model.names = names
3.15 开始训练
# Start training
t0 = time.time() # 记录起始时间
nw = max(round(hyp['warmup_epochs'] * nb), 1000) # number of warmup iterations, max(3 epochs, 1k iterations)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
maps = np.zeros(nc) # mAP per class
results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
scheduler.last_epoch = start_epoch - 1 # do not move
scaler = amp.GradScaler(enabled=cuda)
'''
这段代码的作用是初始化损失函数计算类ComputeLoss。具体解释如下:
1 ComputeLoss(model):这是一个自定义的损失类,通常在训练过程中用于计算模型的损失值。这个类会接收model
作为参数,并根据模型的输出和目标标签来计算损失。
这段代码的作用是初始化一个用于计算损失的类,并将其赋值给变量,以便在训练过程中使用该类的计算模型的损失值
'''
compute_loss = ComputeLoss(model) # init loss class
logger.info(f'Image sizes {imgsz} train, {imgsz_test} test\n'
f'Using {dataloader.num_workers} dataloader workers\n'
f'Logging results to {save_dir}\n'
f'Starting training for {epochs} epochs...')
# 开始epoch循环训练
# epoch ------------------------------------------------------------------
for epoch in range(start_epoch, epochs):
model.train() # 将模型转换为训练模式 step 1
# Update image weights (optional)
# 类权重系数开关，如果为True那么在每次训练的时候都会遍历数据集，然后计算各类数据的权重，然后开始训练
if opt.image_weights:
# Generate indices
if rank in [-1, 0]:
cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights
iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights
dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx
# Broadcast if DDP
if rank != -1:
indices = (torch.tensor(dataset.indices) if rank == 0 else torch.zeros(dataset.n)).int()
dist.broadcast(indices, 0)
if rank != 0:
dataset.indices = indices.cpu().numpy()
# Update mosaic border
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
mloss = torch.zeros(4, device=device) # mean losses (box, obj, cls, total) 初始化训练时打印的平均损失信
if rank != -1: # rank==0,
dataloader.sampler.set_epoch(epoch) # 加载当前的epoch的sampler
pbar = enumerate(dataloader) # 这行代码的作用是创建一个枚举对象，用于遍历数据加载器中的批次数据
logger.info(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'total', 'labels', 'img_size'))
if rank in [-1, 0]:
pbar = tqdm(pbar, total=nb) # 将pbar对象包含在tqdm进度条中，以便在训练过程中显示进度条，nb通常是数据加载器中的总批次数
'''
用于在每个训练批次开始时将优化器的梯度清零。这是因为在pytorch中,梯度会积累,如果不手动清零,
梯度将会在每个批次中累积,导致不正确的训练结果。
'''
optimizer.zero_grad()
for i, (imgs, targets, paths, _) in pbar: # batch -------------------------------------------------------------
# ni用于计算自训练开始以来的累计批次数,这个变量ni通常用于学习率调整,预热等操作
ni = i + nb * epoch # number integrated batches (since train start)
# 将图像数值从float32转换为0-1的范围,uint8到float32的转换,以及图像的归一化
imgs = imgs.to(device, non_blocking=True).float() / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
'''
Warmup
Warmup是在ResNet论文中提到的一种学习率预热的方法,它在训练开始的时候先选择使用一个较小的学习率，
训练了一些epoches或者steps(比如4个epoches,10000steps),再修改为预先设置的学习来进行训练
'''
if ni <= nw: # 前面nw次需要预热
xi = [0, nw] # x interp
# model.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
accumulate = max(1, np.interp(ni, xi, [1, nbs / total_batch_size]).round())
# 遍历优化器中的所有参数组,并将学习率进行插值,使得在前面nw个批次中学习率逐渐增加到预先设置的学习率
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
if 'momentum' in x:
x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
# Multi-scale 设置多尺度训练，即在训练过程中随机选择图像的尺寸，然后进行训练
# 从image*0.5到image*1.5之间随机选择尺寸，然后进行训练
# imgsz是默认的训练图片大小，
# gs是网络的最大步长，模型最大stride为32，因此最大的尺寸为32*1.5=42
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple)
imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
# Forward 前向传播
with amp.autocast(enabled=cuda):
pred = model(imgs) # forward 将图片送入网络，得到一个预测结果 前向传播 step 2
loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size
if rank != -1:
# 采用DDP训练，平均不同gpu之间的梯度
loss *= opt.world_size # gradient averaged between devices in DDP mode
'''
opt.quad通常是指在训练过程中使用四重(quadruple)数据加载策略。这种策略可以提高数据加载效率,特别是在多个GPU进行训练时
opt.quad是一个bool值,表示是否采用mosaic4数据增强,如果采用,loss也要翻4倍,用于启用或禁用四重数据加载器,
'''
if opt.quad:
loss *= 4. # 如果采用collate_fn4取出mosaic4数据，loss也要翻4倍
# Backward 反向传播, scale使用自动混合精度运算
scaler.scale(loss).backward()
if ni % accumulate == 0:
scaler.step(optimizer) # optimizer.step
scaler.update()
optimizer.zero_grad()
# Exponential Moving Average (EMA) 的方法，中文名叫指数滑动平均。
# 它的意义在于利用滑动平均的参数来提高模型在测试数据上的健壮性
if ema:
ema.update(model)
# Print
if rank in [-1, 0]:
mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0) # (GB)
s = ('%10s' * 2 + '%10.4g' * 6) % ('%g/%g' % (epoch, epochs - 1), mem, *mloss, targets.shape[0], imgs.shape[-1])
pbar.set_description(s)
# Plot
if plots and ni < 3:
f = save_dir / f'train_batch{ni}.jpg' # filename
Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
# if tb_writer:
# tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
# tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), []) # add model graph
elif plots and ni == 10 and wandb_logger.wandb:
wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name) for x in save_dir.glob('train*.jpg') if x.exists()]})
# end batch ------------------------------------------------------------------------------------------------
# end epoch ----------------------------------------------------------------------------------------------------
# Scheduler
lr = [x['lr'] for x in optimizer.param_groups] # for tensorboard
scheduler.step()
# DDP process 0 or single-GPU
if rank in [-1, 0]:
# mAP
ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride', 'class_weights'])
final_epoch = epoch + 1 == epochs
if not opt.notest or final_epoch: # Calculate mAP
wandb_logger.current_epoch = epoch + 1
results, maps, times = test.test(data_dict,
batch_size=batch_size * 2,
imgsz=imgsz_test,
model=ema.ema,
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=save_dir,
verbose=nc < 50 and final_epoch,
plots=plots and final_epoch,
wandb_logger=wandb_logger,
compute_loss=compute_loss,
is_coco=is_coco)
# Write
with open(results_file, 'a') as f:
f.write(s + '%10.4g' * 7 % results + '\n') # append metrics, val_loss
if len(opt.name) and opt.bucket:
os.system('gsutil cp %s gs://%s/results/results%s.txt' % (results_file, opt.bucket, opt.name))
# Log
tags = ['train/box_loss', 'train/obj_loss', 'train/cls_loss', # train loss
'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95',
'val/box_loss', 'val/obj_loss', 'val/cls_loss', # val loss
'x/lr0', 'x/lr1', 'x/lr2'] # params
for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags):
if tb_writer:
tb_writer.add_scalar(tag, x, epoch) # tensorboard
if wandb_logger.wandb:
wandb_logger.log({tag: x}) # W&B
# Update best mAP
fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
if fi > best_fitness:
best_fitness = fi
wandb_logger.end_epoch(best_result=best_fitness == fi)
# Save model
if (not opt.nosave) or (final_epoch and not opt.evolve): # if save
ckpt = {'epoch': epoch,
'best_fitness': best_fitness,
'training_results': results_file.read_text(),
'model': deepcopy(model.module if is_parallel(model) else model).half(),
'ema': deepcopy(ema.ema).half(),
'updates': ema.updates,
'optimizer': optimizer.state_dict(),
'wandb_id': wandb_logger.wandb_run.id if wandb_logger.wandb else None}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
if wandb_logger.wandb:
if ((epoch + 1) % opt.save_period == 0 and not final_epoch) and opt.save_period != -1:
wandb_logger.log_model(
last.parent, opt, epoch, fi, best_model=best_fitness == fi)
del ckpt
# end epoch ----------------------------------------------------------------------------------------------------
# end training
这段代码是一个深度学习训练循环的实现，主要用于训练一个目标检测模型。以下是对代码的详细解释：
- 初始化训练参数：
t0 = time.time()：记录训练开始的时间。nw = max(round(hyp['warmup_epochs'] * nb), 1000)：计算预热迭代次数，最多为1000次。maps = np.zeros(nc)：初始化每个类别的mAP（平均精度）。results = (0, 0, 0, 0, 0, 0, 0)：初始化训练结果，包括精度、召回率、mAP等。scheduler.last_epoch = start_epoch - 1：设置学习率调度器的最后一个epoch。scaler = amp.GradScaler(enabled=cuda)：初始化自动混合精度（AMP）的梯度缩放器。
- 初始化损失函数：
compute_loss = ComputeLoss(model)：初始化损失函数类。
- 打印训练信息：
- 打印图像大小、数据加载器工作线程数、日志保存路径和训练的总epoch数。
- 训练循环：
- 外层循环遍历每个epoch：
model.train()：将模型设置为训练模式。- 如果启用了图像权重，计算类权重和图像权重，并更新数据集索引。
- 初始化平均损失
mloss。 - 如果使用分布式训练，设置数据加载器的epoch。
- 创建一个进度条
pbar来显示训练进度。
- 外层循环遍历每个epoch：
- 优化器和梯度清零：
optimizer.zero_grad()：在每个batch开始前清零梯度。
- 批处理循环：
- 内层循环遍历每个batch：
- 计算当前批次的累计次数
ni。 - 将图像转换为浮点数并归一化。
- 如果处于预热阶段，调整学习率和动量。
- 如果启用了多尺度训练，随机调整图像大小。
- 前向传播：计算预测结果和损失。
- 反向传播：使用梯度缩放器进行反向传播。
- 如果达到累积次数，更新优化器参数并清零梯度。
- 更新指数移动平均（EMA）模型。
- 打印当前批次的训练信息。
- 如果启用了绘图，绘制训练图像。
- 计算当前批次的累计次数
- 内层循环遍历每个batch：
- 结束批处理和epoch：
- 完成每个batch的处理后，更新进度条。
- 完成每个epoch的处理后，打印epoch结束信息。
- 学习率调度：
scheduler.step()：更新学习率。
- 评估和保存模型：
- 如果启用了评估，计算mAP并保存结果。
- 更新最佳mAP。
- 保存模型检查点。
总结来说，这段代码实现了一个完整的训练循环，包括数据加载、前向传播、损失计算、反向传播、优化器更新和学习率调度等步骤。通过进度条和日志记录，可以实时监控训练过程。
model.train()
model.train() 是 PyTorch 中的一个方法，用于将模型设置为训练模式。这意味着模型会启用一些特定的行为，例如：
1、Dropout：在训练模式下，Dropout 层会随机将一些神经元的输出置为零，以防止过拟合。
2、Batch Normalization：在训练模式下，Batch Normalization 层会使用当前批次的统计信息（均值和方差）来标准化输入。
当你调用 model.train() 时，模型会启用这些训练模式下的行为。
相反，如果你调用 model.eval()，模型会切换到评估模式，此时 Dropout 和 Batch Normalization 层会使用固定的统计信息，
而不是当前批次的统计信息。
以下是一个简单的示例，展示了如何在训练过程中使用 model.train()：
import torch
import torch.nn as nn
import torch.optim as optim
# 定义一个简单的模型
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc = nn.Linear(10, 1)
def forward(self, x):
return self.fc(x)
# 创建模型实例
model = SimpleModel()
# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
# 切换到训练模式
model.train()
# 生成一些随机数据
inputs = torch.randn(5, 10)
targets = torch.randn(5, 1)
# 前向传播
outputs = model(inputs)
loss = criterion(outputs, targets)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
AMP
scaler = amp.GradScaler(enabled=cuda)
这段代码使用了PyTorch的自动混合精度(Automatic Mixed Precision, AMP)功能,
通过创建一个GradScaler对象来管理梯度缩放,以提高训练速度并减少显存占用。具体解释如下:
1 amp.GradScaler:这是PyTorch提供的用于自动混合精度训练的类。
自动混合精度训练通过在训练过程中自动选择合适的精度
(例如,在计算损失和梯度时使用半精度浮点数(FP16),而在权重更新时使用单精度浮点数(FP32)),
从而在保持模型精度的同时提高训练速度和减少显存占用。
2 enabled=cuda:这个参数控制GradScaler是否启用。
如果cuda为True,则启用梯度缩放;如果cuda为False,则不启用。
通常情况下,cuda的值会根据当前是否使用GPU来决定,如果使用GPU进行训练,则cuda为True,否则为False。
总结来说,这段代码的作用是创建一个GradScaler对象,用于在训练过程中管理梯度缩放,
以实现自动混合精度训练。通过这种方式，可以在保持模型精度的同时提高训练速度和减少显存占用
EMA
if ni % accumulate == 0:
scaler.step(optimizer) # optimizer.step
scaler.update()
optimizer.zero_grad()
# Exponential Moving Average (EMA) 的方法，中文名叫指数滑动平均。
# 它的意义在于利用滑动平均的参数来提高模型在测试数据上的健壮性
if ema:
ema.update(model)
Optimize 训练过程中更新模型参数的部分
1 条件判断 if ni % accumulate == 0:
1 ni是当前的迭代次数(从训练开始算起的总批次数)
2 accumulate是一个整数, 表示在执行优化步骤之前要累积的梯度数量。
常用于模拟更大的批次数量, 以提高训练的稳定性。
例如, 在8个GPU上训练时, 我们通常希望每批的大小为64, 但在每个GPU上, 我们可能只需要计算一次梯度,
然后将其平均分摊到8个GPU上。这就是为什么我们需要设置accumulate=4。???
?? 梯度累计(Graident Accumulation)的作用是什么?
你提到的这段代码和注释描述了在分布式训练环境中使用梯度累积(Gradient Accumulation)的技术。
梯度累积是一种在有限内存资源下模拟更大批量训练的方法。
具体来说,假设你有8个GPU,并且希望每批的总大小为(即全局批量大小为64),但在每个GPU上只计算一次梯度)
然后将这些梯度平均分摊到8个GPU上进行优化步骤。
以下是对这段代码和注释的详细解释：
1 全局批量大小:你希望每批的总大小为64。
2 局部批量大小:在每个GPU上,批量大小为32/8=4。这意味着每个GPU每次处理4个样本。
3 梯度累积:在每个GPU上计算梯度,但不在每个GPU上立即执行优化步骤。
相反,累积梯度直到累积的批量大小达到全局批量大小。
4 累积次数:为了达到全局批量大小64,每个GPU需要累积4次梯度(因为每个GPU每次处理4个样本,累积4次后达到64个样本)。
这就是为什么需要设置 accumulate=2。
accumulate=(Global_batch)/(Local_batch*num_of_GPUs)=(64)/(4*8)=2
5 梯度同步:当累积的批量大小达到全局批量大小时,将梯度从所有GPU上收集到一个或多个GPU上,
然后在这些GPU上执行优化步骤。
通过这种方式，你可以在有限的内存资源下模拟更大的批量训练，从而提高训练的稳定性和性能。
总结一下，设置 accumulate=4 的目的是为了在每个GPU上累积4次梯度,从而达到全局批量大小32,
然后在4个GPU上进行优化步骤。这样可以有效地利用多个GPU进行分布式训练,同时模拟更大的批量大小。
2 scaler.step(optimizer):
1 scaler 是一个 torch.cuda.amp.GradScaler 对象,用于自动混合精度训练(Automatic Mixed Precision, AMP)。
2 scaler.step(optimizer) 用于在缩放梯度后执行优化器的更新步骤。
3 scaler.update():
更新梯度缩放器的状态，以便在下一次迭代中使用。
4 optimizer.zero_grad():
将优化器中的梯度清零，以便在下一个批次中累积新的梯度。
5 指数滑动平均(EMA)更新:
ema 是一个指数滑动平均对象，用于在训练过程中维护模型参数的滑动平均。
ema.update(model) 更新 EMA 对象的状态，使其包含当前模型的参数。
总结来说，这段代码在满足特定条件（即当前迭代次数是 accumulate 的倍数）时，执行优化器的更新步骤，
更新梯度缩放器的状态，清零优化器的梯度，并更新指数滑动平均对象的状态。这些操作确保了模型参数的正确
更新和滑动平均的维护。
Scheduler
# Scheduler
lr = [x['lr'] for x in optimizer.param_groups] # for tensorboard
scheduler.step()
# DDP process 0 or single-GPU
if rank in [-1, 0]:
# mAP
ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride', 'class_weights'])
final_epoch = epoch + 1 == epochs
if not opt.notest or final_epoch: # Calculate mAP
wandb_logger.current_epoch = epoch + 1
results, maps, times = test.test(data_dict,
batch_size=batch_size * 2,
imgsz=imgsz_test,
model=ema.ema,
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=save_dir,
verbose=nc < 50 and final_epoch,
plots=plots and final_epoch,
wandb_logger=wandb_logger,
compute_loss=compute_loss,
is_coco=is_coco)
3.17 plot, Test
if rank in [-1, 0]:
# Plots
if plots:
plot_results(save_dir=save_dir) # save as results.png
if wandb_logger.wandb:
files = ['results.png', 'confusion_matrix.png', *[f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R')]]
wandb_logger.log({"Results": [wandb_logger.wandb.Image(str(save_dir / f), caption=f) for f in files
if (save_dir / f).exists()]})
# Test best.pt
logger.info('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600))
if opt.data.endswith('coco.yaml') and nc == 80: # if COCO
for m in (last, best) if best.exists() else (last): # speed, mAP tests
results, _, _ = test.test(opt.data,
batch_size=batch_size * 2,
imgsz=imgsz_test,
conf_thres=0.001,
iou_thres=0.7,
model=attempt_load(m, device).half(),
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=save_dir,
save_json=True,
plots=False,
is_coco=is_coco)
# Strip optimizers
final = best if best.exists() else last # final model
for f in last, best:
if f.exists():
strip_optimizer(f) # strip optimizers
if opt.bucket:
os.system(f'gsutil cp {final} gs://{opt.bucket}/weights') # upload
if wandb_logger.wandb and not opt.evolve: # Log the stripped model
wandb_logger.wandb.log_artifact(str(final), type='model',
name='run_' + wandb_logger.wandb_run.id + '_model',
aliases=['last', 'best', 'stripped'])
wandb_logger.finish_run()
else:
dist.destroy_process_group()
torch.cuda.empty_cache()
return results
扫一扫
666
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
