【他山之石】PyTorch 手把手搭建神经网络 (MNIST)

“他山之石，可以攻玉”，站在巨人的肩膀才能看得更高，走得更远。在科研的道路上，更需借助东风才能更快前行。为此，我们特别搜集整理了一些实用的代码链接，数据集，软件，编程技巧等，开辟“他山之石”专栏，助你乘风破浪，一路奋勇向前，敬请关注。

作者：知乎—摸鱼

地址：https://www.zhihu.com/people/sun-xu-sen-30

01

数据集介绍

本文使用的是PyTorch自带的MNIST数据集，该数据集中为PIL Image，size=28x28。数据存储形式为(data, target)，其中data为PIL Image，target该图片表示的数字。如下所示：

(<PIL.Image.Image image mode=L size=28x28 at 0x175EF44F160>, 5)

02

模型搭建

• input：torch.Size([64, 1, 28, 28])
• after conv1: torch.Size([64, 32, 26, 26])
• after conv2: torch.Size([64, 64, 24, 24])
• after max_pool2d: torch.Size([64, 64, 12, 12])
• after flatten: torch.Size([64, 9216])
• after fc1: torch.Size([64, 128])
• after fc2: torch.Size([64, 10])
• after log_softmax: torch.Size([64, 10])

由各层的输出结果可以看出，batch_size为64，网络结构图中tensor在各层的变化与输出一致。

2.1 导入库文件

from __future__ import print_function
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
from torch.optim.lr_scheduler import StepLR

2.2 网络结构定义

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        # print("after conv1: {}".format(x.shape))
        x = F.relu(x)
        x = self.conv2(x)
        # print("after conv2: {}".format(x.shape))
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        # print("after max_pool2d: {}".format(x.shape))
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        # print("after flatten: {}".format(x.shape))
        x = self.fc1(x)
        # print("after fc1: {}".format(x.shape))
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        # print("after fc2: {}".format(x.shape))
        output = F.log_softmax(x, dim=1)
        # print("after log_softmax: {}".format(output.shape))
        return output

2.3 transform究竟是干什么的？

相信大家在读取数据集的时候，经常见到下面这段代码

transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
        ])

别着急，我们一个个来看，首先看transforms.ToTensor()

简单来说，就是把PIL Image或者 numpy.ndarray类型的数据转化为张量tensor。原来的在[0, 255]的具有(HxWxC)形式的PIL Image或numpy.ndarray数据，被转换为[0.0, 1.0]范围并且shape变为(CxHxW)。

前面我们看到数据集的格式为PIL Image，这个函数就是将之前的PIL Image变为tensor，后续才可以对其进行一系列操作。

再看transforms.Normalize()

对一个tensor进行归一化，传入的两个参数为元组形式，分别为mean和std

相当于一个容器，将若干个transforms组合到一起。

2.4 DataLoader

taet1 = datasets.MNIST('../data', train=True, download=True,
                   transform=transform)
train_loader = torch.utils.data.DataLoader(dataset1, **train_kwargs)

简而言之，DataLoader的作用就是对传入的数据集进行采样sample，返回一个可迭代的对象。注意到前面定义的batch_size为64，那么在这里，每次返回的可迭代对象的size就是64。

2.5 优化器定义

# 优化器
optimizer = optim.Adadelta(model.parameters(), lr=args.lr)

# 调整学习率
scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)

其中StepLR是用来调整学习率的，具体的调整策略本文暂且不讨论，在之后的博文中对集中调整学习率的方法统一讨论。在这里只需知道随着epoch的增加，learning rate会改变即可。

03

train

train的流程非常简单，首先设置为train模式，分批次读入数据，优化器梯度置零，数据通过网络，计算损失，反向传播，权重更新。如下

def train(args, model, device, train_loader, optimizer, epoch):
    # 设置为train模式
    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()))
            if args.dry_run:
                break

04

test

test的流程与train略有不同，首先要设置为test模式，这是要防止改变已训练好的权重。接着在with torch.no_grad()中让数据通过网络，计算损失和预测是否正确即可。如下

def test(model, device, test_loader):
    # 设置为test模式
    model.eval()
    test_loss = 0
    correct = 0
    # 不进行计算图的构建，即没有grad_fn属性
    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)))

05

集成训练测试

def main():
    # Training settings
    # 声明一个parser
    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=14, metavar='N',
                        help='number of epochs to train (default: 14)')
    parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
                        help='learning rate (default: 1.0)')
    parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
                        help='Learning rate step gamma (default: 0.7)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--dry-run', action='store_true', default=False,
                        help='quickly check a single pass')
    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")

    train_kwargs = {'batch_size': args.batch_size}
    test_kwargs = {'batch_size': args.test_batch_size}
    if use_cuda:
        cuda_kwargs = {'num_workers': 1,
                       'pin_memory': True,
                       'shuffle': True}
        train_kwargs.update(cuda_kwargs)
        test_kwargs.update(cuda_kwargs)

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

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

    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)
        test(model, device, test_loader)
        scheduler.step()

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

06

模型结果

07

摸鱼建议

摸鱼也是最近接触PyTorch这个框架，在学习过程中遇到了一些问题，也总结了一些经验。下面是摸鱼遇到的一些问题以及解决方式

Q1：为什么网络的结构是这样定义的？有什么理由吗？

A1：其实刚开始摸鱼也不清楚为什么网络要这样设计，后来在Andrew Ng的课上，老师提起过一嘴，说这个没有什么特别的原因，如果非要说一个原因的话那就是它在实验上的表现很好。所以我们在学习的过程中，可以借鉴那些经典的网络结构，以此为基础改进来形成我们自己的网络架构。同样网络中的参数也是一般采用设计者给出的会比较好。

Q2：transform，DataLoader等等的到底是干嘛的？在好多地方看到过但还是比较模糊

A2：确实，在看官方文档的时候，经常看到这两段代码。相信看完本文应该就可以解决这个问题了，至于要如何解决类似的问题，我的一个建议是了解数据的源格式以及你想要的的目的格式。其实transform不难理解，就是进行一个数据格式的转换，但是如果不了解数据的源格式，可能对这块就会比较模糊。

本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。

“他山之石”历史文章

更多他山之石专栏文章，

请点击文章底部“阅读原文”查看

分享、点赞、在看，给个三连击呗！