CIFAR10图像识别

CIFAR10图像识别

# 1.导入包
import torch
import torch.nn as nn
import torchvision
from torchvision import transforms, datasets

# 2.设置transforers
transform = transforms.Compose([
    transforms.ToTensor(),  # numpy-----tensor
    transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))
])

# 3.下载数据集
# 3.1训练集
trainset = datasets.CIFAR10(root = './CIFAR10', train = True, download = True, transform = transform)

# 3.2测试集
testset = datasets.CIFAR10(root = '.CIFAR', train = False, download = True, transform = transform)

# 4.批量读取数据
from torch.utils.data.dataloader import DataLoader
BATCH_SIZE = 4
# train_loader = DataLoader(trainset, batch_size = BATCH_SIZE, shuffle = True, num_workers =  1, pin_memory = True)
#
# test_loader = DataLoader(testset, batch_size = BATCH_SIZE, shuffle = True, num_workers = 1, pin_memory = True)

train_loader = DataLoader(trainset, batch_size = BATCH_SIZE, shuffle = True, pin_memory = True)

test_loader = DataLoader(testset, batch_size = BATCH_SIZE, shuffle = True, pin_memory = True)
# 5.可视化显示
import matplotlib.pyplot as plt
import numpy as np

# 6. 10个类别
classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

def imshow(img):
    img = img / 2 + 0.5
    np_img = img.numpy()  # tensor----numpy
    plt.imshow(np.transpose(np_img, (1,2,0)))  # 改变通道顺序
    plt.show()



# 7.随机获取一批数据
imgs,labs = next(iter(train_loader))  # 一批数据，也就是BATCH_SIZE张图片存放在一维的imgs中
print(imgs.shape)
print(labs.shape)

# 8.调用方法
# make_grid把imgs中的128张图片，合成一张图片显示，并且以坐标的形式显示
imshow(torchvision.utils.make_grid(imgs))


# 9.输出这批图片对应的标签
print(' '.join(' %4s' % classes[labs[i]] for i in range(BATCH_SIZE)))
# print('\nlabs')
# print(' %5s'.format(labs[i] for i in range(BATCH_SIZE)))


# 10.定义网络模型
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16*5*5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


# 11.创建模型
net = Net().to('cpu')


# 12.定义优化器和损失函数
import torch.optim as optim
criterion = nn.CrossEntropyLoss()  # 交叉式损失函数
optimizer = optim.SGD(net.parameters(), lr = 0.001, momentum = 0.9)  # 优化器



# 13.定义轮数
EPOCHS = 1
for epoch in range(EPOCHS):
    train_loss = 0.0
    for i, (datas, labels) in enumerate(train_loader):
        datas, labels = datas.to('cpu'), labels.to('cpu')
        optimizer.zero_grad()
        outputs = net(datas)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()
    print('Epoch : %d, Batch: %5d, Loss: %.3f'%
          (epoch + 1, i + 1, train_loss/len(train_loader.dataset)))


# 14.保存模型
PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)

# 15.加载模型
model = Net()
model.load_state_dict(torch.load(PATH))



# 16.测试
correct = 0
total = 0
flag = True
with torch.no_grad():
    for i, (datas, labels) in enumerate(test_loader):
        outputs = model(datas)  # outputs.data.shape---torch.size([128, 10])
        _, predicted = torch.max(outputs.data, dim = 1)  # 第一个是值的张量，第二个是序号的张量
        # 比较有多少个预测正确
        total += labels.size(0)  # labels.size()---torch.size([128])----labels.size(0)---128
        correct += (predicted == labels).sum()  # 相同为1， 不同为0，利用sum（）求和
    print('在10000张测试集图片上的准确率为： {:.3f}%'.format(correct / total *100))


# 17.显示每一类预测的概率
class_correct = list(0. for i in range(10))
total = list(0. for i in range(10))

with torch.no_grad():
    for (image, labels) in test_loader:
        outputs = model(datas)
        _, predicted = torch.max(outputs, dim = 1)  # 获取到每一行最大值的索引
        c = (predicted == labels).squeeze()  # 去掉0维（默认），unsqueeze（）增加一维
        if labels.shape[0] == 4:
            for i in range(BATCH_SIZE):
                label = labels[i]  # 获取每一个label
                class_correct[label] += c[i].item()  # 累计为True的个数，注意：1+True = 2， 1 + false = 1
                total[label] += 1  # 该类总的个数

# 18.输出正确率
for i in range(10):
    print('正确率{：%5s : %2d %%}'% (classes[i], 100 * class_correct[i] / total[i]))