Pytorch 搭建卷积神经网络CNN和循环神经网络RNN在GPU上预测MNIST数据集

卷积神经网络CNN

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

# 下载MNIST数据集
# 若已有该数据集，需改为DOWNLOAD_MNIST = False
DOWNLOAD_MNIST = True
train_data = torchvision.datasets.MNIST(
    root='./mnist',
    train=True,
    transform=torchvision.transforms.ToTensor(),
    download = DOWNLOAD_MNIST
)
test_data = torchvision.datasets.MNIST(root='./mnist/',train=False)
with torch.no_grad():
    test_x = Variable(torch.unsqueeze(test_data.data, dim=1)).type(torch.FloatTensor).cuda()
test_y = test_data.targets[:10000].cuda()

# 配置批处理
train_loader = Data.DataLoader(
    dataset=train_data,
    batch_size=50,
    shuffle=True,
    num_workers=3
)

# CNN网络
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 第一个卷积层
        self.conv1 = nn.Sequential(
            nn.Conv2d(# 原图(1, 28, 28)
                in_channels=1,
                out_channels=16,
                kernel_size=5,
                stride=1,
                padding=2 # (kernel_size-1)/2
            ),# (1, 28, 28) -> (16, 28, 28)
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2) # (16, 28, 28) -> (16, 14, 14)
        )
        # 第二个卷积层
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )# (16, 14, 14) -> (32, 7, 7)
        self.out = nn.Linear(32 * 7 * 7, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        output = self.out(x)
        return output
cnn = CNN().cuda()

# 配置网络优化器
optimizer = torch.optim.Adam(cnn.parameters(), lr = 0.001)
loss_func = nn.CrossEntropyLoss()

plt.ion()
plt.show()
plt.figure()

steplist = []
losslist = []
accuracylist = []
for epoch in range(1):
    for step, (x, y) in enumerate(train_loader):
        b_x = Variable(x).cuda()
        b_y = Variable(y).cuda()
        output = cnn(b_x)
        loss = loss_func(output, b_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if step % 10 == 0:
            test_output = cnn(test_x[:10000])
            pred_y = torch.max(test_output, 1)[1].cpu().data.numpy().squeeze()
            accuracy = sum(pred_y == test_y.cpu().numpy()) / test_y.size(0)
            print('Epoch:', epoch, ' | train loss:%.4f'%loss.item(), ' | test accuracy:%.2f'%accuracy)
            steplist.append(step)
            losslist.append(loss.item())
            accuracylist.append(accuracy)

            plt.subplot(121)
            plt.title('loss')
            plt.plot(steplist, losslist, 'r-', lw=1)
            plt.subplot(122)
            plt.title('accuracy')
            plt.plot(steplist, accuracylist, 'b-', lw=1)
            plt.pause(0.1)

plt.ioff()
plt.show()

结果：

Epoch: 0  | train loss:1.4696  | test accuracy:0.69
Epoch: 0  | train loss:0.7815  | test accuracy:0.80
Epoch: 0  | train loss:0.5405  | test accuracy:0.79
...
Epoch: 0  | train loss:0.0603  | test accuracy:0.98
Epoch: 0  | train loss:0.0451  | test accuracy:0.98
Epoch: 0  | train loss:0.0594  | test accuracy:0.98

可以看到，CNN对MNIST数据集的预测能达到98%的准确率，这是一个相当好的结果

整个预测的可视化过程：

循环神经网络RNN

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

DOWNLOAD_MNIST = True
train_data = torchvision.datasets.MNIST(
    root='./mnist',
    train=True,
    transform=torchvision.transforms.ToTensor(),
    download = DOWNLOAD_MNIST
)

train_loader = Data.DataLoader(
    dataset=train_data,
    batch_size=64,
    shuffle=True,
)

test_data = torchvision.datasets.MNIST(root='./mnist/',train=False)
with torch.no_grad():
    test_x = (torch.unsqueeze(test_data.data, dim=1).type(torch.FloatTensor)/255.).cuda()
test_y = test_data.targets[:10000].cuda()

class RNN(nn.Module):
    def __init__(self):
        super(RNN, self).__init__()
        self.rnn = nn.LSTM(
            input_size=28,
            hidden_size=64,
            num_layers=1,
            batch_first=True
        )
        self.out = nn.Linear(64, 10)

    def forward(self, x):
        r_out, (h_n, h_c) = self.rnn(x, None)
        out = self.out(r_out[:, -1, :]) # (batch, time_step[-1], input_size)
        return out

rnn = RNN().cuda()

optimizer = torch.optim.Adam(rnn.parameters(), lr = 0.001)
loss_func = nn.CrossEntropyLoss()

plt.ion()
plt.show()
plt.figure()

steplist = []
losslist = []
accuracylist = []
for epoch in range(1):
    for step, (x, y) in enumerate(train_loader):
        b_x = Variable(x.view(-1, 28, 28)).cuda()
        b_y = Variable(y).cuda()
        output = rnn(b_x)
        loss = loss_func(output, b_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if step % 10 == 0:
            test_output = rnn(test_x[:10000].view(-1, 28, 28))
            pred_y = torch.max(test_output, 1)[1].cpu().data.numpy().squeeze()
            accuracy = sum(pred_y == test_y.cpu().numpy()) / test_y.size(0)
            print('Epoch:', epoch, ' | train loss:%.4f'%loss.item(), ' | test accuracy:%.2f'%accuracy)
            steplist.append(step)
            losslist.append(loss.item())
            accuracylist.append(accuracy)

            plt.subplot(121)
            plt.title('loss')
            plt.plot(steplist, losslist, 'r-', lw=1)
            plt.subplot(122)
            plt.title('accuracy')
            plt.plot(steplist, accuracylist, 'b-', lw=1)
            plt.pause(0.1)

plt.ioff()
plt.show()

test_output = rnn(test_x[:10])
pred_y = torch.max(test_output, 1)[1].cpu().data.numpy().squeeze()
print(pred_y, 'prediction number')
print(test_y[:10].cpu().numpy(), 'real number')

结果：

Epoch: 0  | train loss:2.2996  | test accuracy:0.13
Epoch: 0  | train loss:2.2977  | test accuracy:0.15
Epoch: 0  | train loss:2.2768  | test accuracy:0.21
...
Epoch: 0  | train loss:0.1875  | test accuracy:0.94
Epoch: 0  | train loss:0.2653  | test accuracy:0.94
Epoch: 0  | train loss:0.3055  | test accuracy:0.94