像 torchvision 这样的 PyTorch 领域库提供了对常见数据集和模型的便捷访问，可用于快速创建最先进的基线。此外，它们还提供了常见的抽象以减少用户可能不得不反复编写的样板代码。torchvision 0.3 版本带来了多项新功能，包括用于语义分割、对象检测、实例分割和人体关键点检测的模型，以及计算机视觉专用的自定义 C++ / CUDA 运算符。
新功能包括：
参考训练/评估脚本： torchvision 现在在 references/ 文件夹下提供以下任务的训练和评估脚本：分类、语义分割、对象检测、实例分割和人体关键点检测。这些脚本记录了如何训练特定模型，并提供基线训练和评估脚本，以快速启动研究。
torchvision 运算符： torchvision 现在包含自定义 C++ / CUDA 运算符。这些运算符专用于计算机视觉，使构建对象检测模型更加容易。这些运算符目前不支持 PyTorch 脚本模式，但计划在下一版本中支持。支持的一些运算符包括：
- roi_pool（以及模块版本 RoIPool）
- roi_align（以及模块版本 RoIAlign）
- nms，用于边界框的非极大值抑制
- box_iou，用于计算两组边界框之间的交并比指标
- 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 作为骨干网络。ResNet101 骨干网络的预训练权重可用，并在 COCO train2017 的一个子集上进行了训练，该子集包含与 Pascal VOC 相同的 20 个类别。
预训练模型在 COCO val2017 的一个子集上给出了以下结果，该子集包含与 Pascal VOC 中相同的 20 个类别：
| 网络 | 平均 IoU | 全局像素级准确率 |
|---|---|---|
| FCN ResNet101 | 63.7 | 91.9 |
| DeepLabV3 ResNet101 | 67.4 | 92.4 |
检测模型
| 网络 | box AP | mask AP | keypoint AP |
|---|---|---|---|
| Faster R-CNN ResNet-50 FPN 在 COCO 上训练 | 37.0 | ||
| Mask R-CNN ResNet-50 FPN 在 COCO 上训练 | 37.9 | 34.6 | |
| Keypoint R-CNN ResNet-50 FPN 在 COCO 上训练 | 54.6 | 65.0 |
用于对象检测、实例分割和关键点检测的模型实现速度很快，尤其是在训练期间。
在下表中，我们使用 8 个 V100 GPU，搭载 CUDA 10.0 和 CUDNN 7.4 来报告结果。训练期间，每个 GPU 使用批量大小为 2，测试期间使用批量大小为 1。
对于测试时间，我们报告模型评估和后处理（包括图像中的掩码粘贴）的时间，但不包括计算精确率-召回率的时间。
| 网络 | 训练时间（秒/迭代） | 测试时间（秒/迭代） | 内存（GB） |
|---|---|---|---|
| Faster R-CNN ResNet-50 FPN | 0.2288 | 0.0590 | 5.2 |
| Mask R-CNN ResNet-50 FPN | 0.2728 | 0.0903 | 5.4 |
| Keypoint R-CNN ResNet-50 FPN | 0.3789 | 0.1242 | 6.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，提供类别字符串）
- 语义边界数据集
- VisionDataset 作为所有数据集的基类
此外，我们还添加了更多图像转换、一般改进和错误修复，以及改进了文档。
请在此处查看完整的发行说明，并在此处查看 Google Colab 上的入门教程，其中描述了如何在自定义数据集上微调您自己的实例分割模型。
干杯！
PyTorch 团队