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}%)]	Loss: {:.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('
Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)
'.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}%)]	Loss: {:.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('
Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)
'.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}%)]	Loss: {:.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}%)]	Loss: {:.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('
Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)
'.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 和 准确率。