用Pytorch搭建LeNet-5模型且基于CIFAR-10資料集訓練

本文中的執行個體是Pytorch官方的一個demo，這裡我結合最近正在看的書《Python深度學習 基于Pytorch》P125 Pytorch實作CIFAR-10多分類這節，進行了梳理與個人的總結，以加深自己的了解。

目的：

①了解LeNet-5模型

②加深Pytorch的學習

③從頭至尾介紹搭建網絡并進行訓練預測的全過程

目錄

• • 簡單介紹
  • • CIFAR-10資料集
    • LeNet-5模型
    • 查找Pytorch中API的具體内容
  • 執行個體
  • • 1.建構神經網絡
    • 2.下載下傳資料集并檢視部分資料
    • 3.訓練模型
    • 4.在測試集中随機預測四張圖看看效果
    • 5.測試模型在測試集上的準确率及10個類别的準确率
    • 6.存儲訓練好的權重檔案進行預測
  • 整體源代碼

簡單介紹

CIFAR-10資料集

CIFAR-10資料集由10個類的60000個32x32彩色圖像組成，以下是十個類别，每個類有6000個圖像。有50000個訓練圖像和10000個測試圖像。

classes = (‘plane’, ‘car’, ‘bird’, ‘cat’,‘deer’, ‘dog’, ‘frog’, ‘horse’, ‘ship’, ‘truck’)

資料集分為五個訓練批次和一個測試批次，每個批次有10000個圖像。測試批次包含來自每個類的恰好1000個随機選擇的圖像。

訓練批次包含随機順序的圖像，但一些訓練批次可能包含來自一個種類的圖像比另一個類更多。

總的訓練批次包含來自每個類的正好5000張圖像。

因為是彩色圖像，3個通道，是以下面模型的輸入為3×32×32

LeNet-5模型

1.模型架構

LeNet-5模型結構為輸入層–卷積層–池化層–卷積層–池化層–全連接配接層–全連接配接層–輸出，為串聯模式，如上圖所示

2.模型特點

a.每個卷積層包含3個部分：卷積、池化和非線性激活函數。

b.使用卷積提前空間特征。

c.采用降采樣（Subsample）的平均池化層（Average Pooling）。

d.使用雙曲正切（Tanh）的激活函數。

e.最後用MLP作為分類器。

3.模型的網絡介紹

輸入：

灰階圖像，通道為1，尺寸為1×32×32

第一層：卷積層

LeNet-5模型接受的輸入層大小是1×32x32。卷積層的過濾器的尺寸是5x5，深度（卷積核個數）為6，不使用全0填充，步長為1。則這一層的輸出的尺寸為32-5+1=28，深度為6。本層的輸出矩陣大小為6×28×28。

第二層：池化層

這一層的輸入是第一層的輸出，是一個6×28x28=4704的節點矩陣。本層采用的過濾器為2x2的大小，長和寬的步長均為2，是以本層的輸出矩陣大小為6×14x14。

第三層：卷積層

本層的輸入矩陣大小為6×14x14，使用的過濾器大小為5x5，深度為16。本層不使用全0填充，步長為1。本層的輸出矩陣大小為16×10x10。

第四層：池化層

本層的輸入矩陣大小是16×10x10，采用的過濾器大小是2x2，步長為2，本層的輸出矩陣大小為16×5x5。

第五層：全連接配接層

本層的輸入矩陣大小為16×5x5。将此矩陣中的節點拉成一個向量，那麼這就和全連接配接層的輸入一樣了，本層的輸出節點個數為120。

第六層：全連接配接層

本層的輸入節點個數為120個，輸出節點個數為84個。

第七層：全連接配接層

本層的輸入節點為84個，輸出節點個數為10個。

下面的神經網絡模型是以LeNet-5為基準，其中就卷積核的個數不同，我看了書上和網上的模型，這裡的參數都不同，猜測卷積核的個數可以自己定，但要保證各層之間次元的統一。

查找Pytorch中API的具體内容

方法一：在Pycharm中将滑鼠移至函數名上，按住ctrl+滑鼠左鍵

方法二：進入Pytorch官網——Docs——Search Docs中輸入需要查找的API名稱

執行個體

Pytorch中Tensor的通道排序：[batch，channel，height，width]

1.建構神經網絡

定義一個LeNet()類，繼承于nn.Module的父類，這個類中包括兩個函數：初始化函數和前向傳播函數。初始化函數中：模型使用到的網絡層結構；前向傳播函數：實作整個前向傳播的過程。

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet,self).__init__()
        # 卷積層1：輸入圖像深度=3，輸出圖像深度=16，卷積核大小=5*5，卷積步長=1;16表示輸出次元，也表示卷積核個數
        self.conv1 = nn.Conv2d(in_channels=3,out_channels=16,kernel_size=5,stride=1)
        # 池化層1：采用最大池化，區域集大小=2*2.池化步長=2
        self.pool1 = nn.MaxPool2d(kernel_size=2,stride=2)
        # 卷積層2
        self.conv2 = nn.Conv2d(in_channels=16,out_channels=32,kernel_size=5,stride=1)
        # 池化層2
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        # 全連接配接層1：輸入大小=32*5*5，輸出大小=120
        self.fc1 = nn.Linear(32*5*5,120)
        # 全連接配接層2
        self.fc2 = nn.Linear(120,84)
        # 全連接配接層3
        self.fc3 = nn.Linear(84,10)

    def forward(self,x):
        x = F.relu(self.conv1(x))  # input(3, 32, 32) output(16, 28, 28)
        x = self.pool1(x)  # output(16, 14, 14)
        x = F.relu(self.conv2(x))  # output(32, 10, 10)
        x = self.pool2(x)  # output(32, 5, 5)
        x = x.view(-1, 32 * 5 * 5)  # output(32*5*5)
        x = F.relu(self.fc1(x))  # output(120)
        x = F.relu(self.fc2(x))  # output(84)
        x = self.fc3(x)  # output(10)
        return x
           

2.下載下傳資料集并檢視部分資料

import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim # 優化器

'''下載下傳資料集'''
# transforms.Compose()函數将兩個函數拼接起來。
# （ToTensor()：把一個PIL.Image轉換成Tensor，Normalize()：标準化，即減均值，除以标準差）
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 訓練集：下載下傳CIFAR10資料集，如果沒有事先下載下傳該資料集，則将download參數改為True
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,download=False, transform=transform)
# 用DataLoader得到生成器，其中shuffle：是否将資料打亂；
# num_workers表示使用多程序加載的程序數，0代表不使用多程序
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,shuffle=True, num_workers=0)

# 測試集資料下載下傳
testset = torchvision.datasets.CIFAR10(root='./data', train=False,download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,shuffle=False, num_workers=0)

classes = ('plane', 'car', 'bird', 'cat','deer', 'dog', 'frog', 'horse', 'ship', 'truck')


# 顯示圖像
def imshow(img):
    # 因為标準化normalize是：output = (input-0.5)/0.5
    # 則反标準化unnormalize是：input = output*0.5 + 0.5
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    # transpose()會更改多元數組的軸的順序
    # Pytorch中是[channel，height，width],這裡改為圖像的[height，width，channel]
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

# 随機擷取部分訓練資料
dataiter = iter(trainloader)
images, labels = dataiter.next()
# 顯示圖像
# torchvision.utils.make_grid()将多張圖檔拼接在一張圖中
imshow(torchvision.utils.make_grid(images))
# 列印标簽
# str.join(sequence)：用于将序列中的元素以指定的字元str連接配接成一個新的字元串。這裡的str是' '，空格
# %5s：表示輸出字元串至少5個字元，不夠5個的話，左側用空格補
print(' '.join('%5s' % classes[labels[j]] for j in range(4)))
           

輸出：

3.訓練模型

'''訓練模型'''
# 有GPU就用GPU跑，沒有就用CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

net = LeNet()
net=net.to(device)
# print("該網絡共有 {} 個參數".format(sum(x.numel() for x in net.parameters())))
# # 該網絡共有 121182 個參數

# 對于多分類問題，應該使用Softmax函數，這裡CIFAR10資料集屬于多分類，卻使用了交叉熵損失函數，
# 是因為進入CrossEntropyLoss()函數内部就會發現其中包含了Softmax函數
loss_function = nn.CrossEntropyLoss() # 使用交叉熵損失函數
# 優化器選擇Adam，學習率設為0.001
optimizer = optim.Adam(net.parameters(), lr=0.001)
# 列印檢視神經網絡的結構
# print(net)
# # LeNet(
# #   (conv1): Conv2d(3, 16, kernel_size=(5, 5), stride=(1, 1))
# #   (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
# #   (conv2): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1))
# #   (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
# #   (fc1): Linear(in_features=800, out_features=120, bias=True)
# #   (fc2): Linear(in_features=120, out_features=84, bias=True)
# #   (fc3): Linear(in_features=84, out_features=10, bias=True)
# # )

for epoch in range(10): # 整個疊代10輪
    running_loss = 0.0 # 初始化損失函數值loss=0
    for i, data in enumerate(trainloader, start=0):
        # 擷取訓練資料
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device) # 将資料及标簽傳入GPU/CPU

        # 權重參數梯度清零
        optimizer.zero_grad()

        # 正向及反向傳播
        outputs = net(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()

        # 顯示損失值
        running_loss += loss.item()
        if i % 2000 == 1999:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' %(epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

print('Finished Training')
           

輸出：

[1, 2000] loss: 1.884

[1, 4000] loss: 1.602

[1, 6000] loss: 1.506

[1, 8000] loss: 1.445

[1, 10000] loss: 1.405

[1, 12000] loss: 1.377

[2, 2000] loss: 1.274

[2, 4000] loss: 1.248

[2, 6000] loss: 1.239

[2, 8000] loss: 1.226

[2, 10000] loss: 1.213

[2, 12000] loss: 1.185

[3, 2000] loss: 1.106

[3, 4000] loss: 1.124

[3, 6000] loss: 1.137

[3, 8000] loss: 1.109

[3, 10000] loss: 1.096

[3, 12000] loss: 1.089

[4, 2000] loss: 1.013

[4, 4000] loss: 1.015

[4, 6000] loss: 1.028

[4, 8000] loss: 1.037

[4, 10000] loss: 1.056

[4, 12000] loss: 1.009

[5, 2000] loss: 0.935

[5, 4000] loss: 0.964

[5, 6000] loss: 0.980

[5, 8000] loss: 0.962

[5, 10000] loss: 0.975

[5, 12000] loss: 0.993

[6, 2000] loss: 0.901

[6, 4000] loss: 0.908

[6, 6000] loss: 0.920

[6, 8000] loss: 0.924

[6, 10000] loss: 0.905

[6, 12000] loss: 0.944

[7, 2000] loss: 0.855

[7, 4000] loss: 0.864

[7, 6000] loss: 0.904

[7, 8000] loss: 0.896

[7, 10000] loss: 0.898

[7, 12000] loss: 0.890

[8, 2000] loss: 0.819

[8, 4000] loss: 0.828

[8, 6000] loss: 0.844

[8, 8000] loss: 0.880

[8, 10000] loss: 0.854

[8, 12000] loss: 0.863

[9, 2000] loss: 0.802

[9, 4000] loss: 0.795

[9, 6000] loss: 0.833

[9, 8000] loss: 0.848

[9, 10000] loss: 0.824

[9, 12000] loss: 0.834

[10, 2000] loss: 0.775

[10, 4000] loss: 0.758

[10, 6000] loss: 0.781

[10, 8000] loss: 0.806

[10, 10000] loss: 0.806

[10, 12000] loss: 0.829

Finished Training

為什麼每計算一個batch，就需要調用一次optimizer.zero_grad() 将權重參數梯度清零？

是為了進行梯度累加。一定條件下，batchsize越大訓練效果越好，梯度累加則實作了batchsize的變相擴大，如果accumulation_steps為8，則batchsize ‘變相’ 擴大了8倍。

通過這個特性能夠變相實作一個很大batch數值的訓練。

4.在測試集中随機預測四張圖看看效果

'''預測四張圖看看'''
dataiter = iter(testloader)
images, labels = dataiter.next()

imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))

images, labels = images.to(device), labels.to(device)
outputs = net(images)
_, predicted = torch.max(outputs, 1)

print('Predicted: ', ' '.join('%5s' % classes[predicted[j]]for j in range(4)))
           

輸出：

可以發現四張圖隻錯了一張

5.測試模型在測試集上的準确率及10個類别的準确率

'''測試模型'''
correct = 0
total = 0
# with是一個上下文管理器
# with torch.no_grad()表示其包括的内容不需要計算梯度，也不會進行反向傳播，節省記憶體
with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = net(images)
        # torch.max(outputs.data, 1)傳回outputs每一行中最大值的那個元素，且傳回其索引
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))

# 列印10個分類的準确率
class_correct = list(0. for i in range(10)) #class_correct=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
class_total = list(0. for i in range(10)) #class_total=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = net(images) # outputs的次元是：4*10
        # torch.max(outputs.data, 1)傳回outputs每一行中最大值的那個元素，且傳回其索引
        # 此時predicted的次元是：4*1
        _, predicted = torch.max(outputs, 1)
        # 此時c的次元：4将預測值與實際标簽進行比較，且進行降維
        c = (predicted == labels).squeeze()
        for i in range(4):
            label = labels[i]
            class_correct[label] += c[i].item()
            class_total[label] += 1
for i in range(10):
    print('Accuracy of %5s : %2d %%' % (classes[i], 100 * class_correct[i] / class_total[i]))
           

輸出：

6.存儲訓練好的權重檔案進行預測

torch.nn.Module子產品中的state_dict變量存放訓練過程中需要學習的權重和偏執系數，state_dict作為python的字典對象将每一層的參數映射成tensor張量

參考部落格：state_dict詳解

Python中.pth檔案的作用

save_path = './Lenet.pth'
torch.save(net.state_dict(), save_path)
           

上面是已經訓練好了的權重檔案，下面随便用一張自己找的圖檔預測下

重建立一個py檔案，讀取訓練好了的權重并對目标圖像進行預測：

import torch
import torchvision.transforms as transforms
from PIL import Image
from model import LeNet

transform = transforms.Compose([transforms.Resize((32, 32)),transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

classes = ('plane', 'car', 'bird', 'cat','deer', 'dog', 'frog', 'horse', 'ship', 'truck')

net = LeNet()
net.load_state_dict(torch.load('Lenet.pth'))

im = Image.open('1.jpg')
im = transform(im)  # [C, H, W]
# 輸入pytorch網絡中要求的格式是[batch，channel，height，width]，是以這裡增加一個次元
im = torch.unsqueeze(im, dim=0)  # [N, C, H, W]

with torch.no_grad():
    outputs = net(im)
    predict = torch.max(outputs, dim=1)[1].data.numpy() # 索引即classed中的類别
print(classes[int(predict)])

# 直接列印張量的預測結果
with torch.no_grad():
    outputs = net(im)
    predict = torch.softmax(outputs,dim=1) # [batch，channel，height，width],這裡因為對batch不需要處理
print(predict)
           

輸出：

整體源代碼

main.py

import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim # 優化器


'''1.建構神經網絡'''
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet,self).__init__()
        # 卷積層1：輸入圖像深度=3，輸出圖像深度=16，卷積核大小=5*5，卷積步長=1;16表示輸出次元，也表示卷積核個數
        self.conv1 = nn.Conv2d(in_channels=3,out_channels=16,kernel_size=5,stride=1)
        # 池化層1：采用最大池化，區域集大小=2*2.池化步長=2
        self.pool1 = nn.MaxPool2d(kernel_size=2,stride=2)
        # 卷積層2
        self.conv2 = nn.Conv2d(in_channels=16,out_channels=32,kernel_size=5,stride=1)
        # 池化層2
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        # 全連接配接層1：輸入大小=32*5*5，輸出大小=120
        self.fc1 = nn.Linear(32*5*5,120)
        # 全連接配接層2
        self.fc2 = nn.Linear(120,84)
        # 全連接配接層3
        self.fc3 = nn.Linear(84,10)

    def forward(self,x):
        x = F.relu(self.conv1(x))  # input(3, 32, 32) output(16, 28, 28)
        x = self.pool1(x)  # output(16, 14, 14)
        x = F.relu(self.conv2(x))  # output(32, 10, 10)
        x = self.pool2(x)  # output(32, 5, 5)
        x = x.view(-1, 32 * 5 * 5)  # output(32*5*5)
        x = F.relu(self.fc1(x))  # output(120)
        x = F.relu(self.fc2(x))  # output(84)
        x = self.fc3(x)  # output(10)
        return x


'''2.下載下傳資料集并檢視部分資料'''
# transforms.Compose()函數将兩個函數拼接起來。
# （ToTensor()：把一個PIL.Image轉換成Tensor，Normalize()：标準化，即減均值，除以标準差）
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 訓練集：下載下傳CIFAR10資料集，如果沒有事先下載下傳該資料集，則将download參數改為True
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,download=False, transform=transform)
# 用DataLoader得到生成器，其中shuffle：是否将資料打亂；
# num_workers表示使用多程序加載的程序數，0代表不使用多程序
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,shuffle=True, num_workers=0)

# 測試集資料下載下傳
testset = torchvision.datasets.CIFAR10(root='./data', train=False,download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,shuffle=False, num_workers=0)

classes = ('plane', 'car', 'bird', 'cat','deer', 'dog', 'frog', 'horse', 'ship', 'truck')


# 顯示圖像
def imshow(img):
    # 因為标準化normalize是：output = (input-0.5)/0.5
    # 則反标準化unnormalize是：input = output*0.5 + 0.5
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    # transpose()會更改多元數組的軸的順序
    # Pytorch中是[channel，height，width],這裡改為圖像的[height，width，channel]
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

# 随機擷取部分訓練資料
dataiter = iter(trainloader)
images, labels = dataiter.next()
# 顯示圖像
# torchvision.utils.make_grid()将多張圖檔拼接在一張圖中
imshow(torchvision.utils.make_grid(images))
# 列印标簽
# str.join(sequence)：用于将序列中的元素以指定的字元str連接配接成一個新的字元串。這裡的str是' '，空格
# %5s：表示輸出字元串至少5個字元，不夠5個的話，左側用空格補
print(' '.join('%5s' % classes[labels[j]] for j in range(4)))


'''3.訓練模型'''
# 有GPU就用GPU跑，沒有就用CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

net = LeNet()
net=net.to(device)
# print("該網絡共有 {} 個參數".format(sum(x.numel() for x in net.parameters())))
# # 該網絡共有 121182 個參數

# 對于多分類問題，應該使用Softmax函數，這裡CIFAR10資料集屬于多分類，卻使用了交叉熵損失函數，
# 是因為進入CrossEntropyLoss()函數内部就會發現其中包含了Softmax函數
loss_function = nn.CrossEntropyLoss() # 使用交叉熵損失函數
# 優化器選擇Adam，學習率設為0.001
optimizer = optim.Adam(net.parameters(), lr=0.001)
# 列印檢視神經網絡的結構
# print(net)
# # LeNet(
# #   (conv1): Conv2d(3, 16, kernel_size=(5, 5), stride=(1, 1))
# #   (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
# #   (conv2): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1))
# #   (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
# #   (fc1): Linear(in_features=800, out_features=120, bias=True)
# #   (fc2): Linear(in_features=120, out_features=84, bias=True)
# #   (fc3): Linear(in_features=84, out_features=10, bias=True)
# # )

for epoch in range(10): # 整個疊代10輪
    running_loss = 0.0 # 初始化損失函數值loss=0
    for i, data in enumerate(trainloader, start=0):
        # 擷取訓練資料
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device) # 将資料及标簽傳入GPU/CPU

        # 權重參數梯度清零
        optimizer.zero_grad()

        # 正向及反向傳播
        outputs = net(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()

        # 顯示損失值
        running_loss += loss.item()
        if i % 2000 == 1999:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' %(epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

print('Finished Training')


'''4.在測試集中随機預測四張圖看看效果'''
dataiter = iter(testloader)
images, labels = dataiter.next()

imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))

images, labels = images.to(device), labels.to(device)
outputs = net(images)
_, predicted = torch.max(outputs, 1)

print('Predicted: ', ' '.join('%5s' % classes[predicted[j]]for j in range(4)))


'''5.測試模型在測試集上的準确率及10個類别的準确率'''
correct = 0
total = 0
# with是一個上下文管理器
# with torch.no_grad()表示其包括的内容不需要計算梯度，也不會進行反向傳播，節省記憶體
with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = net(images)
        # torch.max(outputs.data, 1)傳回outputs每一行中最大值的那個元素，且傳回其索引
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))

# 列印10個分類的準确率
class_correct = list(0. for i in range(10)) #class_correct=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
class_total = list(0. for i in range(10)) #class_total=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = net(images) # outputs的次元是：4*10
        # torch.max(outputs.data, 1)傳回outputs每一行中最大值的那個元素，且傳回其索引
        # 此時predicted的次元是：4*1
        _, predicted = torch.max(outputs, 1)
        # 此時c的次元：4将預測值與實際标簽進行比較，且進行降維
        c = (predicted == labels).squeeze()
        for i in range(4):
            label = labels[i]
            class_correct[label] += c[i].item()
            class_total[label] += 1
for i in range(10):
    print('Accuracy of %5s : %2d %%' % (classes[i], 100 * class_correct[i] / class_total[i]))


'''6.存儲訓練好的權重檔案'''
save_path = './Lenet.pth'
torch.save(net.state_dict(), save_path)
           

predict.py

import torch
import torchvision.transforms as transforms
from PIL import Image
from model import LeNet

transform = transforms.Compose([transforms.Resize((32, 32)),transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

classes = ('plane', 'car', 'bird', 'cat','deer', 'dog', 'frog', 'horse', 'ship', 'truck')

net = LeNet()
net.load_state_dict(torch.load('Lenet.pth'))

im = Image.open('1.jpg')
im = transform(im)  # [C, H, W]
# 輸入pytorch網絡中要求的格式是[batch，channel，height，width]，是以這裡增加一個次元
im = torch.unsqueeze(im, dim=0)  # [N, C, H, W]

with torch.no_grad():
    outputs = net(im)
    predict = torch.max(outputs, dim=1)[1].data.numpy() # 索引即classed中的類别
print(classes[int(predict)])

# 直接列印張量的預測結果
with torch.no_grad():
    outputs = net(im)
    predict = torch.softmax(outputs,dim=1) # [batch，channel，height，width],這裡因為對batch不需要處理
print(predict)