这篇是最后一篇 讲述train.py和trainer.py的一些代码细节
文章目录
train.py
#程序入口
from __future__ import absolute_import
import cupy as cp
import os
import ipdb
import matplotlib
from tqdm import tqdm
from utils._config import opt
from data.dataset import Dataset,TestDataset,inverse_normalize
from model import faster_rcnn_vgg16
from torch.utils import data as data_
from trainer import FasterRCNNTrainer
from utils import array_tool as at
from utils.vis_tool import visdom_bbox
from utils.eval_tool import eval_detection_voc
import resource
rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (2048, rlimit[1]))
matplotlib.use('agg')
def eval(dataloader,faster_rcnn,test_num=10000):
pred_bboxes,pred_labels,pred_scores=list(),list(),list()
gt_bboxes,gt_labels,gt_difficults=list(),list(),list()
for ii,(imgs,sizes,gt_bboxes_,gt_labels_,gt_difficults_) in tqdm(enumerate(dataloader)):
sizes=[sizes[0][0].item(),sizes[1][0].item]
pred_bboxes_,pred_labels_,pred_scores_=faster_rcnn.predict(imgs,[sizes])
gt_bboxes+=list(gt_bboxes_.numpy())
gt_labels+=list(gt_labels_.numpy())
gt_difficults+=list(gt_difficults_.numpy())
pred_bboxes+=pred_bboxes_
pred_labels+=pred_labels_
pred_scores+=pred_scores_
if ii==test_num:break
result=eval_detection_voc(pred_bboxes,pred_labels,pred_scores,gt_bboxes,gt_labels,gt_difficults,use_07_metric=True)
return result
def train(**kwargs):
opt._parse(kwargs)
dataset=Dataset(opt)
print('load data')
dataloader=data_.DataLoader(dataset,batch_size=1,shuffle=True,num_workers=opt.num_works)
testset=TestDataset(opt)
test_dataloader=data_.DataLoader(testset,batch_size=1,num_workers=opt.test_num_works,shuffle=False,pin_memory=True)
#设置pin_memory=True，则意味着生成的Tensor数据最开始是属
# 于内存中的锁页内存，这样将内存的Tensor转义到GPU的显存就会更快一些。
faster_rcnn=faster_rcnn_vgg16()#定义好模型
print('model construct completed')
trainer=FasterRCNNTrainer(faster_rcnn).cuda()#设置trainer 将模型送入trainer 并用cuda()设置好模型加速
if opt.load_path:#加载模型继续训练或者predict
trainer.load(opt.load_path)
print('load pretrained model from %s'% opt.load_path)
trainer.vis.text(dataset.db.label_names,win='labels')
best_map=0
lr_=opt.lr
for epoch in range(opt.epoch):
trainer.reset_meters()#首先在可视化界面重设所有数据
for ii,(img,bbox_,label_,scale) in tqdm(enumerate(dataloader)):
scale=at.scalar(scale)
img,bbox,label=img.cuda().float(),bbox_.cuda(),label_.cuda()
trainer.train_step(img,bbox,label,scale)#调用trainer.py中的函数trainer.train_step(img,bbox,label,scale)进行一次参数迭代优化过程
if (ii+1)%opt.plot_every==0:# 判断数据读取次数是否能够整除plot_every(是否达到了画图次数)，
# 如果达到判断debug_file是否存在，用ipdb工具设置断点，
# 调用trainer中的trainer.vis.plot_many(trainer.get_meter_data())将训练数据读取并上传完成可视化！
if os.path.exists(opt.debug_file):
ipdb.set_trace()
#plot loss
trainer.vis.plot_many(trainer.get_meter_date())
#plot gt
ori_img=inverse_normalize(at.tonumpy(img[0]))
gt_img=visdom_bbox(ori_img,at.tonumpy(bbox_[0]),at.tonumpy(label_[0]))
trainer.vis.img('gt_img',gt_img)
# 将每次迭代读取的图片用dataset文件里面的inverse_normalize()函数进行预处理，将处理后的图片调用Visdom_bbox
#plot predict bboxes
_bboxes,_labels,_scores=trainer.faster_rcnn.predict([ori_img],visualize=True)
pred_img=visdom_bbox(ori_img,at.tonumpy(bbox_[0]),at.tonumpy(label_[0]).reshape(-1),at.tonumpy(_scores[0]))
trainer.vis.img('pred_img',pred_img)
#rpn confusion matrix(meter)
trainer.vis.text(str(trainer.rpn_cm.value().tolist()),win='rpn_cm')
#roi confusion matrix
trainer.vis.img('roi_cm',at.totensor(trainer.roi_cm.conf,False).float())
eval_result = eval(test_dataloader, faster_rcnn, test_num=opt.test_num)
trainer.vis.plot('test_map', eval_result['map'])
lr_ = trainer.faster_rcnn.optimizer.param_groups[0]['lr']
log_info = 'lr:{}, map:{},loss:{}'.format(str(lr_),
str(eval_result['map']),
str(trainer.get_meter_data()))
trainer.vis.log(log_info) # 将损失学习率以及map等信息及时显示更新
if eval_result['map'] > best_map:
best_map = eval_result['map']
best_path = trainer.save(best_map=best_map)
if epoch == 9:
trainer.load(best_path)
trainer.faster_rcnn.scale_lr(opt.lr_decay)
lr_ = lr_ * opt.lr_decay
if epoch == 13:
break
if __name__=='__main__':
import fire
fire.Fire()#入口 方便运行调试各个函数
trainer.py
from __future__ import absolute_import
import os
from collections import namedtuple
# 定义一个namedtuple类型User，并包含name，sex和age属性。
# 比如User = namedtuple('User', ['name', 'sex', 'age'])
import time
from torch.nn import functional as F
from model.utils.creator_tool import AnchorTargetCreator,ProposalCreator,ProposalTargetCreator
from torch import nn
import torch as t
from utils import array_tool as at
from utils.vis_tool import Visualizer
from utils._config import opt
from torchnet.meter import ConfusionMeter,AverageValueMeter
#AverageValueMeter能够计算所有数的平均值和标准差，这里用来统计一个epoch中损失的平均值。
# confusionmeter用来统计分类问题中的分类情况，是一个比准确率更详细的统计指标。
#用于存储几个loss
LossTuple=namedtuple('LossTuple',['rpn_loc_loss','rpn_cls_loss','roi_loc_loss','roi_cls_loss','total_loss'])
class FasterRCNNTrainer(nn.Module):
def __init__(self,faster_rcnn):
super(FasterRCNNTrainer,self).__init__()#在初始化子类之前先初始化其父类
self.faster_rcnn=faster_rcnn
self.rpn_sigma=opt.rpn_sigma
self.roi_sigma=opt.roi_sigma
self.anchor_target_creator=AnchorTargetCreator()
self.proposal_target_creator=ProposalTargetCreator()
self.loc_normalize_mean=faster_rcnn.loc_normalize_mean
self.loc_normalize_std=faster_rcnn.loc_normalize_std
self.optimizer=self.faster_rcnn.get_optimizer()
self.rpn_cm=ConfusionMeter(2)#前后景误差矩阵
self.roi_cm=ConfusionMeter(21)#21分类误差矩阵
self.meters={k:AverageValueMeter() for k in LossTuple._fields} #字典解析，平均loss AverageValueMeter会自动求均值
def forward(self,imgs,bboxes,labels,scale) :
n=bboxes.shape[0]#获取batch个数
if n!=1:
raise ValueError('only batch size 1 is supported')
_,_,H,W=imgs.shape
img_size=(H,W)
features=self.faster_rcnn.extractor(imgs)#生成featuremap vgg16 conv5_3之前的部分提取图片的特征
rpn_locs,rpn_scores,rois,rois_indices,anchor=self.faster_rcnn.rpn(features,img_size,scale)
#rpn_locs=(hh*ww*9,4) rpn_scores=(hh*ww*9,2) rois=(2000,4) roi_indices 用不到 anchor=(hh*ww*9,4)
#因为bs是1 所以把变量都变为单bs的形式
bbox=bboxes[0]#bbox纬度是(N,R,4)
label=labels[0]#label纬度是(N,R)
rpn_scores=rpn_scores[0]#（hh*ww*9,4）
rpn_loc=rpn_locs[0]
roi=rois
#Sample ROIs and forward
sample_roi,gt_roi_loc,gt_roi_label=self.proposal_target_creator(roi,at.tonumpy(bbox),
at.tonumpy(label) ,
self.loc_normalize_mean,self.loc_normalize_std)
#调用proposal_target_creator函数生成sample roi（128，4）、gt_roi_loc（128，4）、gt_roi_label（128，1）
# ，RoIHead网络利用这sample_roi+featue为输入，输出是分类（21类）和回归（进一步微调bbox）的预测值，
# 那么分类回归的groud truth就是ProposalTargetCreator输出的gt_roi_label和gt_roi_loc。
sample_roi_index=t.zeros(len(sample_roi))
roi_cls_loc,roi_score=self.faster_rcnn.head(features,sample_roi,sample_roi_index)
# roi回归输出的是128*84和128*21，然而真实位置参数是128*4和真实标签128*1
#------------------------RPN Loss--------------------#
#输入20000个anchor和bbox 调用anchor_target_creator函数得到2000个anchor与bbox的偏移量和label
gt_rpn_loc,gt_rpn_label=self.anchor_target_creator(
at.tonumpy(bbox),
anchor,
img_size
)
gt_rpn_label=at.totensor(gt_rpn_label).long()#将tensor投射为长整形
gt_rpn_loc=at.totensor(gt_rpn_loc)
rpn_loc_loss=_fast_rcnn_loc_loss(rpn_loc,gt_rpn_loc,gt_rpn_loc.data,self.rpn_sigma)#rpn loc损失
#rpn_scores 为rpn网络得到的（20000个）
rpn_cls_loss=F.cross_entropy(rpn_scores,gt_rpn_label.cuda(),ignore_index=-1)#交叉商损失 默认的ignore_index=-100
_gt_rpn_label=gt_rpn_label[gt_rpn_label>-1]
_rpn_score=at.tonumpy(rpn_scores)[at.tonumpy(gt_rpn_label)>-1]
self.rpn_cm.add(at.totensor(_rpn_score,False),_gt_rpn_label.data.long())
#---------------------ROI losses---------------------#
n_sample=roi_cls_loc.shape[0]
roi_cls_loc=roi_cls_loc.view(n_sample,-1,4)
roi_loc=roi_cls_loc[t.arange(0,n_sample).long().cuda(),at.totensor(gt_roi_label).long()]
gt_roi_label=at.totensor(gt_roi_label).long()
gt_roi_loc=at.totensor(gt_roi_loc)
roi_loc_loss=_fast_rcnn_loc_loss(roi_loc.contiguous(),gt_roi_loc,gt_roi_label.data,self.roi_sigma)
#is_contiguous直观的解释是Tensor底层一维数组元素的存储顺序与Tensor按行优先一维展开的元素顺序是否一致。
# Tensor多维数组底层实现是使用一块连续内存的1维数组（行优先顺序存储，下文描述），
# Tensor在元信息里保存了多维数组的形状，在访问元素时，通过多维度索引转化成1维数组相对于数组起始位置的偏移量即可找到对应的数据。
# 某些Tensor操作（如transpose、permute、narrow、expand）与原Tensor是共享内存中的数据，不会改变底层数组的存储，
# 但原来在语义上相邻、内存里也相邻的元素在执行这样的操作后，在语义上相邻，但在内存不相邻，即不连续了（is not contiguous）。
# 如果想要变得连续使用contiguous方法，如果Tensor不是连续的，则会重新开辟一块内存空间保证数据是在内存中是连续的，
# 如果Tensor是连续的，则contiguous无操作。
roi_cls_loss=nn.CrossEntropyLoss()(roi_score,gt_roi_label.cuda())
self.roi_cm.add(at.totensor(roi_score,False),gt_roi_label.data.long())
losses=[rpn_loc_loss,rpn_cls_loss,roi_loc_loss,roi_cls_loss]
losses=losses+[sum(losses)]
return LossTuple(*losses)
def train_step(self,imgs,bboxes,labels,scale):
#整个模型的一个step 从0梯度到向前过程，误差反向传播更新参数的过程
self.optimizer.zero_grad()#即将梯度初始化为零（
# 因为一个batch的loss关于weight的导数是所有sample的loss关于weight的导数的累加和）
losses=self.forward(imgs,bboxes,labels,scale)
losses.total_loss.backward()
self.optimizer.step()#这个方法会更新所有的参数
#会将所有损失的数据更新到可视化界面上最后将losses返回
self.update_meters(losses)
return losses
def save(self,save_optimizer=False,save_path=None,**kwargs):#保存模型参数和config
save_dict=dict()
save_dict['model']=self.faster_rcnn.state_dict()
save_dict['other_info']=kwargs
save_dict['config']=opt._state_dict()
if save_optimizer:
save_dict['optimizer']=self.optimizer.state_dict()
if save_path is None:
timestr=time.strftime('%m%d%H%M') #保存在checkpoints#时间戳命名
save_path='checkpoint/fasterrcnn_%s'%timestr
for k_,v_ in kwargs.items():
save_path+='_%s'%v_
save_dir=os.path.dirname(save_path)
#os.path.dirname(__file__) 得到当前文件的绝对路径 去掉文件名
if not os.path.exists(save_dir):
os.makedirs(save_dir)
t.save(save_dict,save_path)
def load(self,path,load_optimizer=True,parse_opt=False):#加载模型参数
state_dict=t.load(path)
if 'model' in state_dict:
self.faster_rcnn.load_state_dict(state_dict['model'])
else:
self.faster_rcnn.load_state_dict(state_dict)
if parse_opt:
opt._parse(state_dict['config'])
if 'optimizer' in state_dict and load_optimizer:
self.optimizer.load_state_dict(state_dict['optimizer'])
return self
def update_meters(self,losses):#更新average loss
loss_d={k:at.scalar(v) for k,v in losses._asdict().items()}
#namedtuple 的用法_asdict()
for key,meter in self.meters.items():
meter.add(loss_d[key])
def reset_meters(self):#重置meters
for key,meter in self.meters.items():
meter.reset()
self.roi_cm.reset()
self.rpn_cm.reset()
def get_meter_data(self):#返回meters
return {k:v.value()[0] for k,v in self.meters.items()}
def _smooth_l1_loss(x,t,in_weight,sigma):
sigma2=sigma**2
diff=in_weight*(x-t)
abs_diff=diff.abs()
flag=(abs_diff.data<(1./sigma2)).float()
y=(flag*(sigma2/2.)*(diff**2)+(1-flag)*(abs_diff-0.5/sigma2))
return y.sum()
def _fast_rcnn_loc_loss(pred_loc,gt_loc,gt_label,sigma):
in_weight=t.zeros(gt_loc.shape).cuda()
in_weight[(gt_label>0).view(-1,1).expand_as(in_weight).cuda()]=1
loc_loss=_smooth_l1_loss(pred_loc,gt_loc,in_weight.detach(),sigma)
loc_loss/=((gt_label>=0).sum().float())
return loc_loss
不熟悉的python，pytorch，numpy操作
fire 模块
fire.Fire()#入口 方便运行调试各个函数
直接在输入 python train.py train --env=‘fasterrcnn-caffe’ --plot-every=100 --caffe-pretrain 即可方便的向函数中输入参数
data_.DataLoader
from torch.utils import data as data_ 后
将之前定义的自定义data传入 这个loader包装起来 才可以用于训练
dataloader=data_.DataLoader(dataset,batch_size=1,shuffle=True,num_workers=opt.num_works)
必要的参数包括 batch_size, shuffle ,num_workers
self.faster_rcnn.state_dict()
类似字典的形式得到模型的所有参数 ，这是继承自nn.module 的
t.save(save_dict,save_path)
将得到的字典形式的参数全都保存进pth文件中
这里的t就是指的torch
self.faster_rcnn.load_state_dict(state_dict[‘model’])
先用state_dict=t.load(path) 读入模型 在使用
load_state_dict 自动的将模型对应的参数载入到构建号的model中 安名字来载入
tqdm
tqdm(enumerate(dataloader))
将enumerate用tqdm包装起来 就可以得到读取的进度条
trainer.train_step
def train_step(self,imgs,bboxes,labels,scale):
#整个模型的一个step 从0梯度到向前过程，误差反向传播更新参数的过程
self.optimizer.zero_grad()#即将梯度初始化为零（
# 因为一个batch的loss关于weight的导数是所有sample的loss关于weight的导数的累加和）
losses=self.forward(imgs,bboxes,labels,scale)
losses.total_loss.backward()
self.optimizer.step()#这个方法会更新所有的参数
#会将所有损失的数据更新到可视化界面上最后将losses返回
self.update_meters(losses)
return losses
梯度归零操作，前传，反传，更新参数，更新损失
扫一扫
1287
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
