圖文詳解如何利用PyTorch實現(xiàn)圖像識別

更新時間：2023年04月04日 14:29:42 作者：阿潤菜菜

這篇文章主要給大家介紹了關于如何利用PyTorch實現(xiàn)圖像識別的相關資料,文中通過圖文以及實例代碼介紹的非常詳細,對大家學習或者使用PyTorch具有一定的參考學習價值,需要的朋友可以參考下

使用torchvision庫的datasets類加載常用的數據集或自定義數據集

圖像識別是計算機視覺中的一個基礎任務，它的目標是讓計算機能夠識別圖像中的物體、場景或者概念，并將它們分配到預定義的類別中。例如，給定一張貓的圖片，圖像識別系統(tǒng)應該能夠輸出“貓”這個類別。

為了訓練和評估圖像識別系統(tǒng)，我們需要有大量的帶有標注的圖像數據集。常用的圖像分類數據集有：

ImageNet：一個包含超過1400萬張圖片和2萬多個類別的大型數據庫，是目前最流行和最具挑戰(zhàn)性的圖像分類基準之一。
CIFAR-10/CIFAR-100：一個包含6萬張32×32大小的彩色圖片和10或100個類別的小型數據庫，適合入門級和快速實驗。
MNIST：一個包含7萬張28×28大小的灰度手寫數字圖片和10個類別的經典數據庫，是深度學習中最常用的測試集之一。
Fashion-MNIST：一個包含7萬張28×28大小的灰度服裝圖片和10個類別的數據庫，是MNIST數據庫在時尚領域上更加復雜和現(xiàn)代化版本。

使用torchvision庫可以方便地加載這些常用數據集或者自定義數據集。torchvision.datasets提供了一些加載數據集或者下載數據集到本地緩存文件夾（默認為./data）并返回Dataset對象（torch.utils.data.Dataset）的函數。Dataset對象可以存儲樣本及其對應標簽，并提供索引方式（dataset[i]）來獲取第i個樣本。例如，要加載CIFAR-10訓練集并進行隨機打亂，可以使用以下代碼：

import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose([transforms.ToTensor()]) # 定義轉換函數，將PIL.Image轉換為torch.Tensor
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) # 加載CIFAR-10訓練集
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True) # 定義DataLoader對象，用于批量加載數據

使用torchvision庫進行數據增強和變換，自定義自己的圖像分類數據集并使用torchvision庫加載它們

數據增強和變換：為了提高模型的泛化能力和數據利用率，我們通常會對圖像數據進行一些隨機的變換，例如裁剪、旋轉、翻轉、縮放、亮度調整等。這些變換可以在一定程度上模擬真實場景中的圖像變化，增加模型對不同視角和光照條件下的物體識別能力。torchvision.transforms提供了一些常用的圖像變換函數，可以組合成一個transform對象，并傳入datasets類中作為參數。例如，要對CIFAR-10訓練集進行隨機水平翻轉和隨機裁剪，并將圖像歸一化到[-1, 1]范圍內，可以使用以下代碼：

import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose([
    transforms.RandomHorizontalFlip(), # 隨機水平翻轉
    transforms.RandomCrop(32, padding=4), # 隨機裁剪到32×32大小，并在邊緣填充4個像素
    transforms.ToTensor(), # 將PIL.Image轉換為torch.Tensor
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # 將RGB三個通道的值歸一化到[-1, 1]范圍內
])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) # 加載CIFAR-10訓練集，并應用上述變換
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True) # 定義DataLoader對象，用于批量加載數據

自定義圖像分類數據集：如果我們有自己的圖像分類數據集，我們可以通過繼承torch.utils.data.Dataset類來自定義一個Dataset對象，并實現(xiàn)__len__和__getitem__兩個方法。__len__方法返回數據集中樣本的數量，__getitem__方法根據給定的索引返回一個樣本及其標簽。例如，假設我們有一個文件夾結構如下：

my_dataset/
├── class_0/
│ ├── image_000.jpg
│ ├── image_001.jpg
│ └── ...
├── class_1/
│ ├── image_000.jpg
│ ├── image_001.jpg
│ └── ...
└── ...

其中每個子文件夾代表一個類別，每個子文件夾中包含該類別對應的圖像文件。我們可以使用以下代碼來自定義一個Dataset對象，并加載這個數據集：

import torch.utils.data as data
from PIL import Image
import os

class MyDataset(data.Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir # 根目錄路徑
        self.transform = transform # 變換函數
        
        self.classes = sorted(os.listdir(root_dir)) # 類別列表（按字母順序排序）
        self.class_to_idx = {c: i for i,c in enumerate(self.classes)} # 類別名到索引的映射
        
        self.images = [] # 圖片路徑列表（相對于根目錄）
        self.labels = [] # 標簽列表（整數）
        
        for c in self.classes:
            c_dir = os.path.join(root_dir, c) # 類別子目錄路徑
            for img_name in sorted(os.listdir(c_dir)): # 遍歷每個圖片文件名（按字母順序排序）
                img_path = os.path.join(c,img_name) # 圖片相對路徑（相對于根目錄）
                label = self.class_to_idx[c] # 圖

使用torchvision庫的models類加載預訓練模型或自定義模型

加載預訓練模型或自定義模型：torchvision.models提供了一些常用的圖像分類模型，例如AlexNet、VGG、ResNet等，并且可以選擇是否加載在ImageNet數據集上預訓練好的權重。這些模型可以直接用于圖像分類任務，也可以作為特征提取器或者微調（fine-tune）的基礎。例如，要加載一個預訓練好的ResNet-18模型，并凍結除最后一層外的所有參數，可以使用以下代碼：

import torchvision.models as models

model = models.resnet18(pretrained=True) # 加載預訓練好的ResNet-18模型
for param in model.parameters(): # 遍歷所有參數
    param.requires_grad = False # 將參數的梯度設置為False，表示不需要更新
num_features = model.fc.in_features # 獲取全連接層（fc）的輸入特征數
model.fc = torch.nn.Linear(num_features, 10) # 替換全連接層為一個新的線性層，輸出特征數為10（假設有10個類別）

如果我們想要自定義自己的圖像分類模型，我們可以通過繼承torch.nn.Module類來實現(xiàn)一個Module對象，并實現(xiàn)__init__和forward兩個方法。__init__方法用于定義模型中需要的各種層和參數，forward方法用于定義前向傳播過程。例如，要自定義一個簡單的卷積神經網絡（CNN）模型，可以使用以下代碼：

import torch.nn as nn

class MyCNN(nn.Module):
    def __init__(self):
        super(MyCNN, self).__init__() # 調用父類構造函數
        self.conv1 = nn.Conv2d(3, 6, 5) # 定義第一個卷積層，輸入通道數為3（RGB），輸出通道數為6，卷積核大小為5×5
        self.pool = nn.MaxPool2d(2, 2) # 定義最大池化層，池化核大小為2×2，步長為2
        self.conv2 = nn.Conv2d(6, 16, 5) # 定義第二個卷積層，輸入通道數為6，輸出通道數為16，卷積核大小為5×5
        self.fc1 = nn.Linear(16 * 5 * 5, 120) # 定義第一個全連接層，輸入特征數為16×5×5（根據卷積和池化后的圖像大小計算得到），輸出特征數為120
        self.fc2 = nn.Linear(120, 84) # 定義第二個全連接層，輸入特征數為120，輸出特征數為84
        self.fc3 = nn.Linear(84, 10) # 定義第三個全連接層，輸入特征數為84，

forward方法

forward方法用于定義前向傳播過程，即如何根據輸入的圖像張量（Tensor）計算出輸出的類別概率分布。我們可以使用定義好的各種層和參數，并結合一些激活函數（如ReLU）和歸一化函數（如softmax）來實現(xiàn)forward方法。例如，要實現(xiàn)上面自定義的CNN模型的forward方法，可以使用以下代碼：

import torch.nn.functional as F

class MyCNN(nn.Module):
    def __init__(self):
        # 省略__init__方法的內容
        ...

    def forward(self, x): # 定義前向傳播過程，x是輸入的圖像張量
        x = self.pool(F.relu(self.conv1(x))) # 將x通過第一個卷積層和ReLU激活函數，然后通過最大池化層
        x = self.pool(F.relu(self.conv2(x))) # 將x通過第二個卷積層和ReLU激活函數，然后通過最大池化層
        x = x.view(-1, 16 * 5 * 5) # 將x展平為一維向量，-1表示自動推斷批量大小
        x = F.relu(self.fc1(x)) # 將x通過第一個全連接層和ReLU激活函數
        x = F.relu(self.fc2(x)) # 將x通過第二個全連接層和ReLU激活函數
        x = self.fc3(x) # 將x通過第三個全連接層
        x = F.softmax(x, dim=1) # 將x通過softmax函數，沿著第一個維度（類別維度）進行歸一化，得到類別概率分布
        return x # 返回輸出的類別概率分布

進行模型訓練和測試，使用matplotlib.pyplot庫可視化結果

模型訓練和測試是機器學習中的重要步驟，它們可以幫助我們評估模型的性能和泛化能力。matplotlib.pyplot是一個Python庫，它可以用來繪制各種類型的圖形，包括曲線圖、散點圖、直方圖等。使用matplotlib.pyplot庫可視化結果的一般步驟如下：

導入matplotlib.pyplot模塊，并設置一些參數，如字體、分辨率等。
創(chuàng)建一個或多個圖形對象（figure），并指定大小、標題等屬性。
在每個圖形對象中創(chuàng)建一個或多個子圖（subplot），并指定位置、坐標軸等屬性。
在每個子圖中繪制數據，使用不同的函數和參數，如plot、scatter、bar等。
添加一些修飾元素，如圖例（legend）、標簽（label）、標題（title）等。
保存或顯示圖形。

例如：使用matplotlib.pyplot庫繪制了一個線性回歸模型的訓練誤差和測試誤差曲線：

# 導入模塊
import matplotlib.pyplot as plt
import numpy as np

# 設置字體和分辨率
plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False
%config InlineBackend.figure_format = "retina"

# 生成數據
x = np.linspace(0, 10, 100)
y = 3 * x + 5 + np.random.randn(100) * 2 # 真實值
w = np.random.randn() # 隨機初始化權重
b = np.random.randn() # 隨機初始化偏置

# 定義損失函數
def loss(y_true, y_pred):
    return ((y_true - y_pred) ** 2).mean()

# 定義梯度下降函數
def gradient_descent(x, y_true, w, b, lr):
    y_pred = w * x + b # 預測值
    dw = -2 * (x * (y_true - y_pred)).mean() # 權重梯度
    db = -2 * (y_true - y_pred).mean() # 偏置梯度
    w = w - lr * dw # 更新權重
    b = b - lr * db # 更新偏置
    return w, b

# 訓練模型，并記錄每輪的訓練誤差和測試誤差
epochs = 20 # 訓練輪數
lr = 0.01 # 學習率
train_loss_list = [] # 訓練誤差列表
test_loss_list = [] # 測試誤差列表

for epoch in range(epochs):
    # 劃分訓練集和測試集（8:2）
    train_index = np.random.choice(100, size=80, replace=False)
    test_index = np.setdiff1d(np.arange(100), train_index)
    x_train, y_train = x[train_index], y[train_index]
    x_test, y_test = x[test_index], y[test_index]

    # 梯度下降更新參數，并計算訓練誤差和測試誤差
    w, b = gradient_descent(x_train, y_train, w, b, lr)
    train_loss = loss(y_train, w * x_train + b)
    test_loss = loss(y_test, w * x_test + b)

    # 打印結果，并將誤差添加到列表中
    print(f"Epoch {epoch+1}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}")
    train_loss_list.append(train_loss)
    test_loss_list.append(test_loss)

# 創(chuàng)建一個圖形對象，并設置大小為8*6英寸    
plt.figure(figsize=(8,6))

# 在圖形對象中創(chuàng)建一個子圖，并設置位置為1行1列的第1個
plt.subplot(1, 1, 1)

# 在子圖中繪制訓練誤差和測試誤差曲線，使用不同的顏色和標簽
plt.plot(np.arange(epochs), train_loss_list, "r", label="Train Loss")
plt.plot(np.arange(epochs), test_loss_list, "b", label="Test Loss")

# 添加圖例、坐標軸標簽和標題
plt.legend()
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Linear Regression Loss Curve")

# 保存或顯示圖形
#plt.savefig("loss_curve.png")
plt.show()

運行后，可以看到如下的圖形：