pytorch實現LetNet5並用其訓練minist手寫訓練集

本文針對本人學習pytorch過程記錄，本文主要是針對CNN中的LetNet5

import os
import torch
import torch.nn as nn
import torch.utils.data as Data
import torchvision
from torch.autograd import  Variable

#由於minist數據集市28*28的單通道圖像，60000個訓練集，10000個測試集
# Hyper Parameters
EPOCH = 5               # 一共訓練5次，條件允許可以多訓練幾次
BATCH_SIZE = 100        #批訓練 一批100個
LR = 0.001              # 學習效率
DOWNLOAD_MNIST = False

if not(os.path.exists('./mnist/')) or not os.listdir('./mnist/'):
    # not mnist dir or mnist is empyt dir
    DOWNLOAD_MNIST = True

# 下載數據
train_data = torchvision.datasets.MNIST(
    root='./mnist/',
    train=True,                                      # True 是訓練集
    transform=torchvision.transforms.ToTensor(),     # 把shape=(H x W x C)的像素值範圍爲[0, 255]的PIL.Image或者numpy.ndarray轉換成
                                                     #shape=(C x H x W)的像素值範圍爲[0.0, 1.0]的torch.FloatTensor，黑白圖是單通道
    download=DOWNLOAD_MNIST,
)

test_data = torchvision.datasets.MNIST(
    root='./mnist/',
    train=False,                                      # False 是測試集
    transform=torchvision.transforms.ToTensor(),
    download=DOWNLOAD_MNIST,
)

#放進DataLoader進行批處理
train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True) #shuffle代表是否打亂順序
test_loader = Data.DataLoader(dataset=test_data, batch_size=BATCH_SIZE, shuffle=True) #shuffle代表是否打亂順序

#LetNet5 是七層的結構
class LetNet5(nn.Module):
    def __init__(self):
        super(LetNet5, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(                      #卷積
                in_channels=1,              # 輸入高度 黑白圖片爲1，RGB爲3
                out_channels=16,            # 輸出高度
                kernel_size=5,              # 卷積核 5*5
                stride=1,                   # 步長
                padding=2,
            ),
            nn.ReLU(),                      # 激勵
            nn.MaxPool2d(kernel_size=2),    # 池化
        )
        self.conv2 = nn.Sequential(         # input shape (16, 14, 14)
            nn.Conv2d(                      #卷積
                in_channels=16,              # 輸入高度 黑白圖片爲1，RGB爲3
                out_channels=32,            # 輸出高度
                kernel_size=5,              # 卷積核 5*5
                stride=1,                   # 步長
                padding=2,
            ),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.connect = nn.Linear(32 * 7 * 7, 50)  #全連接層1
        self.out = nn.Linear(50, 10)  # 分爲10類

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)

        x = x.view(x.size(0), -1)  # 將輸出展平爲（batch_size,32 * 7 * 7）

        tem = self.connect(x)
        output = self.out(tem)

        return output

letnet5 = LetNet5()

optimizer = torch.optim.Adam(letnet5.parameters(), lr=LR)   # 採用的是Adam優化器，也可以使用SGD等
loss_func = nn.CrossEntropyLoss()                           # 分類採用的誤差函數

for epoch in range(EPOCH):
    for step, data in enumerate(train_loader):    #訓練集包括數據和標籤
        inputs, labels = data
        inputs, labels = Variable(inputs), Variable(labels)
        outputs = letnet5(inputs)
        loss = loss_func(outputs, labels)  # 計算誤差
        optimizer.zero_grad()  #每次訓練完清零
        loss.backward()        #反向傳遞
        optimizer.step()       #每步加上優化

print("Finished Training")
print("Beginning Testing")

correct = 0
total = 0
for data in test_loader:
    images, labels = data
    outputs = letnet5(Variable(images))
    predicted = torch.max(outputs, 1)[1].data.numpy()
    total += labels.size(0)
    correct += (predicted == labels.data.numpy()).sum()

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