当前位置：首页 > news >正文

0基础学习PyTorch——时尚分类（Fashion MNIST）训练和推理

news 来源：原创 2024/9/28 3:21:13

大纲

环境准备
安装依赖
下载训练集
训练
- 定义模型
- 训练
- - 加载训练集
  - 定义损失函数和优化器
  - 训练模型
  - 保存模型
  - 完整文件
推理
- 加载模型
- 加载并预处理本地文件
- 推理
- 完整文件
代码地址
参考资料

时尚分类是PyTorch官方文档中推荐的案例。本文将拆解这个案例，进行部署以及测试。

环境准备

基础环境可以参考《0基础学习PyTorch——最小Demo》来进行部署。

安装依赖

torchvision 是 PyTorch 的一个官方库，专门用于计算机视觉任务。它提供了常用的数据集、模型架构和图像处理工具，简化了计算机视觉项目的开发过程。后续我们的数据都来源于该库。

source env.sh install torchvision

在这里插入图片描述

下载训练集

将下列内容保存为download.py。

# download.py
import torchvision# Create datasets for training & validation, download if necessary
training_set = torchvision.datasets.FashionMNIST('./data', download=True)
validation_set = torchvision.datasets.FashionMNIST('./data', download=True)# Report split sizes
print('Training set has {} instances'.format(len(training_set)))
print('Validation set has {} instances'.format(len(validation_set)))

然后运行这个文件

python download.py

在这里插入图片描述
此时目录结构如下

训练

定义模型

将下面内容保存为garmentclassifier.py。该文件会被训练和推理两个环节使用。

# garmentclassifier.py
import torch.nn as nn
import torch.nn.functional as F# 定义一个用于服装分类的卷积神经网络
class GarmentClassifier(nn.Module):def __init__(self):super(GarmentClassifier, self).__init__()# 定义第一个卷积层，输入通道数为1，输出通道数为6，卷积核大小为5x5self.conv1 = nn.Conv2d(1, 6, 5)# 定义最大池化层，池化窗口大小为2x2self.pool = nn.MaxPool2d(2, 2)# 定义第二个卷积层，输入通道数为6，输出通道数为16，卷积核大小为5x5self.conv2 = nn.Conv2d(6, 16, 5)# 定义第一个全连接层，输入大小为16*4*4，输出大小为120self.fc1 = nn.Linear(16 * 4 * 4, 120)# 定义第二个全连接层，输入大小为120，输出大小为84self.fc2 = nn.Linear(120, 84)# 定义第三个全连接层，输入大小为84，输出大小为10（对应10个类别）self.fc3 = nn.Linear(84, 10)def forward(self, x):# 通过第一个卷积层和ReLU激活函数，然后通过最大池化层x = self.pool(F.relu(self.conv1(x)))# 通过第二个卷积层和ReLU激活函数，然后通过最大池化层x = self.pool(F.relu(self.conv2(x)))# 展平张量，从多维张量变为二维张量x = x.view(-1, 16 * 4 * 4)# 通过第一个全连接层和ReLU激活函数x = F.relu(self.fc1(x))# 通过第二个全连接层和ReLU激活函数x = F.relu(self.fc2(x))# 通过第三个全连接层（输出层）x = self.fc3(x)return x

训练

加载训练集

这次我们直接从本地加载训练集，但是需要做归一化处理。

from datetime import datetime
import torch
import torchvision
import torchvision.transforms as transforms
from garmentclassifier import GarmentClassifier# 定义图像转换操作：将图像转换为张量，并进行归一化处理
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))]) # 对图像的每个通道进行标准化，使得每个通道的像素值具有零均值和单位标准差# 加载FashionMNIST训练数据集，并应用定义的图像转换操作
training_set = torchvision.datasets.FashionMNIST('./data', train=True, transform=transform)# 创建数据加载器，用于批量加载训练数据，batch_size为4，数据顺序随机打乱
trainloader = torch.utils.data.DataLoader(training_set, batch_size=4, shuffle=True)

定义损失函数和优化器

# 实例化模型
model = GarmentClassifier()
# 定义损失函数为交叉熵损失
loss_fn = torch.nn.CrossEntropyLoss()
# 定义优化器为随机梯度下降（SGD），学习率为0.001，动量为0.9
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

训练模型

# 训练模型，训练2个epoch
for epoch in range(2):running_loss = 0.0  # 初始化累计损失# 枚举数据加载器中的数据，i是批次索引，data是当前批次的数据for i, data in enumerate(trainloader, 0):inputs, labels = data  # 获取输入数据和对应的标签optimizer.zero_grad()  # 清空梯度outputs = model(inputs)  # 前向传播，计算模型输出loss = loss_fn(outputs, labels)  # 计算损失loss.backward()  # 反向传播，计算梯度optimizer.step()  # 更新模型参数running_loss += loss.item()  # 累加损失# 每2000个批次打印一次平均损失if i % 2000 == 1999:print(f'[{epoch + 1}, {i + 1}] loss: {running_loss / 2000}')running_loss = 0.0  # 重置累计损失

保存模型

# 获取当前时间戳，格式为 'YYYYMMDD_HHMMSS'
timestamp = datetime.now().strftime('%Y%m%d%H%M%S.pth')# 定义模型保存路径，包含时间戳
model_path = 'model_{}'.format(timestamp)      # 保存模型的状态字典到指定路径
torch.save(model.state_dict(), model_path)

完整文件

from datetime import datetime
import torch
import torchvision
import torchvision.transforms as transforms
from garmentclassifier import GarmentClassifier# 定义图像转换操作：将图像转换为张量，并进行归一化处理
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))]) # 对图像的每个通道进行标准化，使得每个通道的像素值具有零均值和单位标准差# 加载FashionMNIST训练数据集，并应用定义的图像转换操作
training_set = torchvision.datasets.FashionMNIST('./data', train=True, transform=transform)# 创建数据加载器，用于批量加载训练数据，batch_size为4，数据顺序随机打乱
trainloader = torch.utils.data.DataLoader(training_set, batch_size=4, shuffle=True)# 实例化模型
model = GarmentClassifier()
# 定义损失函数为交叉熵损失
loss_fn = torch.nn.CrossEntropyLoss()
# 定义优化器为随机梯度下降（SGD），学习率为0.001，动量为0.9
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)# 训练模型，训练2个epoch
for epoch in range(2):running_loss = 0.0  # 初始化累计损失# 枚举数据加载器中的数据，i是批次索引，data是当前批次的数据for i, data in enumerate(trainloader, 0):inputs, labels = data  # 获取输入数据和对应的标签optimizer.zero_grad()  # 清空梯度outputs = model(inputs)  # 前向传播，计算模型输出loss = loss_fn(outputs, labels)  # 计算损失loss.backward()  # 反向传播，计算梯度optimizer.step()  # 更新模型参数running_loss += loss.item()  # 累加损失# 每2000个批次打印一次平均损失if i % 2000 == 1999:print(f'[{epoch + 1}, {i + 1}] loss: {running_loss / 2000}')running_loss = 0.0  # 重置累计损失# 获取当前时间戳，格式为 'YYYYMMDD_HHMMSS'
timestamp = datetime.now().strftime('%Y%m%d%H%M%S.pth')# 定义模型保存路径，包含时间戳
model_path = 'model_{}'.format(timestamp)      # 保存模型的状态字典到指定路径
torch.save(model.state_dict(), model_path)

执行该文件，我们会得到一个后缀为pth的模型文件。

推理

加载模型

我们加载上一步创建的模型。

import os
import glob
import torch
import torchvision.transforms as transforms
from PIL import Image
from datetime import datetime
from garmentclassifier import GarmentClassifierdef get_latest_model_path(directory, pattern="model_*.pth"):# 获取目录下所有符合模式的文件model_files = glob.glob(os.path.join(directory, pattern))if not model_files:raise FileNotFoundError("No model files found in the directory.")# 找到最新的模型文件latest_model_file = max(model_files, key=os.path.getmtime)return latest_model_file# 定义图像转换操作：将图像转换为张量，并进行归一化处理
transform = transforms.Compose([transforms.Resize((28, 28)),  # 调整图像大小为28x28transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))
])# 加载训练好的模型
model = GarmentClassifier()
model_path = get_latest_model_path('./')  # 获取最新的模型文件
model.load_state_dict(torch.load(model_path, weights_only=False)) # 加载模型参数
model.eval()  # 设置模型为评估模式

加载并预处理本地文件

在这里插入图片描述

# 从本地加载图像
image_path = 'shoe.jpg'  # 替换为实际的图像路径
image = Image.open(image_path).convert('L')  # 将图像转换为灰度图# 预处理图像
image = transform(image)
image = image.unsqueeze(0)  # 增加一个批次维度

我们使用transform进行归一化处理。

推理

# 推理（预测）
with torch.no_grad():  # 在推理过程中不需要计算梯度outputs = model(image)  # 前向传播，计算模型输出_, predicted = torch.max(outputs, 1)  # 获取预测结果# 定义类别名称
classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat','Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot')# 打印预测结果
print(f'Predicted label: {classes[predicted.item()]}')

由于推理出来的是索引号，为了方便解读，我们将类型映射打印出来。

完整文件

import os
import glob
import torch
import torchvision.transforms as transforms
from PIL import Image
from datetime import datetime
from garmentclassifier import GarmentClassifierdef get_latest_model_path(directory, pattern="model_*.pth"):# 获取目录下所有符合模式的文件model_files = glob.glob(os.path.join(directory, pattern))if not model_files:raise FileNotFoundError("No model files found in the directory.")# 找到最新的模型文件latest_model_file = max(model_files, key=os.path.getmtime)return latest_model_file# 定义图像转换操作：将图像转换为张量，并进行归一化处理
transform = transforms.Compose([transforms.Resize((28, 28)),  # 调整图像大小为28x28transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))
])# 加载训练好的模型
model = GarmentClassifier()
model_path = get_latest_model_path('./')  # 获取最新的模型文件
model.load_state_dict(torch.load(model_path, weights_only=False)) # 加载模型参数
model.eval()  # 设置模型为评估模式# 从本地加载图像
image_path = 'shoe.jpg'  # 替换为实际的图像路径
image = Image.open(image_path).convert('L')  # 将图像转换为灰度图# 预处理图像
image = transform(image)
image = image.unsqueeze(0)  # 增加一个批次维度# 推理（预测）
with torch.no_grad():  # 在推理过程中不需要计算梯度outputs = model(image)  # 前向传播，计算模型输出_, predicted = torch.max(outputs, 1)  # 获取预测结果# 定义类别名称
classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat','Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot')# 打印预测结果
print(f'Predicted label: {classes[predicted.item()]}')