pytorch CNN识别MNIST数据集

import torch
import torchvision
from torch import nn
from torch import optim
import torch.nn.functional as F
from torch.autograd import Variable
from torchvision import transforms

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=2,shuffle=True)
testsets = torchvision.datasets.MNIST(root='./data',train=False,download=False,transform=transform)
testloader = torch.utils.data.DataLoader(testsets,batch_size=2,shuffle=False)

class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1,16,3),
            nn.BatchNorm2d(16),
            nn.ReLU(inplace=True)
        )
        self.layer2 = nn.MaxPool2d(2,2)
        self.layer3 = nn.Sequential(
            nn.Conv2d(16,32,3),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True)
        )
        self.layer4 = nn.MaxPool2d(2,2)
        self.fc = nn.Sequential(
            nn.Linear(32*5*5,512),
            nn.ReLU(inplace=True),
            nn.Linear(512,128),
            nn.ReLU(inplace=True),
            nn.Linear(128,10)
        )
    def forward(self,x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = x.view(x.size(0),-1)
        return self.fc(x)
if __name__ == '__main__':
    net = Net()
    if torch.cuda.is_available():
        net = net.cuda()
    LOSS = nn.CrossEntropyLoss()
    OPT = 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(-1,1,28,28)
            out = net(data)
            loss = LOSS(out,lable)
            OPT.zero_grad()
            loss.backward()
            OPT.step()
            print(loss.data.item())
    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)