圖像分類---利用pytorch搭建LeNet-5網絡模型

前言

     手寫字型識别模型LeNet5誕生于1994年，是最早的卷積神經網絡之一。LeNet5通過巧妙的設計，利用卷積、參數共享、池化等操作提取特征，避免了大量的計算成本，最後再使用全連接配接神經網絡進行分類識别，這個網絡也是最近大量神經網絡架構的起點。

        卷積或池化輸出圖像尺寸的計算公式如下：

        O=輸出圖像的尺寸；I=輸入圖像的尺寸；K=池化或卷積層的核尺寸；S=移動步長；P =填充數 

1、net.py

import torch 
from torch import nn

# 定義一個網絡模型類
class MyLeNet5(nn.Module):
    # 初始化網絡
    def __init__(self):
        super(MyLeNet5, self).__init__()
        # 輸入大小為32*32，輸出大小為28*28，輸入通道為1，輸出為6，卷積核為5
        self.c1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, padding=2)
        # 使用sigmoid激活函數
        self.Sigmoid = nn.Sigmoid()
        # 使用平均池化
        self.s2 = nn.AvgPool2d(kernel_size=2, stride=2)
        self.c3 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)
        self.s4 = nn.AvgPool2d(kernel_size=2, stride=2)
        self.c5 = nn.Conv2d(in_channels=16, out_channels=120, kernel_size=5)

        self.flatten = nn.Flatten()
        self.f6 = nn.Linear(120, 84)
        self.output = nn.Linear(84, 10)

    def forward(self, x):
        # x輸入為32*32*1， 輸出為28*28*6
        x = self.Sigmoid(self.c1(x))
        # x輸入為28*28*6， 輸出為14*14*6
        x = self.s2(x)
        # x輸入為14*14*6， 輸出為10*10*16
        x = self.Sigmoid(self.c3(x))
        # x輸入為10*10*16， 輸出為5*5*16
        x = self.s4(x)
        # x輸入為5*5*16， 輸出為1*1*120
        x = self.c5(x)
        x = self.flatten(x)
        # x輸入為120， 輸出為84
        x = self.f6(x)
        # x輸入為84， 輸出為10
        x = self.output(x)
        return x

if __name__ == "__main__":
    x = torch.rand([1, 1, 28, 28])
    model = MyLeNet5()
    y = model(x)
           

2、train.py

import torch
from torch import nn
from net import MyLeNet5
from torch.optim import lr_scheduler
from torchvision import datasets, transforms
import os


# 将資料轉化為tensor格式
data_transform = transforms.Compose([
    transforms.ToTensor()
])

# 加載訓練資料集
train_dataset = datasets.MNIST(root='./data', train=True, transform=data_transform, download=True)
# 給訓練集建立一個資料加載器, shuffle=True用于打亂資料集，每次都會以不同的順序傳回。
train_dataloader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=16, shuffle=True)
# 加載訓練資料集
test_dataset = datasets.MNIST(root='./data', train=False, transform=data_transform, download=True)
# 給訓練集建立一個資料加載器, shuffle=True用于打亂資料集，每次都會以不同的順序傳回。
test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=16, shuffle=True)

# 如果顯示卡可用，則用顯示卡進行訓練
device = "cuda" if torch.cuda.is_available() else 'cpu'

# 調用net裡面定義的模型，如果GPU可用則将模型轉到GPU
model = MyLeNet5().to(device)

# 定義損失函數（交叉熵損失）
loss_fn = nn.CrossEntropyLoss()

# 定義優化器,SGD,
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)

# 學習率每隔10epoch變為原來的0.1
lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

# 定義訓練函數
def train(dataloader, model, loss_fn, optimizer):
    loss, current, n = 0.0, 0.0, 0
    # enumerate傳回為資料和标簽還有批次
    for batch, (X, y) in enumerate(dataloader):
        # 前向傳播
        X, y = X.to(device), y.to(device)
        output = model(X)
        cur_loss = loss_fn(output, y)
        # torch.max傳回每行最大的機率和最大機率的索引,由于批次是16，是以傳回16個機率和索引
        _, pred = torch.max(output, axis=1)

        # 計算每批次的準确率， output.shape[0]為該批次的多少
        cur_acc = torch.sum(y == pred) / output.shape[0]
        # print(cur_acc)
        # 反向傳播
        optimizer.zero_grad()
        cur_loss.backward()
        optimizer.step()
        # 取出loss值和精度值
        loss += cur_loss.item()
        current += cur_acc.item()
        n = n + 1

    print('train_loss' + str(loss / n))
    print('train_acc' + str(current / n))

# 定義驗證函數
def val(dataloader, model, loss_fn):
    # 将模型轉為驗證模式
    model.eval()
    loss, current, n = 0.0, 0.0, 0
    # 非訓練，推理期用到（測試時模型參數不用更新， 是以no_grad）
    # print(torch.no_grad)
    with torch.no_grad():
        for batch, (X, y) in enumerate(dataloader):
            X, y = X.to(device), y.to(device)
            output = model(X)
            cur_loss = loss_fn(output, y)
            _, pred = torch.max(output, axis=1)
            cur_acc = torch.sum(y == pred) / output.shape[0]
            loss += cur_loss.item()
            current += cur_acc.item()
            n = n + 1
        print('val_loss' + str(loss / n))
        print('val_acc' + str(current / n))

        return current/n

# 開始訓練
epoch = 50
min_acc = 0
for t in range(epoch):
    lr_scheduler.step()
    print(f"epoch{t+1}\n-------------------")
    train(train_dataloader, model, loss_fn, optimizer)
    a = val(test_dataloader, model, loss_fn)
    # 儲存最好的模型權重檔案
    if a > min_acc:
        folder = 'sava_model'
        if not os.path.exists(folder):
            os.mkdir('sava_model')
        min_acc = a
        print('save best model', )
        torch.save(model.state_dict(), "sava_model/best_model.pth")
    # 儲存最後的權重檔案   
    if t == epoch - 1:
        torch.save(model.state_dict(), "sava_model/last_model.pth")
print('Done！')
           

 3、test.py

import torch
from net import MyLeNet5
from torch.autograd import Variable
from torchvision import datasets, transforms
from torchvision.transforms import ToPILImage

# 将資料轉化為tensor格式
data_transform = transforms.Compose([
    transforms.ToTensor()
])

# 加載訓練資料集
train_dataset = datasets.MNIST(root='./data', train=True, transform=data_transform, download=True)
# 給訓練集建立一個資料加載器, shuffle=True用于打亂資料集，每次都會以不同的順序傳回。
train_dataloader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=16, shuffle=True)
# 加載訓練資料集
test_dataset = datasets.MNIST(root='./data', train=False, transform=data_transform, download=True)
# 給訓練集建立一個資料加載器, shuffle=True用于打亂資料集，每次都會以不同的順序傳回。
test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=16, shuffle=True)

#  如果顯示卡可用，則用顯示卡進行訓練
device = "cuda" if torch.cuda.is_available() else 'cpu'

# 調用net裡面定義的模型，如果GPU可用則将模型轉到GPU
model = MyLeNet5().to(device)

# 加載 train.py 裡訓練好的模型
model.load_state_dict(torch.load("D:/PycharmProjects/pytorch_test/LeNet-5/sava_model/best_model.pth"))

# 擷取預測結果
classes = [
    "0",
    "1",
    "2",
    "3",
    "4",
    "5",
    "6",
    "7",
    "8",
    "9",
]

# 把tensor轉成Image， 友善可視化
show = ToPILImage()

# 進入驗證階段
model.eval()
# 對test_dataset裡10000張手寫數字圖檔進行推理
for i in range(len(test_dataloader)):
    x, y = test_dataset[i][0], test_dataset[i][1]
    # tensor格式資料可視化
    show(x).show()
    # 擴充張量次元為4維
    x = Variable(torch.unsqueeze(x, dim=0).float(), requires_grad=False).to(device)
    with torch.no_grad():
        pred = model(x)
        # 得到預測類别中最高的那一類，再把最高的這一類對應classes中的哪一類标簽
        predicted, actual = classes[torch.argmax(pred[0])], classes[y]
        # 最終輸出的預測值與真實值
        print(f'predicted: "{predicted}", actual:"{actual}"')