视觉工具包torchvision重大更新：支持分割模型、检测模型，还有许多数据集

2019 年 5 月 27 日 极市平台

极市正在推出CVPR2019的专题直播分享会，邀请CVPR2019的论文作者进行线上直播，分享优秀的科研工作和技术干货，也欢迎各位小伙伴自荐或推荐优秀的CVPR论文作者到极市进行技术分享~

本周四（5月30日）晚，帝国理工学院计算机系IBUG组博士生邓健康，将为我们分享：ArcFace 构建高效的人脸识别系统（CVPR2019），公众号回复“42”即可获取直播详情。

栗子发自凹非寺
量子位出品 | 公众号 QbitAI

PyTorch宣布了视觉工具包torchvision的重大更新。

终于来到torchvision 0.3了。

这次，工具包里增加了许多新模型：做语义分割的，做目标检测的，做实例分割的……

也增加了许多数据集，比如ImageNet，CelebA，Caltech 101等等等等。

另外，torchvision还有了不少视觉专用的C++/Cuda算子。

消息一出，小伙伴们纷纷奔走相告。

现在，来仔细观察一下，新的torchvision都有哪里变强了。

哪些功能是新来的？

训练/评估脚本

现在，reference/文件夹地下，提供了训练和评估用的脚本，支持许多任务：

分类、语义分割、目标检测、实例分割，以及人物关键点检测。

这些脚本可以当做log：写着某一个特定的模型要怎样训练，并且提供了基线。有了这份快速指引，便于顺利展开研究。

torchvision算子

就像开头提到的那样，torchvision这次有了定制的C++/CUDA算子，计算机视觉专用。

有了这些算子，搭建目标检测模型就更加轻松了：

· roi_pool (以及模块版本RoIPool)
· roi_align (以及模块版本RoIAlign)
· nms，给边界框做非极大抑制 (Non-Maximum Suppression用的)
· box_iou，用来计算两组边界框之间的交集
· box_area, 用来计算一组边界框的面积

等到下次更新，这些算子就会支持PyTorch脚本模式了。

至于这些算子怎么用，官方给出了一些例子：

1import torch 2import torchvision 3 4# create 10 random boxes 5boxes = torch.rand(10, 4) * 100 6# they need to be in [x0, y0, x1, y1] format 7boxes[:, 2:] += boxes[:, :2] 8# create a random image 9image = torch.rand(1, 3, 200, 200)10# extract regions in `image` defined in `boxes`, rescaling11# them to have a size of 3x312pooled_regions = torchvision.ops.roi_align(image, [boxes], output_size=(3, 3))13# check the size14print(pooled_regions.shape)15# torch.Size([10, 3, 3, 3])1617# or compute the intersection over union between18# all pairs of boxes19print(torchvision.ops.box_iou(boxes, boxes).shape)20# torch.Size([10, 10])

哪些模型是新来的？

既然，目标检测、实例分割，以及人物关键点检测模型的支持，全部是最新加入的。

那么，就来看看增加了哪些模型吧：

分割模型

官方博客写到，torchvision 0.3新加入了FCN和DeepLabV3分割模型，用了ResNet50和ResNet101骨架。

ResNet101有预训练的权重可用，是在COCO train2017数据集的一个子集上训练的，20个类别和Pascal VOC一致：

检测模型

torchvision 0.3新包含了预训练的Faster R-CNN、Mask R-CNN以及Keypoint R-CNN。

官方还提到，各种模型的实现都很快，尤其是训练过程很快。

(团队用了8个V100 GPU，带有CUDA 10.0和CUDNN 7.4的那种。训练中每个GPU的批尺寸是2，测试中的批尺寸是1。)

速度如下，都是毫秒级：

而且，只要写几行代码，就可以直接加载、运行那些预训练的检测模型与分割模型了：

1import torchvision 2 3model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True) 4# set it to evaluation mode, as the model behaves differently 5# during training and during evaluation 6model.eval() 7 8image = PIL.Image.open('/path/to/an/image.jpg') 9image_tensor = torchvision.transforms.functional.to_tensor(image)1011# pass a list of (potentially different sized) tensors12# to the model, in 0-1 range. The model will take care of13# batching them together and normalizing14output = model([image_tensor])15# output is a list of dict, containing the postprocessed predictions