VGG-19
- 导包
import torchimport torch.nn as nnimport torch.optim as optimimport torchvisionfrom torchvision import datasets, transformsfrom torch.utils.data import DataLoaderimport 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_sizetrain_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.vgg19(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}%')
打印轮数、损失和准确率