《基于 PyTorch 的猫狗图像分类模型训练与评估》

目录

一、项目背景

二、环境与库准备

三、设备选择

四、数据预处理

五、数据读取与准备

六、模型构建与优化

七、训练与验证过程

八、模型保存与可视化

九、模型测试

十、前端展示



一、项目背景

随着图像数据的海量增长,对图像进行快速准确的分类成为了一项重要的任务。在本项目中,我们选择了常见的图像分类问题,例如区分猫和狗的图像,旨在通过深度学习技术,特别是基于卷积神经网络的方法,实现对这类图像的自动分类。

考虑到实际应用中的需求,我们需要一个能够在不同设备上运行,并且具有较高准确性和泛化能力的模型。通过使用流行的深度学习框架 PyTorch ,我们能够利用其丰富的功能和优化算法,来构建和训练我们的模型。

我们收集了大量的猫和狗的图像数据集,并对其进行了预处理和划分,分为训练集、验证集和测试集。通过精心设计的图像变换和数据增强策略,提高模型的鲁棒性和泛化能力

二、环境与库准备

首先,我们导入了一系列必要的库,包括 torch 相关的模块用于模型构建和训练,torchvision 用于数据处理和模型加载,以及一些数据处理和可视化的库。这些库为我们提供了构建、训练和评估模型所需的各种功能和工具。

import torch.optim as optim
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models
from torch.autograd import Variable
from tqdm import tqdm

import os
from PIL import Image
import cv2

三、设备选择

接下来,通过以下代码选择计算设备:

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
DEVICE

这行代码会自动检测是否有可用的 GPU,如果有则使用 GPU('cuda')进行计算,否则使用 CPU。

四、数据预处理

定义了 transform111 字典,其中包含了针对训练数据和验证数据的不同预处理操作,如随机裁剪、旋转、翻转、标准化等,以增加数据的多样性和模型的泛化能力。

transform111 = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)),
        transforms.RandomRotation(degrees=15),
        transforms.RandomHorizontalFlip(),
        transforms.CenterCrop(size=224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(size=256),
        transforms.CenterCrop(size=224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])
    ])
}

五、数据读取与准备

读取训练和验证数据的文件夹,并使用 ImageFolder 类进行组织。然后,计算训练和验证数据的大小,创建数据加载器,设置批处理大小和工作线程数,为后续的模型训练和验证做好准备。

import shutil
modellr = 1e-3
BATCH_SIZE = 64
EPOCHS = 4
# 读取数据
dataset_train = datasets.ImageFolder('dataset/train', transform111)
print(dataset_train.imgs)
# 对应文件夹的label
print(dataset_train.class_to_idx)
dataset_test = datasets.ImageFolder('dataset/val', transform111)
# 对应文件夹的label
print(dataset_test.class_to_idx)

dataset = './dataset'
train_directory = os.path.join(dataset, 'train')
valid_directory = os.path.join(dataset, 'val')

batch_size = 32
num_classes = 6
print(train_directory)
data = {
    'train': datasets.ImageFolder(root=train_directory, transform=transform111['train']),
    'val': datasets.ImageFolder(root=valid_directory, transform=transform111['val'])
}

train_data_size = len(data['train'])
valid_data_size = len(data['val'])

train_loader = torch.utils.data.DataLoader(data['train'], batch_size=batch_size, shuffle=True, num_workers=8)
test_loader = torch.utils.data.DataLoader(data['val'], batch_size=batch_size, shuffle=True, num_workers=8)

print(train_data_size, valid_data_size)

六、模型构建与优化

  1. 模型定义

使用 ResNet50 模型结构,但不加载预训练权重。

# 下载预训练模型
model = torchvision.models.resnet50(pretrained=False)  # 不使用训练好的预训练模型
model

    2.全连接层调整

criterion = nn.CrossEntropyLoss()
num_ftrs = model.fc.in_features
model.fc = nn.Sequential(
    nn.Linear(num_ftrs, 256),
    nn.ReLU(),
    nn.Dropout(0.4),
    nn.Linear(256, 6),  # 猫狗是二分类,故是 2
    nn.LogSoftmax(dim=1)
)
model.to(DEVICE)

3.优化器选择与学习率设置

选择 Adam 优化器,并定义了学习率调整函数,根据训练轮数动态降低学习率。

optimizer = optim.Adam(model.parameters(), lr=modellr)

def adjust_learning_rate(optimizer, epoch):
    modellrnew = modellr * (0.1 ** (epoch // 50))
    print("lr:", modellrnew)
    for param_group in optimizer.param_groups:
        param_group['lr'] = modellrnew

七、训练与验证过程

  1. 定义了 train 函数进行模型的训练,在每个批次中计算损失、进行反向传播和参数更新,并定期打印训练信息和计算平均准确率与损失。

    def train(model, device, train_loader, optimizer, epoch):
        model.train()
        sum_loss = 0
        train_acc = 0
        total_num = len(train_loader.dataset)
        for batch_idx, (data, target) in enumerate(tqdm(train_loader)):
            data, target = Variable(data).to(device), Variable(target).to(device)
            im = Variable(data)
            output = model(data)
            loss = criterion(output, target)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            print_loss = loss.data.item()
            sum_loss += print_loss
    
            out_t = output.argmax(dim=1)  # 取出预测的最大值
            num_correct = (out_t == target).sum().item()
            acc = num_correct / im.shape[0]
            train_acc += acc
    
            if (batch_idx + 1) % 50 == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),
                           100. * (batch_idx + 1) / len(train_loader), loss.item()))
    
        ave_loss = sum_loss / len(train_loader)
        ave_acc = train_acc / len(train_loader)
        print('epoch:{}, train_acc: {}, loss:{}'.format(epoch, ave_acc, ave_loss))
        return ave_acc, ave_loss
  2. 定义了 val 函数进行模型的验证,计算验证集上的损失和准确率。
    def val(model, device, test_loader):
        model.eval()
        test_loss = 0
        correct = 0
        total_num = len(test_loader.dataset)
        with torch.no_grad():
            for data, target in test_loader:
                data, target = Variable(data).to(device), Variable(target).to(device)
                output = model(data)
                loss = criterion(output, target)
                _, pred = torch.max(output.data, 1)
                correct += torch.sum(pred == target)
                print_loss = loss.data.item()
                test_loss += print_loss
            correct = correct.data.item()
            acc = correct / total_num
            avgloss = test_loss / len(test_loader)
            print('Val set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
                avgloss, correct, len(test_loader.dataset), 100 * acc))
        return acc, avgloss

在训练过程中,通过循环遍历每个训练轮次,调用学习率调整函数、训练函数和验证函数,记录训练和验证的准确率与损失。

然后通过以下代码进行训练和验证

train_acc_list, train_loss_list, val_acc_list, val_loss_list = [], [], [], []

for epoch in range(1, EPOCHS + 1):
    adjust_learning_rate(optimizer, epoch)
    train_acc, train_loss = train(model, DEVICE, train_loader, optimizer, epoch)
    val_acc, val_loss = val(model, DEVICE, test_loader)
    train_acc_list.append(train_acc)
    val_acc_list.append(val_acc)
    train_loss_list.append(train_loss)
    val_loss_list.append(val_loss)

八、模型保存与可视化

训练完成后,保存训练好的模型,并使用 matplotlib 绘制训练和验证的准确率与损失曲线,以便直观地观察模型的性能。

torch.save(model,'model.pth')
import matplotlib.pyplot as plt

epochs_range = range(EPOCHS)
print(epochs_range, train_acc_list)

plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, train_acc_list, label='Training Accuracy')
plt.plot(epochs_range, val_acc_list, label='Validation Accuracy')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')

plt.subplot(1, 2, 2)
plt.plot(epochs_range, train_loss_list, label='Training Loss')
plt.plot(epochs_range, val_loss_list, label='Validation Loss')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')
plt.savefig('./acc-loss.jpg')

预测准确率随着预测次数而提高

九、模型测试

最后,定义了测试数据的预处理操作,加载保存的模型,对测试集进行预测,并输出每个测试样本的图像名称、真实标签和预测标签。

import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.autograd import Variable
from pathlib import Path

classes = ('cat', 'dog')
transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = torch.load("model.pth")
model.eval()
model.to(DEVICE)

dataset_test = datasets.ImageFolder('dataset/test', transform_test)
print(len(dataset_test))
# 对应文件夹的label

y_true, y_sore = [], []
for index in range(len(dataset_test)):
    item = dataset_test[index]
    img, label = item
    img.unsqueeze_(0)
    data = Variable(img).to(DEVICE)
    output = model(data)
    _, pred = torch.max(output.data, 1)
    y_true.append(label)
    y_sore.append(pred.data.item())
    print('Image Name:{}, label:{}, predict:{}'.format(dataset_test.imgs[index][0], classes[label], classes[pred.data.item()]))
    index += 1

图像分类预测结果的记录

十、前端展示

  • 16
    点赞
  • 15
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值