Implementierung handschriftlicher Ziffernerkennung mithilfe des RNN-Netzwerks

import torch
import torchvision
from torchvision import datasets,transforms
from torch.autograd import Variable
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.5],std=[0.5])])
dataset_train = datasets.MNIST(root="./data",transform=transform,train=True,download=True)
dataset_test = datasets.MNIST(root="./data",train=False,transform=transform,download=True)
train_load = DataLoader(dataset=dataset_train,batch_size=64,shuffle=True)
test_load = DataLoader(dataset=dataset_test,batch_size=64,shuffle=False)
images,label = next(iter(train_load))
images_example = torchvision.utils.make_grid(images)
images_example = images_example.numpy().transpose(1,2,0)
mean = [0.5]
std = [0.5]
images_example = images_example * std + mean
plt.imshow(images_example)

class RNN(torch.nn.Module):
    def __init__(self):
        super(RNN,self).__init__()
        self.rnn = torch.nn.RNN(input_size=28,hidden_size=128,num_layers=1,batch_first=True)
        self.output = torch.nn.Linear(128,10)
    def forward(self,input):
        output,_ = self.rnn(input,None)
        output = self.output(output[:,-1,:])
        return output
model = RNN()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

optimizer = torch.optim.Adam(model.parameters())
loss_f = torch.nn.CrossEntropyLoss()
epoch_n = 10
for epoch in range(epoch_n):
    running_loss = 0
    running_correct = 0
    testing_correct = 0
    print("Epoch {}/{}".format(epoch+1,epoch_n))
    print("-"*10)
    for i,(train,label)  in enumerate(train_load):
        train = train.to(device)
        label = label.to(device)
        x_train = train.view(-1,28,28)
        x_train,label = Variable(x_train),Variable(label)
        y_pred = model(x_train)
        loss = loss_f(y_pred,label)
        _,pred = torch.max(y_pred.data,1)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        running_loss += loss
        running_correct += torch.sum(pred==label)
    for i,(test,test_label)  in enumerate(test_load):
        test = test.to(device)
        test_label = test_label.to(device)
        test = test.view(-1,28,28)
        test,test_label = Variable(test),Variable(test_label)
        outputs = model(test)
        _,pred = torch.max(outputs.data,1)
        testing_correct += torch.sum(pred==test_label)
    print("Loss is:{:.4f},Train Accuracy is:{:.4f}%,Test Accuracy is:{:.4f}".format(running_loss/len(dataset_train),
        100*running_correct/len(dataset_train),100*testing_correct/len(dataset_test)))