點擊上方藍字關注我們

微信公眾號：OpenCV學堂

關注獲取更多計算機視覺與深度學習知識

Hello大家好，這篇文章給大家詳細介紹一下pytorch中最重要的組件torchvision，它包含了常見的數據集、模型架構與預訓練模型權重文件、常見圖像變換、計算機視覺任務訓練。可以是說是pytorch中非常有用的模型遷移學習神器。本文將會介紹如何使用torchvison的預訓練模型ResNet50實現圖像分類。

模型

Torchvision.models包里面包含了常見的各種基礎模型架構，主要包括：

AlexNet
VGG
ResNet
SqueezeNet
DenseNet
Inception v3
GoogLeNet
ShuffleNet v2
MobileNet v2
ResNeXt
Wide ResNet
MNASNet

這里我選擇了ResNet50，基于ImageNet訓練的基礎網絡來實現圖像分類， 網絡模型下載與加載如下：

model = torchvision.models.resnet50(pretrained=True).eval().cuda()tf = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] )])

使用模型實現圖像分類

這里首先需要加載ImageNet的分類標簽，目的是最后顯示分類的文本標簽時候使用。然后對輸入圖像完成預處理，使用ResNet50模型實現分類預測，對預測結果解析之后，顯示標簽文本，完整的代碼演示如下：

1with?open('imagenet_classes.txt')?as?f:
2????labels?=?[line.strip()?for?line?in?f.readlines()]
3
4src?=?cv.imread("D:/images/space_shuttle.jpg")?#?aeroplane.jpg
5image?=?cv.resize(src,?(224,?224))
6image?=?np.float32(image)?/?255.0
7image[:,:,]?-=?(np.float32(0.485),?np.float32(0.456),?np.float32(0.406))
8image[:,:,]?/=?(np.float32(0.229),?np.float32(0.224),?np.float32(0.225))
9image?=?image.transpose((2,?0,?1))
10input_x?=?torch.from_numpy(image).unsqueeze(0)
11print(input_x.size())
12pred?=?model(input_x.cuda())
13pred_index?=?torch.argmax(pred,?1).cpu().detach().numpy()
14print(pred_index)
15print("current?predict?class?name?:?%s"%labels[pred_index[0]])
16cv.putText(src,?labels[pred_index[0]],?(50,?50),?cv.FONT_HERSHEY_SIMPLEX,?1.0,?(0,?0,?255),?2)
17cv.imshow("input",?src)
18cv.waitKey(0)
19cv.destroyAllWindows()

運行結果如下：

轉ONNX支持

在torchvision中的模型基本上都可以轉換為ONNX格式，而且被OpenCV DNN模塊所支持，所以，很方便的可以對torchvision自帶的模型轉為ONNX，實現OpenCV DNN的調用，首先轉為ONNX模型，直接使用torch.onnx.export即可轉換(還不知道怎么轉，快點看前面的例子)。轉換之后使用OpenCV DNN調用的代碼如下：

1with?open('imagenet_classes.txt')?as?f:
2????labels?=?[line.strip()?for?line?in?f.readlines()]
3net?=?cv.dnn.readNetFromONNX("resnet.onnx")
4src?=?cv.imread("D:/images/messi.jpg")??#?aeroplane.jpg
5image?=?cv.resize(src,?(224,?224))
6image?=?np.float32(image)?/?255.0
7image[:,?:,?]?-=?(np.float32(0.485),?np.float32(0.456),?np.float32(0.406))
8image[:,?:,?]?/=?(np.float32(0.229),?np.float32(0.224),?np.float32(0.225))
9blob?=?cv.dnn.blobFromImage(image,?1.0,?(224,?224),?(0,?0,?0),?False)
10net.setInput(blob)
11probs?=?net.forward()
12index?=?np.argmax(probs)
13cv.putText(src,?labels[index],?(50,?50),?cv.FONT_HERSHEY_SIMPLEX,?1.0,?(0,?0,?255),?2)
14cv.imshow("input",?src)
15cv.waitKey(0)
16cv.destroyAllWindows()

?運行結果見上圖，這里就不再貼了。

?推薦閱讀?

輕松學Pytorch–環境搭建與基本語法

Pytorch輕松學-構建淺層神經網絡

輕松學pytorch-構建卷積神經網絡

輕松學Pytorch –構建循環神經網絡

輕松學Pytorch-使用卷積神經網絡實現圖像分類

輕松學Pytorch-自定義數據集制作與使用

輕松學Pytorch-Pytorch可視化

輕松學Pytorch–Visdom可視化

輕松學Pytorch – 全局池化層詳解

輕松學Pytorch – 人臉五點landmark提取網絡訓練與使用

輕松學Pytorch – 年齡與性別預測

輕松學Pytorch –車輛類型與顏色識別

輕松學Pytorch-全卷積神經網絡實現表情識別

使用OpenVINO加速Pytorch表情識別模型

輕松學pytorch – 使用多標簽損失函數訓練卷積網絡

志不強者智不達

言不信者行不果

總結

以上是生活随笔為你收集整理的pytorch argmax_轻松学Pytorch使用ResNet50实现图像分类的全部內容，希望文章能夠幫你解決所遇到的問題。

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