在前几篇文章中,我们学习了如何使用全连接神经网络和卷积神经网络(CNN)解决手写数字识别问题。然而,在实际应用中,我们常常需要处理更复杂的图像分类任务,例如识别猫狗、车辆、植物等。对于这些任务,从头训练一个深度学习模型通常需要大量的数据和计算资源。
为了解决这个问题,本文将介绍 迁移学习(Transfer Learning),并展示如何使用预训练模型快速解决图像分类问题。
一、迁移学习简介
迁移学习是一种将在一个任务上学到的知识迁移到另一个相关任务上的技术。在深度学习中,迁移学习通常通过使用预训练模型来实现。
1. 什么是预训练模型?
预训练模型是在大规模数据集(如 ImageNet)上训练好的模型。这些模型已经学会了提取图像特征的通用能力,可以直接用于新的任务,或者通过微调(Fine-tuning)适应新的数据集。
2. 迁移学习的优势
- 节省时间:无需从头训练模型。
- 节省数据:即使在小数据集上也能取得不错的效果。
- 提高性能:预训练模型通常具有更强的特征提取能力。
3. 迁移学习的常见方法
1.特征提取:将预训练模型作为特征提取器,只训练最后的分类层。
2.微调:解冻预训练模型的部分或全部层,并在新数据集上重新训练。
二、使用迁移学习实现图像分类
我们将使用 PyTorch 提供的预训练模型 ResNet18,并在 CIFAR-10 数据集 上进行迁移学习。
1. 问题描述
CIFAR-10 数据集包含 60,000 张 32x32 的彩色图像,分为 10 个类别(飞机、汽车、鸟、猫、鹿、狗、青蛙、马、船、卡车)。我们的目标是使用迁移学习对这些图像进行分类。
2. 实现步骤
- 加载和预处理数据。
- 加载预训练模型并修改最后的分类层。
- 定义损失函数和优化器。
- 训练模型。
- 测试模型并评估性能。
3. 代码实现
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
import matplotlib.pyplot as plt
# 设置 Matplotlib 支持中文显示
plt.rcParams['font.sans-serif'] = ['SimHei'] # 设置字体为 SimHei(黑体)
plt.rcParams['axes.unicode_minus'] = False # 解决负号显示问题
# 1. 加载和预处理数据
transform = transforms.Compose([
transforms.Resize((224, 224)), # 调整图像大小以适应 ResNet 输入
transforms.ToTensor(), # 将图像转换为张量
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) # 标准化
])
# 下载并加载 CIFAR-10 数据集
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)
# 2. 加载预训练模型并修改最后的分类层
model = models.resnet18(pretrained=True) # 加载预训练的 ResNet18 模型
# 冻结所有卷积层的参数
for param in model.parameters():
param.requires_grad = False
# 替换最后的全连接层
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 10) # CIFAR-10 有 10 个类别
# 3. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss() # 交叉熵损失
optimizer = optim.Adam(model.fc.parameters(), lr=0.001) # 只优化最后的全连接层
# 4. 训练模型
num_epochs = 5
loss_history = []
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_loader):
# 前向传播
outputs = model(images)
loss = criterion(outputs, labels)
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 记录损失
if (i + 1) % 100 == 0:
loss_history.append(loss.item())
print(f"Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{len(train_loader)}], Loss: {loss.item():.4f}")
# 5. 测试模型
model.eval() # 设置模型为评估模式
correct = 0
total = 0
with torch.no_grad():
for images, labels in test_loader:
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f"测试集准确率: {100 * correct / total:.2f}%")
# 6. 可视化损失曲线
plt.plot(loss_history)
plt.xlabel("训练步数")
plt.ylabel("损失值")
plt.title("训练损失曲线")
plt.show()
三、代码解析
1.数据加载与预处理:
- 使用 torchvision.datasets.CIFAR10 加载 CIFAR-10 数据集。
- 使用 transforms.Resize 调整图像大小以适应 ResNet 的输入(224x224)。
- 使用 transforms.Normalize 对图像进行标准化。
2.加载预训练模型:
- 使用 torchvision.models.resnet18 加载预训练的 ResNet18 模型。
- 冻结所有卷积层的参数,只训练最后的全连接层。
3.训练过程:
- 使用交叉熵损失函数和 Adam 优化器。
- 训练 5 个 epoch,并记录损失值。
4.测试过程:
- 在测试集上评估模型性能,计算准确率。
5.可视化:
- 绘制训练损失曲线。
四、运行结果
运行上述代码后,你将看到以下输出:
- 训练过程中每 100 步打印一次损失值。
- 测试集准确率(通常在 80% 以上)。
- 训练损失曲线图。
五、总结
本文介绍了迁移学习的基本概念,并使用 PyTorch 提供的预训练模型 ResNet18 在 CIFAR-10 数据集上实现了图像分类。通过迁移学习,我们能够在小数据集上快速构建高性能的深度学习模型。
在下一篇文章中,我们将学习如何使用生成对抗网络(GAN)生成图像。敬请期待!
代码实例说明:
- 本文代码可以直接在 Jupyter Notebook 或 Python 脚本中运行。
- 如果你有 GPU,可以将模型和数据移动到 GPU 上运行,例如:model = model.to('cuda'),images = images.to('cuda')。
希望这篇文章能帮助你更好地理解迁移学习的应用!如果有任何问题,欢迎在评论区留言讨论。