[Pytorch案例实践005]蚂蚁&蜜蜂图像分类

一、项目介绍

        这个项目的目标是使用 PyTorch 构建一个卷积神经网络(CNN),对蚂蚁和蜜蜂两种昆虫的图片进行分类。数据集由包含这两种昆虫的图片组成,分为训练集和验证集。通过训练过程,模型将学习如何区分这两种昆虫,并最终能够准确地对新图片进行分类。

项目结构

  1. 数据准备

  2. 数据集被划分为训练集和验证集。
  3. 使用了 torchvision.transforms 来定义数据预处理步骤,包括图像缩放、裁剪、归一化等。
  4. 使用 torch.utils.data.DataLoader 来加载数据集,方便批量处理。
  5. 模型定义

  6. 构建了一个简单的卷积神经网络,包含多个卷积层、池化层和全连接层。
  7. 最后一层使用全局平均池化层,以便处理不同大小的输入图像,并确保输出特征的尺寸一致。
  8. 训练过程

  9. 使用交叉熵损失函数作为目标函数。
  10. 使用随机梯度下降(SGD)作为优化器,并设置了学习率衰减策略。
  11. 在每个 epoch 结束时,根据验证集上的性能保存最佳模型。
  12. 评估指标

  13. 计算了每个 epoch 的损失和准确率。
  14. 使用了进度条库 tqdm 来可视化训练进度。
  15. 学习率曲线保存

  16. 记录了每个 epoch 的学习率,并绘制了学习率曲线。

二、数据集介绍

数据集概况

  • 来源:数据集来源于公开的数据源百度飞桨–昆虫分类任务_数据集-飞桨AI Studio星河社区 (baidu.com)。
  • 类别:数据集中包含两类昆虫——蚂蚁和蜜蜂。
  • 图像数量:数据集通常会被细分为训练集、验证集和测试集。训练集用于训练模型;验证集用于调整超参数和防止过拟合;测试集则是在训练过程中未使用的数据,用于最终评估模型的泛化能力。
  • 三、完整代码

    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.optim import lr_scheduler
    from torch.utils.data import DataLoader
    from torchvision import datasets, transforms
    import matplotlib.pyplot as plt
    import time
    import os
    import copy
    from tqdm import tqdm
    
    # 定义数据预处理步骤
    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
        'val': transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
    }
    
    # 假设数据集位于以下目录
    data_dir = 'I:/code/pytorch/ants&bees/datasets'
    image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
                      for x in ['train', 'val']}
    dataloaders = {x: DataLoader(image_datasets[x], batch_size=4,
                                 shuffle=True, num_workers=4)
                   for x in ['train', 'val']}
    dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
    class_names = image_datasets['train'].classes
    
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    
    
    # 构建简单的卷积神经网络
    def make_net():
        model = nn.Sequential(
            nn.Conv2d(3, 6, 5),
            nn.MaxPool2d(2, 2),
            nn.ReLU(True),
            nn.Conv2d(6, 16, 5),
            nn.MaxPool2d(2, 2),
            nn.ReLU(True),
            nn.Conv2d(16, 120, 5),
            nn.ReLU(True),
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(120, 84),
            nn.ReLU(True),
            nn.Linear(84, 2),
        )
        return model.to(device)
    
    
    # 设置训练参数
    # 设置训练参数
    def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
        since = time.time()
    
        best_model_wts = copy.deepcopy(model.state_dict())
        best_acc = 0.0
        lrs = []  # 存储每轮的学习率
    
        for epoch in range(num_epochs):
            print(f'Epoch {epoch}/{num_epochs - 1}')
            print('-' * 10)
    
            # 每个epoch有两个阶段:训练和验证
            for phase in ['train', 'val']:
                if phase == 'train':
                    model.train()  # 设置模型为训练模式
                else:
                    model.eval()  # 设置模型为评估模式
    
                running_loss = 0.0
                running_corrects = 0
    
                # 迭代数据
                for inputs, labels in tqdm(dataloaders[phase], desc=f'{phase} epoch {epoch}'):
                    inputs = inputs.to(device)
                    labels = labels.to(device)
    
                    # 清零梯度
                    optimizer.zero_grad()
    
                    # 前向传播
                    # 在训练时跟踪梯度,在验证时不跟踪
                    with torch.set_grad_enabled(phase == 'train'):
                        outputs = model(inputs)
                        _, preds = torch.max(outputs, 1)
                        loss = criterion(outputs, labels)
    
                        # 只有在训练阶段才反向传播 + 优化
                        if phase == 'train':
                            loss.backward()
                            optimizer.step()
    
                    # 统计
                    running_loss += loss.item() * inputs.size(0)
                    running_corrects += torch.sum(preds == labels.data).to(device)  # 确保张量位于正确的设备上
    
                if phase == 'train':
                    scheduler.step()
                    lrs.append(scheduler.get_last_lr()[0])  # 记录学习率
    
                epoch_loss = running_loss / dataset_sizes[phase]
                epoch_acc = running_corrects.double() / dataset_sizes[phase]  # 使用 double() 转换为 double 精度
    
                print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')
    
                # 深拷贝模型
                if phase == 'val' and epoch_acc > best_acc:
                    best_acc = epoch_acc
                    best_model_wts = copy.deepcopy(model.state_dict())
    
            print()
    
        time_elapsed = time.time() - since
        print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
        print(f'Best val Acc: {best_acc:4f}')
    
        # 加载最佳模型权重
        model.load_state_dict(best_model_wts)
    
        # 绘制学习率曲线
        plot_learning_rate(lrs)
        return model
    
    
    def plot_learning_rate(lrs):
        plt.figure()
        plt.plot(lrs, label='Learning Rate')
        plt.title('Learning Rate Over Time')
        plt.xlabel('Epoch')
        plt.ylabel('Learning Rate')
        plt.legend()
        plt.savefig('learning_rate_curve.png')
        plt.close()
    
    
    def plot_loss_accuracy(train_losses, val_losses, train_accs, val_accs):
        epochs = range(len(train_losses))
    
        plt.figure()
        plt.subplot(2, 1, 1)
        plt.plot(epochs, train_losses, label='Training Loss')
        plt.plot(epochs, val_losses, label='Validation Loss')
        plt.title('Loss Over Time')
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.legend()
    
        plt.subplot(2, 1, 2)
        plt.plot(epochs, train_accs, label='Training Accuracy')
        plt.plot(epochs, val_accs, label='Validation Accuracy')
        plt.title('Accuracy Over Time')
        plt.xlabel('Epoch')
        plt.ylabel('Accuracy')
        plt.legend()
    
        plt.tight_layout()
        plt.savefig('loss_accuracy_curve.png')
        plt.close()
    
    
    # 主程序
    if __name__ == "__main__":
        model = make_net()
    
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
        scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
    
        model = train_model(model, criterion, optimizer, scheduler, num_epochs=100)
    
        torch.save(model.state_dict(), 'best_model.pth')
    

    四、总结

            这个项目展示了一个完整的深度学习流程,包括数据预处理、模型构建、训练、评估。

    作者:Seraphina_Lily

    物联沃分享整理
    物联沃-IOTWORD物联网 » [Pytorch案例实践005]蚂蚁&蜜蜂图像分类

    发表回复