pytorch神经网络训练（VGG-16）

VGG-16

1. 导包

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

1. 数据预处理和增强

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的标准化参数

1. 加载数据集

data_path = "D:\\图像处理、深度学习\\flowers"

dataset = datasets.ImageFolder(data_path, transform=transform)

1. 划分训练集和测试集

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])

1. 创建数据加载器

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

1. 使用预训练的VGG16模型

model = torchvision.models.vgg16(pretrained=True)

1. 修改全连接层以适应新的分类任务

num_classes = len(dataset.classes)

model.classifier[6] = nn.Linear(4096, num_classes)

1. 定义损失函数和优化器

criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(model.parameters())

1. 将模型移动到GPU（如果可用）

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model.to(device)

1. 初始化列表来存储每个epoch的损失和准确率

train_losses = []

train_accuracies = []

1. 训练模型

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}%') 

打印轮数、损失和准确率