Faster-RCNN简易复现

Faster-RCNN复现

数据准备

主要目标是利用pytorch框架简易复现Faster-RCNN，我选了一个比较简单的数据集VOC2012，这个数据集标注用的xml格式。

数据读取如下：

 def loadXml(path):
     '''
     读取原始的XMLlabel文件
     :return:
     '''
     dom = xml.dom.minidom.parse(path)
 ​
     root = dom.documentElement
     #获取图片上所有的目标
     objs = root.getElementsByTagName("object")
     #获取这张图片的长宽
     w, h = root.getElementsByTagName("width")[0].childNodes[0].data, root.getElementsByTagName("height")[0].childNodes[0].data
     infos = []
     #把所有目标的信息放到infos中返回
     for item in objs:
         c = item.getElementsByTagName("name")[0].childNodes[0].data
         xmin = item.getElementsByTagName("xmin")[0].childNodes[0].data
         ymin = item.getElementsByTagName("ymin")[0].childNodes[0].data
         xmax = item.getElementsByTagName("xmax")[0].childNodes[0].data
         ymax = item.getElementsByTagName("ymax")[0].childNodes[0].data
         infos.append((c, int(xmin), int(ymin), int(xmax), int(ymax)))
     return int(w), int(h), infos

模型配置

关于训练的参数，anchor的size比论文上多两个，提高小目标的检测能力，长宽比还是三个。初始学习率为0.0001，每2轮后乘以0.1。 训练的相关参数配置如下：

 class Config():
     # anchor的size和长宽比，这里的格式是由pytorch中自带的基础faster_rcnn的输入决定的，参考  
     # torchvision.models.detection.faster_rcnn
     anchor_size = ((32, 64, 128, 256, 512),)
     aspect_ratios = ((0.5, 1, 2),)
     # 特征网络
     backbone = "mobilenet_V2"
     # 目标类型，0是背景，同样是torchvision.models.detection.faster_rcnn的要求
     cls_label = ['__background__', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car',
                  'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike',
                  'person', 'pottedplant', 'sheep', 'sofa', 'train',
                  'tvmonitor', 'cat'] #VOC所有标签
     class_num = 20
     lr = 0.0001
     epoch = 6
     batch_size = 2
     # 模型保存位置
     FasterRCNN_checkpoints = './checkpoints/FasterRCNN/model.pth'

torchvision中自带的faster_rcnn模块

 from torchvision.models.detection import faster_rcnn
 from torchvision.models.detection.rpn import AnchorGenerator
 from torchvision.ops import MultiScaleRoIAlign
 ​
 ​
 # 生成anchor
 anchor_generator = AnchorGenerator(sizes=cfg.anchor_size, aspect_ratios=cfg.aspect_ratios)
 ​
 # 设置roipooling，采用ROIAlign
 roi_pooler = MultiScaleRoIAlign(featmap_names=['0'], output_size=7, sampling_ratio=2)
 ​
 # 建立FasterRCNN模型
 net = faster_rcnn.FasterRCNN(backbone,
                              num_classes=cfg.class_num+1,
                              rpn_anchor_generator=anchor_generator,
                              box_roi_pool=roi_pooler)

faster_rcnn创建主要需要三块，一个是anchor的生成器，一个是roipooling方式，还有一个是用于提取特征的backbone。

backbone可以用torch中提供的，也可以自己写的，但是，如果是自己写的，最好先进行预训练，我用没有预训练的自定义的backbone去试了下，效果不佳。

关于MultiScaleRoIAlign的第一个参数为啥是['0']，torchvision.models.detection.faster_rcnn.py中有说明。。说是如果backbone返回多个特征图层时，写入需要计算的图层名字，比如['feat1', 'feat3']，但如果只返回一个tensor，也就是只有一个图层返回，就['0']。

faster_rcnn中提供很多参数，可以在建立模型时传入。

 def __init__(self, backbone, num_classes=None,
              # transform parameters
              min_size=800, max_size=1333,
              image_mean=None, image_std=None,
              # RPN parameters
              rpn_anchor_generator=None, rpn_head=None,
              rpn_pre_nms_top_n_train=2000, rpn_pre_nms_top_n_test=1000,
              rpn_post_nms_top_n_train=2000, rpn_post_nms_top_n_test=1000,
              rpn_nms_thresh=0.7,
              rpn_fg_iou_thresh=0.7, rpn_bg_iou_thresh=0.3,
              rpn_batch_size_per_image=256, rpn_positive_fraction=0.5,
              # Box parameters
              box_roi_pool=None, box_head=None, box_predictor=None,
              box_score_thresh=0.05, box_nms_thresh=0.5, box_detections_per_img=100,
              box_fg_iou_thresh=0.5, box_bg_iou_thresh=0.5,
              box_batch_size_per_image=512, box_positive_fraction=0.25,
              bbox_reg_weights=None):

训练部分

 from dataset.MyDataset import MyDataset
 from models.FasterRCNN import FasterRCNN
 from torch.utils.data import DataLoader
 from Config import Config
 import torch.optim as optim
 import numpy as np
 from torch.autograd import Variable
 import torch
 from utils.DataProcess import bar
 from torch.utils.tensorboard import SummaryWriter
 ​
 ​
 def train():
     # 调用tensorboard监视训练过程
     writer = SummaryWriter("runs/summary")
     cfg = Config()
 ​
     net = FasterRCNN().net
 ​
     net.cuda()
 ​
     trainset = MyDataset('train')
 ​
     trainLoader = DataLoader(trainset, batch_size=cfg.batch_size, shuffle=True, drop_last=True)
 ​
     params = [p for p in net.parameters() if p.requires_grad]
 ​
     opt = optim.Adam(params, lr=cfg.lr)
 ​
     num_train = len(trainLoader)
 ​
     for epoch in range(1, cfg.epoch + 1):
         net.train()
         # 每2轮学习率下降一次
         if epoch % 2 == 0:
             for p in opt.param_groups:
                 p['lr'] *= 0.1
         for n, data in enumerate(trainLoader):
             imgs = []
             gtbox = []
             for i in range(cfg.batch_size):
                 img_path = data[0][i]
                 gtbox_path = data[1][i]
                 img = trainset.load_pic(img_path)
                 img = Variable(img)
                 gtbbox = np.load(gtbox_path)
                 gtbbox = torch.from_numpy(gtbbox)
                 gtbbox = Variable(gtbbox)
                 if torch.cuda.is_available():
                     gtbbox = gtbbox.cuda()
                     img = img.cuda()
                 imgs.append(img)
                 gtbox.append({"boxes": gtbbox[:, 1:], "labels": gtbbox[:, 0]})
             opt.zero_grad()
             data = net(imgs, gtbox)
             l_cls = data['loss_classifier']
             l_box = data['loss_box_reg']
             l_obj = data['loss_objectness']
             l_box_rpn = data['loss_rpn_box_reg']
             bar('正在第%d轮训练,loss_cls=%.5f,loss_box=%.5f,loss_obj=%.5f,loss_box_rpn=%.5f' %
                 (epoch, l_cls.data, l_box.data, l_obj.data, l_box_rpn.data), n, num_train)
             loss = l_cls + l_box + l_obj + l_box_rpn
             # 添加监视目标
             writer.add_scalar("loss_classifier", l_cls.data, (epoch - 1) * num_train + n)
             writer.add_scalar("loss_box_reg", l_box.data, (epoch - 1) * num_train + n)
             writer.add_scalar("loss_objectness", l_obj.data, (epoch - 1) * num_train + n)
             writer.add_scalar("loss_rpn_box_reg", l_box_rpn.data, (epoch - 1) * num_train + n)
             loss.backward()
             opt.step()
         torch.save(net, cfg.FasterRCNN_checkpoints)
 ​
 ​
 if __name__ == '__main__':
     train()

训练部分用tensorboard记录了一下训练过程。这里注意一下，torch提供的faster_rcnn输入包含两部分，一个是图片集，注意不需要resize，第二个是标注集，标注集是个list，里面元素是Map，包括坐标信息boxes和标签信息labels。至于具体怎么读取数据就看自己怎么设计数据了。训练返回结果包括四个损失，anchor的坐标损失、anchor前背景损失、目标框的坐标损失以及目标分类损失。

模型校验

 net = torch.load(cfg.FasterRCNN_checkpoints)
 net.eval()
 ***********
 ****省略****
 ***********
 img = Variable(img)
 if torch.cuda.is_available():
     img = img.cuda()
 res = net([img])
 res = res[0]
 boxes = res["boxes"]
 labels = res["labels"]
 scores = res["scores"]
 ***********
 ****省略****
 ***********

校验部分就省略写了，主要两点，第一点，不需要新建一个faster_rcnn再加载数据，就直接torch.load就完事了，一定要记得net.eval()；第二点，输入只需要图片列表。校验时，返回结果为坐标、label和得分。最后根据得分进行一下nms即可。

测试效果

自己用的电脑，显卡比较垃圾，就大概练了练，batch_size最大就到2。大概效果如下

源码：https://github.com/Zou-CM/Faster-RCNN-pytorch