迁移学习和微调 代码实例
时间: 2025-01-01 13:20:06 浏览: 9
### 关于迁移学习和微调的代码实例
在深度学习领域,迁移学习是一种有效的方法来利用预训练模型的知识提升新任务上的表现。下面提供了一个基于 PyTorch 的简单例子,在 CIFAR-10 数据集上使用预训练的 ResNet-18 模型进行图像分类的任务。
#### 加载必要的库
```python
import torch
from torchvision import models, transforms, datasets
import torch.nn as nn
from torch.utils.data import DataLoader
```
#### 设置设备配置
```python
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"Using {device} device")
```
#### 下载并准备数据集
```python
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
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 = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
```
#### 获取预训练模型
这里选择了 ResNet-18 作为基础网络结构,并冻结除了最后一层外的所有参数以便执行特征提取[^2]。
```python
model = models.resnet18(pretrained=True).to(device)
for param in model.parameters():
param.requires_grad = False
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 10).to(device) # 修改全连接层适应新的类别数
```
#### 定义损失函数与优化器
为了简化起见,采用交叉熵损失以及随机梯度下降法来进行优化。
```python
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)
```
#### 训练过程
此部分展示了如何迭代地调整权重以最小化误差。
```python
epochs = 5
for epoch in range(epochs):
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 + 1}/{epochs}] Loss: {running_loss / len(train_loader)}")
print('Finished Training')
```
上述代码片段展示了一种典型的迁移学习流程——先加载一个已经在大规模数据集(如 ImageNet)上预先训练好的卷积神经网络;接着针对特定应用场景的需求修改其顶部几层的设计;最后仅对这些新增加的部分实施反向传播更新操作,从而达到快速收敛的目的。
阅读全文