莫凡Pytorch学习笔记（七）

Pytorch搭建CNN网络对MNIST手写数字数据集进行分类

本篇笔记主要对应于莫凡Pytorch中的4.1和4.2节。主要讲了如何使用Pytorch搭建CNN网络对MNIST手写数字数据集进行分类。

MNIST手写数字数据集

MNIST手写数字数据集是一个经典的数据集，经常被用于各类教程的入门案例。

MINST数据库是由Yann提供的手写数字数据库文件，其官方下载地址http://yann.lecun.com/exdb/mnist/。该数据库的里的图像为28*28像素的灰度图像，每个像素的是一个八位字节（0~255）

在Pytorch中，已经为我们提供了相应的下载和处理接口函数。这里，我们可以使用torchvision.datasets.MNIST进行下载。

相关代码如下：

train_data = torchvision.datasets.MNIST(
    root = './mnist',
    train=True,       # 设置True为训练数据，False为测试数据
    transform = torchvision.transforms.ToTensor(),   # 设置将下载的数据改变为tensor（0-1）的形式，原始数据是numpy array（0-255）的形式
    download=DOWNLOAD_MNIST  # 设置True后就自动下载，下载完成后改为False即可
)

这里的几个参数含义如下。首先root参数设置要存放数据集的根目录。train参数设置True表示下载的是训练数据，False为测试数据。transform表示是否对原始数据进行变换处理，这里设置为torchvision.transforms.ToTensor()表示我们将下载的数据改变为tensor（0-1）的形式，而原始数据是numpy array（0-255）的形式。最后一个download参数表示是否需要从互联网上下载，设置True代表从互联网上下载数据集，并把数据集放在root设置的目录下. 设置False则不重新下载。

这里我们简单展示这个数据集的一些参数。

# plot one example
print(train_data.train_data.size())                 # (60000, 28, 28)
print(train_data.train_labels.size())               # (60000)
plt.imshow(train_data.train_data[0].numpy(), cmap='gray')
plt.title('%i' % train_data.train_labels[0])
plt.show()

得到的输出信息为

torch.Size([60000, 28, 28])
torch.Size([60000])

在此之前，我们需要导入相关的库和设置超参数。

相关的库导入和超参数为

import torch
import torch.utils.data as Data
import torch.nn as nn
import torchvision
import matplotlib.pyplot as plt

LR = 0.01
BATCH_SIZE = 50
EPOCH = 1
DOWNLOAD_MNIST = False #设置True后就自动下载，下载完成后改为False即可

数据划分

我们首先生成训练数据的dataloader

train_loader = Data.DataLoader(
    dataset = train_data,
    batch_size = BATCH_SIZE,
    shuffle = True,
    num_workers = 2
)

然后设置测试数据

test_data = torchvision.datasets.MNIST(
    root = './mnist',
    train=False,       # 设置True为训练数据，False为测试数据
)

这里我们只取前2000个数据进行测试

test_x = (torch.unsqueeze(test_data.test_data, dim=1)).float()[:2000]/255
test_y = test_data.test_labels[:2000]

基本网络搭建

这里我们还是使用继承nn.Module的方式来进行模型搭建。主要实现__init__()和forward()函数。
相关代码如下

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 使用nn.Sequential来搭建卷积层
        self.conv1 = nn.Sequential(
            nn.Conv2d(
                in_channels = 1,      # 输入图像的通道数目
                out_channels = 16,    # 卷积核的高度（即卷积核的个数）
                kernel_size = 5,      # 卷积核的长宽为5*5
                stride = 1,           # 卷积核进行扫描时的跳度（移动的步长）
                padding = 2,          # 边缘的填充，
            ),                             # ->(16, 28, 28)
            nn.ReLU(),
            nn.MaxPool2d(
                kernel_size = 2,           # ->(16, 14, 14)
            ),           
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2),    # ->(32, 14, 14)
            nn.ReLU(),
            nn.MaxPool2d(2),               # ->(32, 7, 7)
        )
        self.out = nn.Linear(32 * 7 * 7, 10)
        
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)           # (batch, 32, 7, 7)
        x = x.view(x.size(0), -1)   # (batch, 32*7*7)  展平
        output = self.out(x)
        return output, x

这里我们使用nn.Sequntial的方式来搭建卷积层，我们的网络结构如下。

CNN(
  (conv1): Sequential(
    (0): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (conv2): Sequential(
    (0): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (out): Linear(in_features=1568, out_features=10, bias=True)
)

其由两个卷积层和一个最后的线性层组成。每个卷积层都使用最大池化。

接着，我们设置优化器和损失函数

optimizer = torch.optim.Adam(cnn.parameters(), lr=LR)
loss_func = nn.CrossEntropyLoss()

最后进行训练与展示

# following function (plot_with_labels) is for visualization, can be ignored if not interested
from matplotlib import cm
try: from sklearn.manifold import TSNE; HAS_SK = True
except: HAS_SK = False; print('Please install sklearn for layer visualization')
def plot_with_labels(lowDWeights, labels):
    plt.cla()
    X, Y = lowDWeights[:, 0], lowDWeights[:, 1]
    for x, y, s in zip(X, Y, labels):
        c = cm.rainbow(int(255 * s / 9)); plt.text(x, y, s, backgroundcolor=c, fontsize=9)
    plt.xlim(X.min(), X.max()); plt.ylim(Y.min(), Y.max()); plt.title('Visualize last layer'); plt.show(); plt.pause(0.01)

plt.ion()

for epoch in range(EPOCH):
    for step, (batch_x, batch_y) in enumerate(train_loader):
        b_x = batch_x
        b_y = batch_y
        
        output = cnn(b_x)[0]
        loss = loss_func(output, b_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if step % 100== 0:
            test_output, last_layer = cnn(test_x)
            pred_y = torch.max(test_output, 1)[1].data.numpy()
            print(type(pred_y), type(test_y))
            accuracy = float((pred_y == test_y.data.numpy()).astype(int).sum()) / float(test_y.size(0))
            print('Epoch: ', epoch, '| train loss: %.4f' % loss.data.numpy(), '| test accuracy: %.2f' % accuracy)
            if HAS_SK:
                # Visualization of trained flatten layer (T-SNE)
                tsne = TSNE(perplexity=30, n_components=2, init='pca', n_iter=5000)
                plot_only = 500
                low_dim_embs = tsne.fit_transform(last_layer.data.numpy()[:plot_only, :])
                labels = test_y.numpy()[:plot_only]
                plot_with_labels(low_dim_embs, labels)
plt.ioff()

输出为：

Epoch:  0 | train loss: 0.0019 | test accuracy: 0.97

最终的分类结果可视化展示如上图所示，可以看到基本各个类别被分开。

参考

1. 莫凡Python：Pytorch动态神经网络，https://mofanpy.com/tutorials/machine-learning/torch/