PyTorch入门------训练图像分类器

前言

        1. 操作步骤

        2. 数据集

一、公共部分

        1.加载并归一化 CIFAR10

        2.定义卷积神经网络

二、训练、保存模型参数部分 train_and_save.py

        3.定义损失函数和优化器

        4.训练网络(使用 CPU 或者 GPU)

        5.保存训练好的模型参数

三、加载模型参数、模型推理部分 load_and_infer.py

        6.加载模型参数

        7.在测试数据集上使用加载的模型进行推理

四、代码及运行结果

        1.train_and_save.py

                1.1 代码

                1.2运行结果

CPU

GPU

        2.load_and_infer.py

                2.1 代码

                2.2 运行结果

---

前言

        参考：训练分类器 — PyTorch 教程 2.7.0+cu126 文档 - PyTorch 深度学习库

        1. 操作步骤

        我们将按顺序执行以下步骤训练图像分类器

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

2. 定义卷积神经网络

3. 定义损失函数

4. 在训练数据上训练网络

5. 保持训练好的模型参数

6. 加载模型参数

7. 在测试数据上测试网络

        2. 数据集

        使用 CIFAR10 数据集。它具有以下类别：“airplane”、“automobile”、“bird”、“cat”、“deer”、“dog”、“frog”、“horse”、“ship”、“truck”。CIFAR-10 中的图像大小为 3x32x32，即 3 通道彩色图像，大小为 32x32 像素。

一、公共部分

        1.加载并归一化 CIFAR10

#----- 1.加载并归一化 CIFAR10 -----#
import torch
import torchvision
import torchvision.transforms as transforms# transforms.Compose(): 组合变换操作的容器, 允许将多个图像转换操作按顺序执行
transform = transforms.Compose(# 将范围在 [0, 255] 的 PIL Image 或 numpy.ndarray (H x W x C) 转换为形状为 (C x H x W)、范围在 [0.0, 1.0] 的 torch.FloatTensor[transforms.ToTensor(),  # 使用均值和标准差归一化张量图像，将每个通道的值标准化为 [-1, 1] 区间。 output[channel] = (input[channel] - mean[channel]) / std[channel] = 2x - 1transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]  
)# 设置批量大小
batch_size = 4# 加载CIFAR-10训练数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,download=True, transform=transform)
# 创建训练数据的DataLoader，按批次加载数据并进行随机打乱
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,shuffle=True, num_workers=2)# 加载CIFAR-10测试数据集
testset = torchvision.datasets.CIFAR10(root='./data', train=False,download=True, transform=transform)
# 创建测试数据的DataLoader，按批次加载数据但不打乱
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,shuffle=False, num_workers=2)# CIFAR-10数据集的分类标签，包含10类
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')# 输出训练集和测试集的大小
train_size = len(trainset)
test_size = len(testset)
print(f'Training set size: {train_size}') # 50000
print(f'Test set size: {test_size}')      # 10000# 展示训练集图像
import matplotlib.pyplot as plt
import numpy as npdef imshow(img, title=None):"""显示图像，并且可以选择性地显示标题。参数：img (Tensor): 要显示的图像。通常是经过预处理后的Tensor图像, 形状为 (C, H, W)。title (str, optional): 要显示的标题。如果提供了标题，则在图像上方显示它。"""img = img / 2 + 0.5     # unnormalize，反标准化，将图像的像素值从[-1, 1]映射回[0, 1]npimg = img.numpy()     # 将PyTorch tensor转换为NumPy数组plt.imshow(np.transpose(npimg, (1, 2, 0)))  # 将tensor的维度从 (C, H, W) 转换为 (H, W, C) 以便绘制# 如果提供了标题，就显示标题if title:plt.title(title)plt.show()# 从训练数据加载器中获取一个批次的数据 (随机的)
print("----------随机展示训练集图像----------")
dataiter = iter(trainloader)     # 创建数据加载器的迭代器
images, labels = next(dataiter)  # 获取一批数据，其中images是图像数据，labels是对应的标签# 创建标题，显示标签名称。5s 表示字符串的宽度为 5，不足的部分会使用空格填充
title = ' '.join(f'{classes[labels[j]]:5s}' for j in range(batch_size)) # join() 方法将生成的类别名称列表（每个类别名称都经过格式化）用空格连接成一个单一的字符串。# 显示图像并且显示标题
imshow(torchvision.utils.make_grid(images), title=title) # 使用make_grid将图像组合成一个网格并显示

        2.定义卷积神经网络

#----- 2.定义卷积神经网络 -----#
import torch.nn as nn
import torch.nn.functional as Fclass 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) # torch.flatten(input, start_dim=0, end_dim=-1) ：将输入张量从 start_dim 到 end_dim 维度展平为一个维度。x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))x = self.fc3(x)return xnet = Net()
print("----------模型结构----------")
print(net)

二、训练、保存模型参数部分 train_and_save.py

        3.定义损失函数和优化器

#----- 3.定义损失函数和优化器 -----#
import torch.optim as optim'''CrossEntropyLoss 通常用于分类问题，计算目标类别与模型预测类别之间的误差。交叉熵(Cross-Entropy)本质上是两个概率分布之间的差异度量，常用于比较：一个真实分布P(通常是 one-hot 编码）一个预测分布Q(通常是 softmax 输出）它的数学定义如下：CCrossEntropy(P,Q)=- ∑ [ Pi*log(Qi) ]i=1其中：C 是类别总数Pi是真实标签中第 i 类的概率（一般为 1 或 0)Qi是模型预测中第 i 类的概率(softmax 输出）交叉熵损失只对正确类别的预测概率进行惩罚，其它类别不参与损失。'''
# 定义损失函数（这里使用的是分类交叉熵损失函数（CrossEntropyLoss））
criterion = nn.CrossEntropyLoss()'''optim.SGD 是 PyTorch 提供的随机梯度下降(SGD)优化器。参数解释：net.parameters()：模型参数，即网络中的所有可训练参数。lr=0.001: 学习率，控制每次参数更新的步长。较小的学习率有助于防止过拟合，但可能导致训练变慢。momentum=0.9: 动量项，用来加速收敛，并帮助跳出局部最小值。动量类似于物理中的惯性，它使用上一次的梯度来调整当前的梯度。
'''
# 定义优化器（这里使用的是带动量的随机梯度下降（SGD with momentum））
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

        4.训练网络(使用 CPU 或者 GPU)

#----- 4.训练网络(使用 CPU 或者 GPU) -----#
import time
use_gpu = False  # 改成 False 就强制使用 CPUif not use_gpu:start_time = time.time()  # 记录开始时间print("----------模型训练: CPU----------")for epoch in range(2):  # 遍历整个数据集多次（训练2个 epoch）running_loss = 0.0  # 初始化每2000个小批次（即处理8000张图片）的累计损失for i, data in enumerate(trainloader, 0):  # 遍历训练数据集，每次获取一个小批次（mini-batch）# 获取输入数据；data 是一个包含 [inputs, labels] 的列表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:    # 每2000个小批次（即处理8000张图片）打印一次信息print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')running_loss = 0.0  # 重置累计损失print('CPU Finished Training')  # 打印训练结束的提示end_time = time.time()  # 记录结束时间elapsed_time = end_time - start_time  # 计算消耗的时间print(f"CPU time: {elapsed_time:.6f} 秒")else:start_time = time.time()  # 记录开始时间print("----------模型训练: GPU----------")# 确保在 GPU 上训练，如果有 CUDA 可用device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')print("device", device)# ①将神经网络 net 的所有参数和缓冲区（buffers）转移到 device 上net.to(device)for epoch in range(2):  # 遍历整个数据集多次（训练2个 epoch）running_loss = 0.0  # 初始化每2000个小批次（即处理8000张图片）的累计损失for i, data in enumerate(trainloader, 0):  # 遍历训练数据集，每次获取一个小批次（mini-batch）# 获取输入数据；data 是一个包含 [inputs, labels] 的列表inputs, labels = data# ②将输入数据（即图像）和标签（即真实标签）转移到 device 上inputs, labels = inputs.to(device), labels.to(device)# 将优化器的梯度清零，避免累积梯度optimizer.zero_grad()# 前向传播 + 反向传播 + 更新优化器outputs = net(inputs)              # 通过网络进行前向传播，得到预测结果loss = criterion(outputs, labels)  # 计算损失函数的值loss.backward()   # 计算梯度（反向传播）optimizer.step()  # 使用优化器更新模型的参数# 打印统计信息running_loss += loss.item()  # 将当前损失加到累计损失中if i % 2000 == 1999:    # 每2000个小批次（即处理8000张图片）打印一次信息print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')running_loss = 0.0  # 重置累计损失print('GPU Finished Training')  # 打印训练结束的提示end_time = time.time()  # 记录结束时间elapsed_time = end_time - start_time  # 计算消耗的时间print(f"GPU time: {elapsed_time:.6f} 秒")"""
模型训练: multi GPU
数据并行教程: https://pytorch.ac.cn/tutorials/beginner/blitz/data_parallel_tutorial.html
"""

        5.保存训练好的模型参数

#----- 5.保存训练好的模型参数 -----#
import osprint("----------保存模型参数----------")
PATH = './model/cifar_net.pth'
os.makedirs('model', exist_ok=True) # 创建一个名为 'model' 的文件夹，如果文件夹已存在则不会报错。
torch.save(net.state_dict(), PATH)

三、加载模型参数、模型推理部分 load_and_infer.py

        6.加载模型参数

#----- 6.加载模型参数 -----#
print("----------加载模型参数----------")
PATH = './model/cifar_net.pth'
net = Net()
net.load_state_dict(torch.load(PATH, weights_only=True))

        7.在测试数据集上使用加载的模型进行推理

#----- 7.在测试数据集上使用加载的模型进行推理 -----#
print("----------加载4张测试集图像----------")
# 展示测试集图像
dataiter = iter(testloader)
images, labels = next(dataiter)# 使用训练好的模型进行推理
print("----------模型推理----------")
outputs = net(images)
_, predicted = torch.max(outputs, 1) # 1 是指按行（类别维度）选取最大值。函数返回每一行的最大值，及其对应的索引（类别）print('Predicted: ', ' '.join(f'{classes[predicted[j]]:5s}'for j in range(4)))# show images
title = ' '.join(f'{classes[labels[j]]:5s}' for j in range(4))
imshow(torchvision.utils.make_grid(images), title=title)print("----------网络在整个测试集上的准确率----------")
correct = 0  # 用于记录预测正确的图像数量
total = 0    # 用于记录测试集中图像的总数with torch.no_grad():  # 禁用梯度计算，减少内存消耗并加速推理过程for data in testloader:    # 遍历测试数据集images, labels = data  # 获取当前小批次的图像和标签# # test# print("标签(labels)为：", labels)      # 示例：tensor([6, 0, 5, 7])# print("labels shape: ", labels.shape) # labels shape:  torch.Size([4])# 使用网络进行推理，计算输出outputs = net(images)# 选择输出中最大值的类别作为预测结果_, predicted = torch.max(outputs, 1)  # 1 是指按行（类别维度）选取最大值。函数返回每一行的最大值，及其对应的索引（类别）total += labels.size(0)  # 累加测试集中的样本数量correct += (predicted == labels).sum().item()  # 统计预测正确的样本数量# 打印测试集上模型的准确率
print(f'Accuracy of the network on the 10000 test images: {100 * correct // total} %')print("----------网络在每个类别上的准确率----------")
# 初始化字典
correct_pred = {classname: 0 for classname in classes}  # 初始化每个类别的正确预测数
total_pred = {classname: 0 for classname in classes}    # 初始化每个类别的总预测数# 由于不需要计算梯度，使用 no_grad 来加速推理过程
with torch.no_grad():  for data in testloader:  # 遍历测试数据集images, labels = data  # 获取当前批次的图像和标签outputs = net(images)  # 使用网络进行前向传播，得到预测输出_, predictions = torch.max(outputs, 1)  # 得到每个图像的预测类别（索引）# 收集每个类别的正确预测数量for label, prediction in zip(labels, predictions):  # 同时遍历标签和预测值total_pred[classes[label]] += 1  # 增加对应类别的总预测数if label == prediction:  # 如果标签与预测相同，表示预测正确correct_pred[classes[label]] += 1  # 增加对应类别的正确预测# 打印每个类别的准确率
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} %')  # 打印每个类别的准确率

四、代码及运行结果

        1.train_and_save.py

                1.1 代码

#----- 1.加载并归一化 CIFAR10 -----#
import torch
import torchvision
import torchvision.transforms as transforms# transforms.Compose(): 组合变换操作的容器, 允许将多个图像转换操作按顺序执行
transform = transforms.Compose(# 将范围在 [0, 255] 的 PIL Image 或 numpy.ndarray (H x W x C) 转换为形状为 (C x H x W)、范围在 [0.0, 1.0] 的 torch.FloatTensor[transforms.ToTensor(),  # 使用均值和标准差归一化张量图像，将每个通道的值标准化为 [-1, 1] 区间。 output[channel] = (input[channel] - mean[channel]) / std[channel] = 2x - 1transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]  
)# 设置批量大小
batch_size = 4# 加载CIFAR-10训练数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,download=True, transform=transform)
# 创建训练数据的DataLoader，按批次加载数据并进行随机打乱
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,shuffle=True, num_workers=2)# 加载CIFAR-10测试数据集
testset = torchvision.datasets.CIFAR10(root='./data', train=False,download=True, transform=transform)
# 创建测试数据的DataLoader，按批次加载数据但不打乱
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,shuffle=False, num_workers=2)# CIFAR-10数据集的分类标签，包含10类
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')# 输出训练集和测试集的大小
train_size = len(trainset)
test_size = len(testset)
print(f'Training set size: {train_size}') # 50000
print(f'Test set size: {test_size}')      # 10000# 展示训练集图像
import matplotlib.pyplot as plt
import numpy as npdef imshow(img, title=None):"""显示图像，并且可以选择性地显示标题。参数：img (Tensor): 要显示的图像。通常是经过预处理后的Tensor图像, 形状为 (C, H, W)。title (str, optional): 要显示的标题。如果提供了标题，则在图像上方显示它。"""img = img / 2 + 0.5     # unnormalize，反标准化，将图像的像素值从[-1, 1]映射回[0, 1]npimg = img.numpy()     # 将PyTorch tensor转换为NumPy数组plt.imshow(np.transpose(npimg, (1, 2, 0)))  # 将tensor的维度从 (C, H, W) 转换为 (H, W, C) 以便绘制# 如果提供了标题，就显示标题if title:plt.title(title)plt.show()# 从训练数据加载器中获取一个批次的数据 (随机的)
print("----------随机展示训练集图像----------")
dataiter = iter(trainloader)     # 创建数据加载器的迭代器
images, labels = next(dataiter)  # 获取一批数据，其中images是图像数据，labels是对应的标签# 创建标题，显示标签名称。5s 表示字符串的宽度为 5，不足的部分会使用空格填充
title = ' '.join(f'{classes[labels[j]]:5s}' for j in range(batch_size)) # join() 方法将生成的类别名称列表（每个类别名称都经过格式化）用空格连接成一个单一的字符串。# 显示图像并且显示标题
imshow(torchvision.utils.make_grid(images), title=title) # 使用make_grid将图像组合成一个网格并显示#----- 2.定义卷积神经网络 -----#
import torch.nn as nn
import torch.nn.functional as Fclass 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) # torch.flatten(input, start_dim=0, end_dim=-1) ：将输入张量从 start_dim 到 end_dim 维度展平为一个维度。x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))x = self.fc3(x)return xnet = Net()
print("----------模型结构----------")
print(net)'''
以上为公共部分
'''#----- 3.定义损失函数和优化器 -----#
import torch.optim as optim'''CrossEntropyLoss 通常用于分类问题，计算目标类别与模型预测类别之间的误差。交叉熵(Cross-Entropy)本质上是两个概率分布之间的差异度量，常用于比较：一个真实分布P(通常是 one-hot 编码）一个预测分布Q(通常是 softmax 输出）它的数学定义如下：CCrossEntropy(P,Q)=- ∑ [ Pi*log(Qi) ]i=1其中：C 是类别总数Pi是真实标签中第 i 类的概率（一般为 1 或 0)Qi是模型预测中第 i 类的概率(softmax 输出）交叉熵损失只对正确类别的预测概率进行惩罚，其它类别不参与损失。'''
# 定义损失函数（这里使用的是分类交叉熵损失函数（CrossEntropyLoss））
criterion = nn.CrossEntropyLoss()'''optim.SGD 是 PyTorch 提供的随机梯度下降(SGD)优化器。参数解释：net.parameters()：模型参数，即网络中的所有可训练参数。lr=0.001: 学习率，控制每次参数更新的步长。较小的学习率有助于防止过拟合，但可能导致训练变慢。momentum=0.9: 动量项，用来加速收敛，并帮助跳出局部最小值。动量类似于物理中的惯性，它使用上一次的梯度来调整当前的梯度。
'''
# 定义优化器（这里使用的是带动量的随机梯度下降（SGD with momentum））
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)#----- 4.训练网络(使用 CPU 或者 GPU) -----#
import time
use_gpu = False  # 改成 False 就强制使用 CPUif not use_gpu:start_time = time.time()  # 记录开始时间print("----------模型训练: CPU----------")for epoch in range(2):  # 遍历整个数据集多次（训练2个 epoch）running_loss = 0.0  # 初始化每2000个小批次（即处理8000张图片）的累计损失for i, data in enumerate(trainloader, 0):  # 遍历训练数据集，每次获取一个小批次（mini-batch）# 获取输入数据；data 是一个包含 [inputs, labels] 的列表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:    # 每2000个小批次（即处理8000张图片）打印一次信息print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')running_loss = 0.0  # 重置累计损失print('CPU Finished Training')  # 打印训练结束的提示end_time = time.time()  # 记录结束时间elapsed_time = end_time - start_time  # 计算消耗的时间print(f"CPU time: {elapsed_time:.6f} 秒")else:start_time = time.time()  # 记录开始时间print("----------模型训练: GPU----------")# 确保在 GPU 上训练，如果有 CUDA 可用device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')print("device", device)# ①将神经网络 net 的所有参数和缓冲区（buffers）转移到 device 上net.to(device)for epoch in range(2):  # 遍历整个数据集多次（训练2个 epoch）running_loss = 0.0  # 初始化每2000个小批次（即处理8000张图片）的累计损失for i, data in enumerate(trainloader, 0):  # 遍历训练数据集，每次获取一个小批次（mini-batch）# 获取输入数据；data 是一个包含 [inputs, labels] 的列表inputs, labels = data# ②将输入数据（即图像）和标签（即真实标签）转移到 device 上inputs, labels = inputs.to(device), labels.to(device)# 将优化器的梯度清零，避免累积梯度optimizer.zero_grad()# 前向传播 + 反向传播 + 更新优化器outputs = net(inputs)              # 通过网络进行前向传播，得到预测结果loss = criterion(outputs, labels)  # 计算损失函数的值loss.backward()   # 计算梯度（反向传播）optimizer.step()  # 使用优化器更新模型的参数# 打印统计信息running_loss += loss.item()  # 将当前损失加到累计损失中if i % 2000 == 1999:    # 每2000个小批次（即处理8000张图片）打印一次信息print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')running_loss = 0.0  # 重置累计损失print('GPU Finished Training')  # 打印训练结束的提示end_time = time.time()  # 记录结束时间elapsed_time = end_time - start_time  # 计算消耗的时间print(f"GPU time: {elapsed_time:.6f} 秒")"""
模型训练: multi GPU
数据并行教程: https://pytorch.ac.cn/tutorials/beginner/blitz/data_parallel_tutorial.html
"""#----- 5.保存训练好的模型参数 -----#
import osprint("----------保存模型参数----------")
PATH = './model/cifar_net.pth'
os.makedirs('model', exist_ok=True) # 创建一个名为 'model' 的文件夹，如果文件夹已存在则不会报错。
torch.save(net.state_dict(), PATH)

                1.2运行结果

• CPU

(yolov5_env) wu@WP:~/yolo/pytorch_practice$ python train_and_save.py 
Files already downloaded and verified
Files already downloaded and verified
Training set size: 50000
Test set size: 10000
----------随机展示训练集图像----------
----------模型结构----------
Net((conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))(pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)(conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))(fc1): Linear(in_features=400, out_features=120, bias=True)(fc2): Linear(in_features=120, out_features=84, bias=True)(fc3): Linear(in_features=84, out_features=10, bias=True)
)
----------模型训练: CPU----------
[1,  2000] loss: 2.237
[1,  4000] loss: 1.895
[1,  6000] loss: 1.648
[1,  8000] loss: 1.555
[1, 10000] loss: 1.495
[1, 12000] loss: 1.454
[2,  2000] loss: 1.394
[2,  4000] loss: 1.343
[2,  6000] loss: 1.349
[2,  8000] loss: 1.292
[2, 10000] loss: 1.289
[2, 12000] loss: 1.273
CPU Finished Training
CPU time: 45.390200 秒
----------保存模型参数----------

• GPU

(yolov5_env) wu@WP:~/yolo/pytorch_practice$ python train_and_save.py 
Files already downloaded and verified
Files already downloaded and verified
Training set size: 50000
Test set size: 10000
----------随机展示训练集图像----------
----------模型结构----------
Net((conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))(pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)(conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))(fc1): Linear(in_features=400, out_features=120, bias=True)(fc2): Linear(in_features=120, out_features=84, bias=True)(fc3): Linear(in_features=84, out_features=10, bias=True)
)
----------模型训练: GPU----------
device cuda:0
[1,  2000] loss: 2.236
[1,  4000] loss: 1.876
[1,  6000] loss: 1.694
[1,  8000] loss: 1.633
[1, 10000] loss: 1.556
[1, 12000] loss: 1.492
[2,  2000] loss: 1.420
[2,  4000] loss: 1.386
[2,  6000] loss: 1.381
[2,  8000] loss: 1.343
[2, 10000] loss: 1.307
[2, 12000] loss: 1.283
GPU Finished Training
GPU time: 25.720243 秒
----------保存模型参数----------

        2.load_and_infer.py

                2.1 代码

#----- 1.加载并归一化 CIFAR10 -----#
import torch
import torchvision
import torchvision.transforms as transforms# transforms.Compose(): 组合变换操作的容器, 允许将多个图像转换操作按顺序执行
transform = transforms.Compose(# 将范围在 [0, 255] 的 PIL Image 或 numpy.ndarray (H x W x C) 转换为形状为 (C x H x W)、范围在 [0.0, 1.0] 的 torch.FloatTensor[transforms.ToTensor(),  # 使用均值和标准差归一化张量图像，将每个通道的值标准化为 [-1, 1] 区间。 output[channel] = (input[channel] - mean[channel]) / std[channel] = 2x - 1transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]  
)# 设置批量大小
batch_size = 4# 加载CIFAR-10训练数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,download=True, transform=transform)
# 创建训练数据的DataLoader，按批次加载数据并进行随机打乱
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,shuffle=True, num_workers=2)# 加载CIFAR-10测试数据集
testset = torchvision.datasets.CIFAR10(root='./data', train=False,download=True, transform=transform)
# 创建测试数据的DataLoader，按批次加载数据但不打乱
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,shuffle=False, num_workers=2)# CIFAR-10数据集的分类标签，包含10类
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')# # 输出训练集和测试集的大小
# train_size = len(trainset)
# test_size = len(testset)
# print(f'Training set size: {train_size}') # 50000
# print(f'Test set size: {test_size}')      # 10000# 展示训练集图像
import matplotlib.pyplot as plt
import numpy as npdef imshow(img, title=None):"""显示图像，并且可以选择性地显示标题。参数：img (Tensor): 要显示的图像。通常是经过预处理后的Tensor图像, 形状为 (C, H, W)。title (str, optional): 要显示的标题。如果提供了标题，则在图像上方显示它。"""img = img / 2 + 0.5     # unnormalize，反标准化，将图像的像素值从[-1, 1]映射回[0, 1]npimg = img.numpy()     # 将PyTorch tensor转换为NumPy数组plt.imshow(np.transpose(npimg, (1, 2, 0)))  # 将tensor的维度从 (C, H, W) 转换为 (H, W, C) 以便绘制# 如果提供了标题，就显示标题if title:plt.title(title)plt.show()# # 从训练数据加载器中获取一个批次的数据 (随机的)
# print("----------随机展示训练集图像----------")
# dataiter = iter(trainloader)     # 创建数据加载器的迭代器
# images, labels = next(dataiter)  # 获取一批数据，其中images是图像数据，labels是对应的标签# # 创建标题，显示标签名称。5s 表示字符串的宽度为 5，不足的部分会使用空格填充
# title = ' '.join(f'{classes[labels[j]]:5s}' for j in range(batch_size)) # join() 方法将生成的类别名称列表（每个类别名称都经过格式化）用空格连接成一个单一的字符串。# # 显示图像并且显示标题
# imshow(torchvision.utils.make_grid(images), title=title) # 使用make_grid将图像组合成一个网格并显示#----- 2.定义卷积神经网络 -----#
import torch.nn as nn
import torch.nn.functional as Fclass 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) # torch.flatten(input, start_dim=0, end_dim=-1) ：将输入张量从 start_dim 到 end_dim 维度展平为一个维度。x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))x = self.fc3(x)return x# net = Net()
# print("----------模型结构----------")
# print(net)'''
以上为公共部分
'''#----- 6.加载模型参数 -----#
print("----------加载模型参数----------")
PATH = './model/cifar_net.pth'
net = Net()
net.load_state_dict(torch.load(PATH, weights_only=True))#----- 7.在测试数据集上使用加载的模型进行推理 -----#
print("----------加载4张测试集图像----------")
# 展示测试集图像
dataiter = iter(testloader)
images, labels = next(dataiter)# 使用训练好的模型进行推理
print("----------模型推理----------")
outputs = net(images)
_, predicted = torch.max(outputs, 1) # 1 是指按行（类别维度）选取最大值。函数返回每一行的最大值，及其对应的索引（类别）print('Predicted: ', ' '.join(f'{classes[predicted[j]]:5s}'for j in range(4)))# show images
title = ' '.join(f'{classes[labels[j]]:5s}' for j in range(4))
imshow(torchvision.utils.make_grid(images), title=title)print("----------网络在整个测试集上的准确率----------")
correct = 0  # 用于记录预测正确的图像数量
total = 0    # 用于记录测试集中图像的总数with torch.no_grad():  # 禁用梯度计算，减少内存消耗并加速推理过程for data in testloader:    # 遍历测试数据集images, labels = data  # 获取当前小批次的图像和标签# # test# print("标签(labels)为：", labels)      # 示例：tensor([6, 0, 5, 7])# print("labels shape: ", labels.shape) # labels shape:  torch.Size([4])# 使用网络进行推理，计算输出outputs = net(images)# 选择输出中最大值的类别作为预测结果_, predicted = torch.max(outputs, 1)  # 1 是指按行（类别维度）选取最大值。函数返回每一行的最大值，及其对应的索引（类别）total += labels.size(0)  # 累加测试集中的样本数量correct += (predicted == labels).sum().item()  # 统计预测正确的样本数量# 打印测试集上模型的准确率
print(f'Accuracy of the network on the 10000 test images: {100 * correct // total} %')print("----------网络在每个类别上的准确率----------")
# 初始化字典
correct_pred = {classname: 0 for classname in classes}  # 初始化每个类别的正确预测数
total_pred = {classname: 0 for classname in classes}    # 初始化每个类别的总预测数# 由于不需要计算梯度，使用 no_grad 来加速推理过程
with torch.no_grad():  for data in testloader:  # 遍历测试数据集images, labels = data  # 获取当前批次的图像和标签outputs = net(images)  # 使用网络进行前向传播，得到预测输出_, predictions = torch.max(outputs, 1)  # 得到每个图像的预测类别（索引）# 收集每个类别的正确预测数量for label, prediction in zip(labels, predictions):  # 同时遍历标签和预测值total_pred[classes[label]] += 1  # 增加对应类别的总预测数if label == prediction:  # 如果标签与预测相同，表示预测正确correct_pred[classes[label]] += 1  # 增加对应类别的正确预测# 打印每个类别的准确率
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} %')  # 打印每个类别的准确率

                2.2 运行结果

(yolov5_env) wu@WP:~/yolo/pytorch_practice$ python load_and_infer.py 
Files already downloaded and verified
Files already downloaded and verified
----------加载模型参数----------
----------加载4张测试集图像----------
----------模型推理----------
Predicted:  bird  car   car   ship 
----------网络在整个测试集上的准确率----------
Accuracy of the network on the 10000 test images: 55 %
----------网络在每个类别上的准确率----------
Accuracy for class: plane is 67.1 %
Accuracy for class: car   is 77.8 %
Accuracy for class: bird  is 37.8 %
Accuracy for class: cat   is 46.6 %
Accuracy for class: deer  is 52.9 %
Accuracy for class: dog   is 36.4 %
Accuracy for class: frog  is 51.7 %
Accuracy for class: horse is 57.0 %
Accuracy for class: ship  is 58.2 %
Accuracy for class: truck is 69.3 %