[Pytorch案例实践007]基于VGG16的蚂蚁&蜜蜂图像分类实战

一、项目介绍

项目概述

        这是一个使用PyTorch框架的深度学习项目，目的是构建一个能够区分蜜蜂和蚂蚁的图像分类模型。项目使用了VGG16作为基础模型，并采用了迁移学习的方法来提高模型的性能。

技术栈

• PyTorch: 主要使用的深度学习框架。
• Torchvision: 提供了用于图像处理的常用工具，包括数据加载和转换。
• Matplotlib: 用于绘制训练过程中的损失和准确率曲线。
• tqdm: 进度条显示库，用于监控训练进度。

数据准备

• 数据集路径:
  • 训练数据集路径: I:\code\pytorch\VGG\datasets\train
  • 验证数据集路径: I:\code\pytorch\VGG\datasets\val
• 数据增强:
  • 将所有图像调整为224x224像素。
  • 转换为张量。
  • 使用ImageNet的均值和标准差进行标准化。

模型定义

• VGG16:
  • 类别数 (num_classes): 2（蜜蜂和蚂蚁）。
  • 预训练 (pretrained): 可选择使用或不使用ImageNet上的预训练权重。
  • 模型实例化后被移动到指定的设备（CPU或GPU）。

训练配置

• 超参数:
  • 批大小 (batch_size): 4。
  • 学习率 (learning_rate): 0.001。
  • 训练轮数 (num_epochs): 30。
• 优化器:
  • 使用随机梯度下降 (SGD) 优化器。
  • 动量设置为0.9。
• 学习率调度器:
  • 使用StepLR策略，在每7个epoch后将学习率乘以0.1。

训练过程

• 训练循环:
  • 在每个epoch中，模型首先在训练集上进行训练，然后在验证集上进行评估。
  • 使用交叉熵损失 (CrossEntropyLoss) 作为损失函数。
  • 在训练过程中记录训练和验证损失及准确率。
  • 每个epoch结束时，保存当前模型状态。
• 可视化:
  • 训练完成后，绘制训练和验证的损失、准确率以及学习率的变化曲线，并保存到指定文件夹。

文件结构

• 输出文件夹 (output_dir): 保存训练好的模型、损失/准确率曲线图和学习率曲线图。

执行

• 主函数 (main):
  • 接受一个布尔参数 use_pretrained 来决定是否使用预训练权重。
  • 调用 train_model 函数进行模型训练。

注意事项

• 确保指定的路径正确无误。
• 如果使用GPU，确保安装了正确的PyTorch版本以及CUDA驱动。
• 可能需要根据实际情况调整超参数以获得最佳性能。

二、数据集介绍

        数据集之前介绍过，在此不再赘述。

三、完整代码

训练代码：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import matplotlib.pyplot as plt
from tqdm import tqdm
from VGG import VGG16  # 导入定义的VGG16网络
import os


def train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, num_epochs, device, output_dir):
    """训练和验证模型，并保存损失、准确率和学习率曲线"""

    train_losses = []
    val_losses = []
    train_accuracies = []
    val_accuracies = []
    learning_rates = []

    for epoch in range(num_epochs):
        print(f"Epoch {epoch + 1}/{num_epochs}")

        # 训练阶段
        model.train()
        running_loss = 0.0
        running_corrects = 0
        total_train = 0

        with tqdm(train_loader, unit="batch") as tepoch:
            for inputs, labels in tepoch:
                tepoch.set_description(f"Training Epoch {epoch + 1}")
                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() * inputs.size(0)
                _, preds = torch.max(outputs, 1)
                running_corrects += torch.sum(preds == labels.data)
                total_train += labels.size(0)

                tepoch.set_postfix(loss=loss.item())

        epoch_loss = running_loss / total_train
        epoch_acc = running_corrects.double() / total_train
        train_losses.append(epoch_loss)
        train_accuracies.append(epoch_acc.item())

        # 验证阶段
        model.eval()
        val_running_loss = 0.0
        val_running_corrects = 0
        total_val = 0

        with torch.no_grad():
            with tqdm(val_loader, unit="batch") as vepoch:
                for inputs, labels in vepoch:
                    vepoch.set_description(f"Validation Epoch {epoch + 1}")
                    inputs, labels = inputs.to(device), labels.to(device)

                    outputs = model(inputs)
                    loss = criterion(outputs, labels)

                    val_running_loss += loss.item() * inputs.size(0)
                    _, preds = torch.max(outputs, 1)
                    val_running_corrects += torch.sum(preds == labels.data)
                    total_val += labels.size(0)

                    vepoch.set_postfix(loss=loss.item())

        val_epoch_loss = val_running_loss / total_val
        val_epoch_acc = val_running_corrects.double() / total_val
        val_losses.append(val_epoch_loss)
        val_accuracies.append(val_epoch_acc.item())

        # 更新学习率
        scheduler.step()
        learning_rates.append(optimizer.param_groups[0]['lr'])

        print(f"Training Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}")
        print(f"Validation Loss: {val_epoch_loss:.4f} Acc: {val_epoch_acc:.4f}")

        # 保存模型
        model_save_path = os.path.join(output_dir, f"vgg16_epoch_{epoch + 1}.pth")
        torch.save(model.state_dict(), model_save_path)

    # 绘制并保存损失和准确率曲线
    plt.figure(figsize=(10, 5))
    plt.subplot(1, 2, 1)
    plt.plot(train_losses, label="Training Loss")
    plt.plot(val_losses, label="Validation Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.legend()
    plt.title("Loss Curve")
    plt.savefig(os.path.join(output_dir, "loss_curve.png"))

    plt.subplot(1, 2, 2)
    plt.plot(train_accuracies, label="Training Accuracy")
    plt.plot(val_accuracies, label="Validation Accuracy")
    plt.xlabel("Epoch")
    plt.ylabel("Accuracy")
    plt.legend()
    plt.title("Accuracy Curve")
    plt.savefig(os.path.join(output_dir, "accuracy_curve.png"))

    # 保存学习率曲线
    plt.figure()
    plt.plot(learning_rates, label="Learning Rate")
    plt.xlabel("Epoch")
    plt.ylabel("Learning Rate")
    plt.legend()
    plt.title("Learning Rate Curve")
    plt.savefig(os.path.join(output_dir, "learning_rate_curve.png"))

    print("Training complete. Curves and model saved.")


def main(use_pretrained=False):
    # 检查是否有可用的GPU设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 设置数据集路径和超参数
    train_data_path = r'I:\code\pytorch\VGG\datasets\train'  # 修改为你的训练数据路径
    val_data_path = r'I:\code\pytorch\VGG\datasets\val'  # 修改为你的验证数据路径
    num_classes = 2  # 根据你的分类任务修改类别数
    batch_size = 4
    learning_rate = 0.001
    num_epochs = 30

    # 定义数据预处理
    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]),
    ])

    # 加载数据集
    train_dataset = datasets.ImageFolder(root=train_data_path, transform=transform)
    val_dataset = datasets.ImageFolder(root=val_data_path, transform=transform)

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

    # 实例化模型，选择是否使用预训练权重
    model = VGG16(num_classes=num_classes, pretrained=use_pretrained).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9)

    # 学习率调度器
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

    # 创建输出目录
    output_dir = "output"
    os.makedirs(output_dir, exist_ok=True)

    # 训练模型
    train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, num_epochs, device, output_dir)


if __name__ == "__main__":
    # 修改为True以使用预训练权重，False表示不使用
    main(use_pretrained=True)

测试代码：

import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image, ImageDraw, ImageFont
import os
from VGG import VGG16

def load_model(model_path, num_classes, device):
    model = VGG16(num_classes=num_classes, pretrained=False)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.to(device)
    model.eval()
    return model

def predict(model, image, device, transform):
    image = transform(image).unsqueeze(0).to(device)
    outputs = model(image)
    _, preds = torch.max(outputs, 1)
    confidence = nn.functional.softmax(outputs, dim=1)[0][preds].item()
    return preds.item(), confidence


def draw_label(image, label, confidence):
    draw = ImageDraw.Draw(image)

    # 使用更大的字体
    try:
        font = ImageFont.truetype("arial", 30)  # 使用 Arial 字体，字号为 36
    except IOError:
        font = ImageFont.load_default()  # 如果 Arial 字体不可用，使用默认字体

    text = f"{label}: {confidence:.2f}"

    # 使用 textbbox 获取文本边界框
    text_bbox = draw.textbbox((20, 20), text, font=font)
    text_width = text_bbox[2] - text_bbox[0]
    text_height = text_bbox[3] - text_bbox[1]

    # 矩形背景框
    position = (20, 20)
    draw.rectangle([position, (position[0] + text_width, position[1] + text_height)], fill="black")
    draw.text(position, text, fill="white", font=font)
    return image


def process_image(model, image_path, output_dir, device, transform, class_names):
    image = Image.open(image_path).convert("RGB")
    pred, confidence = predict(model, image, device, transform)
    label = class_names[pred]
    image_with_label = draw_label(image, label, confidence)
    output_path = os.path.join(output_dir, os.path.basename(image_path))
    image_with_label.save(output_path)
    image_with_label.show()  # 显示处理后的图像

def process_folder(model, folder_path, output_dir, device, transform, class_names):
    os.makedirs(output_dir, exist_ok=True)
    for filename in os.listdir(folder_path):
        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
            image_path = os.path.join(folder_path, filename)
            process_image(model, image_path, output_dir, device, transform, class_names)

def main():
    # 硬编码参数
    model_path = r'I:\code\pytorch\VGG\output\vgg16_epoch_30.pth'  # 模型权重文件路径
    input_path = r'I:\code\pytorch\VGG\datasets\val\bees\1181173278_23c36fac71.jpg'  # 输入图片或文件夹路径
    output_dir = r'I:\code\pytorch\VGG\result'  # 输出保存路径
    num_classes = 2  # 分类任务的类别数
    class_names = ['ants', 'bees']  # 类别名称

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = load_model(model_path, num_classes, device)

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

    if os.path.isfile(input_path):
        process_image(model, input_path, output_dir, device, transform, class_names)
    elif os.path.isdir(input_path):
        process_folder(model, input_path, output_dir, device, transform, class_names)
    else:
        print("Invalid input path. Must be a file or directory.")

if __name__ == "__main__":
    main()

四、总结

        此项目实现了使用VGG16模型对蜜蜂和蚂蚁图像进行分类的任务。通过迁移学习的方法利用预训练权重提高了模型的性能。最终，模型经过训练和验证，并且训练过程被可视化，便于分析模型的表现。