Pytorch训练分类器进行图像分类

本文依照Pytorch官方教程源码 训练分类器 — PyTorch 教程 2.3.0+cu121 文档 进行解读，并对使用到的API进行解释说明。

数据准备

在这个模型中我们采用CIRFAR10数据集，这个数据集含有十类物体，分别是airplane，automobile，bird，cat，deer，dog，frog，horse，ship和truck，标签值分别为0-9。

引入相关包

import torch
# 工具库 包含数据库，模型，图片操作库，其他操作等
import torchvision
# 图片操作库
import torchvision.transforms as transforms

准备数据集

# 整合数据转换模式
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
​
batch_size = 4
​
# 训练集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
# 用于批量加载数据
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                          shuffle=True, num_workers=2)
​
# 测试集
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                         shuffle=False, num_workers=2)
​
# 数据集的标签
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

API解释

transforms.ToTensor ToTensor — Torchvision 0.18 documentation (pytorch.org)

用作图像处理，将PIL（Python Imaging Library）或者numpy数组形式转为Pytorch张量的格式。

具体来说，它将图像的像素值从[0,255]的范围转换到[0,1]范围，并且改变其数据类型为torch.FloatTensor。

transforms.Normalize Normalize — Torchvision 0.18 documentation (pytorch.org)

给出n个通道的均值和标准差，对图像的每个通道进行归一化操作。这有助于加快模型的收敛速度和提高模型的性能。

tranforms.Compose Compose — Torchvision 0.18 documentation (pytorch.org)

组合多个转换方式。

torch.utils.data.DataLoader

这样方便后续的图片遍历训练。

• batch_size指每次批量处理的个数，即将图片打包，每个小包有size个，示例中4张图片为一组，用于批量处理。

• shuffle 数据集是否打乱顺序加载。

• num_workers 线程数。

查看图片

import matplotlib.pyplot as plt  # 画图库
import numpy as np # 数组库
​
# functions to show an image
​
​
def imshow(img):
    # 增强图片对比度
    img = img / 2 + 0.5     # unnormalize
    # 图片转换为nparrary 
    npimg = img.numpy()
    # 图片显示，此处需要转置，因为pytorch与numpy多维顺序结构问题
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()
​
​
# get some random training images
# 迭代器
dataiter = iter(trainloader)
# 图片数组（4张）  标签数组（4个） 
images, labels = dataiter.next()
​
# show images
# 图片合并显示
imshow(torchvision.utils.make_grid(images))
# print labels
print(' '.join('%5s' % classes[labels[j]] for j in range(4)))

输出每张图片的类别：

make_grid 的主要功能是将一个迷你批量（mini-batch）图像组合成一个单一图像网格。

定义卷积神经网络

import torch.nn as nn
import torch.nn.functional as F
​
​
class Net(nn.Module):
    def __init__(self):
        super().__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 = torch.flatten(x, 1) # flatten all dimensions except batch
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
​
​
net = Net()

torch.nn

torch.nn是Pytorch的核心模块，专门用于构建和训练神经网络。它提供了许多预定义的神经网络层、激活函数、损失函数和其他相关工具。

主要功能和模块：

1. 神经网络层：

   • 线性层：如 torch.nn.Linear，实现全连接层（也称为密集层）。

   • 卷积层：如 torch.nn.Conv2d，实现二维卷积层，常用于图像处理任务。

   • 循环层：如 torch.nn.LSTM 和 torch.nn.GRU，用于处理序列数据。

2. 激活函数：

   常见的激活函数如 torch.nn.ReLU、torch.nn.Sigmoid、torch.nn.Tanh 等。

3. 损失函数：

   提供了多种损失函数，如 torch.nn.CrossEntropyLoss、torch.nn.MSELoss、torch.nn.BCELoss 等，用于评估模型的性能。

4. 容器：

   • torch.nn.Sequential：按顺序将多个层组合在一起，形成一个神经网络。

   • torch.nn.ModuleList 和 torch.nn.ModuleDict：用于存储和管理层的列表和字典。

5. 其他实用功能：

   • torch.nn.Dropout：实现 dropout 正则化，用于减少过拟合。

   • torch.nn.BatchNorm2d：实现批归一化，用于加速训练和稳定模型。

定义损失函数

定义loss函数和优化器。

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

torch.optim 是 PyTorch 中的一个模块，专门用于优化算法的实现。它提供了多种优化器，用于更新神经网络的权重，以最小化损失函数。优化器在深度学习中扮演着重要角色，通过调整模型的参数来加速训练过程和提高模型的性能。

常见的优化器

1. SGD（随机梯度下降）：

   • torch.optim.SGD：基本的随机梯度下降优化器。

   • 支持动量（momentum）和权重衰减（weight decay）。

2. Adam（自适应矩估计）：

   • torch.optim.Adam：常用的优化器，结合了动量和自适应学习率的方法。

   • 支持学习率衰减（learning rate decay）和权重衰减。

3. RMSprop：

   • torch.optim.RMSprop：优化器，适合处理非平稳目标。

4. 其他优化器：

   • torch.optim.Adagrad：自适应梯度算法。

   • torch.optim.AdamW：Adam 的变种，具有更好的权重衰减方法。

在训练数据上训练网络

# 训练网络
for epoch in range(2):
    running_loss = 0.0
    for i, data in enumerate(trainLoader, 0):
        inputs, labels = data
        optimizer.zero_grad()   # 清零梯度
        outputs = net(inputs)   # 前向传播
        loss = criterion(outputs, labels)   # 计算损失
        loss.backward() # 反向传播
        optimizer.step()    # 更新参数
​
        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')

因为我没有修改数据集的获取方式，所以模型训练过程时会有

Files already downloaded and verified

不过没有太大影响。

快速保存经过训练的模型：

PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)

在测试数据上测试网络

dataiter = iter(testloader)
images, labels = next(dataiter)
​
# print images
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join(f'{classes[labels[j]]:5s}' for j in range(4)))

输出类别：

如果下一次使用时，可以这样调用已经保存的模型文件：

net = Net()
net.load_state_dict(torch.load(PATH))

 接下来我们让训练好的神经网络模型预测上述图片的类别：

outputs = net(images)

模型输出的是十个类别的能量。一个类的预测能量值越高，神经网络就认为这个图像越属于哪一类，所以我们得到最高能量的指数，输出预测类别：

_, predicted = torch.max(outputs, 1)

print('Predicted: ', ' '.join(f'{classes[predicted[j]]:5s}'
                              for j in range(4)))

查看结果：

仅仅几张图片还不能说明结婚，我们需要查看网络在整个数据集上的表现：

correct = 0	# 正确预测的图片数量
total = 0	# 图片总数

# since we're not training, we don't need to calculate the gradients for our outputs
with torch.no_grad():
    for data in testloader:
        images, labels = data
        # calculate outputs by running images through the network
        outputs = net(images)
        # the class with the highest energy is what we choose as prediction
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

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

以及在各个标签下的准确率：

# prepare to count predictions for each class
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} %')

完整代码

"""
只是为了借助yolov5的环境训练分类器，随时可以删去
"""

import torch
# 工具库 包含数据库，模型，图片操作库，其他操作等
import torchvision
# 图片操作库
import torchvision.transforms as transforms

import matplotlib.pyplot as plt
import numpy as np

import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

"""
定义卷积神经网络
"""


class Net(nn.Module):
    def __init__(self):
        super().__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 = torch.flatten(x, 1)  # flatten all dimensions except batch
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


"""
准备数据集
"""
# 归一化
transforms = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

batch_size = 4

# 读取cifar10训练集后并进行标准化后下载到./data中
trainSet = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transforms)
# # 批量加载数据集
trainLoader = torch.utils.data.DataLoader(trainSet, batch_size=4,
                                          shuffle=True, num_workers=2)

testSet = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transforms)
testLoader = torch.utils.data.DataLoader(testSet, batch_size=4,
                                         shuffle=False, num_workers=2)

# 数据集标签
classes = ("plane", "car", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck")


def imshow(img):
    # 增强图片对比度
    img = img / 2 + 0.5
    npImg = img.numpy()
    # 图片显示，此处需要转置
    plt.imshow(np.transpose(npImg, (1, 2, 0)))
    plt.show()


if __name__ == "__main__":
    dataiter = iter(trainLoader)
    images, labels = next(dataiter)
    # show image
    imshow(torchvision.utils.make_grid(images))
    # print labels
    print(' '.join('%5s' % classes[labels[i]] for i in range(batch_size)))

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

    # 训练网络
    for epoch in range(2):
        running_loss = 0.0
        for i, data in enumerate(trainLoader, 0):
            inputs, labels = data
            optimizer.zero_grad()
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            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')

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

    # 测试
    dataiter = iter(testLoader)
    images, labels = next(dataiter)

    # print images
    imshow(torchvision.utils.make_grid(images))
    print('GroundTruth: ', ' '.join(f'{classes[labels[j]]:5s}' for j in range(4)))

    # 让神经网络预测上述图片的类别
    outputs = net(images)
    _, predicted = torch.max(outputs, 1)

    print('Predicted: ', ' '.join(f'{classes[predicted[j]]:5s}'
                                  for j in range(4)))

    # 查看在整个测试集上的表现
    correct = 0
    total = 0
    # since we're not training, we don't need to calculate the gradients for our outputs
    with torch.no_grad():
        for data in testLoader:
            images, labels = data
            # calculate outputs by running images through the network
            outputs = net(images)
            # the class with the highest energy is what we choose as prediction
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

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

    # prepare to count predictions for each class
    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} %')