Pytorch之CIFAR10

CIFAR10

该数据集共有60000张彩色图像，每张图像是32x32x3，分为10个类，第类6000张图。

其中50000张用于训练，构成5个训练批，每一批10000张图；另外10000用于测试，单独构成一批。测试批的数据 里，取自10类中的每一类，每一类随机取1000张。

10类分别是：1.airplane, 2.automobile, 3.bird, 4.cat, 5.deer, 6.dog, 7.frog, 8.horse, 9.ship, 10.truck

一、导入包库

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms, datasets
import torch.optim as optim
from torch.utils.data.dataloader import DataLoader  # 读取数据

二、数据预处理与超参数设置

# 设置transforms,数据转换
transform = transforms.Compose([
    transforms.ToTensor(),  # numpy --> Tensor
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 归一化，范围[-1, 1]
])

# 下载数据集
# 训练集
train_set = datasets.CIFAR10(root='./CIFAR10', train=True, download=True, transform=transform)
# 测试集
test_set = datasets.CIFAR10(root='./CIFAR10', train=False, download=True, transform=transform)
# 超参数
BATCH_SIZE = 128
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
# num_workers:线程，设置成最大线程;pin_memory:直接加载到CPU或GPU中。不设置这两个属性的话会导致CPU或GPU的利用不起来
train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True, num_workers=16, pin_memory=True)
test_loader = DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=True, num_workers=16, pin_memory=True)

三、定义网络模型

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)  # 输入32x32x3,输出28x28x6,计算(32-5)/1+1=28
        self.pool = nn.MaxPool2d(2, 2)  # input:3,output:6,kernel:5,输出：14x14x6
        self.conv2 = nn.Conv2d(6, 16, 5)  # input:6,putput:16,kernel:5,输出:10x10x16
        self.fc1 = nn.Linear(400, 120)  # input:400(16x5x5),output:120
        self.fc2 = nn.Linear(120, 84)  # input:120,output:84
        self.fc3 = nn.Linear(84, 10)  # input:84,output:10

    def forward(self, x):
        # 32x32x3 --> 28x28x6 --> 14x14x6
        x = self.pool(F.relu(self.conv1(x)))
        # 14x14x6 --> 10x10x16 --> 5x5x16
        x = self.pool(F.relu(self.conv2(x)))
        # 改变shape
        x = x.view(-1, 400)  # 400 = 16x5x5
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

四、模型创建

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

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

# 定义轮数
EPOCHS = 200

for epoch in range(EPOCHS):
    train_loss = 0.0
    for i, (datas, labels) in enumerate(train_loader):
        datas, labels = datas.to('cuda'), labels.to('cuda')
        # 梯度置零
        optimizer.zero_gard()
        # 训练
        outputs = net(datas)
        # 计算损失
        loss = criterion(outputs, labels)
        # 反向传播
        loss.backward()
        # 参数更新
        optimizer.step()
        # 累计损失
        train_loss += loss.item()
    print('Epoch :% , Batch : %5d , Loss : %.3f'%(epoch+1, i+1, train_loss/len(train_loader.dataset)))

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

五、利用模型进行测试

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

# 测试
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))

六、其他(直接加入整个代码即可)

1. 输出图片

# 10个类别
import matplotlib.pyplot as plt
import numpy as np
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

def img_show(img):
    img = img / 2 + 0.5		# 逆正则化
    np_img = img.numpy()	# 数据转换：tensor --> numpy
    plt.img_show(np.transpose(np_img, (1, 2, 0)))	# 改变通道顺序
    plt.show()

# 随机获取一批数据
imgs, labs = next(iter(train_loader))
print(imgs.shape)
print(labs.shape)

# 调用方法
img_show(torchvision.utils.make_grid(imgs))		# .make_grid() 以网格的形式展示图片

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

2. 打印每一类的正确率

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

with torch.no_grad():
    for (images, labels) in test_loader:
        # 输出
        outputs = model(images)
        # 获取到每一行最大值的索引
        _, predicted = torch.max(outputs, dim=1)
        c = (predicted == labels).squeeze()     # squeeze() 去掉0维[默认]， unsqueeze() 增加一维
        if labels.shape[0] == 128:
            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   # 该类总的个数

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