当前位置: 首页>>代码示例 >>用法及示例精选 >>正文


Python PyTorch maskrcnn_resnet50_fpn用法及代码示例


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

用法:

torchvision.models.detection.maskrcnn_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 主干的 Mask R-CNN 模型。

参考:“Mask R-CNN”

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

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

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

  • 框 ( FloatTensor[N, 4] ): [x1, y1, x2, y2] 格式的 ground-truth 框,包含 0 <= x1 < x2 <= W0 <= y1 < y2 <= H

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

  • 掩码(UInt8Tensor[N, H, W]):每个实例的分段二进制掩码

该模型在训练期间返回 Dict[Tensor],包含 RPN 和 R-CNN 的分类和回归损失,以及掩码损失。

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

  • 框 ( FloatTensor[N, 4] ):[x1, y1, x2, y2] 格式的预测框,包含 0 <= x1 < x2 <= W0 <= y1 < y2 <= H

  • labels (Int64Tensor[N]):每个实例的预测标签

  • 分数 (Tensor[N]):每个实例的分数

  • 掩码 ( UInt8Tensor[N, 1, H, W] ):每个实例的预测掩码,在 0-1 范围内。为了获得最终的分割掩模,可以对软掩模进行阈值化,一般值为0.5(mask >= 0.5)

有关输出和如何绘制掩码的更多详细信息,您可以参考实例分割模型。

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

例子:

>>> model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
>>> 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, "mask_rcnn.onnx", opset_version = 11)

使用 maskrcnn_resnet50_fpn 的示例:

相关用法


注:本文由纯净天空筛选整理自pytorch.org大神的英文原创作品 torchvision.models.detection.maskrcnn_resnet50_fpn。非经特殊声明,原始代码版权归原作者所有,本译文未经允许或授权,请勿转载或复制。