栗子 發自 凹非寺
量子位 出品 | 公眾號 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腳本模式了。

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

1 import torch2 import torchvision34 # create 10 random boxes5 boxes = torch.rand(10, 4) * 1006 # they need to be in [x0, y0, x1, y1] format7 boxes[:, 2:] += boxes[:, :2]8 # create a random image9 image = torch.rand(1, 3, 200, 200) 10 # extract regions in `image` defined in `boxes`, rescaling 11 # them to have a size of 3x3 12 pooled_regions = torchvision.ops.roi_align(image, [boxes], output_size=(3, 3)) 13 # check the size 14 print(pooled_regions.shape) 15 # torch.Size([10, 3, 3, 3]) 16 17 # or compute the intersection over union between 18 # all pairs of boxes 19 print(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。)

速度如下，都是毫秒級：

而且，只要寫幾行代碼，就可以直接加載、運行那些預訓練的檢測模型與分割模型了：

1 import torchvision23 model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)4 # set it to evaluation mode, as the model behaves differently5 # during training and during evaluation6 model.eval()78 image = PIL.Image.open('/path/to/an/image.jpg')9 image_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 14 output = 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：

torchvision?pypi.org

Colab教程：

https://colab.research.google.com/github/pytorch/vision/blob/temp-tutorial/tutorials/torchvision_finetuning_instance_segmentation.ipynb?colab.research.google.com

詳細的發布說明：

pytorch/vision?github.com

— 完 —

量子位 · QbitAI

?'?' ? 追蹤AI技術和產品新動態

戳右上角「+關注」獲取最新資訊↗↗

如果喜歡，請分享or點贊吧~比心?

總結

以上是生活随笔為你收集整理的ajax 更新模型数据_PyTorch视觉工具包torchvision重大更新！支持各种检测模型、分割模型，还有许多数据集...的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。