pytorch----手写数字体识别

import torch
import matplotlib.pyplot as plt
import os
#data
import numpy as np
import torchvision.datasets as dataset
import torchvision.transforms as transforms
import torch.utils.data as data_utils
train_data = dataset.MNIST(root="mnist",
                           train=True,
                           transform=transforms.ToTensor(),
                           download=True)
test_data = dataset.MNIST(root="mnist",
                           train=False,
                           transform=transforms.ToTensor(),
                           download=False)
#batchsize
train_loader =data_utils.DataLoader(dataset = train_data,batch_size = 64,shuffle=True)
test_loader = data_utils.DataLoader(dataset = test_data,batch_size = 64,shuffle=True)
#net
class CNN(torch.nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv = torch.nn.Sequential(
            torch.nn.Conv2d(1,32,kernel_size=5,padding=2),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2)
        )
        self.fc = torch.nn.Linear(14*14*32,10)
    def forward(self,x):
        out  = self.conv(x)
        out = out.view(out.size()[0],-1)
        out = self.fc(out)
        return out
cnn=CNN()
cnn = cnn.cuda()
#loss
loss_func = torch.nn.CrossEntropyLoss()
#optimizer
optimizer  = torch.optim.Adam(cnn.parameters(),lr=0.01)
#train
for epoch in range(10):
    for i,(images,label)in enumerate(train_loader):
        images = images.cuda()
        labels  =label.cuda()
        outputs = cnn(images)
        loss = loss_func(outputs ,labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print("eporch is {},{}/{},loss is {}".format(epoch+1,i,len(train_data)//64,loss.item()))

    #eval/test
    loss_test = 0
    loss1 = []
    acc1 = []
    acc = 0
    for i,(images,label)in enumerate(test_loader):
        images = images.cuda()
        labels = label.cuda()
        outputs = cnn(images)
        loss_test += loss_func(outputs, labels)
        _,pred = outputs.max(1)
        acc += (pred ==labels).sum().item()
    loss_test = loss_test/(len(test_data)//64)
    acc = acc/len(test_data)
    print("eporch is {},acc is {},loss is {}".format(epoch+1,acc,loss_test))
    



torch.save(cnn,"model/model.pkl")

#load
import torch
import torchvision.datasets as dataset
import torchvision.transforms as transforms
import torch.utils.data as data_utils

class CNN(torch.nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv = torch.nn.Sequential(
            torch.nn.Conv2d(1,32,kernel_size=5,padding=2),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2)
        )
        self.fc = torch.nn.Linear(14*14*32,10)
    def forward(self,x):
        out  = self.conv(x)
        out = out.view(out.size()[0],-1)
        out = self.fc(out)
        return out
test_data = dataset.MNIST(root="mnist",
                           train=False,
                           transform=transforms.ToTensor(),
                           download=False)
#batchsize
test_loader = data_utils.DataLoader(dataset = test_data,batch_size = 64,shuffle=True)
#net
cnn = torch.load('model/model.pkl')
cnn = cnn.cuda()
#eval/test
loss_test = 0
acc = 0
for i,(images,label)in enumerate(test_loader):
    images = images.cuda()
    labels = label.cuda()
    outputs = cnn(images)
    _,pred = outputs.max(1)
    acc += (pred ==labels).sum().item()
acc = acc/len(test_data)
print("acc is {}".format(acc)

)