如何使用Pytorch迅速写一个Mnist数据分类器
- 2020 年 3 月 8 日
- 笔记
如何使用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 xNet不过是一个具有两个卷积层和两个线性全连层的网络。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 xNet1中最后并没有使用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训练器的训练方法。
