如何使用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训练器的训练方法。
