Pytorch torchvision完成Faster-rcnn目标检测demo及源码详解

Torchvision更新到0.3.0后支持了更多的功能，其中新增模块detection中实现了整个faster-rcnn的功能。本博客主要讲述如何通过torchvision和pytorch使用faster-rcnn，并提供一个demo和对应代码及解析注释。

目录

如果你不想深入了解原理和训练，只想用Faster-rcnn做目标检测，请看这里

torchvision中Faster-rcnn接口

一个demo

使用方法

如果你想深入了解原理，并训练自己的模型

环境搭建

准备训练数据

模型训练

单张图片检测

效果

如果你不想深入了解原理和训练，只想用Faster-rcnn做目标检测，请看这里
torchvision中Faster-rcnn接口
torchvision内部集成了Faster-rcnn的模型，其接口和调用方式野非常简洁，目前官方提供resnet50+rpn在coco上训练的模型，调用该模型只需要几行代码：

>>> import torch
>>> import torchvision
 
// 创建模型，pretrained=True将下载官方提供的coco2017模型
>>> model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
>>> model.eval()
>>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
>>> predictions = model(x)
 
 
 
注意网络的输入x是一个Tensor构成的list，而输出prediction则是一个由dict构成list。prediction的长度和网络输入的list中Tensor个数相同。prediction中的每个dict包含输出的结果：

其中boxes是检测框坐标，labels是类别，scores则是置信度。

>>> predictions[0]
 
{'boxes': tensor([], size=(0, 4), grad_fn=<StackBackward>), 'labels': tensor([], dtype=torch.int64), 'scores': tensor([], grad_fn=<IndexBackward>)}
一个demo
如果你不想自己写读取图片/预处理/后处理，我这里有个写好的demo.py，可以跑在任何安装了pytorch1.1+和torchvision0.3+的环境下，不需要其他依赖，可以用来完成目标检测的任务。

为了能够显示类别标签，我们将coco的所有类别写入coco_names.py

names = {'0': 'background', '1': 'person', '2': 'bicycle', '3': 'car', '4': 'motorcycle', '5': 'airplane', '6': 'bus', '7': 'train', '8': 'truck', '9': 'boat', '10': 'traffic light', '11': 'fire hydrant', '13': 'stop sign', '14': 'parking meter', '15': 'bench', '16': 'bird', '17': 'cat', '18': 'dog', '19