pytorch训练图片分类模型

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder

#判断是否支持cuda并设定device运行环境
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# 定义图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),   # 将图像大小重置为 224x224 像素
    transforms.ToTensor(),  # 将图像转换为 PyTorch 的张量格式
    transforms.Normalize(mean=[0.5], std=[0.5])   # 将图像像素值归一化到 [-1,1] 范围内
])

# 定义数据集
trainset = ImageFolder('./data/train', transform=transform)  # 训练集的路径和预处理操作
testset = ImageFolder('./data/test', transform=transform)   # 测试集的路径和预处理操作

# 定义数据加载器
trainloader = DataLoader(trainset, batch_size=32, shuffle=True)  # 训练集数据加载器，批次大小为32，数据随机打乱
testloader = DataLoader(testset, batch_size=32, shuffle=False)   # 测试集数据加载器，批次大小为32，数据不打乱


# 定义模型结构
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)  # 第一层卷积层，输入通道数为3，输出通道数为32，卷积核大小为3x3，边缘填充1
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)  # 第二层卷积层，输入通道数为32，输出通道数为64，卷积核大小为3x3，边缘填充1
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)  # 最大池化层，池化核大小为2x2，步长为2
        self.fc1 = nn.Linear(64 * 56 * 56, 128)  # 第一层全连接层，输入大小为64*56*56，输出大小为128
        self.fc2 = nn.Linear(128, 2)  # 第二层全连接层，输入大小为128，输出大小为2

    def forward(self, x):
        # 第一层卷积
        x = self.pool(torch.relu(self.conv1(x)))
        # 第二层卷积
        x = self.pool(torch.relu(self.conv2(x)))
        # 全连接层
        x = x.view(-1, 64 * 56 * 56)  # 将特征图展开成一维向量
        x = torch.relu(self.fc1(x))   # 激活函数ReLU
        x = self.fc2(x)  # 全连接层
        return x

# 初始化模型
model = Net().to(device)  # 将模型转移到GPU上

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 交叉熵损失函数
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)  # 随机梯度下降优化器,lr 是学习率，momentum 是动量，用于加速梯度下降的收敛

# 训练模型,迭代 70 次
for epoch in range(70):
    running_loss = 0.0  #初始化当前训练周期的损失函数值
    for i, data in enumerate(trainloader, 0):  #开始对训练集的每个批次进行训练，使用enumerate函数迭代训练集数据
        inputs, labels = data[0].to(device), data[1].to(device)  # 将数据转移到GPU上
        optimizer.zero_grad()  # 将模型梯度清零，防止累积
        outputs = model(inputs)  # 正向传播计算模型输出
        loss = criterion(outputs, labels)  #计算当前批次的损失函数值
        loss.backward()  #反向传播计算每个参数的梯度
        optimizer.step()  #根据梯度更新每个参数的值
        running_loss += loss.item()  #计算当前周期的损失函数值
    #打印当前训练周期的批次数和损失函数值的平均值
    print('批次 %d, 损失函数: %.3f' % (epoch + 1, running_loss / len(trainloader)))

torch.save(model.state_dict(), 'fenlei_rmb_1_10.pth')  #保存模型

# 测试
correct = 0  #初始化正确分类的数量为0
total = 0  #初始化总共处理的测试图片数量为0
with torch.no_grad():  #在不需要计算梯度的上下文中执行以下操作
    for data in testloader:  #遍历测试集数据
        images, labels = data[0].to(device), data[1].to(device)  # 获取当前数据的图像和标签，并将数据转移到GPU上
        outputs = model(images)  #运行模型对当前图像进行预测
        _, predicted = torch.max(outputs.data, 1)  #找到输出中概率最大的类别，并返回其在类别列表中的索引，也就是模型的预测值
        total += labels.size(0)  #将当前批次的测试图片数量加到总数中
        correct += (predicted == labels).sum().item()  #计算模型在当前批次的测试图片中预测正确的数量，并将其加到正确分类的数量中
print('测试图片识别率: %d %%' % (100 * correct / total))  #输出模型在测试集上的准确率（百分比）