PyTorch入门------训练图像分类器
----- 2.定义卷积神经网络 -----#x = torch.flatten(x, 1) # torch.flatten(input, start_dim=0, end_dim=-1) ：将输入张量从 start_dim 到 end_dim 维度展平为一个维度。return xprint("----------模型结构----------")print(net)#----- 3.定义损失函数和
·
二、训练、保存模型参数部分 train_and_save.py
三、加载模型参数、模型推理部分 load_and_infer.py
前言
参考：训练分类器 — PyTorch 教程 2.7.0+cu126 文档 - PyTorch 深度学习库
1. 操作步骤
我们将按顺序执行以下步骤训练图像分类器
-
使用
torchvision加载并归一化 CIFAR10 训练和测试数据集 -
定义卷积神经网络
-
定义损失函数
-
在训练数据上训练网络
-
保持训练好的模型参数
-
加载模型参数
-
在测试数据上测试网络
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 - 1
transforms.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 np
def 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 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) # 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)二、训练、保存模型参数部分 train_and_save.py
3.定义损失函数和优化器
#----- 3.定义损失函数和优化器 -----#
import torch.optim as optim
'''
CrossEntropyLoss 通常用于分类问题，计算目标类别与模型预测类别之间的误差。
交叉熵(Cross-Entropy)本质上是两个概率分布之间的差异度量，常用于比较：
一个真实分布P(通常是 one-hot 编码）
一个预测分布Q(通常是 softmax 输出）
它的数学定义如下：
C
CrossEntropy(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 就强制使用 CPU
if 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 os
print("----------保存模型参数----------")
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 - 1
transforms.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 np
def 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 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) # 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)
'''
以上为公共部分
'''
#----- 3.定义损失函数和优化器 -----#
import torch.optim as optim
'''
CrossEntropyLoss 通常用于分类问题，计算目标类别与模型预测类别之间的误差。
交叉熵(Cross-Entropy)本质上是两个概率分布之间的差异度量，常用于比较：
一个真实分布P(通常是 one-hot 编码）
一个预测分布Q(通常是 softmax 输出）
它的数学定义如下：
C
CrossEntropy(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 就强制使用 CPU
if 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 os
print("----------保存模型参数----------")
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 - 1
transforms.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 np
def 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 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) # 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 %
技术共进，成长同行——讯飞AI开发者社区
更多推荐
3
0- 0
已为社区贡献5条内容
所有评论(0)