迁移学习是一种机器学习方法,允许我们利用在一个任务上获得的知识来加速在另一个任务上的学习。使用迁移学习,我们可以在小样本数据集上训练一个模型,尤其是在计算资源有限或者数据稀缺的情况下。
为什么使用迁移学习?
- 减少训练时间: 预训练的模型已经学习到了一些重要的特征,使用这些特征可以加快训练速度。
- 改善性能: 在数据量较小的情况下,利用预训练模型通常能获得更好的性能。
- 节省计算资源: 从头开始训练一个深度学习模型需要大量的计算资源和时间。
迁移学习的基本步骤
- 选择预训练模型: PyTorch提供了多种预训练模型,例如
ResNet
、VGG
、Inception
等。
- 修改模型结构: 根据新任务的需要,修改最后的几层网络结构(通常是全连接层)。
- 定义损失函数和优化器: 选择适合于新任务的损失函数和优化器。
- 微调模型: 在新数据集上对模型进行训练,通常采用较小的学习率。
- 评估模型: 使用验证集或测试集评估模型性能。
使用 PyTorch 进行迁移学习
1. 导入必要的库
1 2 3 4 5
| import torch import torchvision import torchvision.transforms as transforms import torch.nn as nn import torch.optim as optim
|
2. 准备数据集
假设我们要进行图像分类任务,可以使用 torchvision
来加载数据集。
1 2 3 4 5 6 7 8 9 10 11 12
| transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), ])
train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
test_dataset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=False)
|
3. 加载预训练模型
这里选择 ResNet18
作为预训练模型,并替换最后的全连接层。
1 2 3 4 5 6 7 8 9 10
| model = torchvision.models.resnet18(pretrained=True)
num_ftrs = model.fc.in_features model.fc = nn.Linear(num_ftrs, 10)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device)
|
4. 定义损失函数和优化器
1 2 3
| criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001)
|
5. 训练模型
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23
| num_epochs = 10
for epoch in range(num_epochs): model.train() running_loss = 0.0 for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs) loss = criterion(outputs, labels)
loss.backward() optimizer.step()
running_loss += loss.item() print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(train_loader):.4f}")
|
6. 评估模型
1 2 3 4 5 6 7 8 9 10 11 12 13
| model.eval() correct = 0 total = 0
with torch.no_grad(): for inputs, labels in test_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item()
print(f'Accuracy of the model on the test images: {100 * correct / total:.2f}%')
|
注意事项
- 在进行迁移学习时,推荐使用较小的学习率。
- 可以选择“冷冻”某些层,只训练最后的全连接层(Freeze the earlier layers)。
- 微调时可以从较小的 batch size 开始,根据显存情况调整。
结论
通过上述步骤,我们学习了如何在 PyTorch
中利用迁移学习进行图像分类任务。这种方式特别适用于数据量不足的场景,利用预训练模型能够显著提高模型的性能。