PyTorch：训练MNIST数据集

本文主要描述了如何使用现在热度和关注度比较高的Pytorch(深度学习框架)构建一个简单的卷积神经网络，并对MNIST数据集进行了训练和测试。MNIST数据集是一个28*28的手写数字图片集合，使用测试集来验证训练出的模型对手写数字的识别准确率。

加载MNIST数据集

MNIST 包括6万张28x28的训练样本，1万张测试样本，很多教程都会对它”下手”几乎成为一个 “典范”，可以说它就是计算机视觉里面的Hello World。所以我们这里也会使用MNIST来进行实战。

PyTorch中提供了MNIST，CIFAR，COCO等常用数据集的加载方法。MNIST是torchvision.datasets包中的一个类，负责根据传入的参数加载数据集。如果自己之前没有下载过该数据集，可以将download参数设置为True，会自动下载数据集并解包。如果之前已经下载好了，只需将其路径通过root传入即可。

在加载图像后，我们常常需要对图像进行若干预处理。比如减去RGB通道的均值，或者裁剪或翻转图像实现augmentation等，这些操作可以在torchvision.transforms包中找到对应的操作。在下面的代码中，通过使用transforms.Compose()，我们构造了对数据进行预处理的复合操作序列，ToTensor负责将PIL图像转换为Tensor数据（RGB通道从[0, 255]范围变为[0, 1]）， Normalize负责对图像进行规范化。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import torchvision
from torch.autograd import Variable
from torch.utils.data import DataLoader
import cv2
import matplotlib.pyplot as plt

在这些包中torchvision，提供了许多深度学习的数据集torchvision.dataset（比如说MNIST,CIFAR-10等）以及一些pre-train(预训练）过的深度学习网络torchvision.models（比如说Alexnet,Vggnet），而跟数据预处理相关的操作（比如说ToTensor,Crop等操作）则被存在torchvision.transfroms之中。

设置当前设备：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

加载并处理数据（Dataloader）

3.1定义数据处理方式：

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

在这里

transforms.Compose()

是把多种数据处理的方法集合在一起。

transforms.Totensor()

3.2 定义Dataset与DateLoader

#Data set
train_dataset = torchvision.datasets.MNIST(root='./data', 
                                           train=True, 
                                           transform=transforms.ToTensor(),  
                                           download=True)

test_dataset = torchvision.datasets.MNIST(root='./data', 
                                          train=False, 
                                          transform=transforms.ToTensor())

# Data loader
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, 
                                           batch_size=batch_size, 
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset, 
                                          batch_size=batch_size, 
                                          shuffle=False)

dataloader则是加载dataset,并设置其batch_size（单次训练时送入的样本数目），以及shuffle(是否打乱样本顺序，训练集一定要设置shuffle为True，测试集则无强制性规定)

3.3 查看图片信息（可以省略）

dataiter = iter(trainloader)
images, labels = dataiter.next()

print(images.shape)
print(labels.shape)

在每个batch中，有64张图片，每个图片是28*28像素的，

标签的shape也是64，64 images should have 64 labels respectively。

四.定义网络结构

定义网络结构的相关代码：

class  Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.pool3 = nn.MaxPool2d(2, 2)

        self.fc1 = nn.Linear(128*3*3, 625)
        self.fc2 = nn.Linear(625, 10)

    def forward(self, x):
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool2(F.relu(self.conv2(x)))
        x = self.pool3(F.relu(self.conv3(x)))
        x = x.view(-1, 128*3*3)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x 

在网络的定义中，主要用到了init与forward方法

init方法主要是指名网络中有那些会使用到的层结构。

forward方法则展示了数据在网络中的流动过程（前向传播）。

其中的F表示使用了激活函数（上述代码中使用的是relu）

需要注意的是，在最后一个卷积层与全连接层连接时需要使用

x = x.view(-1,128*3*3)

对高维Tensor进行降维，上面代码中的-1是任意的意思。

对网络进行实例化：

net = Net().to(device)

五.定义损失函数与优化器：

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(),lr=1e-3，momentum = 0.9)

优化器采用SGD（小批量梯度下降）

 

六.训练：

num_epochs = 1000

for epoch in range(num_epochs):
    for i, data in enumerate(train_loader):
        images, labels = data
        images = images.to(device)
        labels = labels.to(device)

        # Forward pass
        outputs = F.softmax(net(images))
        loss = criterion(outputs, labels)

        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

从上面的代码中，可以发现训练神经网络的过程包括三个步骤：

1.加载训练数据与标签，将其放置于GPU上

2.前向过程生成结果

因为在此处，我们使用了NLLLoss，所以对于网络的输出应该采用softmax将其转化为概率分布。最后将概率与Labels做NLLLoss.

（当然也可以使用交叉熵CrossEntropy损失函数，CrossEntropy损失函数中封装了softmax）

3.反向更新权值

首先需要将优化器先前保存的梯度信息，然后对loss使用.backward求导。最后将优化器中的变量使用.step更新。

七.测试

测试代码如下:

net.eval()
correct = 0
total = 0
for data_test in test_loader:
    images, labels = data_test
    images = images.to(device)
    labels = labels.to(device)
    output_test = net(images)
    _, predicted = torch.max(output_test, 1)

    total += labels.size(0)
    correct += (predicted == labels).sum()
print("correct1: ", correct)
print("Test acc: {0}".format(correct.item() / len(test_dataset)))

附完整代码：

import numpy as np
import torch
import torchvision
import matplotlib.pyplot as plt
from time import time
from torchvision import datasets, transforms
from torch import nn, optim
import torch.nn.functional as F

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.1307),(0.3081))
     ])
batch_size = 64
#Data set
train_dataset = torchvision.datasets.MNIST(root='./data',
                                           train=True,
                                           transform=transforms.ToTensor(),
                                           download=True)

test_dataset = torchvision.datasets.MNIST(root='./data',
                                          train=False,
                                          transform=transforms.ToTensor())

# Data loader
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          shuffle=False)
dataiter = iter(train_loader)
images, labels = dataiter.next()

print(images.shape)
print(labels.shape)

class  Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.pool3 = nn.MaxPool2d(2, 2)

        self.fc1 = nn.Linear(128*3*3, 625)
        self.fc2 = nn.Linear(625, 10)

    def forward(self, x):
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool2(F.relu(self.conv2(x)))
        x = self.pool3(F.relu(self.conv3(x)))
        x = x.view(-1, 128*3*3)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x


net = Net().to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

num_epochs = 1000

for epoch in range(num_epochs):
    sum_loss = 0.0
    for i, data in enumerate(train_loader):
        images, labels = data
        images = images.to(device)
        labels = labels.to(device)

        # Forward pass
        outputs = F.softmax(net(images))
        loss = criterion(outputs, labels)

        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # calculate loss
        sum_loss += loss.item()
        if i % 100 == 99:
            print('[%d,%d] loss:%.03f' % (epoch+1, i+1, sum_loss / 100))
            sum_loss = 0.0
net.eval()
correct = 0
total = 0
for data_test in test_loader:
    images, labels = data_test
    images = images.to(device)
    labels = labels.to(device)
    output_test = net(images)
    _, predicted = torch.max(output_test, 1)

    total += labels.size(0)
    correct += (predicted == labels).sum()
print("correct1: ", correct)
print("Test acc: {0}".format(correct.item() / len(test_dataset)))