识别cifar数据集

识别cifar数据集

import torchvision as tv
import torch as t
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
from torch.autograd import Variable

'''加载数据'''
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
trainset = tv.datasets.CIFAR10(
    root='D:/pytorch_study/data',
    train=True,
    download=False,
    transform=transform
)
trainloader = t.utils.data.DataLoader(
    trainset,
    batch_size=10,
    shuffle=True,
    num_workers=2

)
testset = tv.datasets.CIFAR10(
    root='D:/pytorch_study/data/',
    train=False,
    download=False,
    transform=transform
)
testloader = t.utils.data.DataLoader(
    testset,
    batch_size=10,
    shuffle=False,
    num_workers=2)

classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship',
           'truck')


# 展示并保存图片
def imshow(inputs):
    plt.ion()
    inputs = inputs / 2 + 0.5  # 加载数据集的时候使用了正则化，需要恢复一下
    inputs = inputs.numpy().transpose((1, 2, 0))  # pytorch加载图片通道数在前，我们展示图片图片的时候通道数在后
    plt.imshow(inputs)
    plt.pause(1)
    plt.close()

'''定义神经网络'''
class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.con1 = nn.Conv2d(3, 6, 5)
        self.con2 = nn.Conv2d(6, 16, 5)
        self.lin1 = nn.Linear(16 * 5 * 5, 120)
        self.lin2 = nn.Linear(120, 84)
        self.lin3 = nn.Linear(84, 10)


    def forward(self,x):
        x = F.max_pool2d(F.relu(self.con1(x)),(2,2))
        x = F.max_pool2d(F.relu(self.con2(x)),2)
        x=x.view(x.size()[0],-1)
        x = F.relu(self.lin1(x))
        x = F.relu(self.lin2(x))
        x = self.lin3(x)
        return x


net=Net()
'''定义损失函数和优化器'''
criterion=nn.CrossEntropyLoss()
optimizer=optim.SGD(net.parameters(),lr=0.001,momentum=0.9)


'''训练函数'''
def tain(model):
    for i,data in enumerate(trainloader,0):
        inputs,labels=data
        inputs, labels = Variable(inputs), Variable(labels)
        #梯度清0
        optimizer.zero_grad()
        #forward+backward
        outputs = model(inputs)
        #计算机损失
        loss = criterion(outputs, labels)
        #反向传播
        loss.backward()
        #更新参数
        optimizer.step()
    if i % 2000 == 1999:
        print(loss)

#提取训练数据
def restore_net():
    net2=t.load('cifar.plk')
    return net2

#测试函数
def test(modul):
    correct = 0
    total = 0
    for data in testloader:
        images, labels = data
        outputs = modul(Variable(images))
        _,predictsd = t.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predictsd == labels).sum()
    print(correct, total)
    accuracy=t.true_divide(correct,total)
    print(accuracy)

def test2(module):
    dataiter = iter(testloader)
    img, lables = dataiter.next()
    imshow(tv.utils.make_grid(img))
    print('the true label:', ' '.join( \
        '%08s' % classes[lables[j]] for j in range(0, 10)))
    outputs = module(Variable(img))
    _, predictsd = t.max(outputs.data, 1)
    print('the pass label:', '  '.join('%5s' \
                             % classes[predictsd[j]] for j in range(0, 10)))


if __name__ == '__main__':
    '''
    for i in range(10):
        tain(net)
    t.save(net, 'cifar.plk')
    '''
    net2=restore_net()
    test(net2)

训练10遍后的结果