python
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)) #输出模型在测试集上的准确率（百分比）