'''
Author: 365JHWZGo
Description: 15.上一次rnn学习的加工
Date: 2021/10/30 19:45
FilePath: day11-2.py
'''# 导包
import os
import torch
import torchvision
from torch.autograd import Variable
import torch.utils.data as Datatorch.manual_seed(1)# 超参数
BATCH_SIZE = 64         # 一遍中每一次只训练64张图片
EPOCH = 1               #  对2000张照片只训练1遍
TIME_STEP = 28          # 对一张图片进行逐行扫描，一共扫描28行
INPUT_SIZE = 28         # 逐行扫描28个像素点
LR = 0.01               # optimizer的学习效率为0.01
DOWNLOAD_MNIST = False  # 是否要下载MNIST# 判断MNIST数据集是否已经在我的目录里，即是否已经下载好
if not (os.path.exists('../mnist')) or not os.listdir('../mnist'):DOWNLOAD_MNIST = True# 从MNIST中获取train_data
train_data = torchvision.datasets.MNIST(root='../mnist',                                # 在哪里寻找MNISTdownload=DOWNLOAD_MNIST,                        # 是否要下载MNISTtrain=True,                                     # 是否为训练数据transform=torchvision.transforms.ToTensor()     # 将train_data做一些转化【将PIL(python image library)图片转化为某种格式】# 在 MNIST 数据集中的每张图片由 28 x 28 个像素点构成, 每个像素点用一个灰度值表示
)# 建立train_data的数据加载器，方便分批训练
train_loader = Data.DataLoader(dataset=train_data,             # 要加载的数据集num_workers=2,                  # 多线程工作shuffle=True,                   # 是否需要随机打乱数据集batch_size=BATCH_SIZE           # 每次加载的数量
)# 与train_data相同
test_data = torchvision.datasets.MNIST(train=False,root='../mnist'
)# 选取测试数据集中的一部分数据  test_data.test_data->图片数据  test_data.test_labels->标签数据
test_x = test_data.test_data.type(torch.FloatTensor)[:2000] / 255.
# print('test_x',test_x.shape)        # test_x torch.Size([2000, 28, 28])
test_y = test_data.test_labels[:2000]# rnn
class RNN(torch.nn.Module):def __init__(self):super(RNN, self).__init__()             # 继承自torch.nn.Module里的__init__()self.rnn = torch.nn.LSTM(input_size=INPUT_SIZE,              # 输入数据大小，即每次输入一行的28个像素点hidden_size=64,                     # 64个隐藏神经元num_layers=1,                       # 有几层RNN,层数越多耗时越长batch_first=True                    # 是否将batch_size作为第一个维度)self.out = torch.nn.Linear(64, 10)def forward(self, x):r_out, (h_n, h_c) = self.rnn(x, None)   # r_out里边包含了batch_size，time_step，output_sizeout = self.out(r_out[:, -1, :])         # 输出最后一步的隐藏状态return out# 创造实例
rnn = RNN()# 优化器
optimizer = torch.optim.Adam(rnn.parameters(), lr=LR)# 损失函数
loss_func = torch.nn.CrossEntropyLoss()if __name__ == '__main__':# 训练for epoch in range(EPOCH):for step, (x, batch_y) in enumerate(train_loader):# print('x',x.shape)  #x torch.Size([64, 1, 28, 28])    #batch_size,channels,width,height# print('----------------------------------')batch_x = x.view(-1, 28, 28)# print('batch',batch_x.shape)    #batch torch.Size([64, 28, 28])output = rnn(batch_x)loss = loss_func(output, batch_y)optimizer.zero_grad()loss.backward()optimizer.step()if step % 50 == 0:test_out = rnn(test_x)pred_y = torch.max(test_out, 1)[1].data.numpy()# accuracy = sum(pred_y == test_y.data.numpy()) / float(test_y.size(0))accuracy = float(sum((pred_y == test_y.data.numpy()).astype(int))) / float(test_y.size(0))print('Epoch: ', epoch, '| train loss: %.4f' % loss.data.numpy(), '| test accuracy: %.2f' % accuracy)