Pytorch 之 MNIST 數據集實現（代碼講解）

原文鏈接：https://blog.csdn.net/weixin_37589575/article/details/95856667

 

目錄1. 數據集介紹2. 代碼2. 讀代碼（個人喜歡的順序）2.1. 導入模塊部分：2.2. Main 函數：
1. 數據集介紹
一般而言，MNIST 數據集測試就是機器學習和深度學習當中的＂Hello World＂工程。幾乎是所有的教程都會把它放在最開始的地方。這是因為，這個簡單的工程包含了大致的機器學習流程，通過練習這個工程有助於讀者加深理解機器學習或者是深度學習的大致流程。
MNIST（Mixed National Institute of Standards and Technology database）是一個計算機視覺數據集，它包含 70000 張手寫數字的灰度圖片，其中每一張圖片包含 28 X 28 個像素點。可以用一個數字數組來表示這張圖片。
2. 代碼

import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4 * 4 * 50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 4 * 4 * 50)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)


def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))


def test(args, model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))


def main():
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                        help='input batch size for testing (default: 1000)')
    parser.add_argument('--epochs', type=int, default=10, metavar='N',
                        help='number of epochs to train (default: 10)')
    parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                        help='learning rate (default: 0.01)')
    parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
                        help='SGD momentum (default: 0.5)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='how many batches to wait before logging training status')
    parser.add_argument('--save-model', action='store_true', default=False,
                        help='For Saving the current Model')
    
    args = parser.parse_args()
    use_cuda = not args.no_cuda and torch.cuda.is_available()
    torch.manual_seed(args.seed)
    device = torch.device("cuda" if use_cuda else "cpu")
    kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('../data', train=True, download=True,
                       transform=transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])),
        batch_size=args.batch_size, shuffle=True, **kwargs)
    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('../data', train=False, transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])),
        batch_size=args.test_batch_size, shuffle=True, **kwargs)

    model = Net().to(device)
    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)

    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)
        test(args, model, device, test_loader)

    if (args.save_model):
        torch.save(model.state_dict(), "mnist_cnn.pt")

if __name__ == '__main__':

    main()

 

 

代碼案例二

 

import  torch
import  torch.nn as nn
import  torch.nn.functional as F
import  torch.optim as optim
from    torchvision import datasets, transforms


batch_size=200
learning_rate=0.01
epochs=10

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('../data', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ])),
    batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('../data', train=False, transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])),
    batch_size=batch_size, shuffle=True)



class MLP(nn.Module):

    def __init__(self):
        super(MLP, self).__init__()

        self.model = nn.Sequential(
            nn.Linear(784, 200),
            nn.ReLU(inplace=True),
            nn.Linear(200, 200),
            nn.ReLU(inplace=True),
            nn.Linear(200, 10),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        x = self.model(x)

        return x

net = MLP()
optimizer = optim.SGD(net.parameters(), lr=learning_rate)
criteon = nn.CrossEntropyLoss()

for epoch in range(epochs):

    for batch_idx, (data, target) in enumerate(train_loader):
        data = data.view(-1, 28*28)

        logits = net(data)
        loss = criteon(logits, target)

        optimizer.zero_grad()
        loss.backward()
        # print(w1.grad.norm(), w2.grad.norm())
        optimizer.step()

        if batch_idx % 100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))


    test_loss = 0
    correct = 0
    for data, target in test_loader:
        data = data.view(-1, 28 * 28)
        logits = net(data)
        test_loss += criteon(logits, target).item()

        pred = logits.data.max(1)[1]
        correct += pred.eq(target.data).sum()

    test_loss /= len(test_loader.dataset)
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

visdom視圖的可視化

from visdom import Visdom
 viz = Visdom()

#一條線時
    viz.line([0.], [0.], win='train_loss', opts=dict(title='train loss'))
    viz.line([loss.item()], [global_step], win='train_loss', update='append')
#兩條線時
    viz.line([[0.0,0.0]], [0.], win='test', opts=dict(title='test loss&acc.',legend=['loss','acc.']))
    viz.line([[test_loss,correct/len(test_loader.dataset)]], [global_step], win='test', update='append')
#圖像可視化
   viz.images(data.view(-1,1,28,28),win='X')
   viz.text(str(pred.detach().cpu().numpy()),win='pred',opts=dict(title='pred'))

 

 

2. 讀代碼（個人喜歡的順序）

2.1. 導入模塊部分：

import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

 

導入基本的 Python 庫和 Python 包:

argparse 是一個模塊，非常好用，可以用來設置模型的默認參數，也可以允許使用者在命令行模式手動設置參數
torch 和 torchvision 是 pytorch 基本的的庫。torch 自然指的是 pytorch， torchvision 是獨立於 pytorch 的關於圖像操作的一些方便工具庫。torchvision主要包括以下幾個包：
（1）vision.datasets : 幾個常用視覺數據集，可以下載和加載，這里主要的高級用法就是可以看源碼如何自己寫自己的Dataset的子類
（2）vision.models : 流行的模型，例如 AlexNet, VGG, ResNet 和 Densenet 以及 與訓練好的參數。
（3）vision.transforms : 常用的圖像操作，例如：隨機切割，旋轉，數據類型轉換，圖像到tensor ,numpy 數組到tensor , tensor 到 圖像等。
（4）vision.utils : 用於把形似 (3 x H x W) 的張量保存到硬盤中，給一個mini-batch的圖像可以產生一個圖像格網。
torch.nn 參考 pytorch 文檔：https://ptorch.com/docs/1/torch-nn
torch.nn.functional 參考 pytorch 文檔：https://ptorch.com/docs/8/torch.nn.functional
import torch.optim 主要包含常見的優化算法

2.2. Main 函數：

parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                    help='input batch size for training (default: 64)')
parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                    help='input batch size for testing (default: 1000)')
parser.add_argument('--epochs', type=int, default=10, metavar='N',
                    help='number of epochs to train (default: 10)')
parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                    help='learning rate (default: 0.01)')
parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
                    help='SGD momentum (default: 0.5)')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='disables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
                    help='random seed (default: 1)')
parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                    help='how many batches to wait before logging training status')
parser.add_argument('--save-model', action='store_true', default=False,
                    help='For Saving the current Model')
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()

這里就是使用了 argparse 模塊，配置了 batch-size, test-batch-size, epochs, lr, momentum, -no-cuda, seed, log-interval, save-model 等參數，參數值就是 default 里面的默認值，同時我們可以在命令行允許的時候修改。例如

python main.py --epochs 20 --lr 0.008

這樣就可以修改 epochs 為 20， 學習率為 0.008。

batch-size：我們一次訓練如果只訓練一張圖片，顯然效率太低了，很多 GPU， CPU 資源都沒有利用到，因此可以一次訓練多張圖片。可以提高效率，但是如果顯存太小，圖片太大，可能導致顯存不夠用的問題，因此這個值在我們 MNIST 小數據集上設置可以隨意一點，大點的數據集需要好好考量。
test-batch-size：測試時的 batch-size
epochs：訓練的 epochs 數量
lr：學習率
momentum：SGD 算法里面的 momentum（動量）
seed：隨機種子
log-interval：多少個 batch 打印一次訓練狀態，設置為10，即 10*64，一個batch size是 64 ，即每 640 張圖片打印一次訓練結果。
save-model：是否保存模型
cuda：GPU 是否可用

torch.manual_seed(args.seed)
device = torch.device("cuda" if use_cuda else "cpu")

 

設置隨機種子，以及將 GPU 命名為 device

 kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('../data', train=True, download=True,
                       transform=transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])),
        batch_size=args.batch_size, shuffle=True, **kwargs)
    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('../data', train=False, transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])),
        batch_size=args.test_batch_size, shuffle=True, **kwargs)

加載訓練集和測試集，pytorch 自帶有 MNIST 數據集。 並且轉為我們需要的格式。

train_loader 里面就是 6000 個訓練集，包含有圖片和圖片的標簽。
test_loader 里面就是 1000 個訓練集，包含有圖片和圖片的標簽。

kwargs = {'num_workers': 1, 'pin_memory': True} 

num_workers 是多進程的加載數，pin_memory：是否將數據保存在pin memory區，pin memory中的數據轉到GPU會快一些。

model = Net().to(device)

構建我們的模型，並送到 GPU 中加速。
模型為：

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4 * 4 * 50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 4 * 4 * 50)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

 

兩個卷積層，兩個線性層：構建我們只看 init 部分就可以了，forward 部分是將數據送到模型，模型如何處理的過程。

optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)

優化方法是 SGD， 第一個參數是模型的參數，第二個是學習率，第三個是動量。

    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)
        test(args, model, device, test_loader)

一共訓練 10 次，每次訓練完都測試一次。

if (args.save_model):
  torch.save(model.state_dict(), "mnist_cnn.pt")

如果配置為要保持模型，就保存模型到 mnist_cnn.pt 文件里面，如果沒有會新建。

def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))

 

訓練過程，接收配置 args, 模型model, GPU device, 訓練數據train_loader，優化器optimizer和當前訓練周期epoch

model.train()

模型進入訓練模式

for batch_idx, (data, target) in enumerate(train_loader):
    data, target = data.to(device), target.to(device)

加載訓練數據集合，這里一次 64 張圖片以及對應的類別（0 - 9），data 是圖片 shape 為[64, 1, 28, 28]，64張圖片，灰度圖片，像素大小為28*28。 target 是 64 個圖片的類別，每個用 0-9 的圖片表示。

 optimizer.zero_grad()
 output = model(data)
 loss = F.nll_loss(output, target)
 loss.backward()
 optimizer.step()

優化過程開始，然后 data 送到模型，經過模型的 forward 步驟，最后一步 softmax 為這張圖片屬於 10 個類的概率，比如[0.7, 0.1, 0.1, 0.1,0, 0, 0, 0, 0, 0] 就是這張圖片屬於 0 的概率為 0.7， 屬於1. 2. 3的概率為 0.1。損失函數為 NLLLoss 。
NLLLoss 的 輸入 是一個對數概率向量和一個目標標簽. 它不會為我們計算對數概率. 適合網絡的最后一層是log_softmax. 損失函數 nn.CrossEntropyLoss() 與 NLLLoss() 相同, 唯一的不同是它為我們去做 softmax. 我們視為分類問題，最常見的損失函數就是交叉熵，我們這里本質也是交叉熵。因為模型最后一步 forward 是 softmax。

CrossEntropyLoss()=log_softmax() + NLLLoss()+softmax()

然后根據 loss 用 SGD 算法優化參數。

 if batch_idx % args.log_interval == 0:
     print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
     epoch, batch_idx * len(data), len(train_loader.dataset),
     100. * batch_idx / len(train_loader), loss.item()))

每 10 * 64 張圖片打印訓練狀態：

 

 

def test(args, model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

 

測試模型。

    model.eval()
    test_loss = 0
    correct = 0

模型進入測試模式，初始化計數變量。然后加載測試圖片，data 送到模型，然后計算疊加損失 loss，pred 就是找出 softmax 之后數組里面最大值得索引，最大值就是預測最有可能的概率，然后索引就是預測的數字。例如上面 [0.7, 0.1, 0.1, 0.1,0, 0, 0, 0, 0, 0] ， 最大概率值是 0.7，索引就是 0，我們預測這個數字是 0。然后疊加准確的個數，最后 loss 和 准確個數除以測試圖片個數就是平均的 loss 和 准確率。