如何使用Pytorch迅速写一个Mnist数据分类器

如何使用Pytorch迅速写一个Mnist数据分类器

一段时间没有更新博文，想着也该写两篇文章玩玩了。而从一个简单的例子作为开端是一个比较不错的选择。本文章会手把手地教读者构建一个简单的Mnist(Fashion-Mnist同理)的分类器，并且会使用相对完整的Pytorch训练框架，因此对于初学者来说应该会是一个方便入门且便于阅读的文章。本文的代码来源于我刚学Pytorch时的小项目，可能在形式上会有引用一些github上的小代码。同时文风可能会和我之前看的一些外国博客有点相近。

本文适用对象： 刚入门的Pytorch新手，想要用Pytorch来完成作业的鱼干。

那么就开始coding吧。

首先，你需要安装好Python 3+,Pytorch 1.0+,我个人使用的是Pytorch1.4，我想1.0以上的版本都可以使用。

然后在想要的位置，新建一个main.py的文件，然后就可以开始敲键盘了。

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

import argparse

第一步自然是导入相应的包。前面的都是Pytorch的包，最后一句导入的argparse便于用来修改训练的参数，这在Pytorch复现深度学习模型时非常常见。

model_names = ['Net','Net1']

parser = argparse.ArgumentParser(description='PyTorch Mnist Training')
parser.add_argument('-a', '--arch', metavar='ARCH', default='Net', choices=model_names,
                    help='model archtecture: ' + '|'.join(model_names) + '(default:Net)')
parser.add_argument('--epochs', default=5, type=int, metavar='N', help='number of total epochs')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum')
parser.add_argument('-b', '--batchsize', default=32, type=int, metavar='N', help='mini-batch size')
parser.add_argument('--lr', '--learning-rate', default=1e-2, type=float, metavar='LR', help='initial learning rata',
                    dest='lr')
args=vars(parser.parse_args())

第一行的model_names是一个list，用来存储我们之后会实现的两种网络结构的名字。然后我定义了一个argparse的对象，关于argparse可以自寻一些教程观看，大概只需要知道可以从指令行输入参数即可。在parser中又定义了arch(使用的网络)，epochs(迭代轮次)，momentum(梯度动量大小)，batchsize(一次送入的图片量大小)，learning-rate(学习率)参数。之前的model_name也正是用在arch参数中，限定了网络框架将会从此二者中选择其一。

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    #parameter
    batch_size = args["batchsize"]
    lr = args["lr"]
    momentum = args["momentum"]
    num_epochs = args["epochs"]

主函数中，先定义cuda对象，便于使用gpu并行运算。在#parameter中，我们把一些从命令行中获得的参数引入到相应的变量中，以便后续书写。

    #prepare the dataset

    mnist_data = datasets.MNIST("./mnist_data",train=True, download=True,
                                transform=transforms.Compose([
                                    transforms.ToTensor(),
                                    transforms.Normalize(mean=(0.13,),std=(0.31,))
                                ]))
    '''
    mnist_data = datasets.FashionMNIST("./fashion_mnist_data", train=True, download=True,
                                transform=transforms.Compose([
                                    transforms.ToTensor(),
                                    transforms.Normalize(mean=(,), std=(,))
                                ]))
    '''
    train_loader = torch.utils.data.DataLoader(
        mnist_data,batch_size=batch_size,shuffle=True, num_workers=1,pin_memory=True
    )
    test_loader = torch.utils.data.DataLoader(
        mnist_data,batch_size=batch_size,shuffle=True, num_workers=1,pin_memory=True
    )

之后将引入Pytorch中datasets包自带的MNIST集，download参数设置为True，以便于本地没有Mnist数据集时直接下载，之后会在当前目录下创建一个mnist_data的文件夹以存放数据，。transform中的transforms.ToTensor()是用于将图片形式的数据转换成tensor类型，而transforms.Normalize(mean=(0.13,),std=(0.31,))则是将tensor类型的数据进行归一化，这里的0.13和0.31可以直接使用。如果你想要使用注释中的FashionMNIST数据集则需要使用的是注释中的内容，当然，mean和std需要另外求解。

之后，定义train_loader和test_loader，将数据集作为可迭代的对象使用。shuffle=True以实现乱序读取数据，一般都会这么设置。num_workers和pin_memory都会影响到数据读取速度，前者是会在读取时创建多少个进程，后者是影响到数据读入到GPU中，一般来说，对于这个项目前者设置为1已经够用，后者设置为True和False都不影响。在更大型的项目中，如果设备较好，前者可以设置大一些。

    model = Net1().to(device) if args["arch"]=='Net1' else Net().to(device)
    optimizer = torch.optim.SGD(model.parameters(),lr=lr,momentum=momentum)
    criteon = nn.CrossEntropyLoss().to(device)

第一行是model实例化，并且会根据args["arch"]选择是用Net还是Net1，to(device)会将model放置于device上运行。第二行定义了一个优化器，使用的是SGD，并且放入model的参数、学习率和动量大小。criteon定义损失函数，这边使用的是交叉熵函数，这一损失函数在分类问题中十分常用。

    #train
    for epoch in range(num_epochs):
        train(model,device,train_loader,optimizer,epoch,criteon)
        test(model,device,test_loader,criteon)


    torch.save(model.state_dict(), "mnist_{}.pth".format(num_epochs))

这就是训练过程，在其中又使用了train和test两个函数(下面会说)，根据num_epochs数目进行循环。循环结束后，torch.save将会把模型的参数model.state_dict()以mnist_{}.pth的形式存放到当前文件夹下。

def train(model,device,train_loader,optimizer,epoch,criteon):
    class_name = model.__class__.__name__

    model.train()
    loss = 0
    for idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        pred = model(data)
        if class_name == 'Net':
            loss = F.nll_loss(pred, target)
        elif class_name == 'Net1':
            loss = criteon(pred, target)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if idx % 100 == 0:
            print("train epoch: {}, iteration: {}, loss: {}".format(epoch, idx, loss.item()))

这里定义了train函数的训练过程。class_name中存放了当前使用的模型名字。 model.train()开启训练模式。在for idx, (data, target) in enumerate(train_loader):取出当前数据集的idx，data和种类target。循环中，先把data和target放置于device上，pred = model(data)会进行一次前传，获得相应数据的预测种类pred。

对不同的模型，我采用了不同定义损失函数的方式，这里需要结合下面的模型结构来看。optimizer.zero_grad()会将上轮累计的梯度清空，之后loss.backward()梯度反向传播，利用optimizer.step()更新参数。而当if idx % 100 == 0:也就是迭代的数据批次到达100的倍数了，就会输出相关信息。

def test(model,device,test_loader,criteon):
    class_name = model.__class__.__name__
    model.eval()
    total_loss = 0 #caculate total loss
    correct = 0
    with torch.no_grad():
        for idx, (data, target) in enumerate(test_loader):
            data, target = data.to(device), target.to(device)
            pred = model(data)
            if class_name == 'Net':
                total_loss += F.nll_loss(pred, target,reduction="sum").item()
            elif class_name == 'Net1':
                total_loss += criteon(pred, target).item()
            correct += pred.argmax(dim=1).eq(target).sum().item()

    total_loss /= len(test_loader.dataset)
    acc = correct/len(test_loader.dataset)
    print("Test loss: {}, Accuracy: {}%".format(total_loss,acc*100))

test函数总体结构类似，model.eval()将会把模型调整测试模式，with torch.no_grad():来声明测试模式下不需要积累梯度信息。correct += pred.argmax(dim=1).eq(target).sum().item()则是会计算出预测对了的数目，之后通过total_loss计算总误差和acc计算准确率。

class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.conv1 = nn.Conv2d(1, 20, kernel_size=5,stride=1)
        self.conv2 = nn.Conv2d(20,50,kernel_size=5,stride=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)
        x = F.log_softmax(x,dim=1)
        return x

Net不过是一个具有两个卷积层和两个线性全连层的网络。self.conv1 = nn.Conv2d(1, 20, kernel_size=5,stride=1)表示conv1是一个接受1个channel的tensor输出20个channel的tensor，且卷积大小为5，步长为1的卷积层。self.fc1 = nn.Linear(4*4*50, 500)则是接收一个4 * 4 * 50长的一维tensor并且输出长为500的一维tensor。

前传函数forward中，x作为输入的数据，输入后会通过 conv1->relu->pooling->conv2->relu->pooling->view将多维tensor转化成一维tensor->fc1->relu->fc2->log_softmax来获得最终的x的值。这里就需要提train和test函数中的if和elif语句了。使用的时Net时，loss = F.nll_loss(pred, target)，这是因为log_softmax之后使用nll_loss和直接使用 nn.CrossEntropyLoss()是等效的，因此：

class Net1(nn.Module):
    def __init__(self):
        super(Net1,self).__init__()
        self.conv_unit=nn.Sequential(
            nn.Conv2d(1, 20, kernel_size=5,stride=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2,stride=2),
            nn.Conv2d(20,50,kernel_size=5,stride=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.fc_unit=nn.Sequential(
            nn.Linear(4*4*50, 500),
            nn.ReLU(),
            nn.Linear(500, 10)
        )


    def forward(self, x):
        x = self.conv_unit(x)
        x = x.view(-1,4*4*50)
        x = self.fc_unit(x)
        return x

Net1中最后并没有使用log_softmax，是因为直接在train的过程中，使用了nn.CrossEntropyLoss()。此外，Net1和Net不同的地方也就是在结构中使用了nn.Sequential()来单元化卷积层和全连层。

if __name__ == '__main__':
    main()

之后就可以使用了！

在命令行中使用：

$ python main.py

就会按照默认的参数训练一个Mnist分类器了。

第三轮的效果:

train epoch: 2, iteration: 1300, loss: 0.010509848594665527
train epoch: 2, iteration: 1400, loss: 0.0020529627799987793
train epoch: 2, iteration: 1500, loss: 0.0027058571577072144
train epoch: 2, iteration: 1600, loss: 0.010049819946289062
train epoch: 2, iteration: 1700, loss: 0.0352507084608078
train epoch: 2, iteration: 1800, loss: 0.009431719779968262
Test loss: 0.01797709318200747, Accuracy: 99.42833333333333%

如果希望查看参数列表，则可以在命令行使用：

$ python main.py -h

就会出现：

usage: main.py [-h] [-a ARCH] [--epochs N] [--momentum M] [-b N] [--lr LR]

PyTorch Mnist Training

optional arguments:
  -h, --help            show this help message and exit
  -a ARCH, --arch ARCH  model archtecture: Net|Net1(default:Net)
  --epochs N            number of total epochs
  --momentum M          momentum
  -b N, --batchsize N   mini-batch size
  --lr LR, --learning-rate LR
                        initial learning rata

于是如果想要使用Net1，lr为0.001的方式训练，就可以按照这样:

$ python main.py -a Net1 --lr 0.001

第三轮结果:

train epoch: 2, iteration: 1200, loss: 0.03096039593219757
train epoch: 2, iteration: 1300, loss: 0.060124486684799194
train epoch: 2, iteration: 1400, loss: 0.08865253627300262
train epoch: 2, iteration: 1500, loss: 0.13717596232891083
train epoch: 2, iteration: 1600, loss: 0.003894627094268799
train epoch: 2, iteration: 1700, loss: 0.06881710141897202
train epoch: 2, iteration: 1800, loss: 0.03184908628463745
Test loss: 0.0013615453808257978, Accuracy: 98.69666666666667%

至此，你获得了一个Mnist训练器的训练方法。