基于Python的人工智能应用案例系列(10):CNN猫狗分类
在本篇博文中,我们展示了如何使用卷积神经网络(CNN)和预训练的AlexNet模型对猫狗图像进行分类。通过数据增强、模型构建和训练,我们能够有效地从图像中提取特征并进行准确的分类。自定义CNN模型展示了从头训练模型的过程,而使用预训练的AlexNet则展示了如何利用迁移学习加速模型训练并提升性能。通过这些实践,我们不仅展示了卷积神经网络的强大功能,还展示了如何借助预训练模型应对图像分类任务。在实际
在本篇文章中,我们将使用卷积神经网络(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,以进一步提升图像分类的效果。
如果你觉得这篇博文对你有帮助,请点赞、收藏、关注我,并且可以打赏支持我!
欢迎关注我的后续博文,我将分享更多关于人工智能、自然语言处理和计算机视觉的精彩内容。
谢谢大家的支持!
开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!
更多推荐
所有评论(0)