机器学习Pytorch图片识别_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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214