Pytorch学习基础——LSTM从训练到测试

在上一篇Pytorch学习基础——LeNet从训练到测试讲述了简单神经网络LeNet识别MNIST数据集的实例，作为对比，本次将结合LSTM实现对MNIST数据集的识别。

实现过程：

• 导入必要的包并设置超参数：

import torch
import torchvision
from torch import nn
from torch.autograd import Variable
import torchvision.datasets as dsets
import torchvision.transforms as transforms
import matplotlib.pyplot as plt

#define hyperparameter
EPOCH = 1
BATCH_SIZE = 64
TIME_STEP = 28    #time_step / image_height
INPUT_SIZE = 28    #input_step / image_width
LR = 0.01
DOWNLOAD = True

• 下载并加载MNIST数据集（如果已经下载MNIST数据集，设置DOWMLOAD=False即可）

#get the mnist dataset
train_data = dsets.MNIST(root='./', train=True, transform=torchvision.transforms.ToTensor(), download=False)
test_data = dsets.MNIST(root='./', train=False, transform=torchvision.transforms.ToTensor())
test_x = test_data.test_data.type(torch.FloatTensor)[:2000]/255
test_y = test_data.test_labels.numpy()[:2000]
#use dataloader to batch input dateset
train_loader = torch.utils.data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)

注意与CNN数据加载时的区别，LSTM将图片按行排列作为序列，实现循环神经网络的训练和测试

• 定义并实例化LSTM神经网络：

#define the RNN class
class RNN(nn.Module):
    #overload __init__() method
    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)
        
    #overload forward() method
    def forward(self, x):
        r_out, (h_n, h_c) = self.rnn(x, None)
        out = self.out(r_out[: ,-1, :])
        return out
rnn = RNN()
print(rnn)

• 定义优化器和损失函数

#define optimizer with Adam optim
optimizer = torch.optim.Adam(rnn.parameters(), lr=LR)
#define cross entropy loss function
loss_func = nn.CrossEntropyLoss()

• 训练模型

#training and testing
for epoch in range(EPOCH):
    for step, (b_x, b_y) in enumerate(train_loader):
        #recover x as (batch, time_step, input_size)
        b_x = b_x.view(-1, 28, 28)
        
        output = rnn(b_x)
        loss = loss_func(output, b_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if step % 50 == 0:
            #train with rnn
            test_output = rnn(test_x)
            #loss function
            pred_y = torch.max(test_output, 1)[1].data.numpy()
            #accuracy calculate
            acc = float((pred_y == test_y).astype(int).sum()) / float(test_y.size)
            print('Epoch: ', (epoch), 'train loss: %.3f'%loss.data.numpy(), 'test acc: %.3f'%(acc))

• 测试模型

# print 100 predictions from test data
numTest = 100
test_output = rnn(test_x[:numTest].view(-1, 28, 28))
pred_y = torch.max(test_output, 1)[1].data.numpy()
print(pred_y, 'prediction number')
print(test_y[:numTest], 'real number')
ErrorCount = 0.0
for i in pred_y:
	if pred_y[i] != test_y[i]:
		ErrorCount += 1
print('ErrorRate : %.3f'%(ErrorCount / numTest))

实验结果：

可以看到，LSTM网络既可以用于语音处理，同时可以进行图像分类，此时的“图像”被抽象为按行排列的序列，对于MNIST数据集 的测试表明，LSTM可以在较短时间内实现对数字手势的识别。