机器学习Pytorch图片识别

基于Pytorchhandbook中的例子做了修改，基于训练的结果，可以输入自己自定义的照片进行识别
基于CIFAR10数据集进行训练，训练结果保存
代码中有详细注释，这里就不一一介绍了

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
import time
import os
from PIL import Image


def showImage(trainloader):
    # 随机获取训练集图片
    dataiter = iter(trainloader)
    images, labels = dataiter.__next__()

    # 展示图片
    imshow(torchvision.utils.make_grid(images))
    # 打印图片类别标签
    print(' '.join('%5s' % classes[labels[j]] for j in range(4)))


# 展示图片的函数
def imshow(img):
    img = img / 2 + 0.5     # 非归一化
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


"""
构建一个卷积神经网络
标准的神经网络计算流程
定义一个多层的神经网络
对数据集的预处理并准备作为网络的输入
将数据输入到网络
计算网络的损失
反向传播，计算梯度
更新网络的梯度，一个简单的更新规则是 weight = weight - learning_rate * gradient
"""
class Net(nn.Module):
    # 定义一个神经网络，下面是一个 5 层的卷积神经网络，包含两层卷积层和三层全连接层：
    def __init__(self):
        super(Net, self).__init__()
        # nn.Conv2d 是 PyTorch 中的二维卷积层。
        # 它有三个参数：输入通道数（3，表示输入图像是三通道RGB图像 3（通道数）x 32（图像高度）x 32（图像宽度）），输出通道数（6，即卷积核的个数），和卷积核的大小（5x5）。
        # CIFAR-10 数据集中的图像大小为 32x32 像素，并且有三个通道（RGB）
        self.conv1 = nn.Conv2d(3, 6, 5)

        # 第二个卷积层，输入通道数为6（来自上一层的输出通道数），输出通道数为16，卷积核大小为5x5。
        self.conv2 = nn.Conv2d(6, 16, 5)

        # 池化层，使用最大池化操作来缩小特征图的大小。
        # nn.MaxPool2d 是 PyTorch 中的二维最大池化层，有两个参数：池化窗口大小（2x2）和步长（2）。
        self.pool = nn.MaxPool2d(2, 2)

        # 全连接层
        # 第一个全连接层，nn.Linear 是 PyTorch 中的线性变换层，将上一层的输出展平为一维向量，然后进行线性变换。
        # 这里输入大小为 16x5x5（来自上一层输出的特征图大小），输出大小为 120
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        # 第二个全连接层，输入大小为120，输出大小为84。
        self.fc2 = nn.Linear(120, 84)
        # 第三个全连接层，输入大小为84，输出大小为10
        # 对于图像分类任务，输出大小通常与类别数相同，这里输出大小为10，表示有10个类别
        self.fc3 = nn.Linear(84, 10)

    # 前向传播方法，也就是定义了数据在神经网络中的流动方式。输入 x 是输入数据，它会经过卷积层、池化层和全连接层的处理，最后输出网络的预测结果。
    def forward(self, x):
        # 将输入 x 输入到第一个卷积层 self.conv1 中，然后通过激活函数 F.relu() 进行激活，再经过池化层 self.pool 进行池化操作。
        x = self.pool(F.relu(self.conv1(x)))

        x = self.pool(F.relu(self.conv2(x)))
        # 将特征图展平为一维向量，以便输入到全连接层中进行线性变换。
        x = x.view(-1, 16 * 5 * 5)
        # 将展平后的特征向量输入到第一个全连接层 self.fc1 中，然后再次通过激活函数 F.relu() 进行激活。
        x = F.relu(self.fc1(x))
        # 将第一全连接层的输出输入到第二个全连接层 self.fc2 中，再次进行激活。
        x = F.relu(self.fc2(x))
        # 将第二个全连接层的输出输入到第三个全连接层 self.fc3 中，得到网络的预测结果。
        x = self.fc3(x)
        return x


def practiceNet():
    # CIFAR-10 数据集中的图像大小为 32x32 像素，并且有三个通道（RGB）
    # torchvision 库中的一个类，用于加载 CIFAR-10 数据集。CIFAR-10 是一个常用的图像分类数据集，
    # 包含 10 个类别的图像（飞机、汽车、猫、狗等），每个类别有 5000 张训练图像和 1000 张测试图像。
    # transform=transform: 这是一个数据预处理的操作，将输入图像转换为 PyTorch 的张量，并进行归一化等处理。
    trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                            download=True, transform=transform)

    # 用于创建数据加载器的类。数据加载器用于将数据划分成小批量并提供数据加载的功能。
    # trainset: 这是通过 torchvision.datasets.CIFAR10 创建的 CIFAR-10 训练集对象。 batch_size每个批次样本数量
    # shuffle 表示在每个 epoch（一个 epoch 表示训练集中所有样本都被遍历一次）之前，对数据进行洗牌（随机打乱样本顺序），从而增加数据的随机性，有助于模型更好地学习。
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                              shuffle=True, num_workers=2)

    # 定义了损失函数。在多分类任务中，通常使用交叉熵损失函数来度量预测结果与真实标签之间的差异
    # 交叉熵损失函数是一个常用的分类损失函数，特别适用于多类别分类问题。对于每个样本，它计算了 模型预测的概率分布 与 真实标签 的差异，并返回一个标量作为损失值。
    criterion = nn.CrossEntropyLoss()

    """
    定义了优化器。在神经网络训练过程中，我们需要使用优化算法来更新网络的权重，以最小化损失函数
    使用随机梯度下降（Stochastic Gradient Descent，SGD）作为优化算法。optim.SGD 是 PyTorch 中提供的 SGD 优化器类
    net.parameters(): 这里传递了神经网络 net 的参数，即网络中需要被优化的 权重 和 偏置 。优化器将根据 损失函数 的 梯度信息 来更新这些参数，从而使网络逐渐学习到更好的参数值。
    lr=0.001: 这是学习率（Learning Rate），它控制了每次参数更新的步长。学习率越大，参数更新越快，但可能导致不稳定和震荡。
    学习率越小，参数更新越慢，但可能使得收敛速度过慢。需要根据具体任务和网络结构来选择合适的学习率
    momentum=0.9: 这是动量（Momentum），它在参数更新中引入了惯性，有助于加快训练速度和避免局部最优。动量的值通常设置为 0.9，但也可以根据实际情况进行调整
    """
    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

    """
    在神经网络训练过程中，我们会使用上述定义的损失函数 criterion 来计算每个批次的损失值，并使用定义的优化器 optimizer 来更新网络的权重，
    从而让网络逐渐学习到最优的参数，以最小化损失函数。
    这样的训练过程通常会经过多个 epochs，每个 epoch 都会遍历整个训练数据集，直到网络达到满意的性能或训练次数达到预定值为止
    """

    # 训练网络
    start = time.time()
    for epoch in range(2):
        running_loss = 0.0
        # 遍历训练数据集的每个 mini-batch。enumerate(trainloader, 0) 会返回一个包含 mini-batch 索引和数据的迭代器
        for i, data in enumerate(trainloader, 0):
            # 获取输入数据
            inputs, labels = data
            # 清空梯度缓存 确保每个 mini-batch 的梯度计算是独立的
            optimizer.zero_grad()
            # 将输入数据输入到神经网络中进行前向传播，得到网络的输出结果。
            outputs = net(inputs)
            # 计算输出结果与真实标签之间的损失（误差），即预测误差
            loss = criterion(outputs, labels)
            # 通过反向传播算法计算损失函数对网络中所有可训练参数（权重和偏置）的梯度（导数）。
            # 这一步是为了计算出梯度信息，用于更新网络的参数。
            loss.backward()
            # 根据计算得到的梯度信息，通过优化器 optimizer 来更新网络的权重和偏置，以最小化损失函数。这一步是训练网络的关键步骤，它使得网络逐渐学习到最优的参数。
            optimizer.step()

            # 打印统计信息
            running_loss += loss.item()
            if i % 2000 == 1999:
                # 每 2000 次迭代打印一次信息 打印每 2000 个 mini-batch 的平均损失值，以便观察训练过程中损失的变化。
                print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
                running_loss = 0.0
    print('Finished Training! Total cost time: ', time.time() - start)


def testNet(net_t):
    testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                           download=True, transform=transform)
    testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                             shuffle=True, num_workers=2)
    # 网络训练完成进行测试
    dataiter = iter(testloader)
    images, labels = dataiter.__next__()
    # 打印图片
    imshow(torchvision.utils.make_grid(images))
    print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))
    # 网络输出
    outputs = net_t(images)
    # 预测结果
    _, predicted = torch.max(outputs, 1)
    print('Predicted: ', ' '.join('%5s' % classes[predicted[j]] for j in range(4)))
    # torch.save(net, PATH)


def test():
    test_image_dir = "/Users/xxx/personal/PyWorkSpace/test4/image"
    transform1 = transforms.Compose([
        transforms.Resize((32, 32))
    ])
    for filename in os.listdir(test_image_dir):
        # 构建完整的图片路径
        test_image_path = os.path.join(test_image_dir, filename)

        test_image = Image.open(test_image_path)

        test_image_tensor = transform(transform1(test_image))

        # 使用已加载的模型进行前向传播
        outputs = net(test_image_tensor)

        # 解释输出结果
        _, predicted = torch.max(outputs, 1)
        predicted_class = classes[predicted.item()]
        print("Predicted class:", predicted_class)


if __name__ == "__main__":
    # 将图片数据从 [0,1] 归一化为 [-1, 1] 的取值范围
    # transform 是一个变换对象，通常使用 transforms.Compose 来将多个预处理操作组合在一起。
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')


    save_dir = "/Users/xxx/personal/PyWorkSpace/test4/practice_result"
    os.makedirs(save_dir, exist_ok=True)
    PATH = os.path.join(save_dir, "model_weights.pth")

    # showImage(trainloader)
    net = Net()
    # net.load_state_dict(torch.load(PATH))
    net = torch.load(PATH)
    # net.eval()
    # practiceNet()
    testNet(net)
    test()
    run_code = 0