YOLOv5中Pytorch训练步骤代码

每次自己写神经网络的时候，总是在参考Yolo系列的代码，其中，第五版代码写的最为清晰的同时，也为可视化编程提供了新思路。最近也是一直在写一个可视化编程的操作界面，这里也对YOLOv5中的一些代码思路做一下总结。

 

除了本文，同时推荐其他相关博文：

1. YOLOv5代码阅读笔记：

   https://blog.csdn.net/weixin_43541325/article/details/108884363

2. DDP及其在Pytorch中的应用:

   https://blog.csdn.net/cdknight_happy/article/details/108262595

 

Pytorch代码学习小笔记

 

01

日志的打印

未对logging进行level设置时无法在终端打印，因此应进行配置，正确代码如下

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
logger.info(f"{INPUT}")

02

DDP模式

 

DP模式与DDP模式

 

DP模式全称为DataParallel，适用于单机多卡训练。

DDP模式全称为DistributeDataParallel，适用于多机多卡训练。

这两种模式都为数据并行方法（不是模型并行）

 

DP模式修改内容：

model = torch.nn.DataParallel(model)

DDP模式

 

DDP模式较为复杂，且存在较多BUG，使用时需要进行调试。

关键参数如下：

group: 进程组，默认情况下只有一个组
rank: 当前进行的序号，用于进程间的通讯，rank=0的进程为主线程
local_rank: 当前进程编号，根据local_rank设置模型运行在那块GPU上
global_rank: 多卡训练时卡的数量global_rank: 多卡训练时卡的数量
sync_bn: 是否使用SyncBatchNorm
total_batch_size: 对于所有GPUs的batch_size(由于数据并行)
world_size: 执行训练的所有进程数，等同于device.count()
backend: 指定通信所用后端,可以是'nccl'、'gloo'

初始阶段

 

启动多进程分布式训练时，示例代码如下：

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP


local_rank = int(input("local_rank"))
torch.cuda.set_device(local_rank)
dist.init_process_group(backend='nccl', init_method='env://')


device = torch.device("cuda", local_rank)
model = parse_model()


model = DDP(model, device_dis=[local_rank], output_device=local_rank)\

 

在YOLOv5中，具体还设计了select_device()函数，该函数定义如下：

def select_device(device='', batch_size=None)

登录后复制

用于获取当前os.environ['CUDA_VISIBLE_DEVICES']中GPU的数量和性能参数，并判断设置的total_batch_size是否合理。
当网络模型需要使用SyncBatchNorm时，启用相应代码为：

model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)

训练阶段

1.设计了torch_distributed_zeros_first(rank)函数，用于判断当前进程是否为主进程，如果不是主进程则利用torch.distribed.barrier()同步不同进程间数据。
2.设置了nominal batch size(模拟batch_size),accumulate技巧，具体代码为：

# nomial batch size
nbs = 64 
# accumulate loss before optimizing
accumulate = max(round(nbs / total_batch_size), 1)
# scale weight_decay
hyp['weight_decay'] *= total_batch_size * accumulate / nbs

后续，accumulate为对应的周期，当经过accumulate个batch后，网络模型才进行一次更新，相当于变向扩大了网络的batch_size。
3. 加载数据集
yolov5设置了create_dataloader()函数，这一函数定义为：

def create_dataloader(path, imgsz, batch_size, stride, opt, hyp=None, augment=False, cache=False, pad=0.0, rect=False,
                      rank=-1, world_size=1, workers=8, image_weights=False, quad=False, prefix=''):
    with torch_distributed_zero_first(rank):
        dataset = LoadImagesAndLabels()
        batch_size = min(batch_size, len(dataset))
        nw = min([os.cpu_count()  world_size, batch_size if batch_size > 1 else 0, workers])
        sampler = torch.utils.data.distributed.DistributedSampler(dataset) if rank!=-1 else None
        loader = torch.utils.data.DataLoader if image_weights else InfiniteDataLoader
        dataloader = loader(dataset, 
                            batch_size=batch_size,
                            num_workers=nw,
                            sampler=sampler,
                            pin_memory=True,
                            collerate_fn=..)
        return dataloader, dataset

其中，sampler和dataloader中的sampler=sampler，可以创建分布式sampler
并且在每一次epoch中，

if rank != -1:
    dataloader.sampler.set_epoch(epoch)

保存模型
只需要对主进程中的模型保存即可，代码为：

if dist.get_rank() == 0:
    torch.save(model.module.state_dict())

结束迭代后：

dist.destory_process_group()

03

train函数中关键代码举例

设置参数并加载模型

1.创建本次训练文件夹并将hyp和opt参数文件进行存储；

2.init_seeds()，并设置nc，names等参数

3.设置wandb存储参数

4.当不是初次训练时：

读取checkpoint文件并加载模型和参数，代码为:

ckpt = torch.load(weights, map_location=device)
model = Model(ckpt['model'].yaml).to(device)
state_dict = ckpt['model'].float().state_dict()
model.load_state_dict(state_dict, strict=False)

 

当初次训练时:

model = Model(opt.cfg).to(device)

5.设置模型freeze

freeze = []  # 模型层号
for k, v in model.named_parameters():
    v.requries_grad = True
    if any(x in k for x in freeze):
        v.requires_grad = False

 

设置OptimizerSGD优化器和Adam优化器推荐：https://zhuanlan.zhihu.com/p/54192512/根据不同层的内容设置不同的参数，代码为：

import torch.optim as optim


optimizer = optim.Adam(model,
                       lr=hyp["lr0"],
                       betas=(hyp["momentum"], 0.999))
# optimizer = optim.SGD(model, lr=hyp["lr0"],
#                       momentum=hyp['momentum'], 
#                       nesterov=True)

 

设置Scheduler
相关代码为：

import torch.optim.lr_scheduler as lr_scheduler
if opt.linear_lr:
    lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp["lrf"]) 
                   + hyp["lrf"]
else:
    lf = one_cycle(1, hyp["lrf"], epochs)


scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)

 

模型EMA

ema = ModelEMA(model) if rank in [-1, 0] else None

 

设置Scaler

from torch.cuda import amp


scaler = amp.GradScaler(enbaled=cuda)

 

设置损失函数

compute_loss = ComputeLoss(model)

 

这里是Yolov5自建的类，其中包含loss的多个组成部分及target的构建，是学习的重点和关键。
正式训练
训练中关键代码包括：

model.train()
# 初始化梯度
optimizer.zero_grad()
# 正向推理
with amp.autocast(enabled=cuda):
    pred = model(imgs)
    loss, loss_items = compute_loss(pred, targets.to(device))
# 根据loss计算网络梯度
scaler.scale(loss).backward()
# 当达到nbs后更新优化器参数
if ni % accumulate == 0:
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()
    if ema:
       ema.update(model)
# 更新优化器参数
scheduler.step()
# 结束训练，清空缓存
torch.cuda.empty_cache()