pytorch全连接网络识别MNIST数据集

import torch
import torchvision
from torchvision import transforms
from torch import nn
from torch.autograd import Variable
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5],[0.5])
])

trainsets = torchvision.datasets.MNIST(root='./data',train=True,download=False,transform=transform)
trainloader = torch.utils.data.DataLoader(trainsets,batch_size=100,shuffle=True)
testsets = torchvision.datasets.MNIST(root='./data',train=False,download=False,transform=transform)
testloader = torch.utils.data.DataLoader(testsets,batch_size=100,shuffle=False)

class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.linear1 = nn.Linear(784,512)
        self.linear2 = nn.Linear(512,256)
        self.linear3 = nn.Linear(256,128)
        self.linear4 = nn.Linear(128,64)
        self.linear5 = nn.Linear(64,10)
    def forward(self,x):
        x = F.relu(self.linear1(x))
        x = F.relu(self.linear2(x))
        x = F.relu(self.linear3(x))
        x = F.relu(self.linear4(x))
        return self.linear5(x)

if __name__ == '__main__':
    net = Net()
    if torch.cuda.is_available():
        net = net.cuda()
    LOSS = nn.CrossEntropyLoss()
    OPT = torch.optim.Adam(net.parameters(),lr=0.01)
    for i in range(100):
        for data,lable in trainloader:
            if torch.cuda.is_available():
                data = data.cuda()
                lable = lable.cuda()
            else:
                data = Variable(data)
                lable = Variable(lable)
            data = data.view(data.size(0),-1)
            out = net(data)
            loss = LOSS(out,lable)
            OPT.zero_grad()
            loss.backward()
            OPT.step()
    net.eval()
    for data,lable in testloader:
        if torch.cuda.is_available():
            data = data.cuda()
            lable = lable.cuda()
        else:
            data = Variable(data)
            lable = Variable(lable)
        data = data.view(data.size(0),-1)
        out = net(data)
        loss = LOSS(out,lable)
        OPT.zero_grad()
        loss.backward()
        OPT.step()
        print(loss)