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

栏目: Python · 发布时间: 5年前

内容简介:PyTorch宣布了视觉工具包终于来到

栗子 发自 凹非寺

量子位 出品 | 公众号 QbitAI

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

终于来到 torchvision 0.3 了。

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

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

也增加了许多 数据集 ,比如 ImageNetCelebACaltech 101 等等等等。

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

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

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

现在,来仔细观察一下,新的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新加入了 FCNDeepLabV3 分割模型,用了ResNet50和ResNet101骨架。

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

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

检测模型

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

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

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

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

速度如下,都是 毫秒级

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

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

 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重大更新:支持分割模型、检测模型,还有许多数据集

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学习干货

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

加入社群

量子位AI社群开始招募啦,社群矩阵: AI讨论群AI+行业群AI技术群

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

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

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

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

վ'ᴗ' ի 追踪AI技术和产品新动态

喜欢就点「在看」吧 !


以上所述就是小编给大家介绍的《视觉工具包torchvision重大更新:支持分割模型、检测模型,还有许多数据集》,希望对大家有所帮助,如果大家有任何疑问请给我留言,小编会及时回复大家的。在此也非常感谢大家对 码农网 的支持!

查看所有标签

猜你喜欢:

本站部分资源来源于网络,本站转载出于传递更多信息之目的,版权归原作者或者来源机构所有,如转载稿涉及版权问题,请联系我们

马云内部讲话

马云内部讲话

阿里巴巴集团 / 红旗出版社 / 2010-12 / 28.00元

马云的话有什么其妙的地方? 为什么员工会把自己的CEO当作偶像? 世界都处在迷茫期,他如何确立阿里巴巴的价值观? 他怎样给已经是富翁的员工寻找新的激情? 风暴袭来,他怎么克服内心的恐惧? 他在互联网合纵连横的动机何在?一起来看看 《马云内部讲话》 这本书的介绍吧!

JS 压缩/解压工具
JS 压缩/解压工具

在线压缩/解压 JS 代码

XML 在线格式化
XML 在线格式化

在线 XML 格式化压缩工具

UNIX 时间戳转换
UNIX 时间戳转换

UNIX 时间戳转换