? ? ? TorchVision中給出了AlexNet的pretrained模型，模型存放位置為https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth ，可通過models.alexnet函數下載，此函數實現在torchvision/models/alexnet.py中，下載后在Ubuntu上存放在~/.cache/torch/hub/checkpoints目錄下，在Windows上存放在C:\Users\spring\.cache\torch\hub\checkpoints目錄下，其中spring為用戶名。

? ? ? AlexNet的介紹參考：https://blog.csdn.net/fengbingchun/article/details/112709281

? ? ? 在推理(inference)過程中，模型的輸入是一個tensor，shape需要是[1,c,h,w]，原始圖像進行預處理操作包括：

? ? ? (1).resize到短邊為256，長邊等比縮放。

? ? ? (2).在中心裁剪圖像大小到224*224。

? ? ? (3).將數據從numpy.ndarray轉換到tensor；原數據shape為[h,w,c]，轉換后tensor shape為[c,h,w]；原數據值范圍為[0,255]，轉換后值范圍為[0.0,1.0]。

? ? ? (4).使用均值和標準差對tensor圖像進行歸一化。

? ? ? (5).將tensor的shape從[c,h,w]轉換到[1,c,h,w]。

? ? ? 模型是通過ImageNet數據集訓練獲得的，它的圖像分類數是1000，ImageNet數據集的介紹參考：https://blog.csdn.net/fengbingchun/article/details/88606621

? ? ? 以下為測試代碼：

import torch
from torchvision import models
from torchvision import transforms
import cv2
from PIL import Image
import math
import numpy as np#print(dir(models))images_path = "../../data/image/"
images_name = ["5.jpg", "6.jpg", "7.jpg", "8.jpg", "9.jpg", "10.jpg"]
images_data = [] # opencv
tensor_data = [] # pytorch tensordef images_stitch(images, cols=3, name="result.jpg"): # 圖像簡單拼接'''images: list, opencv image data; cols: number of images per line; name: save image result name'''width_total = 660width, height = width_total // cols, width_total // colsnumber = len(images)height_total = height * math.ceil(number / cols)mat1 = np.zeros((height_total, width_total, 3), dtype="uint8") # in Python images are represented as NumPy arraysfor idx in range(number):height_, width_, _ = images[idx].shapeif height_ != width_:if height_ > width_:width_ = math.floor(width_ / height_ * width)height_ = heightelse:height_ = math.floor(height_ / width_ * height)width_ = widthelse:height_, width_ = height, widthmat2 = cv2.resize(images[idx], (width_, height_))offset_y, offset_x = (height - height_) // 2, (width - width_) // 2start_y, start_x = idx // cols * height, idx % cols * widthmat1[start_y + offset_y:start_y + height_+offset_y, start_x + offset_x:start_x + width_+offset_x, :] = mat2cv2.imwrite(images_path+name, mat1)for name in images_name:img = cv2.imread(images_path + name)print(f"name: {images_path+name}, opencv image shape: {img.shape}") # (h,w,c)images_data.append(img)img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)img_pil = Image.fromarray(img)transform = 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])])tensor = transform(img_pil)print(f"tensor shape: {tensor.shape}, max: {torch.max(tensor)}, min: {torch.min(tensor)}") # (c,h,w)tensor = torch.unsqueeze(tensor, 0) # 返回一個新的tensor,對輸入的既定位置插入維度1print(f"tensor shape: {tensor.shape}, max: {torch.max(tensor)}, min: {torch.min(tensor)}") # (1,c,h,w)tensor_data.append(tensor)images_stitch(images_data)model = models.alexnet(pretrained=True) # AlexNet網絡
#print(model) # 可查看模型結構,與torchvision/models/alexnet.py中一致
model.eval() # AlexNet is required to be put in evaluation mode in order to do prediction/evaluationwith open("imagenet_classes.txt") as f:classes = [line.strip() for line in f.readlines()] # the line number specified the class numberfor x in range(len(tensor_data)):prediction = model(tensor_data[x])#print(prediction.shape) # [1,1000]_, index = torch.max(prediction, 1)percentage = torch.nn.functional.softmax(prediction, dim=1)[0] * 100print(f"result: {classes[index[0]]}, {percentage[index[0]].item()}")print("test finish")

? ? ? 執行結果如下：以下原始測試圖像來自網絡，每張圖像僅輸出可信度值最高的一個類別。從上往下，從左往右，每張圖像的分類結果依次是：goldfish(金魚)、hen(母雞)、ostrich(鴕鳥)、African crocodile(非洲鱷魚)、goose(鵝)、hartebeest(羚羊)。

?

? ? ? GitHub：https://github.com/fengbingchun/PyTorch_Test?

總結

以上是生活随笔為你收集整理的TorchVision中通过AlexNet网络进行图像分类的全部內容，希望文章能夠幫你解決所遇到的問題。

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