Pytorch实现循环神经网络（二）、LSTM实现MNIST手写数据集分类

先放上完整的训练测试代码：

# -*- coding: utf-8 -*-
"""
Created on Fri Aug  7 15:10:16 2020

@author: wj
"""
import torch
from torch import nn
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.autograd import Variable
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
# torch.manual_seed(1)    # reproducible


num_epoches = 10
BATCH_SIZE = 128         #批训练的数量
TIME_STEP = 28          # 相当于序列长度
INPUT_SIZE = 28         # 特征向量长度
LR = 0.01               # learning rate

# MNIST数据集下载
train_dataset = datasets.MNIST(
    root='./data', train=True, transform=transforms.ToTensor(), download=True)

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

train_loader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=BATCH_SIZE, shuffle=False)


# 定义网络模型
class RNN(nn.Module):
    def __init__(self):
        super(RNN, self).__init__()
        self.rnn = nn.LSTM(input_size=INPUT_SIZE,  # if use nn.RNN(), it hardly learns
                           hidden_size=64,         # rnn 隐藏单元
                           num_layers=1,           # rnn 层数
                           batch_first=True,       # input & output will has batch size as 1s dimension. e.g. (batch, seq, input_size)
                          )
        self.out = nn.Linear(64, 10)  #10分类

    def forward(self, x):

        r_out, (h_n, h_c) = self.rnn(x, None)   # None represents zero initial hidden state
        # choose r_out at the last time step
        out = self.out(r_out[:, -1, :])
        return out

rnn = RNN() # 实例化
rnn=rnn.cuda()
print(rnn)  # 查看模型结构

optimizer = torch.optim.Adam(rnn.parameters(), lr=LR)   # 选择优化器，optimize all cnn parameters
criterion = nn.CrossEntropyLoss()                       # 定义损失函数，the target label is not one-hotted


for epoch in range(num_epoches):
    print('epoch {}'.format(epoch + 1))
    print('*' * 10)
    running_loss = 0.0
    running_acc = 0.0
    
    rnn.train()
    for imgs, labels in train_loader:    
        imgs = imgs.squeeze(1)#(N,28,28)
        imgs = Variable(imgs.cuda())
        labels = Variable(labels.cuda())
        # 前向传播
        out = rnn(imgs)
        loss = criterion(out, labels)
        running_loss += loss.item() * labels.size(0)
        _, pred = torch.max(out, 1)
        num_correct = (pred == labels).sum()
        running_acc += num_correct.item()
        # 向后传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
       
        
        
    print('Finish {} epoch, Loss: {:.6f}, Acc: {:.6f}'.format(
        epoch + 1, running_loss / (len(train_dataset)), running_acc / (len(train_dataset))))
    
    rnn.eval() 
    eval_loss = 0.0
    eval_acc = 0.0      
    for imgs, labels in test_loader:
        imgs = imgs.squeeze(1)#(N,28,28)
        imgs=imgs.cuda()
        labels=labels.cuda()
        
        out = rnn(imgs)
        loss = criterion(out, labels)
        eval_loss += loss.item() * labels.size(0)
        _, pred = torch.max(out, 1)
        num_correct = (pred == labels).sum()
        eval_acc += num_correct.item()
    print('Test Loss: {:.6f}, Acc: {:.6f}'.format(eval_loss / (len(
        test_dataset)), eval_acc / (len(test_dataset))))

最后放上训练结果：