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].datainfos = []#把所有目标的信息放到infos中返回for item in objs:c = item.getElementsByTagName("name")[0].childNodes[0].dataxmin = item.getElementsByTagName("xmin")[0].childNodes[0].dataymin = item.getElementsByTagName("ymin")[0].childNodes[0].dataxmax = item.getElementsByTagName("xmax")[0].childNodes[0].dataymax = item.getElementsByTagName("ymax")[0].childNodes[0].datainfos.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_rcnnanchor_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 = 20lr = 0.0001epoch = 6batch_size = 2# 模型保存位置FasterRCNN_checkpoints = './checkpoints/FasterRCNN/model.pth'

torchvision中自带的faster_rcnn模块

 from torchvision.models.detection import faster_rcnnfrom torchvision.models.detection.rpn import AnchorGeneratorfrom torchvision.ops import MultiScaleRoIAlign# 生成anchoranchor_generator = AnchorGenerator(sizes=cfg.anchor_size, aspect_ratios=cfg.aspect_ratios)# 设置roipooling，采用ROIAlignroi_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 parametersmin_size=800, max_size=1333,image_mean=None, image_std=None,# RPN parametersrpn_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 parametersbox_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 MyDatasetfrom models.FasterRCNN import FasterRCNNfrom torch.utils.data import DataLoaderfrom Config import Configimport torch.optim as optimimport numpy as npfrom torch.autograd import Variableimport torchfrom utils.DataProcess import barfrom torch.utils.tensorboard import SummaryWriterdef train():# 调用tensorboard监视训练过程writer = SummaryWriter("runs/summary")cfg = Config()net = FasterRCNN().netnet.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.1for 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