YOLOv5中Pytorch训练步骤代码
每次自己写神经网络的时候,总是在参考Yolo系列的代码,其中,第五版代码写的最为清晰的同时,也为可视化编程提供了新思路。最近也是一直在写一个可视化编程的操作界面,这里也对YOLOv5中的一些代码思路做一下总结。
除了本文,同时推荐其他相关博文:
-
YOLOv5代码阅读笔记:
https://blog.csdn.net/weixin_43541325/article/details/108884363
-
DDP及其在Pytorch中的应用:
https://blog.csdn.net/cdknight_happy/article/details/108262595
Pytorch代码学习小笔记
01
日志的打印
![](https://oss-emcsprod-public.modb.pro/wechatSpider/modb_20220802_db9ab140-124b-11ed-a22f-fa163eb4f6be.png)
未对logging进行level设置时无法在终端打印,因此应进行配置,正确代码如下
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
logger.info(f"{INPUT}")
02
DDP模式
![](https://oss-emcsprod-public.modb.pro/wechatSpider/modb_20220802_db9ab140-124b-11ed-a22f-fa163eb4f6be.png)
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)\
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技巧,具体代码为:
后续,accumulate为对应的周期,当经过accumulate个batch后,网络模型才进行一次更新,相当于变向扩大了网络的batch_size。# 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
3. 加载数据集
yolov5设置了create_dataloader()函数,这一函数定义为:
其中,sampler和dataloader中的sampler=sampler,可以创建分布式samplerdef 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
并且在每一次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函数中关键代码举例
![](https://oss-emcsprod-public.modb.pro/wechatSpider/modb_20220802_db9ab140-124b-11ed-a22f-fa163eb4f6be.png)
设置参数并加载模型
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
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)
相关代码为:
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 = ModelEMA(model) if rank in [-1, 0] else None
from torch.cuda import amp
scaler = amp.GradScaler(enbaled=cuda)
compute_loss = ComputeLoss(model)
正式训练
训练中关键代码包括:
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()
本文来自博客园,作者:海_纳百川,转载请注明原文链接:https://www.cnblogs.com/chentiao/p/16790625.html,如有侵权联系删除