在本篇文章中,我们将使用卷积神经网络(CNN)和预训练模型来解决猫狗图片分类问题。通过利用PyTorch的 torchvision 库加载数据集,并应用数据增强和数据预处理,我们将训练一个简单的CNN模型,然后再使用预训练的AlexNet模型来提高分类效果。

1. 数据加载与预处理

        首先,我们定义了数据增强和归一化的变换操作。这些变换包括图像旋转、水平翻转、图像尺寸调整、中心裁剪和归一化。这些操作有助于提高模型的泛化能力。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, models
from torchvision.utils import make_grid
import os
import numpy as np
import matplotlib.pyplot as plt

train_transform = transforms.Compose([
        transforms.RandomRotation(10),
        transforms.RandomHorizontalFlip(),
        transforms.Resize(224),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])
    ])

test_transform = transforms.Compose([
        transforms.Resize(224),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])
    ])

root = 'data/cat_dog'
train_data = datasets.ImageFolder(os.path.join(root, 'train'), transform=train_transform)
test_data  = datasets.ImageFolder(os.path.join(root, 'test'),  transform=test_transform)

train_loader = DataLoader(train_data, batch_size=10, shuffle=True)
test_loader  = DataLoader(test_data, batch_size=10, shuffle=True)

        通过 torchvision.datasets.ImageFolder 加载图像数据集,我们可以方便地组织数据,按类别将图片存储在不同文件夹中。

2. 构建卷积神经网络模型

        我们设计了一个简单的CNN模型,包括两个卷积层、两个最大池化层和三层全连接层。输入图像经过卷积层的处理后,提取出特征信息,最后通过全连接层输出预测结果。

class ConvolutionalNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 3, 1)
        self.conv2 = nn.Conv2d(6, 16, 3, 1)
        self.fc1 = nn.Linear(54*54*16, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 2)

    def forward(self, X):
        X = F.relu(self.conv1(X))
        X = F.max_pool2d(X, 2, 2)
        X = F.relu(self.conv2(X))
        X = F.max_pool2d(X, 2, 2)
        X = X.view(-1, 54*54*16)
        X = F.relu(self.fc1(X))
        X = F.relu(self.fc2(X))
        X = self.fc3(X)
        return X

        我们将使用交叉熵损失函数和Adam优化器来训练这个模型。

CNNmodel  = ConvolutionalNetwork()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(CNNmodel.parameters(), lr=0.001)

3. 模型训练

        为了加快训练速度,我们将限制每个 epoch 处理的图像批次数量,分别处理训练集的8000张图片和测试集的3000张图片。

import time
start_time = time.time()

epochs = 3
train_losses = []
test_losses = []
train_correct = []
test_correct = []

for i in range(epochs):
    trn_corr = 0
    tst_corr = 0
    
    for b, (X_train, y_train) in enumerate(train_loader):
        if b == 800:
            break
        y_pred = CNNmodel(X_train)
        loss = criterion(y_pred, y_train)
        predicted = torch.max(y_pred.data, 1)[1]
        batch_corr = (predicted == y_train).sum()
        trn_corr += batch_corr
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
    train_losses.append(loss)
    train_correct.append(trn_corr)

    with torch.no_grad():
        for b, (X_test, y_test) in enumerate(test_loader):
            if b == 300:
                break
            y_val = CNNmodel(X_test)
            predicted = torch.max(y_val.data, 1)[1] 
            tst_corr += (predicted == y_test).sum()
    
    test_losses.append(loss)
    test_correct.append(tst_corr)

print(f'\nDuration: {time.time() - start_time:.0f} seconds')

4. 使用预训练模型

        除了自定义的卷积神经网络外,我们还可以使用预训练的AlexNet模型。通过冻结其卷积层参数,我们只训练模型的全连接层。

AlexNetmodel = models.alexnet(pretrained=True)

for param in AlexNetmodel.parameters():
    param.requires_grad = False

AlexNetmodel.classifier = nn.Sequential(nn.Linear(9216, 1024),
                                 nn.ReLU(),
                                 nn.Dropout(0.4),
                                 nn.Linear(1024, 2),
                                 nn.LogSoftmax(dim=1))

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(AlexNetmodel.classifier.parameters(), lr=0.001)

        我们只需训练一轮,因为卷积层的参数已经经过ImageNet数据集的预训练。

for i in range(1):
    trn_corr = 0
    tst_corr = 0
    
    for b, (X_train, y_train) in enumerate(train_loader):
        if b == 800:
            break
        y_pred = AlexNetmodel(X_train)
        loss = criterion(y_pred, y_train)
        predicted = torch.max(y_pred.data, 1)[1]
        batch_corr = (predicted == y_train).sum()
        trn_corr += batch_corr
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    with torch.no_grad():
        for b, (X_test, y_test) in enumerate(test_loader):
            if b == 300:
                break
            y_val = AlexNetmodel(X_test)
            predicted = torch.max(y_val.data, 1)[1] 
            tst_corr += (predicted == y_test).sum()

5. 结果评估与模型推理

        训练完成后,我们可以评估模型的性能,并对新图片进行分类。通过可视化训练损失和准确率曲线,我们可以观察模型的训练效果。

plt.plot(train_losses, label='training loss')
plt.plot(test_losses, label='validation loss')
plt.title('Loss at the end of each epoch')
plt.legend();

plt.plot([t/80 for t in train_correct], label='training accuracy')
plt.plot([t/30 for t in test_correct], label='validation accuracy')
plt.title('Accuracy at the end of each epoch')
plt.legend();

        最后,我们可以用训练好的模型对新的图像进行预测:

AlexNetmodel.eval()
with torch.no_grad():
    new_pred = AlexNetmodel(test_data[2019][0].view(1,3,224,224)).argmax()
print(f'Predicted value: {new_pred.item()} {class_names[new_pred.item()]}')

结语

        在本篇博文中,我们展示了如何使用卷积神经网络(CNN)和预训练的AlexNet模型对猫狗图像进行分类。通过数据增强、模型构建和训练,我们能够有效地从图像中提取特征并进行准确的分类。自定义CNN模型展示了从头训练模型的过程,而使用预训练的AlexNet则展示了如何利用迁移学习加速模型训练并提升性能。

        通过这些实践,我们不仅展示了卷积神经网络的强大功能,还展示了如何借助预训练模型应对图像分类任务。在实际应用中,迁移学习尤其适用于数据有限但需要高准确率的场景,它能够大大缩短训练时间并提高模型的精度。

        本次案例为我们提供了深度学习在图像处理中的应用方法与经验,未来可以尝试更多高级模型和技术,如ResNet或Transformer,以进一步提升图像分类的效果。

如果你觉得这篇博文对你有帮助,请点赞、收藏、关注我,并且可以打赏支持我!

欢迎关注我的后续博文,我将分享更多关于人工智能、自然语言处理和计算机视觉的精彩内容。

谢谢大家的支持!

Logo

开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!

更多推荐