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

最新推荐文章于 2026-07-20 12:50:14 发布

原创

最新推荐文章于 2026-07-20 12:50:14 发布 · 置顶 · 2.9w 阅读

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

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

torchvision中Faster-rcnn接口

如果你不想深入了解原理和训练，只想用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': 'horse', '20': 'sheep', '21': 'cow', '22': 'elephant', '23': 'bear', '24': 'zebra', '25': 'giraffe', '27': 'backpack', '28': 'umbrella', '31': 'handbag', '32': 'tie', '33': 'suitcase', '34': 'frisbee', '35': 'skis', '36': 'snowboard', '37': 'sports ball', '38': 'kite', '39': 'baseball bat', '40': 'baseball glove', '41': 'skateboard', '42': 'surfboard', '43': 'tennis racket', '44': 'bottle', '46': 'wine glass', '47': 'cup', '48': 'fork', '49': 'knife', '50': 'spoon', '51': 'bowl', '52': 'banana', '53': 'apple', '54': 'sandwich', '55': 'orange', '56': 'broccoli', '57': 'carrot', '58': 'hot dog', '59': 'pizza', '60': 'donut', '61': 'cake', '62': 'chair', '63': 'couch', '64': 'potted plant', '65': 'bed', '67': 'dining table', '70': 'toilet', '72': 'tv', '73': 'laptop', '74':