Pytorch实战-Mnist数据集

上篇理解完官方的代码后，根据自己的理解，写了一份自己风格的CPU版本代码，并加上了详细的注释，方便像我一样的新手小白理解每块代码都在做些什么事情。

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

#parameters
batch_size = 32
lr = 0.01
momentum = 0.5

#defined a Net
class Net(nn.Module):
    def __init__(self):                                     #初始化函数
        super(Net, self).__init__()                     #初始化
        self.conv1 = nn.Sequential(                 #第一部分，一个卷积层，然后用relu函数激活，最后pooling
            nn.Conv2d(1, 20, 5, 1),                      #in_channels = 1,out_channels = 20,kernel_size = 5,stride =1
            nn.ReLU(),
            nn.MaxPool2d(2)                            #kernel_size = 2
        )
        self.conv2 = nn.Sequential(                 #同上，in_channels = 20,out_channels =50
            nn.Conv2d(20, 50, 5, 1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.fc1 = nn.Linear(4 * 4 * 50, 500)        #fully connection,线性函数，
        self.out = nn.Linear(500, 10)                   #ouptut层，最后输出结果有10种可能

    def forward(self, x):                                    #foward函数，往output层传播
       x = self.conv1(x)                                     #通过两层复合卷积层
       x = self.conv2(x)
       x = x.view(-1, 4 * 4 * 50)                         #用view()函数把卷积结果展平
       x = self.fc1(x)                                          #通过fc1线性函数还原
       x = self.out(x)                                          #通过一个线性函数输出
       return F.log_softmax(x, dim=1)                #注意这里return的值要是数值类型，这里用log_softmax优化，模仿官方的做法

def train(model, train_loader, optimizer, epoch):                     #训练函数
    model.train()                                                                      #进入训练状态
    for batch_idx, (data, target) in enumerate(train_loader):       #对train_loader进行批数据训练
        optimizer.zero_grad()                                                     #梯度清零
        output = model(data)                                                      #调用网络
        loss = F.nll_loss(output, target)                                        #计算损失函数，把output的训练结果和target进行比较
        loss.backward()                                                               #计算梯度
        optimizer.step()                                                              #更新参数
        if batch_idx % 50 == 0:                                                    #每训练50批打印，epoch，当前进度，当前损失值
            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(model, test_loader):                                             #测试函数
    model.eval()                                                                #进入test状态
    test_loss = 0                                                                 #初始化
    correct = 0
    with torch.no_grad():                                                       #test不需要梯度，猜想这个函数是想节约内存
        for data, target in test_loader:
            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()      #比较pred和target的值，equal为1

    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():
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('../mnist', train=True, download=True,             #下载并创建数据集
                       transform=transforms.Compose([                           #对数据集进行处理
                           transforms.ToTensor(),                                       #转化为tensor类型
                           transforms.Normalize((0.1307,), (0.3081,))           #按照mean和std处理数据
                       ])),
        batch_size=batch_size, shuffle=True)
    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('../mnist', train=False, transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])),
        batch_size=batch_size, shuffle=True)

    model = Net()                  #建模
    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)  #用SGD优化模型

    for epoch in range(1):                                                                          #遍历全部数据集一次
        train(model, train_loader, optimizer, epoch)                                       #训练
        test(model,test_loader)                                                                     #测试

main()

下面是我的测试结果，比官方的略微好了一点点，我的batch_size定义的比较小，是32，官方64。