Pytorch入门实战（3）：使用简单CNN实现物体分类

代码地址

本文涉及知识点

PytorchVision Transforms的基本使用

Pytorch nn.Linear的基本用法

nn.Conv2d官方文档

本文内容

使用自己随便构造的一个简单的CNN网络，使用CIFAR10数据集（包含10种类别，图片大小为32x32）训练一个分类网络。

-----开冲-----

数据预处理

首先导入需要的包：

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

import matplotlib.pyplot as plt
import numpy as np

打印下版本，我这使用的是pytorch-1.10.2

torch.__version__

'1.10.2'

定义transform，表明图片的处理方式和batch_size：

# 如果你的内存不够的话，可以减小batch_size
# 一般batch_size越大，模型越稳定，训练速度越快。（但也不是越大越好）
batch_size = 16

transform = transforms.Compose(
    [transforms.ToTensor(), # 将图片转为Tensor类型
     # 对图片进行正则化。第一个参数为mean(均值)，第二个为std(方差)。每个参数之所以有三个0.5，是因为有RGB三个通道。
     # 综上，这句就是把图片的RGB三个通道都正则化到均值为0.5，方差为0.5的分布上。
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

之后开始准备数据集，这里直接使用官方提供的数据集。如果你的网够快，可以把download=False改为True，或者通过百度网盘下载，然后解压到data目录下

trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=False, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                          shuffle=True)
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                         shuffle=False)

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

接下来随便打印几个看看：

def imshow(img):
	# 因为之前正则化了，所以显示图片前要恢复
    img = img / 2 + 0.5     # unnormalize
    # 转成Numpy.ndarray格式，plt不认tensor
    npimg = img.numpy()
    # plt.imshow接受的图片shape为(h, w, c)
    # 即通道在最后一维，而传过来的img是(c, h, w)
    # 所以要用transpose改一下
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

# 获取一些图片及其对应的标签
dataiter = iter(trainloader)
"""
这里images是一个tensor，shape为(16, 3, 32, 32)
	其中16为batch_size，即16张图片，3为RGB通道，图片大小为32x32
labels也是tensor，shape为(16)，为每个图片的对应的标签。
"""
images, labels = dataiter.next()

"""
make_grid：制作表格。即把多张图片拼到一张中去。
		   nrow=8，表示生成表格的列数。
从下面的输出可以看到，images的16张图片被make_grid
按照2x8的表格拼成了一张大图片。该方法方便人们进行
图像展示。
"""
imshow(torchvision.utils.make_grid(images, nrow=8))
# print labels
print(' '.join(f'{classes[labels[j]]:5s}' for j in range(batch_size)))

truck cat   truck dog   deer  horse plane horse truck horse frog  ship  truck bird  frog  deer

定义CNN分类模型

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        """
      	定义卷积层, nn.Conv2d官方文档：https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html
      	nn.Conv2d包含三个重要参数：
      		in_channels: 输入的通道数
      		out_channels: 输出的通道数
      		kernel_size: 卷积核的大小
      		stride: 步长，默认为1
      		padding: 填充，默认为0，即不进行填充
      		
      	补充：这里卷积层2d的意思是数据是“2维的”，
      		  例如图片就是2维数据（长×宽）。同理也有Conv1d，
      		  是针对于文本、信号等1维数据，也有Conv3d
      		  是针对视频等这种3维数据。
      	"""
        self.classifier = nn.Sequential(
            nn.Conv2d(3, 6, 5),
            # 激活函数
            nn.ReLU(),
            # 使用MaxPool进行下采样。
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            # 当完成卷积后，使用flatten将数据展开
            # 即将tensor的shape从(batch_size, c, h, w)变成(batch_size, c*h*w)，这样才能送给全连接层
            nn.Flatten(),
            # 最后接全连接层。
            # 计算方式可以参考：https://blog.csdn.net/zhaohongfei_358/article/details/123269313
            nn.Linear(16 * 5 * 5, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, 10)
            # 注意这里并没有调用Softmax，也不能调Softmax
            # 这是因为Softmax被包含在了CrossEntropyLoss损失函数中
            # 如果这里调用的话，就会调用两遍，最后网络啥都学不着
        )

    def forward(self, x):
        return self.classifier(x)

net = Net()
# 使用简单的CrossEntorpyLoss作为损失函数，一般多分类问题都用这个
criterion = nn.CrossEntropyLoss()
# 使用简单的SGD作为优化器
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

训练网络

开始训练网络，由于网络较小，这里直接用cpu进行训练

# 把所有训练样本看过一遍称为1个epoch
# 简单起见，这里只训练20个epochs
epochs = 20 

for epoch in range(epochs):
	# 记录一下损失
    running_loss = 0.0
    for i, data in enumerate(trainloader):
		# trainloader返回的是tuple，第一个是图片数，第二个对应的labels
        inputs, labels = data

        # 清除之前的梯度
        optimizer.zero_grad()

        # 进行前向传播
        outputs = net(inputs)
        # 计算损失
        loss = criterion(outputs, labels)
        # 反向传播
        loss.backward()
        # 更新参数
        optimizer.step()

        # 记录损失，每2000次打印一次损失
        running_loss += loss.item()
        if i % 2000 == 1999:   
            print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
            running_loss = 0.0

print('Finished Training')

[1,  2000] loss: 2.162
[2,  2000] loss: 1.603
[3,  2000] loss: 1.397
。。。
[18,  2000] loss: 0.704
[19,  2000] loss: 0.670
[20,  2000] loss: 0.648
Finished Training

在训练了20个epoch后，损失降到了0.648。可以看到损失还在下降，大家可以尝试多训练一段时间。

测试模型

测试下模型总的精准率：

correct = 0 # 记录正确的数量
total = 0 # 记录总数

# torch.no_grad表示不需要计算梯度。
with torch.no_grad():
    for data in testloader:
        images, labels = data
		# 前向传播
        outputs = net(images)
		"""
		outputs.shape为(16, 10)，batch_size为16, 10为类别
		output这16张图片的各个类别的可能性（未经Softmax处理）
		所以通过torch.max找到最大的那个。
		torch.max接受两个参数，第一个是tensor，第二个是dim（维度）
				 这里传1，意思是在类别这个维度上取最大的
		torch.max有两个输出，values和indexes，
				 values就是最大的数是什么，
				 indexes是这些最大的数的index是什么
		这里我们只需要index即可，所以忽略第一个参数
		"""
        _, predicted = torch.max(outputs, 1)
        # 记录总数量
        total += labels.size(0)
        # 计算正确数量
        correct += (predicted == labels).sum().item()

print(f'Accuracy of the network on the 10000 test images: {100 * correct // total} %')

Accuracy of the network on the 10000 test images: 64 %

接下来计算下每个类别精准率：

# 统计每个类别的正确数量和总数量
correct_pred = {classname: 0 for classname in classes}
total_pred = {classname: 0 for classname in classes}

# again no gradients needed
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predictions = torch.max(outputs, 1)
        # collect the correct predictions for each class
        for label, prediction in zip(labels, predictions):
            if label == prediction:
                correct_pred[classes[label]] += 1
            total_pred[classes[label]] += 1


# print accuracy for each class
for classname, correct_count in correct_pred.items():
    accuracy = 100 * float(correct_count) / total_pred[classname]
    print(f'Accuracy for class: {classname:5s} is {accuracy:.1f} %')

Accuracy for class: plane is 70.1 %
Accuracy for class: car   is 73.1 %
Accuracy for class: bird  is 49.7 %
Accuracy for class: cat   is 45.1 %
Accuracy for class: deer  is 57.6 %
Accuracy for class: dog   is 52.3 %
Accuracy for class: frog  is 79.3 %
Accuracy for class: horse is 65.2 %
Accuracy for class: ship  is 81.7 %
Accuracy for class: truck is 71.3 %

参考文献

pytorch官方CNN分类样例: https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html