Python PyTorch fasterrcnn_resnet50_fpn用法及代码示例

本文简要介绍python语言中 torchvision.models.detection.fasterrcnn_resnet50_fpn 的用法。

用法:

torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, progress=True, num_classes=91, pretrained_backbone=True, trainable_backbone_layers=None, **kwargs)

参数：

• pretrained(bool) -如果为 True，则返回在 COCO train2017 上预训练的模型

• progress(bool) -如果为 True，则显示下载到 stderr 的进度条

• num_classes(int) -模型的输出类数(包括背景)

• pretrained_backbone(bool) -如果为 True，则返回一个在 Imagenet 上预训练过主干的模型

• trainable_backbone_layers(int) -从最终块开始的可训练(未冻结)resnet 层数。有效值介于 0 和 5 之间，其中 5 表示所有主干层都是可训练的。

使用 ResNet-50-FPN 主干构建 Faster R-CNN 模型。

参考：“更快的R-CNN：通过区域提议网络实现实时目标检测”.

模型的输入应该是一个张量列表，每个形状为 [C, H, W] ，每个图像一个，并且应该在 0-1 范围内。不同的图像可以有不同的尺寸。

模型的行为取决于它是处于训练模式还是评估模式。

在训练期间，模型需要输入张量以及目标(字典列表)，其中包含：

• 框 ( FloatTensor[N, 4] )： [x1, y1, x2, y2] 格式的 ground-truth 框，包含 0 <= x1 < x2 <= W 和 0 <= y1 < y2 <= H 。

• labels (Int64Tensor[N])：每个ground-truth框的类标签

该模型在训练期间返回 Dict[Tensor]，其中包含 RPN 和 R-CNN 的分类和回归损失。

在推理过程中，模型只需要输入张量，并将后处理的预测作为 List[Dict[Tensor]] 返回，每个输入图像一个。 Dict的字段如下，其中N为检测次数：

• 框 ( FloatTensor[N, 4] )：[x1, y1, x2, y2] 格式的预测框，包含 0 <= x1 < x2 <= W 和 0 <= y1 < y2 <= H 。

• labels (Int64Tensor[N])：每次检测的预测标签

• 分数(Tensor[N])：每次检测的分数

有关输出的更多详细信息，您可以参考实例分割模型。

Faster R-CNN 可导出到 ONNX 以用于固定批量大小，输入图像大小固定。

例子：

>>> model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
>>> # For training
>>> images, boxes = torch.rand(4, 3, 600, 1200), torch.rand(4, 11, 4)
>>> labels = torch.randint(1, 91, (4, 11))
>>> images = list(image for image in images)
>>> targets = []
>>> for i in range(len(images)):
>>>     d = {}
>>>     d['boxes'] = boxes[i]
>>>     d['labels'] = labels[i]
>>>     targets.append(d)
>>> output = model(images, targets)
>>> # For inference
>>> model.eval()
>>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
>>> predictions = model(x)
>>>
>>> # optionally, if you want to export the model to ONNX:
>>> torch.onnx.export(model, x, "faster_rcnn.onnx", opset_version = 11)

使用 fasterrcnn_resnet50_fpn 的示例：

将蒙版重新用于边界框

可视化实用程序