操作系统

Ubuntu（Linux）

程序语言

Python（3.11.4）

第三方依赖

torch, torchvision, matplotlib

基本参数

参数值

Batch_size

32

激活函数

ReLU

优化器

Adam

学习率

0.001

训练epoch

30

基本参数

参数值

Batch_size

32

激活函数

ReLU

优化器

Adam

学习率

0.001

训练epoch

30

# 训练集的转换：包含数据增强

train_transform = transforms.Compose([

    transforms.Resize((224, 224)),

    transforms.RandomCrop(224, padding=4),  # 添加随机裁剪

    transforms.ToTensor(),

    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

])

# 测试集/验证集的转换：不包含数据增强

test_transform = transforms.Compose([

    transforms.Resize((224, 224)),

    transforms.ToTensor(),

    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

])

dataroot = "/home/ubuntu/apple"

# 加载数据集时应用相应的转换

train_dataset_full = torchvision.datasets.ImageFolder(root=dataroot, transform=train_transform)

test_dataset_full = torchvision.datasets.ImageFolder(root=dataroot, transform=test_transform)

# 划分 train 和 test 数据集

train_size = int(0.8 * len(train_dataset_full))

test_size = len(train_dataset_full) - train_size

# 这里使用相同的索引划分方法来确保训练集和测试集的一致性

train_indices = torch.arange(0, train_size)

test_indices = torch.arange(train_size, train_size + test_size)

train_dataset = torch.utils.data.Subset(train_dataset_full, train_indices)

val_dataset = torch.utils.data.Subset(test_dataset_full, test_indices)

# 创建数据加载器

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

val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)

dataset = test_dataset_full

import torch

import torchvision

import torchvision.transforms as transforms

from torchvision import models

import torch.nn as nn

import torch.optim as optim

from torch.utils.data import DataLoader, random_split

import torch.nn.functional as F

import matplotlib.pyplot as plt

# visualization part

def plot_loss(train_losses, val_losses):

    plt.figure(figsize=(10, 5))

    plt.plot(train_losses, label='Training Loss')

    plt.plot(val_losses, label='Validation Loss')

    plt.xlabel('Epoch')

    plt.ylabel('Loss')

    plt.title('Training and Validation Loss')

    plt.legend()

    # plt.show()

    plt.savefig('loss-vgg.png')

def plot_accuracy(train_accuracies, val_accuracies):

    plt.figure(figsize=(10, 5))

    plt.plot(train_accuracies, label='Training Accuracy')

    plt.plot(val_accuracies, label='Validation Accuracy')

    plt.xlabel('Epoch')

    plt.ylabel('Accuracy')

    plt.title('Training and Validation Accuracy')

    plt.legend()

    # plt.show()

    plt.savefig('acc-vgg.png')

# Data Preparation

transform = transforms.Compose([

    transforms.Resize((224, 224)),

    transforms.ToTensor(),

    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

])

dataroot = r"/home/ubuntu/apple"

dataset = torchvision.datasets.ImageFolder(root=dataroot, transform=transform)

# 划分 train 和 test 数据集

train_size = int(0.8 * len(dataset))

test_size = len(dataset) - train_size

train_dataset, val_dataset = random_split(dataset, [train_size, test_size])

# 创建数据加载器

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

val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)

# 实例化模型

# model = VGG16(num_classes=6)

model = models.vgg16(pretrained=True)

# 修改分类器部分

model.classifier = nn.Sequential(

    nn.Flatten(),

    nn.LayerNorm(25088),

    nn.Linear(25088, len(dataset.classes))

)

# Training and Evaluation

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

model.to(device)

criterion = nn.CrossEntropyLoss()

# optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for i, layer in enumerate(model.features):

    if isinstance(layer, nn.ReLU):

        model.features[i] = nn.LeakyReLU(negative_slope=0.01)

for i, layer in enumerate(model.classifier):

    if isinstance(layer, nn.ReLU):

        model.classifier[i] = nn.LeakyReLU(negative_slope=0.01)

# Training loop

num_epochs = 30

train_losses = []

val_losses = []

train_accuracies = []

val_accuracies = []

for epoch in range(num_epochs):

    # train

    model.train()

    running_loss = 0.0

    correct = 0

    total = 0

    test_correct = 0

    test_total = 0

    for i, (inputs, labels) in enumerate(train_loader):

        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()

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

    train_losses.append(running_loss / len(train_loader))

    train_accuracies.append(100 * correct / total)

    # evaluate

    model.eval()

    with torch.no_grad():

        val_loss = 0.0

        val_correct = 0

        val_total = 0

        for inputs, labels in val_loader:

            inputs, labels = inputs.to(device), labels.to(device)

            outputs = model(inputs)

            loss = criterion(outputs, labels)

            val_loss += loss.item()

            _, predicted = torch.max(outputs.data, 1)

            val_total += labels.size(0)

            val_correct += (predicted == labels).sum().item()

    val_losses.append(val_loss / len(val_loader))

    val_accuracies.append(100 * val_correct / val_total)

    print(f"Epoch {epoch+1}, Loss: {train_losses[-1]}, Accuracy: {train_accuracies[-1]}%, Validation Loss: {val_losses[-1]}, Validation Accuracy: {val_accuracies[-1]}%")

# 绘制损失曲线和准确率曲线

plot_loss(train_losses, val_losses)

plot_accuracy(train_accuracies, val_accuracies)

print('Finished')