迁移学习是一种常见的机器学习技术,它利用在一个任务上训练好的模型,在另一个相关任务上进行微调。微调是指在新任务上对预训练模型进行一定程度的调整,以适应新任务的特点。在这个教程中,我们将学习如何使用PyTorch进行迁移学习和微调。
pip install torch torchvision
import torch
import torchvision.models as models
model = models.resnet18(pretrained=True)
num_classes = 10 # 新任务的类别数
model.fc = torch.nn.Linear(model.fc.in_features, num_classes)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
# 微调模型
model.train()
for epoch in range(num_epochs):
for images, labels in train_loader:
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
model.eval()
correct = 0
total = 0
with torch.no_grad():
for images, labels in val_loader:
outputs = model(images)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
accuracy = correct / total
print('Validation accuracy: {:.2f}%'.format(100 * accuracy))
通过以上步骤,我们完成了使用PyTorch进行迁移学习和微调的教程。希朇对你有所帮助!