博客

torchvision 0.3:分割、检测模型、新数据集等等……

作者: 2019年5月22日2024年11月16日暂无评论

像 torchvision 这样的 PyTorch 领域库提供了对常用数据集和模型的便捷访问,可用于快速创建最先进的基准模型。此外,它们还提供了通用的抽象层,减少了用户原本必须反复编写的样板代码。torchvision 0.3 版本带来了多项新功能,包括用于语义分割、目标检测、实例分割和人体关键点检测的模型,以及专门针对计算机视觉的自定义 C++ / CUDA 算子。

新功能包括:

参考训练/评估脚本: torchvision 现在在 references/ 文件夹下提供了用于以下任务的训练和评估脚本:分类、语义分割、目标检测、实例分割和人体关键点检测。这些脚本记录了如何训练特定模型,并提供了基准训练和评估脚本,以帮助快速启动研究工作。

torchvision 算子: torchvision 现在包含自定义的 C++ / CUDA 算子。这些算子专用于计算机视觉,能够更轻松地构建目标检测模型。这些算子目前不支持 PyTorch 脚本模式(script mode),但计划在下一个版本中加入支持。目前支持的部分算子包括:

  • roi_pool(以及模块版本 RoIPool)
  • roi_align(以及模块版本 RoIAlign)
  • nms,用于边界框的非极大值抑制(non-maximum suppression)
  • box_iou,用于计算两组边界框之间的交并比(Intersection over Union)指标
  • box_area,用于计算一组边界框的面积

以下是使用 torchvision 算子的几个示例

import torch
import torchvision

# create 10 random boxes
boxes = torch.rand(10, 4) * 100
# they need to be in [x0, y0, x1, y1] format
boxes[:, 2:] += boxes[:, :2]
# create a random image
image = torch.rand(1, 3, 200, 200)
# extract regions in `image` defined in `boxes`, rescaling
# them to have a size of 3x3
pooled_regions = torchvision.ops.roi_align(image, [boxes], output_size=(3, 3))
# check the size
print(pooled_regions.shape)
# torch.Size([10, 3, 3, 3])

# or compute the intersection over union between
# all pairs of boxes
print(torchvision.ops.box_iou(boxes, boxes).shape)
# torch.Size([10, 10])

新模型和数据集: torchvision 现在增加了对目标检测、实例分割和人体关键点检测模型的支持。此外,还添加了几个流行的数据集。注意:该 API 目前处于实验阶段,可能会在 torchvision 的未来版本中发生变化。新模型包括

分割模型

0.3 版本还包含用于图像密集像素预测的模型。它增加了 FCN 和 DeepLabV3 分割模型,并使用了 ResNet50 和 ResNet101 作为主干网络(backbone)。我们提供了基于 ResNet101 主干的预训练权重,这些权重是在 COCO train2017 的一个子集上进行训练的,该子集包含了与 Pascal VOC 相同的 20 个类别。

在包含与 Pascal VOC 相同 20 个类别的 COCO val2017 子集上,预训练模型的表现如下:

网络平均交并比 (mean IoU)全局像素准确率 (global pixelwise acc)
FCN ResNet10163.791.9
DeepLabV3 ResNet10167.492.4

检测模型

网络框 AP (box AP)掩码 AP (mask AP)关键点 AP (keypoint AP)
在 COCO 上训练的 Faster R-CNN ResNet-50 FPN37.0  
在 COCO 上训练的 Mask R-CNN ResNet-50 FPN37.934.6 
在 COCO 上训练的 Keypoint R-CNN ResNet-50 FPN54.6 65.0

这些目标检测、实例分割和关键点检测模型的实现速度很快,特别是在训练期间。

在下表中,我们使用 8 张 V100 GPU、CUDA 10.0 和 CUDNN 7.4 来汇报结果。在训练期间,每张 GPU 使用的 batch size 为 2;在测试期间,batch size 为 1。

对于测试时间,我们汇报了模型评估和后处理(包括将掩码粘贴到图像中)的时间,但不包括计算精确率-召回率(precision-recall)的时间。

网络训练时间 (秒/次)测试时间 (秒/次)显存占用 (GB)
Faster R-CNN ResNet-50 FPN0.22880.05905.2
Mask R-CNN ResNet-50 FPN0.27280.09035.4
Keypoint R-CNN ResNet-50 FPN0.37890.12426.8

您只需几行代码即可加载并使用预训练的检测和分割模型

import torchvision

model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
# set it to evaluation mode, as the model behaves differently
# during training and during evaluation
model.eval()

image = PIL.Image.open('/path/to/an/image.jpg')
image_tensor = torchvision.transforms.functional.to_tensor(image)

# pass a list of (potentially different sized) tensors
# to the model, in 0-1 range. The model will take care of
# batching them together and normalizing
output = model([image_tensor])
# output is a list of dict, containing the postprocessed predictions

分类模型

增加了以下分类模型:

  • GoogLeNet (Inception v1)
  • MobileNet V2
  • ShuffleNet v2
  • ResNeXt-50 32x4d 和 ResNeXt-101 32x8d

数据集

增加了以下数据集:

  • Caltech101, Caltech256 和 CelebA
  • ImageNet 数据集(改进了 ImageFolder,提供类名字符串)
  • 语义边界数据集 (Semantic Boundaries Dataset)
  • 作为所有数据集基类的 VisionDataset

此外,我们还增加了更多图像变换方法、进行了常规改进和错误修复,并改进了文档。

查看完整的发行说明请点击这里,以及这份入门教程在 Google Colab 上的链接,该教程介绍了如何在自定义数据集上微调您自己的实例分割模型。

干杯!

PyTorch 团队