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

栗子 发自 凹非寺
量子位 出品 | 公众号 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`,?rescaling
11#?them?to?have?a?size?of?3x3
12pooled_regions?=?torchvision.ops.roi_align(image,?[boxes],?output_size=(3,?3))
13#?check?the?size
14print(pooled_regions.shape)
15#?torch.Size([10,?3,?3,?3])
16
17#?or?compute?the?intersection?over?union?between
18#?all?pairs?of?boxes
19print(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)
10
11#?pass?a?list?of?(potentially?different?sized)?tensors
12#?to?the?model,?in?0-1?range.?The?model?will?take?care?of
13#?batching?them?together?and?normalizing
14output?=?model([image_tensor])
15#?output?is?a?list?of?dict,?containing?the?postprocessed?predictions

分类模型

新加入的分类模型有这些：

· GoogLeNet (Inception v1)
· MobiliNet V2
· ShuffleNet v2
· ResNeXt-50 32x4d，以及ResNeXt-101 32x8d

哪些数据集是新来的？

先看人脸数据集，加入了这些：

Caltech101，Caltech256，以及CelebA

然后，ImageNet也加入了。

还有，Semantic Boundaries Dataset (语义边界数据集) 。

以及，VisionDataset，作为各种数据集的基类 (base class) 。

快去用一下

如果你还有什么疑问，可以观察更详细的发布说明。

如果你不知道怎么用，可以服用Colab教程。

如果你已经准备好了，就开始用新的torchvision去做项目吧。

torchvision 0.3.0
https://pypi.org/project/torchvision/

Colab教程：
https://colab.research.google.com/github/pytorch/vision/blob/temp-tutorial/tutorials/torchvision_finetuning_instance_segmentation.ipynb

详细的发布说明：
https://github.com/pytorch/vision/releases

作者系网易新闻·网易号“各有态度”签约作者

— 完 —

小程序|get更多AI学习干货

加入社群

量子位AI社群开始招募啦，社群矩阵：AI讨论群、AI+行业群、AI技术群；

目前已有4万AI行业从业者、爱好者加入，AI技术群更有来自海内外各大高校实验室大牛、各明星AI公司工程师等。自由互相交流AI发展现状及趋势。

欢迎对AI感兴趣的同学，在量子位公众号（QbitAI）对话界面回复关键字“微信群”，获取入群方式。（技术群与AI+行业群需经过审核，审核较严，敬请谅解）

量子位?QbitAI · 头条号签约作者

?'?' ? 追踪AI技术和产品新动态

喜欢就点「在看」吧 !

关注公众号：拾黑（shiheibook）了解更多

[广告]赞助链接：

四季很好，只要有你，文娱排行榜：https://www.yaopaiming.com/
让资讯触达的更精准有趣：https://www.0xu.cn/