以前看CNN训练代码的时候,往往代码比较易懂,基本很快就能知道各个模块功能,但到了后面很多出来的网络中,由于加入了大量的trick,导致很多人看不懂代码,代码下载以后无从下手。训练参数和利用yaml定义网络详细过程可以看我另外的文章,都有写清楚。
其实不管什么网络,训练部分大体都分几个部分:
本文主要会将上述几个部分代码列出来,致于其中的trick部分这里暂不解释【后续我会再写有关内容】,只是为了方便大家先了解训练过程。
注:这里并不是完整的代码,仅仅是对训练代码中几个重要的部分进行梳理、
1.网络的定义
加载权重一般分两种情况:1)有预权重;2)无预权重
1)有预权重情况下的网络定义;
if pretrained: # 加载预权重 ckpt = torch.load(weights, map_location=device) # 加载模型 # 模型的定义 model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create exclude = ['anchor'] if (opt.cfg or hyp.get('anchors'))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
ckpt是加载的预权重【通过torch.load就能看出来】,weights是预权重路径,device用cpu还是gpu。
state_dict是获取的预权重的权值,通过state_dict()可以看出来。
intersect_dicts就是将预权重和模型本身默认的权值进行对比后再赋值给预权重【比如对比层的shape或keys】
model.load_state_dict就是将预权重加载到网络重。
2)无预权重
无预权重就可以直接使用模型默认权值训练,或者也可以自己去初始化一下。
model = Model(opt.cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
2.数据集的处理与加载
数据集的处理的第一步当然是读取数据集路径。这个路径一般在yaml文件中,比如data/coco.yaml.路径文件是txt格式的【后面我会写一篇训练自己的数据集】
获得路径:这个data_dict就是我们读取的yaml文件
train_path = data_dict['train'] # 训练集路径,加载的txt test_path = data_dict['val'] # 测试集路径,加载的txt
要想定位到数据集的处理,只需要找到两个函数,dataset和Dataloader,因为数据集的处理都需要继承torc提供的这两个函数。
dataset用来处理,dataloader用来加载,这两个函数以及处理数据集在我博客中都有写。
yolov7中是实现了一个create_dataloader函数进行数据集的处理【具体实现可以看我yolov7专栏中的数据集处理】
训练集处理:
# Trainloader 训练数据集的处理 dataloader, dataset = create_dataloader(train_path, imgsz, 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, quad=opt.quad, prefix=colorstr('train: '))
验证集处理:
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]
取出图像(后面要放入model中)
pbar = enumerate(dataloader) pbar = tqdm(pbar, total=nb) # progress bar for i, (imgs, targets, paths, _) in pbar: # batch ------------------------------------------------------------- ni = i + nb * epoch # number integrated batches (since train start) imgs = imgs.to(device, non_blocking=True).float() / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
3.训练超参数的定义与初始化
v7中的超参数,比如学习率,优化器类型等是有个hyp的yaml定义的。
然后在代码中又建立了pg0,pg1,pg2列表用来存储网络中的参数【这些参数就是可导的,也就是我们要的最终网络权重】。
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups,pg0放BN层weight,implicit,pg1放卷积weight,pg2放bias
for k, v in model.named_modules(): # =model.modules() if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): pg2.append(v.bias) # biases if isinstance(v, nn.BatchNorm2d): pg0.append(v.weight) # no decay elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter): pg1.append(v.weight) # apply decay if hasattr(v, 'im'): if hasattr(v.im, 'implicit'): pg0.append(v.im.implicit) else: for iv in v.im: pg0.append(iv.implicit) if hasattr(v, 'imc'): if hasattr(v.imc, 'implicit'): pg0.append(v.imc.implicit) else: for iv in v.imc: pg0.append(iv.implicit) if hasattr(v, 'imb'): if hasattr(v.imb, 'implicit'): pg0.append(v.imb.implicit) else: for iv in v.imb: pg0.append(iv.implicit) if hasattr(v, 'imo'): if hasattr(v.imo, 'implicit'): pg0.append(v.imo.implicit) else: for iv in v.imo: pg0.append(iv.implicit) if hasattr(v, 'ia'): if hasattr(v.ia, 'implicit'): pg0.append(v.ia.implicit) else: for iv in v.ia: pg0.append(iv.implicit) if hasattr(v, 'attn'): if hasattr(v.attn, 'logit_scale'): pg0.append(v.attn.logit_scale) if hasattr(v.attn, 'q_bias'): pg0.append(v.attn.q_bias) if hasattr(v.attn, 'v_bias'): pg0.append(v.attn.v_bias) if hasattr(v.attn, 'relative_position_bias_table'): pg0.append(v.attn.relative_position_bias_table) if hasattr(v, 'rbr_dense'): if hasattr(v.rbr_dense, 'weight_rbr_origin'): pg0.append(v.rbr_dense.weight_rbr_origin) if hasattr(v.rbr_dense, 'weight_rbr_avg_conv'): pg0.append(v.rbr_dense.weight_rbr_avg_conv) if hasattr(v.rbr_dense, 'weight_rbr_pfir_conv'): pg0.append(v.rbr_dense.weight_rbr_pfir_conv) if hasattr(v.rbr_dense, 'weight_rbr_1x1_kxk_idconv1'): pg0.append(v.rbr_dense.weight_rbr_1x1_kxk_idconv1) if hasattr(v.rbr_dense, 'weight_rbr_1x1_kxk_conv2'): pg0.append(v.rbr_dense.weight_rbr_1x1_kxk_conv2) if hasattr(v.rbr_dense, 'weight_rbr_gconv_dw'): pg0.append(v.rbr_dense.weight_rbr_gconv_dw) if hasattr(v.rbr_dense, 'weight_rbr_gconv_pw'): pg0.append(v.rbr_dense.weight_rbr_gconv_pw) if hasattr(v.rbr_dense, 'vector'): pg0.append(v.rbr_dense.vector)
调用优化器是要到 optim这个库:
from torch import optim
这里是可以选用Adam和SGD两种。
# 梯度优化方法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)
4.损失函数的定义
网络有了,数据集有了,优化器有了,现在就还需要损失函数了。
compute_loss_ota = ComputeLossOTA(model) # init loss class compute_loss = ComputeLoss(model) # init loss class
5.训练
训练阶段,我们需要定义训练的epoch,即利用for训练不停的将image喂入model,将获得的output和真实值之间建立损失函数进行训练。
训练阶段有个很明显的地方就是model.train()【看到这里就知道已经进入了训练阶段】,测试阶段就是model.eval()。
model.train()
for epoch in range(start_epoch, epochs): # epoch model.train() if rank != -1: dataloader.sampler.set_epoch(epoch) pbar = enumerate(dataloader) if rank in [-1, 0]: pbar = tqdm(pbar, total=nb) # progress bar optimizer.zero_grad() for i, (imgs, targets, paths, _) in pbar: # batch ni = i + nb * epoch # number integrated batches (since train start) imgs = imgs.to(device, non_blocking=True).float() / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
pbar就是从dataloader中读取数据,并将imgs放入device后除以255.0【像素的归一化】
5.1.1前向传播
pred就是得到的output,然后与ground truth签求loss。
# Forward with amp.autocast(enabled=cuda): pred = model(imgs) # forward loss, loss_items = compute_loss_ota(pred, targets.to(device), imgs) # loss scaled by batch_size
5.1.2反向传播
看到backward()这个地方就是梯度反向传播的过程
# Backward scaler.scale(loss).backward()
5.1.3梯度更新
看到update()和zero_grad()就是梯度更新的地方,每次更新后在进行清零操作。
# Optimize if ni % accumulate == 0: scaler.step(optimizer) # optimizer.step scaler.update() optimizer.zero_grad()
5.1.4模型保存
模型的保存会用到torch.save(),所以只要看到这个地方就知道是模型训练好的模型。保存的ckpt就是我们的训练模型,只不过里面还保存了训练的epoch,best_fitness等。
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': None} # Save last, best and delete torch.save(ckpt, last) if best_fitness == fi: torch.save(ckpt, best) if (best_fitness == fi) and (epoch >= 200): torch.save(ckpt, wdir / 'best_{:03d}.pt'.format(epoch)) if epoch == 0: torch.save(ckpt, wdir / 'epoch_{:03d}.pt'.format(epoch)) elif ((epoch+1) % 25) == 0: torch.save(ckpt, wdir / 'epoch_{:03d}.pt'.format(epoch)) elif epoch >= (epochs-5): torch.save(ckpt, wdir / 'epoch_{:03d}.pt'.format(epoch)) del ckpt