Faster RCNN模型如何自定义损失函数
Faster RCNN模型如何自定义损失函数
- 1. 代码分析:
-
- 1.1 _fasterRCNN类:最基础的模型类
- 1.2 resnet类:继承_fasterRCNN类
- 1.3 vgg类:继承_fasterRCNN类
- 1.4 如果要修改loss,一共需要修改哪些内容:
-
- 1.4.1 模型代码:/lib/model/faster_rcnn/faster_rcnn.py,/lib/model/faster_rcnn/resnet.py,/lib/model/faster_rcnn/vgg16.py
- 1.4.2 训练代码:trainval_net.py
- 1.4.3 测试代码:test_net.py
- 1.4.4 运行参数代码:\\lib\\model\\utils\\parser_func.py
- 2. 修改模型代码:(以添加confidence loss为例)
-
- 2.1 修改_fasterRCNN类:
- 2.2 修改resnet类:
- 2.3 修改vgg16类:
- 3. 修改训练代码:
- 4. 修改测试代码:
- 5. 修改运行参数代码:
1. 代码分析:
1.1 _fasterRCNN类:最基础的模型类
/lib/model/faster_rcnn/faster_rcnn.py
class _fasterRCNN(nn.Module):
1.2 resnet类:继承_fasterRCNN类
/lib/model/faster_rcnn/resnet.py
class resnet(_fasterRCNN):
1.3 vgg类:继承_fasterRCNN类
与1.2同理
1.4 如果要修改loss,一共需要修改哪些内容:
1.4.1 模型代码:/lib/model/faster_rcnn/faster_rcnn.py,/lib/model/faster_rcnn/resnet.py,/lib/model/faster_rcnn/vgg16.py
1.4.2 训练代码:trainval_net.py
1.4.3 测试代码:test_net.py
1.4.4 运行参数代码:\\lib\\model\\utils\\parser_func.py
2. 修改模型代码:(以添加confidence loss为例)
2.1 修改_fasterRCNN类:
修改/lib/model/faster_rcnn/faster_rcnn.py
class _fasterRCNN(nn.Module):# 旧的__init__:# def __init__(self, classes, class_agnostic):# 新的__init__:def __init__(self, classes, class_agnostic, conf):......# 旧的初始化loss:# self.RCNN_loss_cls = 0# self.RCNN_loss_bbox = 0# 新的初始化loss,这里添加自己定义的损失,如:self.RCNN_loss_cls = 0self.RCNN_loss_bbox = 0self.conf = conf......# 旧的forward:# def forward(self, im_data, im_info, gt_boxes, num_boxes):# 新的forward:def forward(self, im_data,im_info,gt_boxes,num_boxes,target=False,test=False,eta=1.0,hints=False,):...# 旧的compute object classification probability:# cls_score = self.RCNN_cls_score(pooled_feat)# cls_prob = F.softmax(cls_score, 1)# 新的compute object classification probability:cls_score = self.RCNN_cls_score(pooled_feat)cls_prob = F.softmax(cls_score, 1)if self.conf:pooled_feat_conf = pooled_feat.detach()# confidenceconfidence = F.sigmoid(self.netD_confidence(pooled_feat_conf))# Make sure we don't have any numerical instabilityeps = 1e-12pred_original = torch.clamp(cls_prob, 0.0 + eps, 1.0 - eps)confidence = torch.clamp(confidence, 0.0 + eps, 1.0 - eps)confidence_loss = -torch.log(confidence)...# 旧的返回值:# return rois, cls_prob, bbox_pred, rpn_loss_cls, rpn_loss_bbox, RCNN_loss_cls, RCNN_loss_bbox, rois_label# 新的返回值:if self.conf:return (rois,cls_prob,bbox_pred,rpn_loss_cls,rpn_loss_bbox,RCNN_loss_cls,RCNN_loss_bbox,rois_label,None,None,confidence_loss,confidence,)else:return (rois,cls_prob,bbox_pred,rpn_loss_cls,rpn_loss_bbox,RCNN_loss_cls,RCNN_loss_bbox,rois_label,None,None,None,None,)
2.2 修改resnet类:
/lib/model/faster_rcnn/resnet.py
# 整个netD_confidence都是新加的,是用来计算confidence loss的层结构:
class netD_confidence(nn.Module):def __init__(self, feat_d):super(netD_confidence, self).__init__()self.fc1 = nn.Linear(feat_d, 128)self.bn = nn.BatchNorm1d(128)self.fc2 = nn.Linear(128, 1)def forward(self, x):x = self.fc2(F.dropout(self.bn(self.fc1(x))))return xclass resnet(_fasterRCNN):# 旧的 __init__:#def __init__(self, classes, num_layers=101, pretrained=False, class_agnostic=False):# 新的 __init__:多个conf=Nonedef __init__( self, classes, num_layers=101, pretrained=False, class_agnostic=False, conf=None):self.model_path = 'data/pretrained_model/resnet101_caffe.pth'self.dout_base_model = 1024self.pretrained = pretrainedself.class_agnostic = class_agnostic# 多加一行conf的初始化:self.conf = conf...# 旧的_fasterRCNN.__init__:# _fasterRCNN.__init__(self, classes, class_agnostic)# 新的_fasterRCNN.__init__:多一个conf_fasterRCNN.__init__(self, classes, class_agnostic, conf)...def _init_modules(self):...# 旧的build resnet:# self.RCNN_base = nn.Sequential(resnet.conv1, resnet.bn1,resnet.relu,# resnet.maxpool,resnet.layer1,resnet.layer2,resnet.layer3)# self.RCNN_top = nn.Sequential(resnet.layer4)# self.RCNN_cls_score = nn.Linear(2048, self.n_classes)# if self.class_agnostic:# self.RCNN_bbox_pred = nn.Linear(2048, 4)# else:# self.RCNN_bbox_pred = nn.Linear(2048, 4 * self.n_classes)# 新的build resnet:self.RCNN_base = nn.Sequential(resnet.conv1, resnet.bn1,resnet.relu,resnet.maxpool,resnet.layer1,resnet.layer2,resnet.layer3)self.RCNN_top = nn.Sequential(resnet.layer4)feat_d = 2048self.RCNN_cls_score = nn.Linear(feat_d, self.n_classes)self.netD_confidence = netD_confidence(feat_d)if self.class_agnostic:self.RCNN_bbox_pred = nn.Linear(feat_d, 4)else:self.RCNN_bbox_pred = nn.Linear(feat_d, 4 * self.n_classes)...
2.3 修改vgg16类:
同2.2
3. 修改训练代码:
trainval_net.py
...
if __name__ == '__main__':...# 旧的initilize the network:# if args.net == 'vgg16':# fasterRCNN = vgg16(imdb.classes, pretrained=True, class_agnostic=args.class_agnostic)#elif args.net == 'res101':# fasterRCNN = resnet(imdb.classes, 101, pretrained=True, class_agnostic=args.class_agnostic)#elif args.net == 'res50':# fasterRCNN = resnet(imdb.classes, 50, pretrained=True, class_agnostic=args.class_agnostic)#elif args.net == 'res152':# fasterRCNN = resnet(imdb.classes, 152, pretrained=True, class_agnostic=args.class_agnostic)#else:# print("network is not defined")# pdb.set_trace()# 新的initilize the network:if args.net == 'vgg16':fasterRCNN = vgg16(imdb.classes, pretrained=True, class_agnostic=args.class_agnostic, conf=args.conf)elif args.net == 'res101':fasterRCNN = resnet(imdb.classes, 101, pretrained=True, class_agnostic=args.class_agnostic, conf=args.conf)elif args.net == 'res50':fasterRCNN = resnet(imdb.classes, 50, pretrained=True, class_agnostic=args.class_agnostic, conf=args.conf)elif args.net == 'res152':fasterRCNN = resnet(imdb.classes, 152, pretrained=True, class_agnostic=args.class_agnostic, conf=args.conf)else:print("network is not defined")pdb.set_trace()...# 添加一行conf_gamma,即conf损失在总损失函数中的权重:conf_gamma = args.conf_gamma...fasterRCNN.zero_grad()(rois,cls_prob,bbox_pred,rpn_loss_cls,rpn_loss_box,RCNN_loss_cls,RCNN_loss_bbox,rois_label,out_d_pixel, # 新加的out_d, # 新加的confidence_loss, # 新加的_, # 新加的) = fasterRCNN(im_data, im_info, gt_boxes, num_boxes, hints=True)loss = rpn_loss_cls.mean() + rpn_loss_box.mean() + RCNN_loss_cls.mean() + RCNN_loss_bbox.mean()# 加上conf_loss:if args.conf:conf_loss = confidence_loss.mean()...if args.mGPUs:...else:loss_rpn_cls = rpn_loss_cls.item()loss_rpn_box = rpn_loss_box.item()loss_rcnn_cls = RCNN_loss_cls.item()loss_rcnn_box = RCNN_loss_bbox.item()fg_cnt = torch.sum(rois_label.data.ne(0))bg_cnt = rois_label.data.numel() - fg_cnt# 加上conf:if args.conf:loss_conf = conf_loss.item()
4. 修改测试代码:
test_net.py
在这里插入代码片
5. 修改运行参数代码:
\\lib\\model\\utils\\parser_func.py
def parse_args():...# 添加conf和conf_gamma:parser.add_argument("--ga",dest="conf_gamma",help="the weight for confidence loss.",default=0.1,type=float,)parser.add_argument("--conf",dest="conf",help="whether use the confidence branch to help teaching process",action="store_false",)...