lstm实现minist手写字符识别

import torch
from torch import nn
import torchvision.datasets as dsets
import torchvision.transforms as transforms
from torch.autograd import Variable
import matplotlib.pyplot as plt
import numpy as np
torch.manual_seed(1)    # reproducible

# Hyper Parameters
EPOCH = 1           # 训练整批数据多少次, 为了节约时间, 我们只训练一次
BATCH_SIZE = 64
TIME_STEP = 28      # rnn 时间步数 / 图片高度
INPUT_SIZE = 28     # rnn 每步输入值 / 图片每行像素
LR = 0.01           # learning rate
# DOWNLOAD_MNIST = False  # 如果你已经下载好了mnist数据就写上 Fasle
class RNN(nn.Module):
    def __init__(self):
        super(RNN, self).__init__()

        self.rnn = nn.LSTM(     # LSTM 效果要比 nn.RNN() 好多了
            input_size=28,      # 图片每行的数据28像素点
            hidden_size=64,     # rnn hidden unit
            num_layers=1,       # 有几层 RNN layers
            batch_first=True,   # input & output 会是以 batch size 为第一维度的特征集 e.g. (batch, time_step, input_size)
        )

        self.out = nn.Linear(64, 10)    # 输出层

    def forward(self, x):
        # x shape (batch, time_step, input_size)
        # r_out shape (batch, time_step, output_size)
        # h_n shape (n_layers, batch, hidden_size)   LSTM 有两个 hidden states, h_n 是分线, h_c 是主线
        # h_c shape (n_layers, batch, hidden_size)
        r_out, (h_n, h_c) = self.rnn(x, None)   # None 表示 hidden state 会用全0的 state


        # 选取最后一个时间点的 r_out 输出
        # 这里 r_out[:, -1, :] 的值也是 h_n 的值
        out = self.out(r_out[:, -1, :])
        return out
rnn = RNN()
print(rnn)
optimizer = torch.optim.Adam(rnn.parameters(), lr=LR)   # optimize all parameters
loss_func = nn.CrossEntropyLoss()   # the target label is not one-hotted
#以下是读测试数据：+++++++++++++++++++++++++++++++++++++++++++++++++++++
data_file = open("C:/Users/蟋蟀/Desktop/CNN-master/mnist_test.csv", 'r')
wude_data = data_file.readlines()
data_file.close()
ahh_data=[]
ahh_data0=[]
for i in range(2000):
  zhongjian_data = wude_data[i].split(',')
  zhongjian_data0 = zhongjian_data[0]
  zhongjian_data = np.asfarray(zhongjian_data[1:])/255
  zhongjian_data = zhongjian_data.tolist()
  ahh_data.append(zhongjian_data)
  ahh_data0.append(zhongjian_data0)
# print("zhongjian_data:")
# print(ahh_data)
ahh_data = np.array(ahh_data).reshape((2000,1,28,28))
ahh_data0 =np.asfarray(ahh_data0)
# print("after zhongjian_data:")
# print(ahh_data)
ahh_data=torch.Tensor(ahh_data)
print("读测试数据完成!")
data_file = open("C:/Users/蟋蟀/Desktop/CNN-master/mnist_train.csv", 'r')
try_train_data = data_file.readlines()
data_file.close()
train_x=[]
train_y=[]
for i in range(60000):
    zhongjian_data = try_train_data[i].split(',')
    zhongjian_data0 = zhongjian_data[0]
    zhongjian_data = np.asfarray(zhongjian_data[1:]) / 255
    zhongjian_data = zhongjian_data.tolist()
    train_x.append(zhongjian_data)
    train_y.append(zhongjian_data0)
train_x = np.array(train_x).reshape((-1,60,28,28))
train_y =np.asfarray(train_y)
train_y=train_y.reshape((-1,60))
train_x=torch.Tensor(train_x)
train_y=torch.tensor(train_y,dtype=torch.long)
print("读训练数据完成！")
print("进入训练：")
loss_print=[]
for i in range(1000) :
    x=train_x[i].view(-1,28,28)
    y=train_y[i]
    output = rnn(x)
    loss = loss_func(output, y)  # cross entropy loss
    optimizer.zero_grad()  # clear gradients for this training step
    loss.backward()  # backpropagation, compute gradients
    optimizer.step()
    if i%100==0:
       loss_print.append(loss.data)
plt.plot(loss_print)
plt.show()
test_output = rnn(ahh_data[10:20].view(-1, 28, 28))
pred_y = torch.max(test_output, 1)[1].data.numpy().squeeze()
print(pred_y, 'prediction number')
print(ahh_data0[10:20], 'real number')