VGG-16
- 导包
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt- 数据预处理和增强
transform = transforms.Compose([
transforms.Resize((224, 224)), 调整图像大小为 224x224像素，符合VGG16输入
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])使用ImageNet的标准化参数
- 加载数据集
data_path = "D:\\图像处理、深度学习\\flowers"
dataset = datasets.ImageFolder(data_path, transform=transform)- 划分训练集和测试集
train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])- 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)- 使用预训练的VGG16模型
model = torchvision.models.vgg16(pretrained=True)- 修改全连接层以适应新的分类任务
num_classes = len(dataset.classes)
model.classifier[6] = nn.Linear(4096, num_classes)- 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())- 将模型移动到GPU（如果可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)- 初始化列表来存储每个epoch的损失和准确率
train_losses = []
train_accuracies = []
- 训练模型
num_epochs = 30 定义训练的轮数为30轮
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
correct = 0
total = 0 开始训练循环，将模型设置为训练模式初始化累计损失为0.0，初始化正确预测的数量为0，初始化总样本数量为0
for inputs, labels in train_loader: 遍历训练数据加载器中的每个批次
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad() 将输入和标签移动到定义的设备上（如GPU），清空模型的梯度
outputs = model(inputs)
loss = criterion(outputs, labels) 将输入传递给模型以获得输出，计算输出和标签之间的损失
loss.backward()
optimizer.step()
running_loss += loss.item() 反向传播损失以计算梯度，更新模型的参数，累计批次损失
_, predicted = torch.max(outputs.data, 1) 获取每个样本的最高预测值
total += labels.size(0)
correct += (predicted == labels).sum().item() 累计总样本数量，累计正确预测的样本数量
epoch_loss = running_loss / len(train_loader) 计算平均损失
epoch_accuracy = 100 * correct / total 计算准确率
train_losses.append(epoch_loss)
train_accuracies.append(epoch_accuracy) 将损失添加到训练损失列表中，将准确率添加到训练准确率列表中
print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {epoch_loss}, Accuracy: {epoch_accuracy}%') 打印轮数、损失和准确率
