PyTorch入门-CIFAR10图像分类

CIFAR10数据集下载

CIFAR10数据集包含10个类别，图像尺寸为 3×32×32

官方下载地址很慢，这里给一个百度云：

https://pan.baidu.com/s/1oTvW8wNa-VOjhn0WE5Vmiw 提取码: me8s

下载后在项目目录新建一个data目录解压进去

导入相关包

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import time
import copy

MINI_BATCH = 8      # 数据集的图片数量很大，无法一次性加载所有数据，所以一次加载一个mini-batch的图片
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')     # GPU可用则使用GPU

使用torchvision加载并且归一化训练和测试数据集

CIFAR10数据集的输出是范围在[0,1]之间的PILImage，我们将它转换并归一化范围在[-1,1]之间的Tensor：

# ToTensor(): 将ndarrray格式的图像转换为Tensor张量
# Normalize(mean, std) mean：每个通道颜色平均值，这里的平均值为0.5，私人数据集自己计算；std：每个通道颜色标准偏差，(原始数据 - mean) / std 得到归一化后的数据
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

数据加载器：

# 训练数据加载
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=False, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=MINI_BATCH, shuffle=True, num_workers=4)
# 测试数据加载
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=MINI_BATCH, shuffle=False, num_workers=4)

定义卷积神经网络

我们实现一个简单的神经网络 LeNet-5来进行分类：

这个网络具有两个卷积层，两个池化层，三个全连接层，原网络用于手写数字识别，输入为灰度图，这里我们输入图像是RGB所以修改输入数据为 3×32×32 的Tensorr数据，输出数据维度为 1*10 ，表示图片属于10个类别的概率，图中数据维度变化说明：

• 二维卷积层输出大小 out = (in - F + 2P) / S + 1 ，其中：
  
  F： 卷积核大小 F×F
  
  P： Padding，默认为0
  
  S： 步长Stride，默认为1
  
  如图中第一层卷积层 (32 - 5) / 1 + 1 = 28
• 池化层输出大小 out = (in - F) / S + 1 ，其中：
  
  F： 池化窗口大小 F×F
  
  S： 池化窗口移动的步长Stride，默认和池化窗口维度相同
  
  如图中第二层池化层 (28 - 2) / 2 + 1 = 14

这部分可以写成一个独立的文件，在训练代码中引入此文件中的网络结构：

# net.py
import torch
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)				# 卷积层：3通道到6通道，卷积5*5
        self.conv2 = nn.Conv2d(6, 16, 5)			# 卷积层：6通道到16通道，卷积5*5

        self.pool = nn.MaxPool2d(2, 2)				# 池化层，在2*2窗口上进行下采样

		# 三个全连接层 ：16*5*5 -> 120 -> 84 -> 10
        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 = F.relu(self.conv1(x))        # F.relu 是一个常用的激活函数
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        x = self.pool(x)

        x = x.view(-1, 16 * 5 * 5)			# 变换数据维度为 1*(16*5*5)，-1表示根据后面推测

        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)

        return x

定义一个通用的训练函数，得到最优参数

def train(model, criterion, optimizer, epochs):
    since = time.time()

    best_acc = 0.0      # 记录模型测试时的最高准确率
    best_model_wts = copy.deepcopy(model.state_dict())  # 记录模型测试出的最佳参数

    for epoch in range(epochs):
        print('-' * 30)
        print('Epoch {}/{}'.format(epoch+1, epochs))

        # 训练模型
        running_loss = 0.0
        for i, data in enumerate(trainloader):
            inputs, labels = data
            inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)

            # 前向传播，计算损失
            outputs = net(inputs)
            loss = criterion(outputs, labels)

            # 反向传播+优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

            # 每1000批图片打印训练数据
            if (i != 0) and (i % 1000 == 0):
                print('step: {:d},  loss: {:.3f}'.format(i, running_loss/1000))
                running_loss = 0.0

        # 每个epoch以测试数据的整体准确率为标准测试一下模型
        correct = 0
        total = 0
        with torch.no_grad():
            for data in testloader:
                images, labels = data
                images, labels = images.to(DEVICE), labels.to(DEVICE)

                outputs = net(images)
                _, predicted = torch.max(outputs.data, 1)

                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        acc = correct / total
        if acc > best_acc:      # 当前准确率更高时更新
            best_acc = acc
            best_model_wts = copy.deepcopy(model.state_dict())

    time_elapsed = time.time() - since
    print('-' * 30)
    print('训练用时： {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
    print('最高准确率: {}%'.format(100 * best_acc))

    # 返回测试出的最佳模型
    model.load_state_dict(best_model_wts)
    return model

定义好损失函数和优化器后训练模型

from net import Net

net = Net()
net.to(DEVICE)

# 使用分类交叉熵 Cross-Entropy 作损失函数，动量SGD做优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# 训练10个epoch
net = train(net, criterion, optimizer, 10)
# 保存模型参数
torch.save(net.state_dict(), 'net_dict.pt')

测试模型

# 图像类别
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

net = Net()
net.load_state_dict(torch.load('net_dict.pt'))  # 加载各层参数
net.to(DEVICE)

# 整体正确率
correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(DEVICE), labels.to(DEVICE)

        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('整体准确率: {}%'.format(100 * correct / total))

print('=' * 30)

# 每一个类别的正确率
class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
    for data in testloader:
        images, labels = data
        if torch.cuda.is_available():
            images, labels = images.cuda(), labels.cuda()
        outputs = net(images)
        _, predicted = torch.max(outputs, 1)
        c = (predicted == labels).squeeze()
        for i in range(labels.size(0)):
            label = labels[i]
            class_correct[label] += c[i].item()
            class_total[label] += 1

for i in range(10):
    print('{}的准确率 : {:.2f}%'.format(classes[i], 100 * class_correct[i] / class_total[i]))

模型对测试集图片的一些预测结果

import matplotlib.pyplot as plt
import numpy as np

# 定义一个显示图片的函数
def imshow(img):
    # 输入数据：torch.tensor[c, h, w]
    img = img * 0.5 + 0.5     # 反归一
    npimg = np.transpose(img.numpy(), (1, 2, 0))    # [c, h, w] -> [h, w, c]
    plt.imshow(npimg)
    plt.show()

# 取一批图片
testdata = iter(testloader)
images, labels = testdata.next()
imshow(torchvision.utils.make_grid(images))
print('真实类别: ', ' '.join('{}'.format(classes[labels[j]]) for j in range(labels.size(0))))

# 预测是10个标签的权重，一个类别的权重越大，神经网络越认为它是这个类别，所以输出最高权重的标签。
outputs = net(images)
_, predicted = torch.max(outputs, 1)
print('预测结果: ', ' '.join('{}'.format(classes[predicted[j]]) for j in range(labels.size(0))))