PyTorch 視覺工具包 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
5boxes = torch.rand(10, 4) * 100
6
7boxes[:, 2:] += boxes[:, :2]
8
9image = torch.rand(1, 3, 200, 200)
10
11
12pooled_regions = torchvision.ops.roi_align(image, [boxes], output_size=(3, 3))
13
14print(pooled_regions.shape)
15
16
17
18
19print(torchvision.ops.box_iou(boxes, boxes).shape)
20

哪些模型是新來的？

既然，目標檢測、實例分割，以及人物關鍵點檢測模型的支持，全部是最新加入的。

那麼，就來看看增加了哪些模型吧：

分割模型

官方博客寫到，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
5
6model.eval()
7
8image = PIL.Image.open('/path/to/an/image.jpg')
9image_tensor = torchvision.transforms.functional.to_tensor(image)
10
11
12
13
14output = model([image_tensor])
15

分類模型

新加入的分類模型有這些：

· 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

作者系網易新聞·網易號「各有態度」籤約作者

覺得這篇文章有幫助？請轉發給更多人

關注 AI有道 加星標，獲取最新 AI 乾貨

最新 AI 乾貨，我在看❤️